Update app.py

app.py

@@ -19,8 +19,15 @@ for model_name in model_list:
 @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)["stabilityai/stable-diffusion-xl-base-1.0"][0]
+    output = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
+    print("All keys in the output:", output.keys())
+    if "stabilityai/stable-diffusion-xl-base-1.0" in output:
+        image = output["stabilityai/stable-diffusion-xl-base-1.0"][0]
+    else:
+        print("The key 'stabilityai/stable-diffusion-xl-base-1.0' does not exist in the output.")
+        image = None
     return image
+
     
 
 # Create the Gradio blocks