Update app.py

app.py

@@ -7,37 +7,40 @@ from diffusers import DiffusionPipeline
 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
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
 else:
     torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
+def infer(model_id: Optional[str] = "CompVis/stable-diffusion-v1-4",
+          prompt: str = "",
+          negative_prompt: str = "",
+          seed: Optional[int] = 42,
+          randomize_seed: bool = True,
+          width: int = 1024,
+          height: int = 1024,
+          guidance_scale: Optional[float] = 7,
+          num_inference_steps: Optional[int] = 20,
+          progress = gr.Progress(track_tqdm=True)):
+    if model_id:
+        model_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
 
+    pipe = DiffusionPipeline.from_pretrained(
+        pretrained_model_name_or_path=model_id,
+        torch_dtype=torch_dtype)
+    pipe = pipe.to(device)
+
     image = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -66,8 +69,17 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+        gr.Markdown(" # Text-to-Image Gradio Form")
 
+        with gr.Row():
+            model_id = gr.Dropdown(
+                choices=["stabilityai/sdxl-turbo", "CompVis/stable-diffusion-v1-4"],
+                multiselect=False,
+                allow_custom_value=True,
+                label="Model",
+                info="Choose model ID",
+            )
+        
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -138,6 +150,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
+            model_id,
             prompt,
             negative_prompt,
             seed,