Match standalone probe: PIL 256 ONNX feed and 960px LUT apply.

Co-authored-by: Cursor <cursoragent@cursor.com>

magicmatch_core/image_ops.py

@@ -0,0 +1,34 @@
+"""Image sizing aligned with scripts/color_match_probe.py (standalone parity)."""
+
+from __future__ import annotations
+
+import numpy as np
+
+# Same as MAX_DISPLAY_EDGE in color_match_probe.py
+PROBE_MAX_EDGE = 960
+
+
+def downscale_hwc_max_edge(hwc: np.ndarray, max_edge: int = PROBE_MAX_EDGE) -> np.ndarray:
+    """Bilinear downscale H×W×3 float RGB so max(w,h) <= max_edge (probe fit_for_preview)."""
+    hwc = np.asarray(hwc, dtype=np.float32)
+    h, w, _ = hwc.shape
+    scale = min(1.0, max_edge / max(h, w))
+    if scale >= 1.0:
+        return hwc.copy()
+    nw = max(1, int(w * scale))
+    nh = max(1, int(h * scale))
+    try:
+        from PIL import Image
+
+        arr = (np.clip(hwc, 0, 1) * 255.0).astype(np.uint8)
+        pil = Image.fromarray(arr, "RGB").resize((nw, nh), Image.Resampling.BILINEAR)
+        return np.asarray(pil, dtype=np.float32) / 255.0
+    except ImportError:
+        ys = (np.linspace(0, h - 1, nh)).astype(np.int32)
+        xs = (np.linspace(0, w - 1, nw)).astype(np.int32)
+        return hwc[np.ix_(ys, xs)]
+
+
+def prepare_apply_source(source_hwc: np.ndarray) -> np.ndarray:
+    """Source pixels passed to apply_merged_lut_preview (probe source_rgb)."""
+    return downscale_hwc_max_edge(source_hwc, PROBE_MAX_EDGE)

magicmatch_core/inference.py

@@ -26,24 +26,15 @@ _SESSION: ort.InferenceSession | None = None
 
 
 def _resize_hwc_to_nhwc(image_hwc: np.ndarray, size: int = IMAGE_SIZE) -> np.ndarray:
-    """H×W×3 float [0,1] → [1, size, size, 3] bilinear."""
+    """H×W×3 float [0,1] → [1, size, size, 3] PIL bilinear (same as probe load_rgb_nhwc)."""
+    from PIL import Image
+
     image_hwc = np.asarray(image_hwc, dtype=np.float32)
-    try:
-        import torch
-        import torch.nn.functional as F
-
-        t = torch.from_numpy(image_hwc).permute(2, 0, 1).unsqueeze(0)
-        t = F.interpolate(t, size=(size, size), mode="bilinear", align_corners=False)
-        out = t.permute(0, 2, 3, 1).numpy().astype(np.float32)
-        return np.clip(out, 0.0, 1.0)
-    except ImportError:
-        from PIL import Image
-
-        arr = (np.clip(image_hwc, 0, 1) * 255.0).astype(np.uint8)
-        pil = Image.fromarray(arr, "RGB")
-        pil = pil.resize((size, size), Image.Resampling.BILINEAR)
-        out = np.asarray(pil, dtype=np.float32) / 255.0
-        return out[np.newaxis, ...]
+    arr = (np.clip(image_hwc, 0, 1) * 255.0).astype(np.uint8)
+    pil = Image.fromarray(arr, "RGB")
+    pil = pil.resize((size, size), Image.Resampling.BILINEAR)
+    out = np.asarray(pil, dtype=np.float32) / 255.0
+    return np.clip(out[np.newaxis, ...], 0.0, 1.0)
 
 
 def get_session(onnx_path: str | Path | None = None) -> ort.InferenceSession:

magicmatch_core/live_cache.py

@@ -7,35 +7,14 @@ import io
 
 import numpy as np
 
-LIVE_PREVIEW_MAX_EDGE = 2048
-
-
-def _downscale_hwc(hwc: np.ndarray, max_edge: int) -> np.ndarray:
-    hwc = np.asarray(hwc, dtype=np.float32)
-    h, w, _ = hwc.shape
-    scale = min(1.0, max_edge / max(h, w))
-    if scale >= 1.0:
-        return hwc.copy()
-    nw = max(1, int(w * scale))
-    nh = max(1, int(h * scale))
-    try:
-        from PIL import Image
-
-        arr = (np.clip(hwc, 0, 1) * 255.0).astype(np.uint8)
-        pil = Image.fromarray(arr, "RGB").resize((nw, nh), Image.Resampling.BILINEAR)
-        return np.asarray(pil, dtype=np.float32) / 255.0
-    except ImportError:
-        # Nearest stride fallback
-        ys = (np.linspace(0, h - 1, nh)).astype(np.int32)
-        xs = (np.linspace(0, w - 1, nw)).astype(np.int32)
-        return hwc[np.ix_(ys, xs)]
+from .image_ops import PROBE_MAX_EDGE, downscale_hwc_max_edge
 
 
 def pack_live_cache(source_hwc: np.ndarray, merged_lut: np.ndarray) -> dict:
     """JSON-serializable bundle for the Comfy front-end (PNG + float32 LUT)."""
     from PIL import Image
 
-    small = _downscale_hwc(source_hwc, LIVE_PREVIEW_MAX_EDGE)
+    small = downscale_hwc_max_edge(source_hwc, PROBE_MAX_EDGE)
     merged = np.asarray(merged_lut, dtype=np.float32).reshape(-1)
 
     arr_u8 = (np.clip(small, 0, 1) * 255.0).astype(np.uint8)

nodes.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import numpy as np
 import torch
 
+from .magicmatch_core.image_ops import prepare_apply_source
 from .magicmatch_core.inference import build_merged_lut
 from .magicmatch_core.live_cache import pack_live_cache
 from .magicmatch_core.lut import apply_merged_lut_preview
@@ -55,7 +56,9 @@ class MagicMatchBuild:
     RETURN_NAMES = ("lut",)
     FUNCTION = "build"
     CATEGORY = "MAGICMATCH"
-    DESCRIPTION = "Build merged color-match LUT from source + reference (run once per image pair)."
+    DESCRIPTION = (
+        "Build merged LUT (ONNX on full-res pair, PIL 256 — matches standalone probe)."
+    )
 
     @classmethod
     def IS_CHANGED(cls, source, reference):
@@ -102,8 +105,8 @@ class MagicMatchPreview:
     FUNCTION = "preview"
     CATEGORY = "MAGICMATCH"
     DESCRIPTION = (
-        "In-node live preview after first run. Drag strength without re-queue; "
-        "queue again to bake the same strength into the output image."
+        "Apply cached LUT at probe parity (960px max edge, PIL 256 net). "
+        "Live slider after first run; queue again to export."
     )
 
     @classmethod
@@ -121,10 +124,10 @@ class MagicMatchPreview:
         lut: MagicMatchLUT,
         strength: float,
     ) -> dict:
-        src = _image_batch_to_hwc(source)
+        src_apply = prepare_apply_source(_image_batch_to_hwc(source))
         strength = float(np.clip(strength, 0.0, 1.0))
-        out = apply_merged_lut_preview(src, lut.merged_lut, strength)
-        cache = pack_live_cache(src, lut.merged_lut)
+        out = apply_merged_lut_preview(src_apply, lut.merged_lut, strength)
+        cache = pack_live_cache(src_apply, lut.merged_lut)
         return {
             "ui": {"magicmatch_live": [cache]},
             "result": (_hwc_to_image(out),),
@@ -157,7 +160,9 @@ class MagicMatch:
     RETURN_NAMES = ("image",)
     FUNCTION = "match"
     CATEGORY = "MAGICMATCH"
-    DESCRIPTION = "Single-node color match (re-runs neural net when strength changes)."
+    DESCRIPTION = (
+        "One-shot color match; output at 960px max edge like the standalone probe."
+    )
 
     def match(
         self,
@@ -168,8 +173,9 @@ class MagicMatch:
         src = _image_batch_to_hwc(source)
         ref = _image_batch_to_hwc(reference)
         merged = build_merged_lut(src, ref)
+        src_apply = prepare_apply_source(src)
         strength = float(np.clip(strength, 0.0, 1.0))
-        out = apply_merged_lut_preview(src, merged, strength)
+        out = apply_merged_lut_preview(src_apply, merged, strength)
         return (_hwc_to_image(out),)