Update app.py

app.py

@@ -16,6 +16,12 @@ for model_name in model_list:
 
 @spaces.GPU
 def generate_images(model_name, prompt, negative_prompt, num_inference_steps, guidance_scale, num_images=4):
+    if model_name not in models:
+        try:
+            models[model_name] = DiffusionPipeline.from_pretrained(model_name).to("cuda")
+        except Exception as e:
+            print(f"Error loading model {model_name}: {e}")
+            return []
     pipe = models[model_name]
     outputs = []
     for _ in range(num_images):
@@ -23,6 +29,7 @@ def generate_images(model_name, prompt, negative_prompt, num_inference_steps, gu
         outputs.append(output)
     return outputs
 
+
 # Create the Gradio blocks
 with gr.Blocks() as demo:
     with gr.Row(equal_height=False):
@@ -40,6 +47,6 @@ with gr.Blocks() as demo:
         with gr.Column():
             output_gallery = gr.Gallery(label="Generated Images", height=480, scale=1)
 
-generate_btn.click(generate_images, inputs=[model_dropdown, prompt, negative_prompt, num_inference_steps, guidance_scale, num_images], outputs=output_gallery)
+    generate_btn.click(generate_images, inputs=[model_dropdown, prompt, negative_prompt, num_inference_steps, guidance_scale, num_images], outputs=output_gallery)
 
 demo.launch()