Update app.py

app.py

@@ -1,24 +1,47 @@
-import gradio as gr
-from diffusers import DiffusionPipeline
-
-# Load the pre-trained model pipeline
-pipe = DiffusionPipeline.from_pretrained("ImageInception/ArtifyAI")
-
-# Define a function to generate the image based on user input text
-def generate_image(prompt):
-    # Generate the image from the prompt
-    image = pipe(prompt).images[0]
-    return image
-
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=generate_image, 
-    inputs="text", 
-    outputs="image", 
-    title="ArtifyAI Generator",
-    description="Enter a prompt and generate an image using the ArtifyAI model."
-)
-
-# Launch the Gradio app
-if __name__ == "__main__":
-    iface.launch()
+import gradio as gr
+from diffusers import DiffusionPipeline
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Pre-load the models but do not load them yet
+models = {
+    "ArtifyAI v1.1": "ImageInception/ArtifyAI-v1.1",
+    "ArtifyAI v1.0": "ImageInception/ArtifyAI-v1.0"
+}
+
+# Function to load the selected model
+def load_model(model_name):
+    return DiffusionPipeline.from_pretrained(models[model_name], torch_dtype=torch.float16 if device == "cuda" else torch.float32).to(device)
+
+# Initially load the first model
+pipe = load_model("ArtifyAI v1.1")
+
+def generate_image(prompt, selected_model):
+    global pipe
+    
+    # Load the selected model if it's not already loaded
+    if selected_model in models:
+        pipe = load_model(selected_model)
+    
+    if device == "cuda":
+        with torch.cuda.amp.autocast():
+            image = pipe(prompt).images[0]
+    else:
+        image = pipe(prompt).images[0]
+    return image
+
+# Gradio interface
+demo = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Enter your prompt"),
+        gr.Dropdown(choices=list(models.keys()), label="Select Model", value="ArtifyAI v1.1")
+    ],
+    outputs=gr.Image(type="pil"),
+    title="ArtifyAI",
+    description="Generate images using the selected model."
+)
+
+if __name__ == "__main__":
+    demo.launch()