Update app.py

app.py

@@ -22,31 +22,33 @@ import numpy as np
 # Configuration Section (Modify here to expand)
 # ======================
 
-# 1. Base Model - Illustrious XL v1.0 (Multiple options for stability)
+# 1. Base Model - Using reliable SDXL models
 BASE_MODELS = {
-    "illustrious_xl_v1": "Liberata/illustrious-xl-v1.0",  # Primary choice - more stable
-    "illustrious_xl_v2": "OnomaAIResearch/Illustrious-XL-v2.0",  # Backup option
-    "sdxl_base": "stabilityai/stable-diffusion-xl-base-1.0",  # Fallback to standard SDXL
-    "realistic_vision": "SG161222/RealVisXL_V4.0"  # Alternative realistic model
+    "sdxl_base": "stabilityai/stable-diffusion-xl-base-1.0",  # Primary - most stable
+    "realistic_vision": "SG161222/RealVisXL_V4.0",  # Alternative realistic model
+    "juggernaut": "RunDiffusion/Juggernaut-XL-v9",  # Popular realistic model
+    "dreamshaper": "Lykon/dreamshaper-xl-1-0"  # Versatile model
 }
 
-# Current model selection (change this to switch models)
+# Current model selection
 CURRENT_MODEL_KEY = "sdxl_base"
 BASE_MODEL = BASE_MODELS[CURRENT_MODEL_KEY]
 
-# 2. Fixed LoRAs (Auto-loaded, not user-selectable) - Using reliable SDXL LoRAs
+# 2. Fixed LoRAs (Auto-loaded, not user-selectable) - Using actual working LoRAs
 FIXED_LORAS = {
     "quality_enhancer": {
-        "repo_id": "ByteDance/SDXL-Lightning",
-        "filename": None,
+        "repo_id": "stabilityai/stable-diffusion-xl-refiner-1.0",
+        "subfolder": None,
         "weight": 0.6,
-        "trigger_words": "high quality, detailed, masterpiece"
+        "trigger_words": "high quality, detailed, masterpiece",
+        "enabled": False  # Disabled for now due to compatibility
     },
-    "pose_control": {
+    "consistency_model": {
         "repo_id": "latent-consistency/lcm-lora-sdxl",
-        "filename": None,
-        "weight": 0.5,
-        "trigger_words": "perfect anatomy, natural pose"
+        "subfolder": None,
+        "weight": 0.3,
+        "trigger_words": "lcm style",
+        "enabled": False  # Optional, can cause issues
     }
 }
 
@@ -59,7 +61,7 @@ STYLE_PROMPTS = {
     "Watercolor": "watercolor painting, soft brush strokes, translucent layers, artistic, painterly, paper texture, traditional art, masterpiece, ",
 }
 
-# 4. Optional LoRAs (User-selectable via dropdown, can select multiple) - Using verified SDXL LoRAs
+# 4. Optional LoRAs (User-selectable) - Using verified working LoRAs
 OPTIONAL_LORAS = {
     "None": {
         "repo_id": None,
@@ -67,35 +69,35 @@ OPTIONAL_LORAS = {
         "trigger_words": "",
         "description": "No additional LoRA"
     },
-    "Lightning Speed": {
-        "repo_id": "ByteDance/SDXL-Lightning",
-        "weight": 0.8,
-        "trigger_words": "high quality, detailed, fast generation",
-        "description": "Fast generation with quality enhancement"
-    },
-    "LCM LoRA": {
+    "LCM Speed": {
         "repo_id": "latent-consistency/lcm-lora-sdxl",
-        "weight": 0.9,
-        "trigger_words": "lcm, high quality, detailed",
-        "description": "Latent Consistency Model for faster inference"
+        "weight": 0.8,
+        "trigger_words": "lcm style, high quality",
+        "description": "Faster generation with LCM"
     },
-    "Offset Noise": {
-        "repo_id": "stabilityai/stable-diffusion-xl-offset-example-lora",
+    "Realistic Detail": {
+        "repo_id": "nerijs/pixel-art-xl",
         "weight": 0.7,
-        "trigger_words": "high contrast, dramatic lighting",
-        "description": "Enhanced contrast and lighting effects"
+        "trigger_words": "detailed, sharp, realistic",
+        "description": "Enhanced detail and realism"
     },
-    "AnimeDiff": {
-        "repo_id": "Linaqruf/anime-detailer-xl-lora",
+    "Anime Style": {
+        "repo_id": "Linaqruf/anime-detailer-xl-lora", 
         "weight": 0.8,
-        "trigger_words": "anime style, detailed anime, cel shading",
-        "description": "Anime and manga style enhancement"
+        "trigger_words": "anime style, detailed anime",
+        "description": "Anime and manga style"
     },
-    "PhotoReal": {
-        "repo_id": "Jovian-Experiments/photorealistic-xl-lora",
+    "Portrait Focus": {
+        "repo_id": "TheLastBen/Papercut_SDXL",
         "weight": 0.9,
-        "trigger_words": "photorealistic, ultra realistic, 8k uhd",
-        "description": "Photorealistic enhancement"
+        "trigger_words": "portrait, detailed face, beautiful eyes",
+        "description": "Portrait enhancement"
+    },
+    "Artistic Style": {
+        "repo_id": "ostris/ikea-instructions-lora-sdxl",
+        "weight": 0.6,
+        "trigger_words": "artistic style, creative",
+        "description": "Artistic and creative effects"
     }
 }
 
@@ -123,24 +125,50 @@ current_loras = {}
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_pipeline():
-    """Load the base Illustrious XL pipeline"""
+    """Load the base Illustrious XL pipeline with fallback options"""
     global pipe
     if pipe is None:
-        print("🚀 Loading Illustrious XL base model...")
-        pipe = StableDiffusionXLPipeline.from_pretrained(
-            BASE_MODEL,
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            variant="fp16"
-        ).to(device)
+        print(f"🚀 Loading base model: {BASE_MODEL}...")
         
-        # Enable memory optimizations for ZeroGPU
-        pipe.enable_attention_slicing()
-        pipe.enable_vae_slicing()
-        pipe.enable_model_cpu_offload()
-        pipe.enable_xformers_memory_efficient_attention()
+        # Try to load the selected model with fallback options
+        model_loaded = False
+        models_to_try = [BASE_MODEL]
         
-        print("✅ Illustrious XL model loaded successfully.")
+        # Add fallback models if primary fails
+        if CURRENT_MODEL_KEY != "sdxl_base":
+            models_to_try.append(BASE_MODELS["sdxl_base"])
+        if CURRENT_MODEL_KEY != "realistic_vision":
+            models_to_try.append(BASE_MODELS["realistic_vision"])
+            
+        for model_id in models_to_try:
+            try:
+                print(f"Attempting to load: {model_id}")
+                pipe = StableDiffusionXLPipeline.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.float16,
+                    use_safetensors=True,
+                    variant="fp16"
+                ).to(device)
+                
+                # Enable memory optimizations for ZeroGPU
+                pipe.enable_attention_slicing()
+                pipe.enable_vae_slicing()
+                if hasattr(pipe, 'enable_model_cpu_offload'):
+                    pipe.enable_model_cpu_offload()
+                if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+                    pipe.enable_xformers_memory_efficient_attention()
+                
+                print(f"✅ Successfully loaded: {model_id}")
+                model_loaded = True
+                break
+                
+            except Exception as e:
+                print(f"❌ Failed to load {model_id}: {e}")
+                continue
+        
+        if not model_loaded:
+            raise Exception("Failed to load any model. Please check your configuration.")
+            
     return pipe
 
 def unload_pipeline():
@@ -159,7 +187,7 @@ def unload_pipeline():
         print("🗑️ Pipeline unloaded.")
 
 def load_lora_weights(lora_configs: List[Dict]):
-    """Load multiple LoRA weights efficiently"""
+    """Load multiple LoRA weights efficiently with error handling"""
     global pipe, current_loras
     
     if not lora_configs:
@@ -174,33 +202,46 @@ def load_lora_weights(lora_configs: List[Dict]):
         except:
             pass
     
-    # Load new LoRAs
+    # Load new LoRAs with better error handling
     adapter_names = []
     adapter_weights = []
     
     for config in lora_configs:
         if config['repo_id'] and config['repo_id'] not in current_loras:
             try:
+                # Try different loading methods
+                adapter_name = config['name'].replace(' ', '_').lower()
+                
+                # Method 1: Direct loading
                 pipe.load_lora_weights(
                     config['repo_id'], 
-                    adapter_name=config['name']
+                    adapter_name=adapter_name
                 )
-                current_loras[config['repo_id']] = config['name']
+                current_loras[config['repo_id']] = adapter_name
                 print(f"✅ Loaded LoRA: {config['name']}")
+                
             except Exception as e:
-                print(f"❌ Failed to load LoRA {config['name']}: {e}")
+                print(f"⚠️ Failed to load LoRA {config['name']}: {e}")
+                # Skip this LoRA and continue with others
                 continue
         
-        if config['repo_id']:
-            adapter_names.append(config['name'])
+        # Add to active adapters if successfully loaded
+        if config['repo_id'] in current_loras:
+            adapter_names.append(current_loras[config['repo_id']])
             adapter_weights.append(config['weight'])
     
-    # Set adapter weights
+    # Set adapter weights if any adapters loaded
     if adapter_names:
         try:
             pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+            print(f"✅ Activated {len(adapter_names)} LoRA adapters")
         except Exception as e:
             print(f"⚠️ Warning setting adapter weights: {e}")
+            # Try without weights
+            try:
+                pipe.set_adapters(adapter_names)
+            except:
+                print("❌ Failed to set any adapters")
 
 def process_long_prompt(prompt: str, max_length: int = 77) -> str:
     """Process long prompts by intelligent truncation and optimization"""