Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+from gradio_imageslider import ImageSlider
 import math
 import numpy as np
 import random
@@ -7,7 +8,7 @@ import spaces
 
 from PIL import Image
 from diffusers import QwenImageEditPlusPipeline
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Any
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -42,39 +43,16 @@ _VAE_IMAGE_SIZE = 1024 * 1024
 
 
 def calculate_vae_gen_size(image: Image.Image) -> tuple:
-    """
-    Return (gen_w, gen_h) that exactly matches the pipeline's internal VAE
-    conditioning scale for this image.
-
-    The pipeline always resizes every input image to VAE_IMAGE_SIZE (~1MP) before
-    VAE-encoding it into image_latents, using:
-        vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, w / h)
-
-    img_shapes (used for 2-D RoPE) is built from BOTH the output size (height/width)
-    AND the conditioning sizes (vae_width, vae_height).  When they differ, the RoPE
-    coordinate systems are misaligned → huge pixel shift.
-
-    Passing gen_h/gen_w = the same 1MP-equivalent makes the output tokens and Image 1
-    conditioning tokens share an identical coordinate system → no shift.
-    This is exactly what ComfyUI’s ImageScaleToTotalPixels (megapixels=1.0) achieves.
-    """
     W, H = image.size
     ratio = W / H
     gen_w = math.sqrt(_VAE_IMAGE_SIZE * ratio)
     gen_h = gen_w / ratio
-    # pipeline rounds to multiples of 32 (also satisfies the ÷16 divisibility requirement)
     gen_w = round(gen_w / 32) * 32
     gen_h = round(gen_h / 32) * 32
     return int(gen_w), int(gen_h)
 
 
-
 def update_dimensions_on_upload(image: Optional[Image.Image]) -> Image.Image:
-    """
-    Cap longest side to 1328px, snap to multiples of 16.
-    Pipeline requires divisibility by vae_scale_factor * 2 = 8 * 2 = 16.
-    Never upscales.
-    """
     if image is None:
         return image
 
@@ -83,7 +61,6 @@ def update_dimensions_on_upload(image: Optional[Image.Image]) -> Image.Image:
     original_width, original_height = image.size
     scale = min(MAX_SIDE / original_width, MAX_SIDE / original_height, 1.0)
 
-    # Must be multiples of 16 (vae_scale_factor * 2)
     new_width  = (int(original_width  * scale) // 16) * 16
     new_height = (int(original_height * scale) // 16) * 16
 
@@ -102,10 +79,7 @@ def infer(
     true_guidance_scale: float = 1.0,
     num_inference_steps: int = 4,
     progress=gr.Progress(track_tqdm=True)
-) -> Tuple[Image.Image, int]:
-    """
-    Transfer color grading from a reference image onto a source image.
-    """
+): 
     if source_image is None:
         raise gr.Error("Please upload a source image (Image 1).")
     if reference_image is None:
@@ -118,16 +92,8 @@ def infer(
     src_img = source_image.convert("RGB")
     ref_img = reference_image.convert("RGB")
 
-    # Original size — used to resize the output back at the end
     out_w, out_h = src_img.size
 
-    # Generate at the 1MP-equivalent of Image 1’s aspect ratio.
-    # The pipeline internally scales ALL input images to VAE_IMAGE_SIZE (~1MP) before
-    # VAE-encoding them as conditioning latents.  img_shapes (for 2-D RoPE) combines
-    # the output size (height/width) with those conditioning sizes.  If they differ,
-    # the RoPE coordinate systems are misaligned → huge pixel shift.
-    # Using the same 1MP formula as the pipeline eliminates the mismatch.
-    # (ComfyUI achieves this via ImageScaleToTotalPixels at megapixels=1.0.)
     gen_w, gen_h = calculate_vae_gen_size(src_img)
 
     result = pipe(
@@ -141,11 +107,22 @@ def infer(
         num_images_per_prompt=1,
     ).images[0]
 
-    # Resize output back to the original image dimensions
-    # if result.size != (out_w, out_h):
-    #     result = result.resize((out_w, out_h), Image.LANCZOS)
+    # Return updates to make sure the standard image shows up first, and the toggle button appears
+    return (
+        gr.update(value=result, visible=True),             # Show single image
+        gr.update(value=(src_img, result), visible=False), # Hide slider
+        seed,
+        gr.update(visible=True, value="🔍 Compare Before & After") # Show toggle button
+    )
 
-    return result, seed
+def toggle_compare_view(btn_text):
+    """Toggles visibility between the single image and the comparison slider."""
+    if "Compare" in btn_text:
+        # Hide image, Show Slider, Update Button text
+        return gr.update(visible=False), gr.update(visible=True), gr.update(value="🖼️ Show Only Result")
+    else:
+        # Show image, Hide Slider, Update Button text
+        return gr.update(visible=True), gr.update(visible=False), gr.update(value="🔍 Compare Before & After")
 
 
 # --- UI ---
@@ -208,7 +185,12 @@ with gr.Blocks() as demo:
                     )
 
             with gr.Column():
-                result = gr.Image(label="Color Graded Output", interactive=False)
+                # Both components exist, but only one is visible at a time
+                result_image = gr.Image(label="Final Color-Graded Output", interactive=False)
+                compare_slider = ImageSlider(label="Before & After Comparison", interactive=False, visible=False)
+                
+                # Hidden button that only appears AFTER the image is generated
+                compare_btn = gr.Button("🔍 Compare Before & After", visible=False, variant="secondary")
 
         gr.Examples(
             examples=[
@@ -216,22 +198,28 @@ with gr.Blocks() as demo:
                 ["images/image2.jpeg","images/image1.jpg"],
             ],
             inputs=[source_image, reference_image],
-            outputs=[result, seed],
+            outputs=[result_image, compare_slider, seed, compare_btn],
             fn=infer,
-            cache_examples=True,
-            cache_mode="lazy",
+            cache_examples=False, # Disabled caching to ensure gr.update() functions flawlessly
             elem_id="examples"
         )
 
+    # Trigger Generation
     inputs = [
         source_image, reference_image,
         seed, randomize_seed, true_guidance_scale,
         num_inference_steps,
     ]
-    outputs = [result, seed]
-
+    outputs = [result_image, compare_slider, seed, compare_btn]
     run_btn.click(fn=infer, inputs=inputs, outputs=outputs)
 
+    # Trigger Toggle View
+    compare_btn.click(
+        fn=toggle_compare_view,
+        inputs=[compare_btn],
+        outputs=[result_image, compare_slider, compare_btn]
+    )
+
     source_image.upload(
         fn=update_dimensions_on_upload,
         inputs=[source_image],