Update app.py

app.py

@@ -7,10 +7,10 @@ import spaces
 model_list = [model.strip() for model in os.environ.get("MODELS").split(",")]
 lora_list = [model.strip() for model in os.environ.get("LORAS").split(",")]
 
+# Load all models upfront
 models = {}
 for model_name in model_list:
     try:
-        print("Loading " + model_name)
         models[model_name] = DiffusionPipeline.from_pretrained(model_name).to("cuda")
     except Exception as e:
         print(f"Error loading model {model_name}: {e}")
@@ -18,18 +18,16 @@ 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 []
+        return []
+
     pipe = models[model_name]
     outputs = []
+
     for _ in range(num_images):
         output = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)["images"][0]
         outputs.append(output)
-    return outputs
 
+    return outputs
 
 # Create the Gradio blocks
 with gr.Blocks() as demo:
@@ -43,11 +41,12 @@ with gr.Blocks() as demo:
                     num_inference_steps = gr.Slider(minimum=10, maximum=50, step=1, value=25, label="Number of Inference Steps")
                     guidance_scale = gr.Slider(minimum=1, maximum=20, step=0.5, value=7.5, label="Guidance Scale")
                     num_images = gr.Slider(minimum=1, maximum=4, step=1, value=4, label="Number of Images")
+
                 generate_btn = gr.Button("Generate Image")
-                
-        with gr.Column():
-            output_gallery = gr.Gallery(label="Generated Images", height=480, scale=1)
+
+            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)
 
-demo.launch()
+demo.launch()