Update app.py

app.py

@@ -2,27 +2,50 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load the Stable Diffusion pipeline
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
-)
-pipe = pipe.to("cuda")  # Use GPU for faster generation
+# Load different models
+models = {
+    "Stable Diffusion v1.5": "runwayml/stable-diffusion-v1-5",
+    "Stable Diffusion v2.1": "stabilityai/stable-diffusion-2-1",
+    "Anime Diffusion": "hakurei/waifu-diffusion-v1-4",
+}
 
-def generate_image(prompt):
-    # Generate an image based on the text prompt
-    image = pipe(prompt).images[0]
+# Function to load the selected model
+def load_model(model_name):
+    model_id = models[model_name]
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id, torch_dtype=torch.float16
+    )
+    pipe = pipe.to("cuda")  # Use GPU
+    return pipe
+
+# Load the default model
+current_pipe = load_model("Stable Diffusion v1.5")
+
+# Function to generate image
+def generate_image(prompt, model_name):
+    global current_pipe
+    # Reload pipeline if the model changes
+    if model_name not in current_pipe.config["_name_or_path"]:
+        current_pipe = load_model(model_name)
+    # Generate the image
+    image = current_pipe(prompt).images[0]
     return image
 
-# Create the Gradio interface
+# Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("### Text-to-Image Generator")
+    gr.Markdown("### Multi-Model Text-to-Image Generator")
     with gr.Row():
         with gr.Column():
-            text_input = gr.Textbox(label="Enter a text prompt")
-            submit_button = gr.Button("Generate Image")
+            text_input = gr.Textbox(label="Enter a text prompt", placeholder="Describe the image you want...")
+            model_selector = gr.Dropdown(
+                label="Select Model", choices=list(models.keys()), value="Stable Diffusion v1.5"
+            )
+            generate_button = gr.Button("Generate Image")
         with gr.Column():
             output_image = gr.Image(label="Generated Image")
-    
-    submit_button.click(generate_image, inputs=text_input, outputs=output_image)
+
+    generate_button.click(
+        generate_image, inputs=[text_input, model_selector], outputs=output_image
+    )
 
 demo.launch()