Update app.py

app.py

@@ -1,12 +1,7 @@
 import gradio as gr
 import torch
 import functools
-
-from diffusers import (
-    StableDiffusionPipeline,
-    StableDiffusionXLPipeline,
-    DPMSolverMultistepScheduler,
-)
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
 MODEL_OPTS = {
     "SD v1.5 (base)":        "runwayml/stable-diffusion-v1-5",
@@ -22,50 +17,32 @@ DTYPE = torch.float16 if DEVICE != "cpu" else torch.float32
 
 @functools.lru_cache(maxsize=len(MODEL_OPTS))
 def get_pipeline(model_id: str):
-    # Use the SDXL pipeline class whenever the model_id is the XL Base
-    if "stable-diffusion-xl-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)
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=DTYPE,
+        safety_checker=None
+    ).to(DEVICE)
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
     return pipe
 
 def generate(prompt, steps, guidance, seed, model_name):
     model_id = MODEL_OPTS[model_name]
-    # ultra-fast “Turbo” runs only 4 steps
     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))
-
-    # ALWAYS pass dicts so SDXL's UNet doesn't see None
-    output = pipe(
+    imgs = pipe(
         prompt,
         num_inference_steps=int(steps),
         guidance_scale=float(guidance),
-        generator=generator,
-        cross_attention_kwargs={},   # never None
-        added_cond_kwargs={},        # never None for SDXL
-    )
-    return output.images
+        generator=generator
+    ).images
+    return imgs
 
 with gr.Blocks() as demo:
-    gr.Markdown("Model-Switcher Stable Diffusion 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")
+    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(0, label="Seed (0=random)")
@@ -75,9 +52,8 @@ with gr.Blocks() as demo:
     btn.click(
         fn=generate,
         inputs=[prompt, steps, guidance, seed, checkpoint],
-        outputs=gallery,
+        outputs=gallery
     )
 
 if __name__ == "__main__":
     demo.launch()
-