Update app.py

app.py

@@ -21,13 +21,18 @@ pipe = pipe.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def load_model(model_repo_id):
+    global pipe
+    pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype, use_cache=False)
+    pipe = pipe.to(device)
+
 
 @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
-    # model_id: Optional[str] = "CompVis/stable-diffusion-v1-4",
+    model_id = "CompVis/stable-diffusion-v1-4",
     negative_prompt = 'cat, dog',
-    prompt: str = 'cute_animal',
-    seed: int = 42,
+    prompt = 'cute_animal',
+    seed = 42,
     randomize_seed = False,
     width = 128,
     height = 128,
@@ -35,6 +40,8 @@ def infer(
     num_inference_steps = 20,
     progress=gr.Progress(track_tqdm=True),
 ):
+    load_model(model_id) 
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -70,6 +77,13 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # Text-to-Image Gradio Template")
 
+        with gr.Row():
+            model_id = gr.Dropdown(
+                label="Model",
+                choices=["CompVis/stable-diffusion-v1-4", "stabilityai/sdxl-turbo"],
+                value="CompVis/stable-diffusion-v1-4",
+            )
+        
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -140,6 +154,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
+            model_id,
             prompt,
             negative_prompt,
             seed,