Update app.py

app.py

@@ -14,7 +14,7 @@ sd_dreambooth_model_1.diffusion_model.load_weights("dreambooth_riffusion_model_c
 
 
 def generate_images(prompt: str, num_steps: int, unconditional_guidance_scale: int, temperature: int):
-    generated_img = sd_dreambooth_model_1.generate(
+    img = sd_dreambooth_model_1.generate(
         prompt, 
         num_steps=num_steps,
         unconditional_guidance_scale=unconditional_guidance_scale,
@@ -22,9 +22,9 @@ def generate_images(prompt: str, num_steps: int, unconditional_guidance_scale: i
         batch_size=1,
     )
 
-    img = Image.fromarray(generated_img)
+    pil_img = Image.fromarray(img[0])
 
-    return img
+    return pil_img
 
 
 # pass function, input type for prompt, the output for multiple images