Update app.py

app.py

@@ -1,71 +1,74 @@
 import gradio as gr
 import torch
 import functools
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
-# Model options for switching between Stable Diffusion variants
+from diffusers import (
+    StableDiffusionPipeline,
+    StableDiffusionXLPipeline,
+    DPMSolverMultistepScheduler,
+)
+
 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
+    # Choose the correct pipeline class for SDXL
+    if "sdxl-base" in model_id:
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            model_id,
+            torch_dtype=DTYPE,
+            safety_checker=None,
+        )
+    else:
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=DTYPE,
+            safety_checker=None,
+        )
+
+    pipe = pipe.to(DEVICE)
+    # swap in the faster DPMSolver scheduler everywhere
     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 int(seed) == 0 else torch.manual_seed(int(seed))
+    generator = None if seed == 0 else torch.manual_seed(int(seed))
 
-    # Ensure added_cond_kwargs is always a dict to avoid NoneType errors
-    output = pipe(
-        prompt=prompt,
+    # explicitly pass empty dicts so SDXL's UNet never sees None
+    out = pipe(
+        prompt,
         num_inference_steps=int(steps),
         guidance_scale=float(guidance),
         generator=generator,
-        cross_attention_kwargs={},
-        added_cond_kwargs={},
+        cross_attention_kwargs={},       # always a dict
+        added_cond_kwargs={},            # always a dict (for SDXL)
     )
-    return output.images
+    return out.images
 
-# Build the Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## ✨ Model-Switcher Stable Diffusion Demo")
-    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")
+    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")
     steps      = gr.Slider(1, 50, value=30, label="Inference Steps")
     guidance   = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
-    seed       = gr.Number(value=0, label="Seed (0 = random)")
+    seed       = gr.Number(0, label="Seed (0=random)")
     btn        = gr.Button("Generate")
     gallery    = gr.Gallery(label="Gallery", columns=2, height="auto")
 
@@ -77,3 +80,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.launch()
+