Single image output

app.py

@@ -145,7 +145,7 @@ pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype
 MAX_SEED = np.iinfo(np.int32).max
 
 # --- Main Inference Function (with hardcoded negative prompt) ---
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=120)
 def infer(
     image,
     prompt,
@@ -154,7 +154,6 @@ def infer(
     true_guidance_scale=1.0,
     num_inference_steps=50,
     rewrite_prompt=True,
-    num_images_per_prompt=1,
     progress=gr.Progress(track_tqdm=True),
 ):
     """
@@ -177,17 +176,17 @@ def infer(
         print(f"Rewritten Prompt: {prompt}")
 
     # Generate the image
-    image = pipe(
+    images = pipe(
         image,
         prompt=prompt,
         negative_prompt=negative_prompt,
         num_inference_steps=num_inference_steps,
         generator=generator,
         true_cfg_scale=true_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt
+        num_images_per_prompt=1
     ).images
-
-    return image, seed
+    
+    return images[0], seed
 
 # --- Examples and UI Layout ---
 examples = []
@@ -208,8 +207,7 @@ with gr.Blocks(css=css) as demo:
             with gr.Column():
                 input_image = gr.Image(label="Input Image", show_label=False, type="pil")
 
-            # result = gr.Image(label="Result", show_label=False, type="pil")
-            result = gr.Gallery(label="Result", show_label=False, type="pil")
+            result = gr.Image(label="Result", show_label=False, type="pil")
         with gr.Row():
             prompt = gr.Text(
                     label="Prompt",
@@ -250,14 +248,6 @@ with gr.Blocks(css=css) as demo:
                     value=50,
                 )
                 
-                num_images_per_prompt = gr.Slider(
-                    label="Number of images per prompt",
-                    minimum=1,
-                    maximum=4,
-                    step=1,
-                    value=1,
-                )
-                
                 rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=True)
 
         # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
@@ -273,7 +263,6 @@ with gr.Blocks(css=css) as demo:
             true_guidance_scale,
             num_inference_steps,
             rewrite_prompt,
-            num_images_per_prompt,
         ],
         outputs=[result, seed],
     )