Create app.py

app.py

@@ -0,0 +1,246 @@
+import gradio as gr
+import math
+import numpy as np
+import random
+import torch
+import spaces
+
+from PIL import Image
+from diffusers import QwenImageEditPlusPipeline
+from typing import Optional, Tuple
+
+MAX_SEED = np.iinfo(np.int32).max
+
+# --- Model Loading ---
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = QwenImageEditPlusPipeline.from_pretrained(
+    "Qwen/Qwen-Image-Edit-2511",
+    torch_dtype=dtype
+).to(device)
+
+# Load the lightning LoRA for fast inference
+pipe.load_lora_weights(
+    "lightx2v/Qwen-Image-Edit-2511-Lightning",
+    weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors",
+    adapter_name="lightning"
+)
+
+# Load the color grade transfer LoRA
+pipe.load_lora_weights(
+    "ovi054/QIE-2511-Color-Grade-Transfer-LoRA",
+    weight_name="QIE-2511-Color-Grade-Transfer-LoRA.safetensors",
+    adapter_name="color"
+)
+
+pipe.set_adapters(["lightning", "color"], adapter_weights=[1.0, 1.0])
+
+
+# VAE_IMAGE_SIZE must match the pipeline constant (pipeline_qwenimage_edit_plus.py line 67)
+_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
+
+    MAX_SIDE = 1328
+
+    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
+
+    if (new_width, new_height) == (original_width, original_height):
+        return image
+
+    return image.resize((new_width, new_height), Image.LANCZOS)
+
+
+@spaces.GPU
+def infer(
+    source_image: Optional[Image.Image] = None,
+    reference_image: Optional[Image.Image] = None,
+    seed: int = 0,
+    randomize_seed: bool = True,
+    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:
+        raise gr.Error("Please upload a color grade reference image (Image 2).")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    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(
+        image=[src_img, ref_img],
+        prompt="Transfer ONLY the color grading from Image 2 onto Image 1",
+        height=gen_h,
+        width=gen_w,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        true_cfg_scale=true_guidance_scale,
+        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 result, seed
+
+
+# --- UI ---
+css = '''
+#col-container { max-width: 1000px; margin: 0 auto; }
+.dark .progress-text { color: white !important }
+#examples { max-width: 1000px; margin: 0 auto; }
+.image-container { min-height: 300px; }
+'''
+
+with gr.Blocks() as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown("## 🎨 Color Grade Transfer - Qwen Image Edit + LoRA")
+        gr.Markdown("""
+            Transfer color grading and tones from a reference image onto your source image ✨
+            Using my [ovi054/Color-Grade-Transfer-LoRA](https://huggingface.co/ovi054/QIE-2511-Color-Grade-Transfer-LoRA) and 4 step inference
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    source_image = gr.Image(
+                        label="Image 1 (Source — content to preserve)",
+                        type="pil",
+                        elem_classes="image-container"
+                    )
+                    reference_image = gr.Image(
+                        label="Image 2 (Color Grade Reference)",
+                        type="pil",
+                        elem_classes="image-container"
+                    )
+
+                run_btn = gr.Button("🎨 Transfer Color Grade", variant="primary", size="lg")
+
+                with gr.Accordion("Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0
+                    )
+                    randomize_seed = gr.Checkbox(
+                        label="Randomize Seed",
+                        value=True
+                    )
+                    true_guidance_scale = gr.Slider(
+                        label="True Guidance Scale",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.1,
+                        value=1.0
+                    )
+                    num_inference_steps = gr.Slider(
+                        label="Inference Steps",
+                        minimum=1,
+                        maximum=40,
+                        step=1,
+                        value=4
+                    )
+
+            with gr.Column():
+                result = gr.Image(label="Color Graded Output", interactive=False)
+
+        gr.Examples(
+            examples=[
+                ["images/image1.jpg", "images/image2.jpeg"],
+                ["images/image2.jpeg","images/image1.jpg"],
+            ],
+            inputs=[source_image, reference_image],
+            outputs=[result, seed],
+            fn=infer,
+            cache_examples=True,
+            cache_mode="lazy",
+            elem_id="examples"
+        )
+
+    inputs = [
+        source_image, reference_image,
+        seed, randomize_seed, true_guidance_scale,
+        num_inference_steps,
+    ]
+    outputs = [result, seed]
+
+    run_btn.click(fn=infer, inputs=inputs, outputs=outputs)
+
+    source_image.upload(
+        fn=update_dimensions_on_upload,
+        inputs=[source_image],
+        outputs=[source_image]
+    )
+    reference_image.upload(
+        fn=update_dimensions_on_upload,
+        inputs=[reference_image],
+        outputs=[reference_image]
+    )
+
+demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=css, footer_links=["api", "gradio", "settings"])