Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 import torch
 
 # --- Configuration ---
-HF_REPO_ID = "aanchal77/Final-One" 
+HF_REPO_ID = "aanchal77/Final-One"
 BASE_MODEL_ID = "runwayml/stable-diffusion-v1-5"
 
 # --- Define Available LoRAs ---
@@ -31,14 +31,20 @@ dtype = torch.float16 if device == "cuda" else torch.float32
 print(f"Using device: {device}")
 
 print(f"🎨 Loading base model: {BASE_MODEL_ID}")
-pipe = StableDiffusionPipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=dtype).to(device)
+pipe = StableDiffusionPipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=dtype)
+
+# 🔧 Replace the fragile PNDM scheduler with a robust one to avoid index/NoneType errors
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+
+pipe = pipe.to(device)
 if device == "cpu":
     pipe.enable_attention_slicing()
 
 # --- The Core Generation Function ---
 def generate(prompt, quality, lora_choice):
+    # Reset to base weights
     pipe.unload_lora_weights()
-    
+
     lora_subfolder = AVAILABLE_LORAS.get(lora_choice)
 
     if lora_subfolder:
@@ -47,7 +53,7 @@ def generate(prompt, quality, lora_choice):
             pipe.load_lora_weights(
                 HF_REPO_ID,
                 subfolder=lora_subfolder,
-                weight_name="adapter_model.safetensors"  # ✅ Explicit LoRA file
+                weight_name="adapter_model.safetensors"  # ensure exact file
             )
         except Exception as e:
             print(f"❌ Failed to load LoRA from Hub '{HF_REPO_ID}/{lora_subfolder}': {e}")
@@ -56,16 +62,23 @@ def generate(prompt, quality, lora_choice):
 
     steps = 25 if quality == "Fast" else 40
     guidance_scale = 7.5
-    
+
     print(f"🚀 Generating with prompt: '{prompt}'")
     with torch.no_grad():
-        image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
-        
+        image = pipe(
+            prompt,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale
+        ).images[0]
+
     return image
 
 # --- Build the Gradio UI ---
 title = f"🎨 Stable Diffusion Gallery from {HF_REPO_ID}"
-description = "Select a trained LoRA model from your Hugging Face repository to apply its style. The first time you select a LoRA, it may take a moment to download."
+description = (
+    "Select a trained LoRA model from your Hugging Face repository to apply its style. "
+    "The first time you select a LoRA, it may take a moment to download."
+)
 
 demo = gr.Interface(
     fn=generate,
@@ -74,7 +87,7 @@ demo = gr.Interface(
         gr.Dropdown(["Fast", "High Quality"], value="Fast", label="Generation Quality"),
         gr.Dropdown(
             choices=list(AVAILABLE_LORAS.keys()),
-            value="None (Base Model)", 
+            value="None (Base Model)",
             label="Select a Trained LoRA Model"
         )
     ],
@@ -84,7 +97,8 @@ demo = gr.Interface(
     examples=[
         ["A portrait of an astronaut, cinematic lighting, by vincent van gogh", "Fast", "Artist: Vincent van Gogh"],
         ["A peaceful village in the mountains, impressionism style", "High Quality", "Style: Impressionism"],
-    ]
+    ],
+    cache_examples=False,   # 🔒 prevent startup 500s if an example errors
 )
 
 demo.launch()