Update app.py

app.py

@@ -6,9 +6,35 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
 from typing import Tuple
 
+# Updated function to handle different models
+def load_and_prepare_model(model_choice: str):
+    if model_choice == "RealVisXL_V5.0_Lightning":
+        model_id = "SG161222/RealVisXL_V5.0_Lightning"
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            use_safetensors=True,
+            add_watermarker=False,
+        ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
+    else:  # Default is CompVis/stable-diffusion-v-1-4-original
+        model_id = "CompVis/stable-diffusion-v-1-4-original"
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            use_safetensors=True,
+            add_watermarker=False,
+        ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
+    
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    return pipe
+
+# Initial model setup
+model_choice = "CompVis/stable-diffusion-v-1-4-original"
+model = load_and_prepare_model(model_choice)
+
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     styles = {
         "3840 x 2160": (
@@ -22,24 +48,6 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + negative
 
-def load_and_prepare_model(model_choice: str):
-    if model_choice == "RealVisXL_V5.0_Lightning":
-        model_id = "SG161222/RealVisXL_V5.0_Lightning"
-    elif model_choice == "Stable Diffusion 2":
-        model_id = "stabilityai/stable-diffusion-2"
-    else:
-        raise ValueError("Invalid model choice")
-
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        use_safetensors=True,
-        add_watermarker=False,
-    ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
-    
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    return pipe
-
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, np.iinfo(np.int32).max)
@@ -61,9 +69,11 @@ def generate(
         num_inference_steps: int = 25,
         randomize_seed: bool = False,
 ):
-    # Load the chosen model
-    model = load_and_prepare_model(model_choice)
+    global model
 
+    # Load the selected model dynamically
+    model = load_and_prepare_model(model_choice)
+    
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=model.device).manual_seed(seed)
 
@@ -90,28 +100,53 @@ with gr.Blocks(theme="soft") as demo:
     
     with gr.Row():
         with gr.Column(scale=3):
-            # Dropdown to select model
-            model_choice = gr.Dropdown(
-                choices=["RealVisXL_V5.0_Lightning", "Stable Diffusion 2"],
-                label="Select Model",
-                value="RealVisXL_V5.0_Lightning",  # Default value
+            # Dropdown for selecting model
+            model_dropdown = gr.Dropdown(
+                label="Choose Model",
+                choices=["RealVisXL_V5.0_Lightning", "CompVis/stable-diffusion-v-1-4-original"],
+                value="CompVis/stable-diffusion-v-1-4-original",  # Default model
             )
-            
+
+            # Input Prompt Textbox
             prompt = gr.Textbox(
                 label="Input Prompt",
                 placeholder="Describe the image you want to create",
                 lines=2,
             )
-            run_button = gr.Button("Generate Image")
             
-            gr.Markdown("Developed using the RealVisXL_V5.0_Lightning model.", elem_id="model_info")
+            # Generate Image Button
+            run_button = gr.Button("Generate Image")
+
+            # Information about model
+            gr.Markdown("Developed using the RealVisXL_V5.0_Lightning or CompVis/stable-diffusion-v-1-4-original model.", elem_id="model_info")
+        
         with gr.Column(scale=7):
-            result = gr.Image(label="Generated Image", type="filepath")  # Corrected type to 'filepath'
+            # Display Generated Images (a single image in this case)
+            result = gr.Image(label="Generated Image", type="file")
 
+    # Set the callback for the button
     run_button.click(
         fn=generate,
-        inputs=[model_choice, prompt],
+        inputs=[model_dropdown, prompt],
         outputs=[result],
     )
 
+    # Footer with custom style and text
+    gr.Markdown("""
+    <style>
+    .footer {
+        position: relative;
+        left: 0;
+        bottom: 0;
+        width: 100%;
+        background-color: white;
+        color: black;
+        text-align: center;
+    }
+    </style>
+    <div class="footer">
+    <p>Developed with ❤ by Aklavya(Bucky)</p>
+    </div>
+    """)
+
 demo.launch()