Update app.py

app.py

@@ -4,10 +4,10 @@ import gradio as gr
 import os
 import spaces
 
+# Load the models outside of the generate_images function
 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:
@@ -17,12 +17,11 @@ 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:
+    pipe = models.get(model_name)
+    if pipe is None:
         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)
@@ -41,12 +40,10 @@ 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)
+            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()
+demo.launch()
+