Update app.py

app.py

@@ -30,6 +30,7 @@ def infer(
     height=1024,
     guidance_scale=3,
     num_inference_steps=20,
+    num_images=1,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -37,7 +38,7 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
+    images = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
@@ -45,10 +46,10 @@ def infer(
         width=width,
         height=height,
         generator=generator,
-    ).images[0]
-
-    return image, seed
+        num_images_per_prompt=num_images,
+    ).images
 
+    return images, seed
 
 examples = [
         "A capybara wearing a suit holding a sign that reads Hello World",
@@ -63,7 +64,7 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # [Stable Diffusion 3.5 Medium Turbo")
+        gr.Markdown(" # [Stable Diffusion 3.5 Medium Turbo]")
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -75,7 +76,7 @@ with gr.Blocks(css=css) as demo:
 
             run_button = gr.Button("Run", scale=0, variant="primary")
 
-        result = gr.Image(label="Result", show_label=False)
+        result = gr.Gallery(label="Results").style(grid=[4], preview=True)
 
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
@@ -129,6 +130,14 @@ with gr.Blocks(css=css) as demo:
                     value=20,
                 )
 
+            num_images = gr.Slider(
+                label="Number of Images",
+                minimum=1,
+                maximum=4,
+                step=1,
+                value=1,
+            )
+
         gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True, cache_mode="lazy")
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -142,6 +151,7 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            num_images,
         ],
         outputs=[result, seed],
     )