Update app.py

app.py

@@ -19,14 +19,8 @@ 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]
-    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
+    output = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)["images"][0]
+    return output