update app.pyp

app.py

@@ -1,7 +1,272 @@
+"""ICTone Hugging Face Spaces demo optimized for ZeroGPU.
+
+Space setup:
+1. Select ZeroGPU hardware in the Space settings.
+2. Add a Space secret named HF_TOKEN. The token owner must have accepted
+   the access conditions for black-forest-labs/FLUX.1-Fill-dev.
+3. Keep inference.py in the same directory as this file.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+
 import gradio as gr
+import numpy as np
+import spaces
+import torch
+from diffusers import FluxFillPipeline
+from PIL import Image
+
+from inference import (
+    DEFAULT_INSTANCE_PROMPT,
+    apply_lut,
+    estimate_lut,
+    run_one,
+)
+
+
+MAX_SEED = np.iinfo(np.int32).max
+
+FLUX_PATH = os.getenv(
+    "FLUX_PATH",
+    "black-forest-labs/FLUX.1-Fill-dev",
+)
+LORA_PATH = os.getenv(
+    "LORA_PATH",
+    "ToneStyle/ICTone-Fill-LoRA",
+)
+IMAGE_SIZE = int(os.getenv("IMAGE_SIZE", "512"))
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+
+def load_pipeline() -> FluxFillPipeline:
+    """Load FluxFill + ICTone LoRA once at Space startup.
+
+    ZeroGPU recommends placing the model on CUDA at module scope. During Space
+    startup this uses ZeroGPU's CUDA emulation; a real GPU is attached only
+    while a @spaces.GPU function is running.
+    """
+    print(f"[ICTone] Loading base model: {FLUX_PATH}")
+    print(f"[ICTone] Loading LoRA: {LORA_PATH}")
+
+    load_kwargs = {
+        "torch_dtype": torch.bfloat16,
+    }
+    if HF_TOKEN:
+        load_kwargs["token"] = HF_TOKEN
+
+    pipe = FluxFillPipeline.from_pretrained(
+        FLUX_PATH,
+        **load_kwargs,
+    )
+    pipe.load_lora_weights(LORA_PATH)
+
+    # Required placement pattern for ZeroGPU. Do not lazy-load/move the model
+    # inside infer().
+    pipe.to("cuda")
+
+    print("[ICTone] Pipeline ready.")
+    return pipe
+
+
+# Load once at module scope for efficient ZeroGPU model placement.
+pipe = load_pipeline()
+
+
+@spaces.GPU(duration=60)
+def infer(
+    content: Image.Image,
+    reference: Image.Image,
+    seed: int,
+    randomize_seed: bool,
+    guidance_scale: float,
+    num_inference_steps: int,
+    lut_size: int,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Run ICTone and reconstruct the result at the original content resolution."""
+    if content is None or reference is None:
+        raise gr.Error("Please upload both a content image and a reference image.")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    seed = int(seed)
+    guidance_scale = float(guidance_scale)
+    num_inference_steps = int(num_inference_steps)
+    lut_size = int(lut_size)
+
+    content_rgb = content.convert("RGB")
+    reference_rgb = reference.convert("RGB")
+
+    with torch.inference_mode():
+        pred, panel, _, _ = run_one(
+            pipe,
+            content_rgb,
+            reference_rgb,
+            size=IMAGE_SIZE,
+            prompt=DEFAULT_INSTANCE_PROMPT,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            seed=seed,
+            generator_device="cuda",
+        )
+
+        # Lift the low-resolution Flux prediction back to the original content
+        # resolution using ICTone's fitted 3D LUT.
+        if lut_size > 1:
+            before = np.asarray(content_rgb)
+            after = np.asarray(
+                pred.resize(content_rgb.size, Image.Resampling.BILINEAR)
+            )
+
+            before_flat = before.reshape(-1, 3)
+            after_flat = after.reshape(-1, 3)
+
+            # Bound LUT fitting cost for very large uploaded images.
+            max_samples = 500_000
+            if len(before_flat) > max_samples:
+                rng = np.random.default_rng(0)
+                selected = rng.choice(
+                    len(before_flat),
+                    max_samples,
+                    replace=False,
+                )
+                before_flat = before_flat[selected]
+                after_flat = after_flat[selected]
+
+            lut = estimate_lut(
+                before_flat,
+                after_flat,
+                size=lut_size,
+                device="cuda",
+            )
+            output = Image.fromarray(
+                apply_lut(
+                    before,
+                    lut,
+                    device="cuda",
+                )
+            )
+        else:
+            output = pred
+
+    return output, panel, seed
+
+
+with gr.Blocks(title="ICTone · In-Context Tone Style Transfer") as demo:
+    gr.Markdown(
+        """
+# ICTone
+
+**In-Context Tone Style Transfer**
+
+Upload a **content image** and a **reference image**. ICTone transfers the
+reference color, contrast, and photographic tone while preserving the content
+of the source image.
+
+The demo uses **FLUX.1-Fill-dev** with the **ICTone LoRA** and runs on
+Hugging Face **ZeroGPU**. A short queue may appear when shared GPUs are busy.
+"""
+    )
+
+    with gr.Row():
+        content = gr.Image(
+            label="Content image",
+            type="pil",
+        )
+        reference = gr.Image(
+            label="Reference image",
+            type="pil",
+        )
+
+    with gr.Accordion("Generation settings", open=False):
+        with gr.Row():
+            seed = gr.Number(
+                label="Seed",
+                value=666,
+                precision=0,
+            )
+            randomize_seed = gr.Checkbox(
+                label="Randomize seed",
+                value=False,
+            )
+        with gr.Row():
+            guidance = gr.Slider(
+                label="Guidance scale",
+                minimum=1,
+                maximum=100,
+                value=50,
+                step=1,
+            )
+            steps = gr.Slider(
+                label="Inference steps",
+                minimum=1,
+                maximum=28,
+                value=4,
+                step=1,
+            )
+            lut = gr.Slider(
+                label="LUT size",
+                minimum=0,
+                maximum=33,
+                value=33,
+                step=1,
+            )
+
+    run = gr.Button(
+        "Transfer tone",
+        variant="primary",
+    )
+
+    with gr.Row():
+        output = gr.Image(
+            label="Result",
+            type="pil",
+        )
+        preview = gr.Image(
+            label="Content | Reference | Result",
+            type="pil",
+        )
+
+    used_seed = gr.Number(
+        label="Used seed",
+        precision=0,
+    )
+
+    run.click(
+        fn=infer,
+        inputs=[
+            content,
+            reference,
+            seed,
+            randomize_seed,
+            guidance,
+            steps,
+            lut,
+        ],
+        outputs=[
+            output,
+            preview,
+            used_seed,
+        ],
+        show_progress="full",
+    )
+
+    gr.Markdown(
+        """
+**Models:** `black-forest-labs/FLUX.1-Fill-dev` +
+`ToneStyle/ICTone-Fill-LoRA`
+
+FLUX.1-Fill-dev is subject to the FLUX.1 [dev] license and access conditions.
+"""
+    )
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    demo.queue().launch(
+        server_name="0.0.0.0",
+        server_port=int(os.getenv("PORT", "7860")),
+    )

inference.py

@@ -0,0 +1,753 @@
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from diffusers import FluxFillPipeline
+from PIL import Image
+
+
+DEFAULT_INSTANCE_PROMPT = (
+    "A side-by-side triptych. Left: source photo. "
+    "Middle: a color and tone reference photo. "
+    "Right: the same scene as the left, re-graded so its colors, "
+    "contrast, and film look match the middle reference, while "
+    "preserving the left's content and details."
+)
+
+
+def _resize_to_wh(pil: Image.Image, width: int, height: int) -> Image.Image:
+    """Aspect-preserving resize-cover then center-crop to (width, height)."""
+    if pil.mode != "RGB":
+        pil = pil.convert("RGB")
+    w, h = pil.size
+    s = max(width / w, height / h)
+    nw = max(int(round(w * s)), width)
+    nh = max(int(round(h * s)), height)
+    pil = pil.resize((nw, nh), Image.LANCZOS)
+    left = (nw - width) // 2
+    top = (nh - height) // 2
+    return pil.crop((left, top, left + width, top + height))
+
+
+def _resize_width_keep_aspect(
+    pil: Image.Image, width: int, height_multiple: int = 16
+) -> Image.Image:
+    """Resize so output width == ``width``, height scales by original aspect.
+
+    Height is rounded to the nearest positive multiple of ``height_multiple``
+    (FluxFill's VAE + patch stride, so the pipeline accepts it without extra
+    padding). No cropping is performed.
+    """
+    if pil.mode != "RGB":
+        pil = pil.convert("RGB")
+    w, h = pil.size
+    new_h = int(round(h * width / w))
+    if height_multiple > 1:
+        new_h = int(round(new_h / height_multiple)) * height_multiple
+        new_h = max(height_multiple, new_h)
+    else:
+        new_h = max(1, new_h)
+    return pil.resize((width, new_h), Image.LANCZOS)
+
+
+def _build_lut_laplacian(size: int):
+    """Sparse 3D graph-Laplacian on the (size, size, size) LUT grid.
+
+    Each vertex has up to 6 axis-aligned neighbors. ``L L^T`` acts as a
+    curvature penalty so uncovered / sparsely-covered LUT cells extrapolate
+    smoothly from their neighbors instead of collapsing to a fixed anchor.
+    """
+    from scipy.sparse import coo_matrix
+    V = size ** 3
+    idx = np.arange(V, dtype=np.int64).reshape(size, size, size)
+    rows: list[np.ndarray] = []
+    cols: list[np.ndarray] = []
+    data: list[np.ndarray] = []
+    for axis in range(3):
+        # Pair each vertex with its +1-neighbor along ``axis`` (no wrap-around).
+        take = [slice(None)] * 3
+        take[axis] = slice(None, -1)
+        src = idx[tuple(take)].ravel()
+        take[axis] = slice(1, None)
+        dst = idx[tuple(take)].ravel()
+        # Edge (src -> dst): L[src] - L[dst] = 0 (finite difference row).
+        n = src.shape[0]
+        edge_rows = np.arange(n, dtype=np.int64) + sum(r.shape[0] for r in rows) // 2 * 0  # local re-baseline below
+        # We'll assemble a single edge-per-row Laplacian by concatenating below.
+        rows.append(src)
+        cols.append(src)
+        data.append(np.ones(n, dtype=np.float32))
+        rows.append(src)
+        cols.append(dst)
+        data.append(-np.ones(n, dtype=np.float32))
+        rows.append(dst)
+        cols.append(dst)
+        data.append(np.ones(n, dtype=np.float32))
+        rows.append(dst)
+        cols.append(src)
+        data.append(-np.ones(n, dtype=np.float32))
+    r = np.concatenate(rows)
+    c = np.concatenate(cols)
+    d = np.concatenate(data)
+    # This is a graph Laplacian L in vertex-index form; ``LtL`` == L (since L is
+    # symmetric PSD). Return it as the penalty operator directly.
+    return coo_matrix((d, (r, c)), shape=(V, V)).tocsr()
+
+
+def estimate_lut(
+    before,
+    after,
+    size: int = 33,
+    lam_smooth: float = 0.05,
+    lam_anchor: float = 1e-5,
+    device=None,
+):
+    """Estimate a size^3 3D LUT mapping ``before`` colors to ``after`` colors.
+
+    Trilinear-consistent sparse least squares with a **Laplacian smoothness
+    prior** and a tiny identity anchor to break gauge invariance:
+
+        min_L  ||W L - after||^2
+               + lam_smooth * L^T Δ L                  (3D grid smoothness)
+               + lam_anchor * ||L - identity||^2       (gauge fix)
+
+    The smoothness prior is the key fix for tone-migration artifacts (e.g.
+    saturated colors like lips): sparsely-covered LUT cells extrapolate from
+    neighboring *covered* cells that carry the correct color transform,
+    instead of being pulled back toward the input color by an identity anchor.
+
+    ``before`` / ``after`` may be HxWx3 arrays or flat (N, 3) arrays with the
+    same length. ``device`` picks the solver backend:
+      - None / "cpu": scipy sparse direct solve (spsolve).
+      - "cuda" / torch.device("cuda"): torch sparse block-CG on GPU (~90x
+        faster on hi-res pairs).
+    Falls back to nearest-cell averaging if SciPy is unavailable on CPU path.
+    """
+    if device is not None and str(device) != "cpu":
+        return _estimate_lut_gpu(
+            before, after, size, lam_smooth, lam_anchor, device=device
+        )
+    return _estimate_lut_cpu(before, after, size, lam_smooth, lam_anchor)
+
+
+def _estimate_lut_cpu(before, after, size, lam_smooth, lam_anchor):
+    before = np.asarray(before, dtype=np.float32).reshape(-1, 3)
+    after = np.asarray(after, dtype=np.float32).reshape(-1, 3)
+    assert before.shape == after.shape, "before/after must have same shape"
+
+    V = size ** 3
+
+    # Identity LUT anchor (used by both the LS solve and the fallback).
+    grid = np.linspace(0, 255, size, dtype=np.float32)
+    R, G, B = np.meshgrid(grid, grid, grid, indexing='ij')
+    identity_flat = np.stack([R, G, B], axis=-1).reshape(V, 3)
+
+    try:
+        from scipy.sparse import coo_matrix, eye as sp_eye
+        from scipy.sparse.linalg import spsolve
+    except ImportError:
+        return _estimate_lut_nearest_fallback(
+            before, after, size, identity_flat
+        ).reshape(size, size, size, 3)
+
+    # Continuous LUT-grid coordinates per pixel.
+    pos = before / 255.0 * (size - 1)
+    i0 = np.clip(np.floor(pos).astype(np.int64), 0, size - 1)
+    i1 = np.clip(i0 + 1, 0, size - 1)
+    f = pos - i0
+    fr, fg, fb = f[:, 0], f[:, 1], f[:, 2]
+    ofr, ofg, ofb = 1.0 - fr, 1.0 - fg, 1.0 - fb
+    r0, g0, b0 = i0[:, 0], i0[:, 1], i0[:, 2]
+    r1, g1, b1 = i1[:, 0], i1[:, 1], i1[:, 2]
+
+    def _vidx(r, g, b):
+        return (r * size + g) * size + b
+
+    verts = np.stack([
+        _vidx(r0, g0, b0), _vidx(r0, g0, b1),
+        _vidx(r0, g1, b0), _vidx(r0, g1, b1),
+        _vidx(r1, g0, b0), _vidx(r1, g0, b1),
+        _vidx(r1, g1, b0), _vidx(r1, g1, b1),
+    ], axis=1)
+    wts = np.stack([
+        ofr * ofg * ofb, ofr * ofg * fb,
+        ofr * fg * ofb, ofr * fg * fb,
+        fr * ofg * ofb, fr * ofg * fb,
+        fr * fg * ofb, fr * fg * fb,
+    ], axis=1).astype(np.float32)
+
+    N = before.shape[0]
+    rows = np.repeat(np.arange(N, dtype=np.int64), 8)
+    cols = verts.reshape(-1)
+    data = wts.reshape(-1)
+    W = coo_matrix((data, (rows, cols)), shape=(N, V)).tocsr()
+
+    WtW = (W.T @ W).tocsc()
+    scale = max(float(WtW.diagonal().mean()), 1.0)
+    lap = _build_lut_laplacian(size).tocsc()
+    A = (WtW + (lam_smooth * scale) * lap + lam_anchor * sp_eye(V, format='csc')).tocsc()
+    rhs = W.T @ after + lam_anchor * identity_flat
+
+    lut_flat = np.empty((V, 3), dtype=np.float32)
+    for c in range(3):
+        lut_flat[:, c] = spsolve(A, rhs[:, c])
+
+    return lut_flat.reshape(size, size, size, 3)
+
+
+def _estimate_lut_gpu(
+    before, after, size, lam_smooth, lam_anchor, device,
+    tol: float = 1e-4, max_iter: int = 200,
+):
+    """GPU LS solve via preconditioned block conjugate gradient.
+
+    Solves the same system as ``_estimate_lut_cpu`` but on the specified CUDA
+    device using ``torch.sparse.mm`` for matvecs. Returns a numpy array shaped
+    ``(size, size, size, 3)`` matching the CPU path.
+    """
+    b = torch.as_tensor(before, dtype=torch.float32, device=device).reshape(-1, 3)
+    y = torch.as_tensor(after, dtype=torch.float32, device=device).reshape(-1, 3)
+    N = b.shape[0]
+    V = size ** 3
+
+    pos = b / 255.0 * (size - 1)
+    i0 = pos.floor().clamp(0, size - 1).long()
+    i1 = (i0 + 1).clamp(0, size - 1)
+    f = pos - i0.float()
+    fr, fg, fb = f[:, 0], f[:, 1], f[:, 2]
+    ofr, ofg, ofb = 1 - fr, 1 - fg, 1 - fb
+    r0, g0, b0 = i0[:, 0], i0[:, 1], i0[:, 2]
+    r1, g1, b1 = i1[:, 0], i1[:, 1], i1[:, 2]
+
+    def _vidx(r, g, c):
+        return (r * size + g) * size + c
+
+    verts = torch.stack([
+        _vidx(r0, g0, b0), _vidx(r0, g0, b1),
+        _vidx(r0, g1, b0), _vidx(r0, g1, b1),
+        _vidx(r1, g0, b0), _vidx(r1, g0, b1),
+        _vidx(r1, g1, b0), _vidx(r1, g1, b1),
+    ], dim=1)
+    wts = torch.stack([
+        ofr * ofg * ofb, ofr * ofg * fb,
+        ofr * fg * ofb, ofr * fg * fb,
+        fr * ofg * ofb, fr * ofg * fb,
+        fr * fg * ofb, fr * fg * fb,
+    ], dim=1)
+
+    rows = torch.arange(N, device=device).repeat_interleave(8)
+    cols = verts.reshape(-1)
+    vals = wts.reshape(-1)
+    W = torch.sparse_coo_tensor(torch.stack([rows, cols]), vals, (N, V)).coalesce()
+    Wt = W.transpose(0, 1).coalesce()
+
+    # 3D graph Laplacian on the LUT grid.
+    idx3 = torch.arange(V, device=device).view(size, size, size)
+    L_rows: list[torch.Tensor] = []
+    L_cols: list[torch.Tensor] = []
+    L_vals: list[torch.Tensor] = []
+    for ax in range(3):
+        sl = [slice(None)] * 3
+        sl[ax] = slice(None, -1)
+        src = idx3[tuple(sl)].reshape(-1)
+        sl[ax] = slice(1, None)
+        dst = idx3[tuple(sl)].reshape(-1)
+        ones_s = torch.ones_like(src, dtype=torch.float32)
+        ones_d = torch.ones_like(dst, dtype=torch.float32)
+        L_rows += [src, dst, src, dst]
+        L_cols += [src, dst, dst, src]
+        L_vals += [ones_s, ones_d, -ones_s, -ones_d]
+    Lop = torch.sparse_coo_tensor(
+        torch.stack([torch.cat(L_rows), torch.cat(L_cols)]),
+        torch.cat(L_vals), (V, V)).coalesce()
+
+    # WᵀW diag + Laplacian diag → Jacobi preconditioner.
+    WtW_diag = torch.zeros(V, device=device).scatter_add_(
+        0, verts.reshape(-1), wts.reshape(-1) ** 2)
+    scale = max(float(WtW_diag.mean().item()), 1.0)
+    lam_s = lam_smooth * scale
+    lam_a = lam_anchor
+    L_diag = torch.zeros(V, device=device)
+    same = Lop.indices()[0] == Lop.indices()[1]
+    L_diag.scatter_add_(0, Lop.indices()[0][same], Lop.values()[same])
+    Minv = 1.0 / (WtW_diag + lam_s * L_diag + lam_a)
+
+    grid = torch.linspace(0, 255, size, device=device)
+    R, G, B = torch.meshgrid(grid, grid, grid, indexing='ij')
+    identity = torch.stack([R, G, B], dim=-1).reshape(V, 3)
+
+    def A_matvec(X):
+        return (torch.sparse.mm(Wt, torch.sparse.mm(W, X))
+                + lam_s * torch.sparse.mm(Lop, X)
+                + lam_a * X)
+
+    rhs = torch.sparse.mm(Wt, y) + lam_a * identity
+
+    # Preconditioned block Conjugate Gradient (per-channel in parallel).
+    X = torch.zeros_like(rhs)
+    R_ = rhs - A_matvec(X)
+    Z = Minv.unsqueeze(1) * R_
+    P = Z.clone()
+    rz_old = (R_ * Z).sum(dim=0)
+    b_norm = rhs.norm(dim=0).clamp_min(1e-30)
+    for _ in range(max_iter):
+        AP = A_matvec(P)
+        alpha = rz_old / ((P * AP).sum(dim=0) + 1e-30)
+        X = X + alpha.unsqueeze(0) * P
+        R_ = R_ - alpha.unsqueeze(0) * AP
+        if (R_.norm(dim=0) / b_norm).max().item() < tol:
+            break
+        Z = Minv.unsqueeze(1) * R_
+        rz_new = (R_ * Z).sum(dim=0)
+        P = Z + (rz_new / rz_old).unsqueeze(0) * P
+        rz_old = rz_new
+
+    return X.reshape(size, size, size, 3).detach().cpu().numpy()
+
+
+def _estimate_lut_nearest_fallback(before, after, size, identity_flat):
+    """Old nearest-cell averaging path, used only if SciPy is missing."""
+    idx = np.clip(
+        np.round(before / 255.0 * (size - 1)).astype(np.int32),
+        0, size - 1)
+    flat = (idx[:, 0] * size + idx[:, 1]) * size + idx[:, 2]
+
+    V = size ** 3
+    lut = np.zeros((V, 3), dtype=np.float32)
+    counts = np.zeros(V, dtype=np.int64)
+    np.add.at(lut, flat, after)
+    np.add.at(counts, flat, 1)
+    filled = counts > 0
+    lut[filled] /= counts[filled, None]
+    lut[~filled] = identity_flat[~filled]
+    return lut
+
+
+def apply_lut(content, lut, device=None):
+    """Apply a 3D LUT to an RGB image (any resolution) with trilinear interp.
+
+    ``device=None`` runs the numpy path (portable); a CUDA device runs it via
+    ``torch.nn.functional.grid_sample`` (~50x faster on hi-res images).
+    """
+    if device is not None and str(device) != "cpu":
+        return _apply_lut_gpu(content, lut, device)
+    return _apply_lut_cpu(content, lut)
+
+
+def _apply_lut_cpu(content, lut):
+    size = lut.shape[0]
+    img = np.asarray(content, dtype=np.float32) / 255.0 * (size - 1)
+
+    i0 = np.floor(img).astype(np.int32)
+    i0 = np.clip(i0, 0, size - 1)
+    i1 = np.clip(i0 + 1, 0, size - 1)
+    f = img - i0
+
+    r0, g0, b0 = i0[..., 0], i0[..., 1], i0[..., 2]
+    r1, g1, b1 = i1[..., 0], i1[..., 1], i1[..., 2]
+    fr = f[..., 0:1]
+    fg = f[..., 1:2]
+    fb = f[..., 2:3]
+
+    c000 = lut[r0, g0, b0]
+    c001 = lut[r0, g0, b1]
+    c010 = lut[r0, g1, b0]
+    c011 = lut[r0, g1, b1]
+    c100 = lut[r1, g0, b0]
+    c101 = lut[r1, g0, b1]
+    c110 = lut[r1, g1, b0]
+    c111 = lut[r1, g1, b1]
+
+    c00 = c000 * (1 - fb) + c001 * fb
+    c01 = c010 * (1 - fb) + c011 * fb
+    c10 = c100 * (1 - fb) + c101 * fb
+    c11 = c110 * (1 - fb) + c111 * fb
+    c0 = c00 * (1 - fg) + c01 * fg
+    c1 = c10 * (1 - fg) + c11 * fg
+    out = c0 * (1 - fr) + c1 * fr
+
+    return np.clip(out, 0, 255).astype(np.uint8)
+
+
+def _apply_lut_gpu(content, lut, device):
+    """GPU trilinear LUT application via ``F.grid_sample`` (5D volume)."""
+    import torch.nn.functional as F
+    size = lut.shape[0]
+    lut_np = np.asarray(lut, dtype=np.float32)
+    lut_t = torch.from_numpy(lut_np).permute(3, 0, 1, 2).unsqueeze(0).to(device)
+    # LUT dim layout after permute: (N=1, C=3, D=R, H=G, W=B).
+
+    img = torch.from_numpy(np.asarray(content, dtype=np.float32)).to(device)
+    coord = img / 255.0 * (size - 1)
+    coord = 2 * coord / (size - 1) - 1                # → [-1, 1] w.r.t. (R, G, B)
+    # grid_sample expects last-dim order (x, y, z) == (W, H, D) == (B, G, R).
+    coord = coord[..., [2, 1, 0]].unsqueeze(0).unsqueeze(0)   # (1, 1, H, W, 3)
+    out = F.grid_sample(
+        lut_t, coord, mode="bilinear", align_corners=True, padding_mode="border"
+    )
+    out = out.squeeze(0).squeeze(1).permute(1, 2, 0)
+    return out.clamp(0, 255).byte().cpu().numpy()
+
+
+def build_triptych(
+    content: Image.Image, reference: Image.Image, panel_w: int, panel_h: int
+):
+    """Return (triptych_pil, mask_pil) both at (3*panel_w, panel_h).
+
+    - Triptych: ``[content | reference | content]`` — the right third is a
+      copy of content, acting as an identity prior for prepare_latents
+      (matches the training/validation convention).
+    - Mask: L-mode PIL, 0 elsewhere and 255 on the right third (fill region).
+    """
+    canvas_w = panel_w * 3
+    canvas_h = panel_h
+
+    tri = Image.new("RGB", (canvas_w, canvas_h))
+    tri.paste(content, (0, 0))
+    tri.paste(reference, (panel_w, 0))
+    tri.paste(content, (2 * panel_w, 0))
+
+    mask_arr = np.zeros((canvas_h, canvas_w), dtype=np.uint8)
+    mask_arr[:, 2 * panel_w:] = 255
+    mask = Image.fromarray(mask_arr, mode="L")
+    return tri, mask
+
+
+def run_one(
+    pipe: FluxFillPipeline,
+    content_pil: Image.Image,
+    reference_pil: Image.Image,
+    *,
+    size: int,
+    prompt: str,
+    guidance_scale: float,
+    num_inference_steps: int,
+    seed: int,
+    generator_device: str,
+) -> tuple[Image.Image, Image.Image, Image.Image, Image.Image]:
+    """Run one triptych inference. Returns (pred_pil, panel_pil, tri_pil, content_sq).
+
+    - ``pred_pil``: just the right-third region — the migrated result.
+    - ``panel_pil``: horizontal panel ``[content | reference | pred]``.
+    - ``tri_pil``: the full pipeline output. Useful for debugging.
+    - ``content_sq``: the resized content actually fed to the pipeline. Kept
+      so callers can pair it with ``pred_pil`` for LUT fitting.
+
+    Content's width is resized to ``size`` while its height keeps the original
+    aspect ratio (rounded to a multiple of 16 for the VAE). Reference is
+    force-resized to the same ``(size, H)`` so the triptych panels line up.
+    Canvas is ``3*size x H``.
+    """
+    content = _resize_width_keep_aspect(content_pil, size, height_multiple=16)
+    panel_w, panel_h = content.size
+    # Stretch reference directly to content's (W, H) — no aspect preservation,
+    # no cropping. Panels line up by construction.
+    if reference_pil.mode != "RGB":
+        reference_pil = reference_pil.convert("RGB")
+    reference = reference_pil.resize((panel_w, panel_h), Image.LANCZOS)
+
+    tri, mask = build_triptych(content, reference, panel_w, panel_h)
+
+    generator = torch.Generator(device=generator_device).manual_seed(seed)
+    result = pipe(
+        prompt=prompt,
+        image=tri,
+        mask_image=mask,
+        height=panel_h,
+        width=panel_w * 3,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        max_sequence_length=512,
+        generator=generator,
+    ).images[0]
+
+    pred = result.crop((panel_w * 2, 0, panel_w * 3, panel_h))
+    panel = Image.new("RGB", (panel_w * 3, panel_h))
+    panel.paste(content, (0, 0))
+    panel.paste(reference, (panel_w, 0))
+    panel.paste(pred, (panel_w * 2, 0))
+    return pred, panel, result, content
+
+
+def main():
+    parser = argparse.ArgumentParser("Zero-shot tone migration with FLUX.1-Fill")
+    # Single-pair inputs
+    parser.add_argument("--content", type=str, default=None,
+                        help="Path to the content (source) image.")
+    parser.add_argument("--reference", type=str, default=None,
+                        help="Path to the reference (style/tone) image.")
+    # Batch mode over TST2K
+    parser.add_argument("--tst2k-dir", type=str, default=None,
+                        help="Root of TST2K-style eval set. Each subdir must contain "
+                             "content.png + reference.png. If set, --content/--reference "
+                             "are ignored and up to --tst2k-num subdirs are processed.")
+    parser.add_argument("--triplet-list", type=str, default=None,
+                        help="Text file with 3 whitespace-separated columns per row: "
+                             "<content_path> <reference_path> <gt_path>. Comment lines "
+                             "starting with '#' are skipped. Overrides --tst2k-dir and "
+                             "--content/--reference. Row index (0-based) becomes the "
+                             "output stem (NNNN.png).")
+    parser.add_argument("--tst2k-num", type=int, default=50,
+                        help="Max number of samples (subdirs or list rows) to iterate.")
+
+    parser.add_argument("--output-dir", type=str, default="./tone_out",
+                        help="Directory for batch outputs. Ignored when --output-file is set.")
+    parser.add_argument("--output-file", type=str, default=None,
+                        help="Exact output image path for single-pair inference. "
+                             "When set, --output-dir is not used.")
+    parser.add_argument("--flux-path", type=str,
+                        default="ckpt/FLUX-Fill")
+    parser.add_argument("--lora-path", type=str, default=None,
+                        help="Optional LoRA path (dir or .safetensors). Omit to test "
+                             "the base FluxFill model with no fine-tuning.")
+
+    parser.add_argument("--image-size", type=int, default=512,
+                        help="Content width; height keeps the source aspect "
+                             "ratio (rounded to a multiple of 16 for the VAE). "
+                             "Canvas width = 3 * this.")
+    parser.add_argument("--lut-size", type=int, default=33,
+                        help="3D LUT grid size per channel used to lift the "
+                             "512-res migration back onto the original hi-res "
+                             "content. Set to 0 to skip hi-res reconstruction.")
+    parser.add_argument("--lut-device", type=str, default="auto",
+                        choices=["auto", "cpu", "cuda"],
+                        help="Solver device for LUT estimate + apply. 'auto' "
+                             "uses CUDA when available (~90x faster estimate, "
+                             "~50x faster apply on hi-res images).")
+    parser.add_argument("--num-inference-steps", type=int, default=28)
+    parser.add_argument("--guidance-scale", type=float, default=30.0,
+                        help="FluxFill's guidance-distilled embed value. ICEdit's "
+                             "reference inference.py uses 50; 30 is FluxFill default.")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--prompt", type=str, default=None,
+                        help="Text prompt. Omit to use the built-in triptych "
+                             "instruction (identical to the training default).")
+
+    parser.add_argument("--enable-model-cpu-offload", action="store_true")
+    parser.add_argument("--dtype", type=str, default="bfloat16",
+                        choices=["bfloat16", "float16", "float32"])
+    parser.add_argument("--generator-device", type=str, default="cpu",
+                        choices=["cpu", "cuda"],
+                        help="Where the noise generator lives. ICEdit uses 'cpu'; "
+                             "'cuda' silences the diffusers 'passed generator was "
+                             "created on cpu' warning at the cost of slightly "
+                             "different bit-exact noise across restarts.")
+
+    # Data-parallel sharding across independent processes (one per GPU).
+    # Each shard iterates the same global job list but only processes indices
+    # ``i`` where ``i % num_shards == shard_index``. The per-sample seed is
+    # derived from the *global* index so results are identical to a
+    # single-process run.
+    parser.add_argument("--num-shards", type=int, default=1,
+                        help="Total number of parallel shards (processes).")
+    parser.add_argument("--shard-index", type=int, default=0,
+                        help="This shard's index in [0, num_shards).")
+
+    args = parser.parse_args()
+
+    if args.triplet_list is None and args.tst2k_dir is None \
+            and (args.content is None or args.reference is None):
+        parser.error(
+            "Provide one of: --triplet-list, --tst2k-dir, or both --content and --reference."
+        )
+
+    if args.num_shards < 1 or not (0 <= args.shard_index < args.num_shards):
+        parser.error(
+            f"Invalid sharding: shard_index={args.shard_index}, "
+            f"num_shards={args.num_shards}."
+        )
+
+    if args.output_file is not None and (
+        args.triplet_list is not None or args.tst2k_dir is not None
+    ):
+        parser.error("--output-file is only supported for single-pair inference.")
+
+    torch_dtype = {
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+        "float32": torch.float32,
+    }[args.dtype]
+
+    print(f"[load] FluxFill from {args.flux_path}  dtype={args.dtype}")
+    pipe = FluxFillPipeline.from_pretrained(args.flux_path, torch_dtype=torch_dtype)
+
+    if args.lora_path:
+        print(f"[load] LoRA weights from {args.lora_path}")
+        pipe.load_lora_weights(args.lora_path)
+    else:
+        print("[load] no LoRA — testing base FluxFill zero-shot")
+
+    if args.enable_model_cpu_offload:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe = pipe.to("cuda")
+
+    prompt = args.prompt if args.prompt is not None else DEFAULT_INSTANCE_PROMPT
+    print(f"[prompt] {prompt}")
+
+    if args.lut_device == "auto":
+        lut_device = "cuda" if torch.cuda.is_available() else None
+    elif args.lut_device == "cuda":
+        lut_device = "cuda"
+    else:
+        lut_device = None
+    print(f"[lut] solver device: {lut_device or 'cpu'}")
+
+    output_file = Path(args.output_file) if args.output_file is not None else None
+    if output_file is not None:
+        output_file.parent.mkdir(parents=True, exist_ok=True)
+        dir_out = str(output_file.parent)
+    else:
+        os.makedirs(args.output_dir, exist_ok=True)
+        dir_out = os.path.join(args.output_dir, "outputs")
+        os.makedirs(dir_out, exist_ok=True)
+    S = int(args.image_size)
+
+    # ---- Assemble list of (stem, content_path, reference_path, gt_path) ----
+    jobs: list[tuple[str, Path, Path, Path | None]] = []
+    if args.triplet_list is not None:
+        list_path = Path(args.triplet_list)
+        list_base = list_path.resolve().parent
+
+        def resolve_list_path(value: str) -> Path:
+            path = Path(value)
+            return path if path.is_absolute() else list_base / path
+
+        with open(list_path) as f:
+            for raw in f:
+                line = raw.strip()
+                if not line or line.startswith("#"):
+                    continue
+                parts = line.split()
+                if len(parts) < 2:
+                    continue
+                c = resolve_list_path(parts[0])
+                r = resolve_list_path(parts[1])
+                g = resolve_list_path(parts[2]) if len(parts) >= 3 else None
+                if not (c.exists() and r.exists()):
+                    continue
+                if g is not None and not g.exists():
+                    g = None
+                stem = f"{len(jobs):04d}"
+                jobs.append((stem, c, r, g))
+                if len(jobs) >= args.tst2k_num:
+                    break
+        if not jobs:
+            raise SystemExit(f"[error] no valid rows in {list_path}")
+        print(f"[batch] {len(jobs)} samples from {list_path}")
+    elif args.tst2k_dir is not None:
+        root = Path(args.tst2k_dir)
+        subs = sorted([p for p in root.iterdir() if p.is_dir()])
+        for sub in subs:
+            c = sub / "content.png"
+            r = sub / "reference.png"
+            if not (c.exists() and r.exists()):
+                continue
+            g = sub / "gt.png"
+            jobs.append((sub.name, c, r, g if g.exists() else None))
+            if len(jobs) >= args.tst2k_num:
+                break
+        if not jobs:
+            raise SystemExit(f"[error] no valid subdirs found under {root}")
+        print(f"[batch] {len(jobs)} samples from {root}")
+    else:
+        c = Path(args.content)
+        r = Path(args.reference)
+        jobs.append((c.stem, c, r, None))
+
+    # ---- Filter jobs for this shard while keeping the global index ----
+    # ``global_i`` is the index into the full (unsharded) job list. It drives
+    # both the per-sample seed (``args.seed + global_i``) and the output file
+    # prefix, so different shards write disjoint filenames and any single
+    # sample gets the same seed regardless of shard configuration.
+    total_jobs = len(jobs)
+    sharded = [
+        (gi, job) for gi, job in enumerate(jobs)
+        if gi % args.num_shards == args.shard_index
+    ]
+    if args.num_shards > 1:
+        print(
+            f"[shard] {args.shard_index+1}/{args.num_shards}: "
+            f"{len(sharded)}/{total_jobs} samples"
+        )
+
+    # ---- Run inference ----
+    for local_i, (global_i, (stem, cpath, rpath, gpath)) in enumerate(sharded):
+        content_pil = Image.open(cpath).convert("RGB")
+        reference_pil = Image.open(rpath).convert("RGB")
+
+        pred, panel, full, content_sq = run_one(
+            pipe,
+            content_pil,
+            reference_pil,
+            size=S,
+            prompt=prompt,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            seed=args.seed + global_i,
+            generator_device=args.generator_device,
+        )
+
+        # ---- Hi-res reconstruction via 3D LUT ----
+        # The pipeline works at S=512, but ``content_pil`` is usually
+        # higher-resolution. Fit a 3D LUT so its trilinear evaluation on the
+        # hi-res content reproduces the diffusion output.
+        #
+        # Training pair: hi-res content pixels paired with pred bilinearly
+        # upsampled to hi-res. This gives the LUT the exact color distribution
+        # it will be applied to (important for saturated regions like lips,
+        # whose peak reds get lost when content is first downsampled). We
+        # subsample to keep the sparse LS problem small.
+        pred_hires = None
+        if args.lut_size and args.lut_size > 1:
+            content_arr = np.asarray(content_pil)
+            pred_up = np.asarray(
+                pred.resize(content_pil.size, Image.BILINEAR)
+            )
+            flat_before = content_arr.reshape(-1, 3)
+            flat_after = pred_up.reshape(-1, 3)
+            max_samples = 500_000
+            if flat_before.shape[0] > max_samples:
+                rng = np.random.default_rng(0)
+                sel = rng.choice(
+                    flat_before.shape[0], size=max_samples, replace=False
+                )
+                flat_before = flat_before[sel]
+                flat_after = flat_after[sel]
+            lut = estimate_lut(
+                flat_before,
+                flat_after,
+                size=int(args.lut_size),
+                device=lut_device,
+            )
+            hires_arr = apply_lut(content_arr, lut, device=lut_device)
+            pred_hires = Image.fromarray(hires_arr)
+
+        # Content-resolution reconstruction goes to ``outputs/``.
+        # Filename is just the (1-based) global index, zero-padded to 4 digits,
+        # so shards never collide and results sort naturally.
+        code = f"{global_i :04d}"
+        name_out = str(output_file) if output_file is not None else str(Path(dir_out) / code)
+        if pred_hires is not None:
+            pred_hires.save(
+                str(output_file) if output_file is not None else f"{name_out}.png"
+            )
+
+        print(
+            f"[done] shard {args.shard_index+1}/{args.num_shards}  "
+            f"{local_i+1}/{len(sharded)}  (global {global_i+1}/{total_jobs})  "
+            f"{stem}  →  {name_out}.png"
+        )
+
+    print(
+        f"[all done] shard {args.shard_index+1}/{args.num_shards} — "
+        f"results under {os.path.abspath(args.output_dir)}"
+    )
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,18 @@
+# ZeroGPU-compatible inference environment
+torch==2.8.0
+
+diffusers==0.39.0
+transformers==4.57.6
+accelerate==1.12.0
+peft==0.18.0
+safetensors==0.8.0
+huggingface-hub==0.36.0
+sentencepiece==0.2.1
+protobuf==7.35.1
+
+numpy==2.1.3
+scipy==1.15.2
+Pillow==11.1.0
+
+gradio>=4,<7
+spaces