From 2f7dc90c88ed1dcce65bb2117779a4108b824104 Mon Sep 17 00:00:00 2001 From: poxular Date: Mon, 1 Jun 2026 17:35:29 +0200 Subject: [PATCH] Match standalone probe: PIL 256 ONNX feed and 960px LUT apply. Co-authored-by: Cursor --- magicmatch_core/image_ops.py | 34 ++++++++++++++++++++++++++++++++++ magicmatch_core/inference.py | 25 ++++++++----------------- magicmatch_core/live_cache.py | 25 ++----------------------- nodes.py | 22 ++++++++++++++-------- 4 files changed, 58 insertions(+), 48 deletions(-) create mode 100644 magicmatch_core/image_ops.py diff --git a/magicmatch_core/image_ops.py b/magicmatch_core/image_ops.py new file mode 100644 index 0000000..e920caa --- /dev/null +++ b/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) diff --git a/magicmatch_core/inference.py b/magicmatch_core/inference.py index b4d7c3f..4db638b 100644 --- a/magicmatch_core/inference.py +++ b/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: diff --git a/magicmatch_core/live_cache.py b/magicmatch_core/live_cache.py index f0a9364..715b1eb 100644 --- a/magicmatch_core/live_cache.py +++ b/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) diff --git a/nodes.py b/nodes.py index 2aa7906..2468d45 100644 --- a/nodes.py +++ b/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),)