Update app.py

app.py

@@ -7,16 +7,15 @@ import spaces
 model_list = os.environ.get("MODELS").split(",")
 lora_list = os.environ.get("LORAS") # Not in use
 
-@spaces.GPU
-# Load the available models and their pipelines
-models = []
+models = {}
 for model_name in model_list:
     try:
         models[model_name] = pipeline("text-to-image", model=model_name, torch_dtype=torch.float16).to("cuda")
     except Exception as e:
         print(f"Error loading model {model_name}: {e}")
 
-# Define the function to generate the image
+
+@spaces.GPU
 def generate_image(model_name, prompt, negative_prompt, num_inference_steps, guidance_scale):
     pipe = models[model_name]
     image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)["sample"][0]