Update app.py

app.py

@@ -103,13 +103,13 @@ extras_b.sampling_configs.update({
     't_start': 1.0
 })
 
-@spaces.GPU
-def generate_image(prompt, height, width, seed):
+@spaces.GPU(duration=240)
+def generate_images(prompt, height, width, seed, num_images):
     torch.manual_seed(seed)
     random.seed(seed)
     np.random.seed(seed)
 
-    batch_size = 1
+    batch_size = num_images
     height_lr, width_lr = get_target_lr_size(height / width, std_size=32)
     stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)
     stage_c_latent_shape_lr, stage_b_latent_shape_lr = calculate_latent_sizes(height_lr, width_lr, batch_size=batch_size)
@@ -139,17 +139,18 @@ def generate_image(prompt, height, width, seed):
         
         torch.cuda.empty_cache()
         imgs = show_images(sampled)
-        return imgs[0]
+        return imgs
 
 iface = gr.Interface(
-    fn=generate_image,
+    fn=generate_images,
     inputs=[
         gr.Textbox(label="Prompt"),
         gr.Slider(minimum=256, maximum=2560, step=32, label="Height", value=1024),
         gr.Slider(minimum=256, maximum=5120, step=32, label="Width", value=1024),
-        gr.Number(label="Seed", value=42)
+        gr.Number(label="Seed", value=42),
+        gr.Slider(minimum=1, maximum=10, step=1, label="Number of Images", value=1)
     ],
-    outputs=gr.Image(type="pil"),
+    outputs=gr.Gallery(label="Generated Images", columns=5, rows=2),
     title="UltraPixel Image Generation",
     description="Generate high-resolution images using UltraPixel model.",
     theme='bethecloud/storj_theme'