Update app.py

app.py

@@ -6,23 +6,28 @@ import random
 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
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+available_models = [
+    "stabilityai/sdxl-turbo",
+    "CompVis/stable-diffusion-v1-4",
+    "runwayml/stable-diffusion-v1-5",
+]
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def init_model(model_repo_id):
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    return pipe.to(device)
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
+    model_repo_id, 
     prompt,
     negative_prompt,
     seed,
@@ -33,6 +38,8 @@ def infer(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    pipe = init_model(model_repo_id)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -82,6 +89,11 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
+            model_repo_id = gr.Dropdown(available_models, 
+                                        multiselect=False,
+                                        label="Model",
+                                        info="Choose models for generation")
+            
             negative_prompt = gr.Text(
                 label="Negative prompt",
                 max_lines=1,
@@ -138,6 +150,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
+            model_repo_id,
             prompt,
             negative_prompt,
             seed,