Update app.py

app.py

@@ -42,16 +42,39 @@ _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
 
@@ -60,6 +83,7 @@ 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
 
@@ -78,7 +102,10 @@ 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:
@@ -91,8 +118,16 @@ 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(
@@ -106,19 +141,11 @@ def infer(
         num_images_per_prompt=1,
     ).images[0]
 
-    # Return updates: Show single image by default, update slider in the background
-    return (
-        gr.update(value=result, visible=True),             
-        gr.update(value=(src_img, result), visible=False), 
-        seed,
-        gr.update(visible=True, value="🔍 Compare Before & After") 
-    )
+    # Resize output back to the original image dimensions
+    # if result.size != (out_w, out_h):
+    #     result = result.resize((out_w, out_h), Image.LANCZOS)
 
-def toggle_compare_view(btn_text):
-    if "Compare" in btn_text:
-        return gr.update(visible=False), gr.update(visible=True), gr.update(value="🖼️ Show Only Result")
-    else:
-        return gr.update(visible=True), gr.update(visible=False), gr.update(value="🔍 Compare Before & After")
+    return (src_img, result), seed
 
 
 # --- UI ---
@@ -181,14 +208,7 @@ with gr.Blocks() as demo:
                     )
 
             with gr.Column():
-                # Clean single image layout as the default output view
-                result_image = gr.Image(label="Final Color-Graded Output", interactive=False)
-                
-                # Using the core native gr.ImageSlider component (hidden initially)
-                compare_slider = gr.ImageSlider(label="Before & After Comparison", interactive=False, visible=False)
-                
-                # Subtle secondary toggle button
-                compare_btn = gr.Button("🔍 Compare Before & After", visible=False, variant="secondary")
+                result = gr.ImageSlider(label="Color Graded Output", interactive=False)
 
         gr.Examples(
             examples=[
@@ -196,9 +216,10 @@ with gr.Blocks() as demo:
                 ["images/image2.jpeg","images/image1.jpg"],
             ],
             inputs=[source_image, reference_image],
-            outputs=[result_image, compare_slider, seed, compare_btn],
+            outputs=[result, seed],
             fn=infer,
-            cache_examples=False, 
+            cache_examples=True,
+            cache_mode="lazy",
             elem_id="examples"
         )
 
@@ -207,14 +228,9 @@ with gr.Blocks() as demo:
         seed, randomize_seed, true_guidance_scale,
         num_inference_steps,
     ]
-    outputs = [result_image, compare_slider, seed, compare_btn]
-    run_btn.click(fn=infer, inputs=inputs, outputs=outputs)
+    outputs = [result, seed]
 
-    compare_btn.click(
-        fn=toggle_compare_view,
-        inputs=[compare_btn],
-        outputs=[result_image, compare_slider, compare_btn]
-    )
+    run_btn.click(fn=infer, inputs=inputs, outputs=outputs)
 
     source_image.upload(
         fn=update_dimensions_on_upload,