Update app.py

app.py

@@ -81,11 +81,12 @@ flux_pipeline = FluxKontextPipeline.from_pretrained(
 flux_pipeline.to(device_gen)
 
 TASK_CHOICES = [
+    "Prompt Tuning - Image Editing",
+    "Prompt Tuning - T2I Generation",
     "Pointwise - Image Editing",
     "Pointwise - T2I Generation",
     "Pairwise - Image Editing",
     "Pairwise - T2I Generation",
-    "Prompt Tuning - Image Editing",
 ]
 
 # ============================================================
@@ -355,13 +356,30 @@ def create_instruction(prompt, task_type):
         raise ValueError(f"Unknown task type: {task_type}")
 
 def update_ui_for_task(task_type):
-    """Update image component visibility/labels and instruction label based on selected task type."""
-    if task_type == "Pointwise - Image Editing":
+    """Update image component visibility/labels, instruction label, and generate button based on selected task type."""
+    if task_type == "Prompt Tuning - Image Editing":
+        return (
+            gr.update(visible=True, label="Source Image"),
+            gr.update(visible=True, label="Generated Image", interactive=False, value=None),
+            gr.update(visible=False, label="(unused)", value=None),
+            gr.update(label="Editing Instruction", placeholder="Enter the instruction for editing..."),
+            gr.update(visible=True),   # generate_btn visible
+        )
+    elif task_type == "Prompt Tuning - T2I Generation":
+        return (
+            gr.update(visible=True, label="Generated Image"),
+            gr.update(visible=False, label="(unused)", value=None),
+            gr.update(visible=False, label="(unused)", value=None),
+            gr.update(label="T2I Prompt", placeholder="Enter the text-to-image generation prompt…"),
+            gr.update(visible=False),  # generate_btn hidden (no T2I pipeline)
+        )
+    elif task_type == "Pointwise - Image Editing":
         return (
             gr.update(visible=True, label="Source Image"),
             gr.update(visible=True, label="Edited Image"),
             gr.update(visible=False, label="Image B", value=None),
             gr.update(label="Editing Instruction", placeholder="Describe the edit that was applied to the source image…"),
+            gr.update(visible=False),  # generate_btn hidden
         )
     elif task_type == "Pointwise - T2I Generation":
         return (
@@ -369,6 +387,7 @@ def update_ui_for_task(task_type):
             gr.update(visible=False, label="(unused)", value=None),
             gr.update(visible=False, label="(unused)", value=None),
             gr.update(label="Text-to-Image Prompt", placeholder="Enter the text-to-image generation prompt…"),
+            gr.update(visible=False),  # generate_btn hidden
         )
     elif task_type == "Pairwise - Image Editing":
         return (
@@ -376,6 +395,7 @@ def update_ui_for_task(task_type):
             gr.update(visible=True, label="Image A"),
             gr.update(visible=True, label="Image B"),
             gr.update(label="Editing Instruction", placeholder="Describe the edit that was applied to the source image…"),
+            gr.update(visible=False),  # generate_btn hidden
         )
     elif task_type == "Pairwise - T2I Generation":
         return (
@@ -383,13 +403,7 @@ def update_ui_for_task(task_type):
             gr.update(visible=True, label="Image B"),
             gr.update(visible=False, label="(unused)", value=None),
             gr.update(label="Text-to-Image Prompt", placeholder="Enter the text-to-image generation prompt…"),
-        )
-    elif task_type == "Prompt Tuning - Image Editing":
-        return (
-            gr.update(visible=True, label="Source Image"),
-            gr.update(visible=True, label="Generated Image", interactive=False, value=None),
-            gr.update(visible=False, label="(unused)", value=None),
-            gr.update(label="Instruction", placeholder="Enter the instruction for editing..."),
+            gr.update(visible=False),  # generate_btn hidden
         )
     else:
         raise ValueError(f"Unknown task type: {task_type}")
@@ -440,76 +454,67 @@ def run_vlm_evaluation(messages, loaded_images):
         buffer += new_text
         yield buffer
 
+def generate_image(task_type, instruction_text, image1):
+    """Generate an edited image using Flux (Kontext) and return the path."""
+    if task_type != "Prompt Tuning - Image Editing":
+        raise gr.Error("Generate is only available for Prompt Tuning – Image Editing.")
+    if not image1:
+        raise gr.Error("Please upload the Source Image first.")
+    if not instruction_text:
+        raise gr.Error("Please enter an editing instruction first.")
+
+    source_img = load_image(image1)
+    width, height = source_img.size
+
+    generated_image = run_flux_generation(instruction_text, source_img, width, height)
+
+    timestamp = int(time.time())
+    os.makedirs("generated_images", exist_ok=True)
+    generated_image_path = f"generated_images/flux_edit_{timestamp}.png"
+    generated_image.save(generated_image_path)
+
+    return generated_image_path
+
 @spaces.GPU(duration=300)
 def model_inference(task_type, instruction_text, image1, image2, image3, progress=gr.Progress()):
-    """Run model inference based on the selected task type and uploaded images."""
-    
+    """Run VLM evaluation based on the selected task type and uploaded images."""
+
     loaded_images = []
     task_for_template = task_type
-    generated_image_path = None
 
     # Validate inputs and collect images based on task
-    if task_type == "Pointwise - Image Editing":
+    if task_type in ("Pointwise - Image Editing", "Prompt Tuning - Image Editing"):
         if not image1 or not image2:
-            yield "Error: Please upload both Source Image and Edited Image.", None
+            yield "Error: Please upload Source Image and Edited/Generated Image."
             return
         files = [image1, image2]
         loaded_images = [load_image(img) for img in files]
+        task_for_template = "Pointwise - Image Editing"
 
-    elif task_type == "Pointwise - T2I Generation":
+    elif task_type in ("Pointwise - T2I Generation", "Prompt Tuning - T2I Generation"):
         if not image1:
-            yield "Error: Please upload the Generated Image.", None
+            yield "Error: Please upload the Generated Image."
             return
         files = [image1]
         loaded_images = [load_image(img) for img in files]
+        task_for_template = "Pointwise - T2I Generation"
 
     elif task_type == "Pairwise - Image Editing":
         if not image1 or not image2 or not image3:
-            yield "Error: Please upload Source Image, Image A, and Image B.", None
+            yield "Error: Please upload Source Image, Image A, and Image B."
             return
         files = [image1, image2, image3]
         loaded_images = [load_image(img) for img in files]
 
     elif task_type == "Pairwise - T2I Generation":
         if not image1 or not image2:
-            yield "Error: Please upload both Image A and Image B.", None
+            yield "Error: Please upload both Image A and Image B."
             return
         files = [image1, image2]
         loaded_images = [load_image(img) for img in files]
 
-    elif task_type == "Prompt Tuning - Image Editing":
-        if not image1:
-            yield "Error: Please upload the Source Image.", None
-            return
-        
-        progress(0, desc="Starting generation...")
-        yield "Generating edited image with Flux... (This may take a minute)", None
-        
-        # Load source image
-        try:
-            source_img = load_image(image1)
-            width, height = source_img.size
-
-            # Run Flux generation on GPU
-            generated_image = run_flux_generation(instruction_text, source_img, width, height)
-
-            # Save generated image
-            timestamp = int(time.time())
-            os.makedirs("generated_images", exist_ok=True)
-            generated_image_path = f"generated_images/flux_edit_{timestamp}.png"
-            generated_image.save(generated_image_path)
-            
-        except Exception as e:
-            yield f"Error generating image: {str(e)}", None
-            return
-
-        yield "Image generated! Evaluating...", generated_image_path
-        
-        loaded_images = [source_img, generated_image]
-        task_for_template = "Pointwise - Image Editing"
-
     else:
-        yield "Error: Unknown task type selected.", None
+        yield "Error: Unknown task type selected."
         return
 
     # Build instruction with <image> placeholders
@@ -528,7 +533,7 @@ def model_inference(task_type, instruction_text, image1, image2, image3, progres
     # Run VLM evaluation on GPU (streaming)
     progress(0.9, desc="Evaluating...")
     for text in run_vlm_evaluation(messages, loaded_images):
-        yield text, gr.update(value=generated_image_path) if generated_image_path else gr.update()
+        yield text
 
 # ============================================================
 # Gradio UI
@@ -537,17 +542,18 @@ def model_inference(task_type, instruction_text, image1, image2, image3, progres
 OVERVIEW_MD = """
 ### 📋 Task Overview
 
-This demo supports **four evaluation tasks**. Select one to get started:
+This demo supports **six tasks**. Select one to get started:
 
 | Task | Description |
 |------|-------------|
+| **Prompt Tuning – Image Editing** | Generate an edit using Flux (Kontext) from a source image and instruction, then evaluate it. Use the refinement to tune your prompt. |
+| **Prompt Tuning – T2I Generation** | Upload a generated image and a text-to-image prompt, then evaluate it. Use the refinement to iteratively improve your prompt. |
 | **Pointwise – Image Editing** | Rate a single edited image against its source image and the editing instruction. Produces per-aspect scores and a refined request. |
 | **Pointwise – T2I Generation** | Rate a single generated image against a text-to-image prompt. Produces per-aspect scores and a refined prompt. |
 | **Pairwise – Image Editing** | Compare two edited images (A vs B) given a source image and editing instruction. Determines which edit is better per aspect. |
 | **Pairwise – T2I Generation** | Compare two generated images (A vs B) given a text-to-image prompt. Determines which generation is better per aspect. |
-| **Prompt Tuning – Image Editing** | Generate an edit using Flux (Kontext) from a source image and instruction, then evaluate it. Use the refinement to tune your prompt. |
 
-**Try the examples below - they're basically begging to be clicked! 🎯**
+**Try the examples on the right - they're basically begging to be clicked! 🎯**
 """
 
 with gr.Blocks(css="""
@@ -559,11 +565,11 @@ with gr.Blocks(css="""
     gr.Markdown(OVERVIEW_MD)
 
     with gr.Row(equal_height=True):
-        # ============ LEFT COLUMN – all inputs (scrollable) ============
+        # ============ LEFT COLUMN – inputs ============
         with gr.Column(scale=1, elem_id="input-panel"):
             task_selector = gr.Radio(
                 choices=TASK_CHOICES,
-                value="Pointwise - Image Editing",
+                value="Prompt Tuning - Image Editing",
                 label="Task Type",
                 info="Select the evaluation task",
             )
@@ -578,9 +584,10 @@ with gr.Blocks(css="""
                     )
                 with gr.Column(scale=1, min_width=160):
                     image2 = gr.Image(
-                        label="Edited Image",
+                        label="Generated Image",
                         type="filepath",
                         sources=["upload", "clipboard"],
+                        interactive=False,
                     )
                 with gr.Column(scale=1, min_width=160):
                     image3 = gr.Image(
@@ -590,40 +597,49 @@ with gr.Blocks(css="""
                         visible=False,
                     )
 
-            # ---- Instruction + Evaluate ----
+            # ---- Instruction + Buttons ----
             instruction = gr.Textbox(
                 label="Editing Instruction",
                 lines=3,
-                placeholder="Describe the edit that was applied to the source image…",
+                placeholder="Enter the instruction for editing...",
             )
-            submit_btn = gr.Button("Evaluate", variant="primary")
+            with gr.Row():
+                generate_btn = gr.Button("Generate Image", variant="secondary", visible=True)
+                submit_btn = gr.Button("Evaluate", variant="primary")
 
-            # ---- Examples ----
+        # ============ RIGHT COLUMN – examples ============
+        with gr.Column(scale=1):
             gr.Examples(
                 examples=[
+                    ["Prompt Tuning - Image Editing", "Remove the arrows from the blue sign and add the text of Detour ahead, no right turns.", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_source.png", None, None],
                     ["Pointwise - Image Editing", "Remove the arrows from the blue sign and add the text of Detour ahead, no right turns.", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_source.png", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_ovis_u1_Image A.png", None],
                     ["Pairwise - Image Editing", "Remove the arrows from the blue sign and add the text of Detour ahead, no right turns.", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_source.png", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_ovis_u1_Image A.png", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_ovis_u1_Image A.png"],
-                    ["Prompt Tuning - Image Editing", "Remove the arrows from the blue sign and add the text of Detour ahead, no right turns.", "example_images/0016cb70b187efe39969766dc4b3f9ed_b63ed6db519f685c33b860b511879cfe2fa7351059a17ebe5eafa83213e222fb_13_source.png", None, None],
                 ],
                 inputs=[task_selector, instruction, image1, image2, image3],
             )
 
-        # ============ RIGHT COLUMN – evaluation output ============
-        with gr.Column(scale=1):
-            output = gr.Textbox(label="Evaluation Result", lines=30)
+    # ---- Evaluation result (full width, below the input/examples row) ----
+    output = gr.Textbox(label="Evaluation Result", lines=20)
 
-    # ---- Wire task selector to update image visibility/labels + instruction label ----
+    # ---- Wire task selector to update image visibility/labels, instruction label, and generate button ----
     task_selector.change(
         fn=update_ui_for_task,
         inputs=[task_selector],
-        outputs=[image1, image2, image3, instruction],
+        outputs=[image1, image2, image3, instruction, generate_btn],
+    )
+
+    # ---- Wire generate button (Prompt Tuning – Image Editing only) ----
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[task_selector, instruction, image1],
+        outputs=[image2],
     )
 
     # ---- Wire evaluate button ----
     submit_btn.click(
         fn=model_inference,
         inputs=[task_selector, instruction, image1, image2, image3],
-        outputs=[output, image2],
+        outputs=[output],
     )
 
     gr.Markdown(tos_markdown)