Update app.py

app.py

@@ -60,7 +60,7 @@ def generate_image(prompt: str, negative_prompt: str = "", inference_steps: int
             guidance_scale=float(guidance_scale),
             jit=True,
         ).images
-    return image
+    return image[0]
 
 examples = [
     ["A watercolor painting of a bird"],
@@ -271,9 +271,9 @@ with block as demo:
                     )
                 btn = gr.Button("Generate image", label="Primary Button", variant="primary")
 
-        gallery = gr.Gallery(
+        gallery = gr.Image(
             label="Generated images", show_label=False, elem_id="gallery"
-        ).style(height="auto")
+        )
 
 
         with gr.Row():