Update app.py

app.py

@@ -3,30 +3,31 @@ from diffusers import DiffusionPipeline
 import torch
 import spaces
 
-# Load the model
-# The model will be downloaded and cached the first time the app runs.
-pipe = DiffusionPipeline.from_pretrained(
-    "Spestly/OdysseyXL-V2.5", 
-    torch_dtype=torch.float16, 
-    use_safetensors=True
-)
-# Move the pipeline to the GPU
-pipe.to("cuda")
+model_options = {
+    "OdysseyXL V2.5": "Spestly/OdysseyXL-V2.5",
+    "OdysseyXL V2": "Spestly/OdysseyXL-V2",
+    "OdysseyXL V1": "Spestly/OdysseyXL-V1",
+    "OdysseyXL Origin": "Spestly/OdysseyXL-Origin"
+}
 
-@spaces.GPU
-def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps):
-    """
-    Generates an image from a text prompt using the OdysseyXL V2.5 model.
+# Cache for loaded pipelines
+loaded_pipelines = {}
 
-    Args:
-        prompt (str): The text prompt to generate the image from.
-        negative_prompt (str): The negative text prompt.
-        guidance_scale (float): The guidance scale for the generation.
-        num_inference_steps (int): The number of inference steps.
+def load_model(model_name):
+    model_id = model_options[model_name]
+    if model_name not in loaded_pipelines:
+        pipe = DiffusionPipeline.from_pretrained(
+            model_id, 
+            torch_dtype=torch.float16,
+            use_safetensors=True
+        )
+        pipe.to("cuda")
+        loaded_pipelines[model_name] = pipe
+    return loaded_pipelines[model_name]
 
-    Returns:
-        PIL.Image.Image: The generated image.
-    """
+@spaces.GPU
+def generate_image(model_name, prompt, negative_prompt, guidance_scale, num_inference_steps):
+    pipe = load_model(model_name)
     image = pipe(
         prompt=prompt, 
         negative_prompt=negative_prompt,
@@ -35,14 +36,18 @@ def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps)
     ).images[0]
     return image
 
-# --- Gradio Interface ---
-
+# --- Gradio UI ---
 with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("# 🎨 OdysseyXL V2.5 Image Generation")
-    gr.Markdown("A Gradio UI for the [Spestly/OdysseyXL V2.5](https://huggingface.co/Spestly/OdysseyXL%20V2.5) SDXL model, optimized for ZeroGPU.")
-    
+    gr.Markdown("# 🎨 OdysseyXL Image Generation")
+    gr.Markdown("Choose from multiple OdysseyXL model versions running on ZeroGPU (H200).")
+
     with gr.Row():
         with gr.Column(scale=2):
+            model_name = gr.Dropdown(
+                label="Model Version",
+                choices=list(model_options.keys()),
+                value="OdysseyXL V2.5"
+            )
             prompt = gr.Textbox(
                 label="Prompt", 
                 show_label=False, 
@@ -70,27 +75,28 @@ with gr.Blocks(css="style.css") as demo:
                     value=30
                 )
             run_button = gr.Button("Generate Image", variant="primary")
+
         with gr.Column(scale=1):
             image_output = gr.Image(label="Generated Image", show_label=False)
 
     gr.Examples(
         examples=[
-            ["A futuristic cityscape, vibrant neon colors, ultra-realistic, 8K", "blurry, low quality", 7.5, 30],
-            ["A majestic lion with a crown of stars, cosmic background, fantasy art", "cartoon, sketch", 8.0, 40],
-            ["An enchanted forest at night, glowing mushrooms, fireflies, mystical atmosphere", "daytime, bright", 7.0, 35],
-            ["A delicious-looking gourmet burger on a wooden table, hyperrealistic food photography", "messy, unappetizing", 7.5, 25]
+            ["OdysseyXL V2.5", "A futuristic cityscape, vibrant neon colors, ultra-realistic, 8K", "blurry, low quality", 7.5, 30],
+            ["OdysseyXL V2.1", "A majestic lion with a crown of stars, cosmic background, fantasy art", "cartoon, sketch", 8.0, 40],
+            ["OdysseyXL V1.9", "An enchanted forest at night, glowing mushrooms, fireflies, mystical atmosphere", "daytime, bright", 7.0, 35],
+            ["OdysseyXL V1.5", "A delicious-looking gourmet burger on a wooden table, hyperrealistic food photography", "messy, unappetizing", 7.5, 25]
         ],
-        inputs=[prompt, negative_prompt, guidance_scale, num_inference_steps],
+        inputs=[model_name, prompt, negative_prompt, guidance_scale, num_inference_steps],
         outputs=image_output,
         fn=generate_image,
-        cache_examples=True,
+        cache_examples=True
     )
 
     run_button.click(
-        fn=generate_image, 
-        inputs=[prompt, negative_prompt, guidance_scale, num_inference_steps],
+        fn=generate_image,
+        inputs=[model_name, prompt, negative_prompt, guidance_scale, num_inference_steps],
         outputs=image_output
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()