Update app.py

app.py

@@ -3,50 +3,68 @@ import torch
 import functools
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
+# Model options for switching between Stable Diffusion variants
 MODEL_OPTS = {
     "SD v1.5 (base)":        "runwayml/stable-diffusion-v1-5",
     "SDXL Base 1.0":         "stabilityai/stable-diffusion-xl-base-1.0",
     "SD-Turbo (ultra-fast)": "stabilityai/sd-turbo"
 }
 
+# Auto-detect compute device (GPU/CPU)
 DEVICE = (
     "mps"  if torch.backends.mps.is_available() else
     "cuda" if torch.cuda.is_available()     else
     "cpu"
 )
+# Use float16 on GPU/MPS for performance, float32 on CPU
 DTYPE = torch.float16 if DEVICE != "cpu" else torch.float32
 
 @functools.lru_cache(maxsize=len(MODEL_OPTS))
 def get_pipeline(model_id: str):
+    """
+    Lazily load and cache the Stable Diffusion pipeline for a given model.
+    """
     pipe = StableDiffusionPipeline.from_pretrained(
         model_id,
         torch_dtype=DTYPE,
         safety_checker=None
     ).to(DEVICE)
+    # Switch to a faster sampler
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
     return pipe
 
+
 def generate(prompt, steps, guidance, seed, model_name):
+    """
+    Generate images using the specified Stable Diffusion model.
+    """
     model_id = MODEL_OPTS[model_name]
+    # Limit steps for turbo model
     if "Turbo" in model_name:
         steps = min(int(steps), 4)
     pipe = get_pipeline(model_id)
-    generator = None if seed == 0 else torch.manual_seed(int(seed))
-    imgs = pipe(
-        prompt,
+    generator = None if int(seed) == 0 else torch.manual_seed(int(seed))
+
+    # Ensure cross_attention_kwargs is always a dict (avoids NoneType errors)
+    output = pipe(
+        prompt=prompt,
         num_inference_steps=int(steps),
         guidance_scale=float(guidance),
-        generator=generator
-    ).images
-    return imgs
+        generator=generator,
+        cross_attention_kwargs={},
+    )
+    return output.images
 
+# Build the Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## Model-Switcher Stable Diffusion Demo")
-    prompt     = gr.Textbox("Retro robot in neon city", label="Prompt")
-    checkpoint = gr.Dropdown(list(MODEL_OPTS.keys()), value="SD v1.5 (base)", label="Checkpoint")
+    prompt     = gr.Textbox(label="Prompt", value="Retro robot in neon city")
+    checkpoint = gr.Dropdown(choices=list(MODEL_OPTS.keys()),
+                            value="SD v1.5 (base)",
+                            label="Checkpoint")
     steps      = gr.Slider(1, 50, value=30, label="Inference Steps")
     guidance   = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
-    seed       = gr.Number(0, label="Seed (0=random)")
+    seed       = gr.Number(value=0, label="Seed (0 = random)")
     btn        = gr.Button("Generate")
     gallery    = gr.Gallery(label="Gallery", columns=2, height="auto")