Update app.py

app.py

@@ -20,19 +20,24 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 @spaces.GPU(duration=190)
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, num_images=1, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-        prompt=prompt, 
-        width=width,
-        height=height,
-        num_inference_steps=num_inference_steps, 
-        generator=generator,
-        guidance_scale=guidance_scale
-    ).images[0] 
-    return image, seed
+    
+    images = []
+    for _ in range(num_images):
+        image = pipe(
+            prompt=prompt, 
+            width=width,
+            height=height,
+            num_inference_steps=num_inference_steps, 
+            generator=generator,
+            guidance_scale=guidance_scale
+        ).images[0]
+        images.append(image)
+    
+    return images, seed
  
 examples = [
     "a white husky knocking everything down in a living room",
@@ -67,7 +72,7 @@ with gr.Blocks(css=css) as demo:
             
             run_button = gr.Button("Run", scale=0)
         
-        result = gr.Image(label="Result", show_label=False)
+        result = gr.Gallery(label="Result", show_label=False)
 
         # Display all settings directly without the Accordion
         seed = gr.Slider(
@@ -116,6 +121,14 @@ with gr.Blocks(css=css) as demo:
                 value=28,
             )
         
+        num_images = gr.Slider(
+            label="Number of images",
+            minimum=1,
+            maximum=4,
+            step=1,
+            value=1,
+        )
+        
         gr.Examples(
             examples=examples,
             fn=infer,
@@ -127,7 +140,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images],
         outputs=[result, seed]
     )