:tada: init

app.py

@@ -2,6 +2,15 @@ import gradio as gr
 import spaces
 import torch
 from diffusers import DiffusionPipeline
+from PIL import Image
+
+# Global Variables
+current_model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+pipe = DiffusionPipeline.from_pretrained(
+    current_model_id,
+    torch_dtype=torch.float16,
+).to(device)
 
 
 @spaces.GPU()
@@ -11,13 +20,18 @@ def inference(
     prompt: str,
     negative_prompt: str = "",
     progress=gr.Progress(track_tqdm=True),
-):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+) -> Image.Image:
+    global current_model_id, pipe
 
-    pipe = DiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-    ).to(device)
+    if model_id != current_model_id:
+        try:
+            pipe = DiffusionPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.float16,
+            ).to(device)
+            current_model_id = model_id
+        except Exception as e:
+            raise gr.Error(str(e))
 
     image = pipe(
         prompt,
@@ -31,8 +45,14 @@ if __name__ == "__main__":
     demo = gr.Interface(
         fn=inference,
         inputs=[
-            gr.Text(
+            gr.Dropdown(
                 label="Model ID",
+                choices=[
+                    "stabilityai/stable-diffusion-3-medium-diffusers",
+                    "stabilityai/stable-diffusion-xl-base-1.0",
+                    "stabilityai/stable-diffusion-2-1",
+                    "runwayml/stable-diffusion-v1-5",
+                ],
                 value="stabilityai/stable-diffusion-3-medium-diffusers",
             ),
             gr.Text(label="Prompt", value=""),
@@ -42,4 +62,4 @@ if __name__ == "__main__":
             gr.Image(label="Image", type="pil"),
         ],
     )
-    demo.launch()
+    demo.queue().launch()