Update app.py

app.py

@@ -1,25 +1,39 @@
 import spaces
 import gradio as gr
+from diffusers import FluxKontextPipeline
 from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
 from transformers.image_utils import load_image
 from threading import Thread
 import torch
+import os
 
 from serve_constants import html_header, bibtext, learn_more_markdown, tos_markdown
 
-MODEL_ID = "JasperHaozhe/RationalRewards-Both-Demo"  
+MODEL_ID = "JasperHaozhe/RationalRewards-Both-Demo"
+FLUX_MODEL_ID = "black-forest-labs/FLUX.1-Kontext-dev"
+
 processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForImageTextToText.from_pretrained(
     MODEL_ID,
     trust_remote_code=True,
     torch_dtype=torch.bfloat16
-).to("cuda").eval()
+).to("cpu").eval()
+
+# Load Flux Pipeline
+flux_pipeline = FluxKontextPipeline.from_pretrained(
+    FLUX_MODEL_ID,
+    torch_dtype=torch.bfloat16
+)
+# Fix VAE precision for Flux to avoid artifacts
+flux_pipeline.vae.to(dtype=torch.float32)
+flux_pipeline.to("cpu")
 
 TASK_CHOICES = [
     "Pointwise - Image Editing",
     "Pointwise - T2I Generation",
     "Pairwise - Image Editing",
     "Pairwise - T2I Generation",
+    "Prompt Tuning - Image Editing",
 ]
 
 # ============================================================
@@ -318,40 +332,107 @@ def update_ui_for_task(task_type):
             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..."),
+        )
+    else:
+        raise ValueError(f"Unknown task type: {task_type}")
 
 @spaces.GPU
 def model_inference(task_type, instruction_text, image1, image2, image3):
     """Run model inference based on the selected task type and uploaded images."""
+    
+    loaded_images = []
+    task_for_template = task_type
+    generated_image = None
+
     # Validate inputs and collect images based on task
     if task_type == "Pointwise - Image Editing":
         if not image1 or not image2:
-            yield "Error: Please upload both Source Image and Edited Image."
+            yield "Error: Please upload both Source Image and Edited Image.", None
             return
         files = [image1, image2]
+        loaded_images = [load_image(img) for img in files]
+
     elif task_type == "Pointwise - T2I Generation":
         if not image1:
-            yield "Error: Please upload the Generated Image."
+            yield "Error: Please upload the Generated Image.", None
             return
         files = [image1]
+        loaded_images = [load_image(img) for img in files]
+
     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."
+            yield "Error: Please upload Source Image, Image A, and Image B.", None
             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."
+            yield "Error: Please upload both Image A and Image B.", None
             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
+        
+        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
+            
+            # Ensure model is offloaded to CPU to make space for Flux
+            model.to("cpu")
+            torch.cuda.empty_cache()
+            
+            # Move Flux to CUDA
+            flux_pipeline.to("cuda")
+            
+            # Run Flux
+            generator = torch.Generator("cuda").manual_seed(42)
+            with torch.no_grad():
+                generated_image = flux_pipeline(
+                    prompt=instruction_text,
+                    image=source_img,
+                    guidance_scale=3.5,
+                    num_inference_steps=28,
+                    width=width,
+                    height=height,
+                    generator=generator,
+                ).images[0]
+            
+            # Move Flux back to CPU
+            flux_pipeline.to("cpu")
+            torch.cuda.empty_cache()
+            
+        except Exception as e:
+            # Attempt to recover state
+            flux_pipeline.to("cpu")
+            torch.cuda.empty_cache()
+            yield f"Error generating image: {str(e)}", None
+            return
+
+        yield "Image generated! Evaluating...", generated_image
+        
+        loaded_images = [source_img, generated_image]
+        task_for_template = "Pointwise - Image Editing"
+
     else:
-        yield "Error: Unknown task type selected."
+        yield "Error: Unknown task type selected.", None
         return
 
-    # Load images
-    loaded_images = [load_image(image) for image in files]
-
     # Build instruction with <image> placeholders
-    instruction = create_instruction(instruction_text, task_type)
+    instruction = create_instruction(instruction_text, task_for_template)
 
     # Interleave images into the <image> placeholders
     content = []
@@ -363,6 +444,9 @@ def model_inference(task_type, instruction_text, image1, image2, image3):
 
     messages = [{"role": "user", "content": content}]
 
+    # Ensure model is on CUDA for evaluation
+    model.to("cuda")
+
     # Generate and stream text
     prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = processor(
@@ -373,15 +457,19 @@ def model_inference(task_type, instruction_text, image1, image2, image3):
     ).to("cuda")
 
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048, temperature=0.3)
 
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
     buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
+    try:
+        for new_text in streamer:
+            buffer += new_text
+            yield buffer, gr.update()
+    finally:
+
+        pass
 
 # ============================================================
 # Gradio UI
@@ -398,6 +486,9 @@ This demo supports **four evaluation tasks**. Select one to get started:
 | **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! 🎯**
 """
 
 with gr.Blocks(css="""
@@ -408,7 +499,7 @@ with gr.Blocks(css="""
     # ---- Overview ----
     gr.Markdown(OVERVIEW_MD)
 
-    with gr.Row(equal_height=False):
+    with gr.Row(equal_height=True):
         # ============ LEFT COLUMN – all inputs (scrollable) ============
         with gr.Column(scale=1, elem_id="input-panel"):
             task_selector = gr.Radio(
@@ -452,6 +543,8 @@ with gr.Blocks(css="""
             gr.Examples(
                 examples=[
                     ["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],
             )
@@ -471,7 +564,7 @@ with gr.Blocks(css="""
     submit_btn.click(
         fn=model_inference,
         inputs=[task_selector, instruction, image1, image2, image3],
-        outputs=output,
+        outputs=[output, image2],
     )
 
     gr.Markdown(tos_markdown)