Update app.py

app.py

@@ -2,10 +2,27 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load the Stable Diffusion model
-model_id = "runwayml/stable-diffusion-v1-5"
-pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipeline = pipeline.to("cuda")
+# Function to load the selected Stable Diffusion model
+def load_model(model_id):
+    """Load the specified Stable Diffusion model."""
+    pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+    pipeline = pipeline.to("cpu")
+    return pipeline
+
+# Initial model
+model_options = [
+    "CompVis/stable-diffusion-v1-1",
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    "stabilityai/stable-diffusion-2-1"
+]
+current_model_id = model_options[0]
+pipeline = load_model(current_model_id)
+
+def switch_model(selected_model):
+    """Switch the Stable Diffusion model."""
+    global pipeline
+    pipeline = load_model(selected_model)
+    return f"Model switched to: {selected_model}"
 
 def generate_image(prompt, num_inference_steps=50, guidance_scale=7.5):
     """Generate an image from a text prompt using Stable Diffusion."""
@@ -14,10 +31,17 @@ def generate_image(prompt, num_inference_steps=50, guidance_scale=7.5):
 
 # Define the Gradio app layout
 with gr.Blocks() as app:
-    gr.Markdown("# Stable Diffusion v1.5 Gradio App\nGenerate stunning images from text prompts!")
+    gr.Markdown("# Stable Diffusion Gradio App\nGenerate stunning images from text prompts and switch between models!")
 
     with gr.Row():
         with gr.Column():
+            model_selector = gr.Dropdown(
+                label="Select Model",
+                choices=model_options,
+                value=current_model_id,
+                interactive=True
+            )
+            model_switch_status = gr.Textbox(label="Model Status", value=f"Current model: {current_model_id}", interactive=False)
             prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here", lines=2)
             num_inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=50, step=5)
             guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, value=7.5, step=0.5)
@@ -26,6 +50,12 @@ with gr.Blocks() as app:
         with gr.Column():
             output_image = gr.Image(label="Generated Image", type="pil")
 
+    model_selector.change(
+        switch_model, 
+        inputs=[model_selector], 
+        outputs=model_switch_status
+    )
+
     generate_btn.click(
         generate_image, 
         inputs=[prompt, num_inference_steps, guidance_scale],