Multiple models

app.py

@@ -9,6 +9,12 @@ import torch
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
 
+MODEL_CHOICES = [
+    "stabilityai/sdxl-turbo",
+    "stabilityai/sd-turbo",
+    "CompVis/stable-diffusion-v1-4",
+]
+
 if torch.cuda.is_available():
     torch_dtype = torch.float16
 else:
@@ -17,11 +23,13 @@ else:
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
+current_model_id = model_repo_id
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
+@spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     prompt,
     negative_prompt,
@@ -77,6 +85,12 @@ with gr.Blocks(css=css) as demo:
                 container=False,
             )
 
+            model_dropdown = gr.Dropdown(
+            label="Diffusion model",
+            choices=MODEL_CHOICES,
+            value=model_repo_id,
+            )
+
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)