Upload 2 files

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+import torch
+import functools
+from diffusers import StableDiffusionPipeline, 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"
+}
+
+DEVICE = (
+    "mps"  if torch.backends.mps.is_available() else
+    "cuda" if torch.cuda.is_available()     else
+    "cpu"
+)
+DTYPE = torch.float16 if DEVICE != "cpu" else torch.float32
+
+@functools.lru_cache(maxsize=len(MODEL_OPTS))
+def get_pipeline(model_id: str):
+    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]
+    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,
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
+        generator=generator
+    ).images
+    return imgs
+
+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")
+    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)")
+    btn        = gr.Button("Generate")
+    gallery    = gr.Gallery(label="Gallery", columns=2, height="auto")
+
+    btn.click(
+        fn=generate,
+        inputs=[prompt, steps, guidance, seed, checkpoint],
+        outputs=gallery
+    )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.2
+diffusers>=0.28
+transformers>=4.42
+accelerate>=0.29
+safetensors
+gradio>=4.32