commit 4d2d3d42369dc04538f0f17a176fd75add7fc84b Author: poxular Date: Mon Jun 1 16:51:00 2026 +0200 Initial release: MAGICMATCH ComfyUI color-match nodes with bundled ONNX model. Co-authored-by: Cursor diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b1b448b --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +.pytest_cache/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..75e7f3c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 MAGICMATCH contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..0c3f0ed --- /dev/null +++ b/README.md @@ -0,0 +1,65 @@ +# MAGICMATCH + +ComfyUI custom nodes for **neural color match**: two images in, graded source out. Tune strength on a preview node before you wire the result to save/export. + +## Install + +1. Clone into ComfyUI `custom_nodes`: + + ```bash + cd ComfyUI/custom_nodes + git clone https://github.com/YOUR_USER/MAGICMATCH.git + ``` + +2. Install Python deps (ComfyUI’s environment): + + ```bash + pip install -r MAGICMATCH/requirements.txt + ``` + +3. Confirm the model file is present: + + ```text + MAGICMATCH/models/color_match.onnx + ``` + +4. Restart ComfyUI. Nodes appear under **MAGICMATCH**. + +## Workflow (strength before export) + +Use **Build** + **Preview** so changing the slider only re-applies the LUT (fast): + +```text +[Source] ──┬──► MagicMatch Build LUT ◄── [Reference] + │ │ + │ ▼ lut + └──────► MagicMatch Preview (strength) ──► Preview Image + │ + └──► Save / export when ready +``` + +| Step | Node | +|------|------| +| 1 | **MagicMatch Build LUT** — run once per source/reference pair | +| 2 | **MagicMatch Preview (strength)** — `0` = unchanged source, `1` = full match; re-queue after moving the slider | +| 3 | Connect preview output to your save/export chain | + +**MagicMatch (one-shot)** combines build + apply in one node (simpler, slower when you change strength). + +## Requirements + +- ComfyUI with standard `IMAGE` tensors (batch size **1**) +- `onnxruntime`, `numpy` (see `requirements.txt`) +- CPU inference by default (~5–15s for Build on first run per pair) + +## Nodes + +| Class | Display name | +|-------|----------------| +| `MagicMatchBuild` | MagicMatch Build LUT | +| `MagicMatchPreview` | MagicMatch Preview (strength) | +| `MagicMatch` | MagicMatch (one-shot) | + +## License + +See [LICENSE](LICENSE). The bundled `color_match.onnx` is part of this package; use responsibly and in compliance with applicable terms for any upstream model you obtained. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8758186 --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +""" +MAGICMATCH — ComfyUI custom nodes for neural color match (ONNX). +""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/magicmatch_core/__init__.py b/magicmatch_core/__init__.py new file mode 100644 index 0000000..5673bbf --- /dev/null +++ b/magicmatch_core/__init__.py @@ -0,0 +1,12 @@ +"""MAGICMATCH core — ONNX inference and merged LUT apply.""" + +from .inference import build_merged_lut, run_inference_from_images +from .lut import apply_merged_lut_preview, get_merged_lut, nn_outputs_to_flat + +__all__ = [ + "build_merged_lut", + "run_inference_from_images", + "get_merged_lut", + "nn_outputs_to_flat", + "apply_merged_lut_preview", +] diff --git a/magicmatch_core/inference.py b/magicmatch_core/inference.py new file mode 100644 index 0000000..b4d7c3f --- /dev/null +++ b/magicmatch_core/inference.py @@ -0,0 +1,84 @@ +""" +ONNX inference for MAGICMATCH color-match graph (Comfy IMAGE tensors). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np + +try: + import onnxruntime as ort +except ImportError as e: + raise ImportError("MAGICMATCH: pip install onnxruntime") from e + +PACKAGE_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_ONNX = PACKAGE_ROOT / "models" / "color_match.onnx" +IMAGE_SIZE = 256 + +INPUT_NAMES = ("input_img:0", "ref_img:0") +OUTPUT_LUT3D = "3D_LUT" +OUTPUT_LUT1D = "1D_LUT" + +_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.""" + 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, ...] + + +def get_session(onnx_path: str | Path | None = None) -> ort.InferenceSession: + global _SESSION + path = Path(onnx_path or DEFAULT_ONNX) + if not path.is_file(): + raise FileNotFoundError( + f"MAGICMATCH ONNX not found: {path}\n" + "Expected: MAGICMATCH/models/color_match.onnx in this custom node folder." + ) + if _SESSION is None: + _SESSION = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + return _SESSION + + +def run_inference_from_images( + source_hwc: np.ndarray, + reference_hwc: np.ndarray, + *, + session: ort.InferenceSession | None = None, +) -> tuple[np.ndarray, np.ndarray]: + """Run graph on H×W×3 RGB (float 0–1). Returns lut3d, lut1d.""" + sess = session or get_session() + feeds = { + INPUT_NAMES[0]: _resize_hwc_to_nhwc(source_hwc), + INPUT_NAMES[1]: _resize_hwc_to_nhwc(reference_hwc), + } + out_map = {o.name: o for o in sess.get_outputs()} + names = [out_map[OUTPUT_LUT3D].name, out_map[OUTPUT_LUT1D].name] + lut3d, lut1d = sess.run(names, feeds) + return np.asarray(lut3d, dtype=np.float32), np.asarray(lut1d, dtype=np.float32) + + +def build_merged_lut(source_hwc: np.ndarray, reference_hwc: np.ndarray) -> np.ndarray: + from .lut import get_merged_lut, nn_outputs_to_flat + + lut3d, lut1d = run_inference_from_images(source_hwc, reference_hwc) + return get_merged_lut(*nn_outputs_to_flat(lut1d, lut3d)) diff --git a/magicmatch_core/lut.py b/magicmatch_core/lut.py new file mode 100644 index 0000000..194310c --- /dev/null +++ b/magicmatch_core/lut.py @@ -0,0 +1,402 @@ +""" +1D/3D LUT apply + merged 25³ cube (17³ NN outputs → 25³ merged LUT). +""" + +from __future__ import annotations + +import numpy as np + +MERGED_LUT_SIZE = 25 +NN_LUT_SIZE = 17 + + +def _clamp01(x: np.ndarray) -> np.ndarray: + return np.clip(x, 0.0, 1.0) + + +def apply_1d_lut(image: np.ndarray, lut1d: np.ndarray) -> np.ndarray: + """ + image: (N, 3) or flat length N*3 in [0,1] + lut1d: length 3*M (channel-major: R slice, G slice, B slice) + """ + flat = np.asarray(image, dtype=np.float32).reshape(-1) + result = np.empty_like(flat) + m = lut1d.size // 3 + scale = m - 1 + m_minus_2 = m - 2 + r_base, g_base, b_base = 0, m, m * 2 + + for idx in range(0, flat.size, 3): + for ch, base in enumerate((r_base, g_base, b_base)): + scaled = float(flat[idx + ch]) * scale + i_low = min(int(np.floor(scaled)), m_minus_2) + frac = scaled - i_low + low = float(lut1d[base + i_low]) + high = float(lut1d[base + i_low + 1]) + result[idx + ch] = np.clip(low * (1.0 - frac) + high * frac, 0.0, 1.0) + return result + + +def apply_3d_lut(image: np.ndarray, lut3d: np.ndarray, lut_size: int) -> np.ndarray: + """ + image: flat RGB + lut3d: length 3*M^3, layout c * M^3 + x * M^2 + y * M + z + where x=B, y=G, z=R. + """ + flat = np.asarray(image, dtype=np.float32).reshape(-1) + result = np.empty_like(flat) + m = lut_size + scale = m - 1 + m_minus_2 = m - 2 + m2 = m * m + m3 = m * m * m + r_base, g_base, b_base = 0, m3, m3 * 2 + + for idx in range(0, flat.size, 3): + x_scaled = float(flat[idx + 2]) * scale # B -> X + y_scaled = float(flat[idx + 1]) * scale # G -> Y + z_scaled = float(flat[idx + 0]) * scale # R -> Z + + x0 = min(int(np.floor(x_scaled)), m_minus_2) + y0 = min(int(np.floor(y_scaled)), m_minus_2) + z0 = min(int(np.floor(z_scaled)), m_minus_2) + x1, y1, z1 = x0 + 1, y0 + 1, z0 + 1 + fx, fy, fz = x_scaled - x0, y_scaled - y0, z_scaled - z0 + + i000 = x0 * m2 + y0 * m + z0 + i001 = x0 * m2 + y0 * m + z1 + i010 = x0 * m2 + y1 * m + z0 + i011 = x0 * m2 + y1 * m + z1 + i100 = x1 * m2 + y0 * m + z0 + i101 = x1 * m2 + y0 * m + z1 + i110 = x1 * m2 + y1 * m + z0 + i111 = x1 * m2 + y1 * m + z1 + + for ch, base in enumerate((r_base, g_base, b_base)): + c000 = lut3d[base + i000] + c001 = lut3d[base + i001] + c010 = lut3d[base + i010] + c011 = lut3d[base + i011] + c100 = lut3d[base + i100] + c101 = lut3d[base + i101] + c110 = lut3d[base + i110] + c111 = lut3d[base + i111] + + c00 = c000 * (1 - fx) + c100 * fx + c01 = c001 * (1 - fx) + c101 * fx + c10 = c010 * (1 - fx) + c110 * fx + c11 = c011 * (1 - fx) + c111 * fx + c0 = c00 * (1 - fy) + c10 * fy + c1 = c01 * (1 - fy) + c11 * fy + result[idx + ch] = np.clip(c0 * (1 - fz) + c1 * fz, 0.0, 1.0) + + return result + + +def transform_image( + image: np.ndarray, + lut1d: np.ndarray, + lut3d: np.ndarray, +) -> np.ndarray: + """1D LUT then 3D LUT (colorStrength=1 path in original).""" + intermediate = apply_1d_lut(image, lut1d) + m = lut1d.size // 3 + return apply_3d_lut(intermediate, lut3d, m) + + +def get_merged_lut( + lut1d: np.ndarray, + lut3d: np.ndarray, + merged_size: int = MERGED_LUT_SIZE, +) -> np.ndarray: + """ + Build 25³ merged LUT (RGBc interleaved: index (r,g,b)*3+c). + Builds RGBc interleaved 25³ merged LUT from 1D + 3D NN outputs. + """ + lut1d = np.asarray(lut1d, dtype=np.float32).reshape(-1) + lut3d = np.asarray(lut3d, dtype=np.float32).reshape(-1) + + hald = np.empty(merged_size * merged_size * merged_size * 3, dtype=np.float32) + interval = 1.0 / (merged_size - 1) + for r in range(merged_size): + for g in range(merged_size): + for b in range(merged_size): + index = (r * merged_size + g) * merged_size + b + hald[index * 3] = r * interval + hald[index * 3 + 1] = g * interval + hald[index * 3 + 2] = b * interval + + return transform_image(hald, lut1d, lut3d) + + +def reshape_merged_lut_for_apply(merged_lut3d: np.ndarray, size: int = MERGED_LUT_SIZE) -> np.ndarray: + """RGBc merged cube → cBGR flat layout for apply_3d_lut (transformImageWithMergedLut).""" + merged = np.asarray(merged_lut3d, dtype=np.float32).reshape(-1) + r = np.arange(size, dtype=np.int32) + g = np.arange(size, dtype=np.int32) + b = np.arange(size, dtype=np.int32) + c = np.arange(3, dtype=np.int32) + rr, gg, bb, cc = np.meshgrid(r, g, b, c, indexing="ij") + src = (rr * size * size + gg * size + bb) * 3 + cc + dst = cc * size**3 + bb * size**2 + gg * size + rr + reshaped = np.empty_like(merged) + reshaped[dst.ravel()] = merged[src.ravel()] + return reshaped + + +def _lut_cube_from_flat(lut_flat: np.ndarray, lut_size: int) -> np.ndarray: + """cBGR flat → (3, B, G, R) float32.""" + return np.asarray(lut_flat, dtype=np.float32).reshape(3, lut_size, lut_size, lut_size) + + +def apply_3d_lut_vectorized(image_hw3: np.ndarray, lut_flat: np.ndarray, lut_size: int) -> np.ndarray: + """Vectorized trilinear 3D LUT (same indexing as apply_3d_lut).""" + h, w, _ = image_hw3.shape + lut = _lut_cube_from_flat(lut_flat, lut_size) + scale = float(lut_size - 1) + rgb = np.asarray(image_hw3, dtype=np.float32).reshape(-1, 3) + r, g, b = rgb[:, 0], rgb[:, 1], rgb[:, 2] + + x = np.clip(b * scale, 0.0, lut_size - 1 - 1e-5) + y = np.clip(g * scale, 0.0, lut_size - 1 - 1e-5) + z = np.clip(r * scale, 0.0, lut_size - 1 - 1e-5) + + x0 = np.minimum(np.floor(x).astype(np.int32), lut_size - 2) + y0 = np.minimum(np.floor(y).astype(np.int32), lut_size - 2) + z0 = np.minimum(np.floor(z).astype(np.int32), lut_size - 2) + x1, y1, z1 = x0 + 1, y0 + 1, z0 + 1 + fx, fy, fz = (x - x0), (y - y0), (z - z0) + + out = np.empty_like(rgb) + for c in range(3): + c000 = lut[c, x0, y0, z0] + c001 = lut[c, x0, y0, z1] + c010 = lut[c, x0, y1, z0] + c011 = lut[c, x0, y1, z1] + c100 = lut[c, x1, y0, z0] + c101 = lut[c, x1, y0, z1] + c110 = lut[c, x1, y1, z0] + c111 = lut[c, x1, y1, z1] + + c00 = c000 * (1.0 - fx) + c100 * fx + c01 = c001 * (1.0 - fx) + c101 * fx + c10 = c010 * (1.0 - fx) + c110 * fx + c11 = c011 * (1.0 - fx) + c111 * fx + c0 = c00 * (1.0 - fy) + c10 * fy + c1 = c01 * (1.0 - fy) + c11 * fy + out[:, c] = np.clip(c0 * (1.0 - fz) + c1 * fz, 0.0, 1.0) + + return out.reshape(h, w, 3) + + +def _rgb_to_hue(r: float, g: float, b: float) -> float: + mx = max(r, g, b) + mn = min(r, g, b) + if mx == mn: + return 0.0 + d = mx - mn + if mx == r: + h = (g - b) / d + (6.0 if g < b else 0.0) + elif mx == g: + h = (b - r) / d + 2.0 + else: + h = (r - g) / d + 4.0 + return (h / 6.0) % 1.0 + + +def _hsv_to_rgb(h: float, s: float, v: float) -> tuple[float, float, float]: + h = float(np.clip(h, 0.0, 1.0)) + s = float(np.clip(s, 0.0, 1.0)) + v = float(np.clip(v, 0.0, 1.0)) + i = int(h * 6.0) + f = h * 6.0 - i + p = v * (1.0 - s) + q = v * (1.0 - f * s) + t = v * (1.0 - (1.0 - f) * s) + i = i % 6 + if i == 0: + r, g, b = v, t, p + elif i == 1: + r, g, b = q, v, p + elif i == 2: + r, g, b = p, v, t + elif i == 3: + r, g, b = p, q, v + elif i == 4: + r, g, b = t, p, v + else: + r, g, b = v, p, q + return r, g, b + + +def _rgb_to_hsv(r: float, g: float, b: float) -> tuple[float, float, float]: + mx = max(r, g, b) + mn = min(r, g, b) + d = mx - mn + if mx == 0.0: + return 0.0, 0.0, 0.0 + s = d / mx + if d == 0.0: + return 0.0, 0.0, mx + if mx == r: + h = (g - b) / d + (6.0 if g < b else 0.0) + elif mx == g: + h = (b - r) / d + 2.0 + else: + h = (r - g) / d + 4.0 + return (h / 6.0) % 1.0, s, mx + + +def _set_hue(r: float, g: float, b: float, hue: float) -> tuple[float, float, float]: + h, s, v = _rgb_to_hsv(r, g, b) + return _hsv_to_rgb(hue, s, v) + + +def _luma_rec709(r: float, g: float, b: float) -> float: + return 0.299 * r + 0.587 * g + 0.114 * b + + +def apply_color_light_strength( + image: np.ndarray, + result: np.ndarray, + color_strength: float, + light_strength: float = 1.0, + original_input: np.ndarray | None = None, +) -> np.ndarray: + """ + Color + luminance strength blend after LUT (lightStrength=1, no originalInput). + """ + src = np.asarray(image, dtype=np.float32).reshape(-1) + res = np.asarray(result, dtype=np.float32).reshape(-1) + out = np.empty_like(src) + inv_color = 1.0 - color_strength + inv_light = 1.0 - light_strength + + for idx in range(0, src.size, 3): + ir, ig, ib = float(src[idx]), float(src[idx + 1]), float(src[idx + 2]) + rr, rg, rb = float(res[idx]), float(res[idx + 1]), float(res[idx + 2]) + + mr = ir * inv_color + rr * color_strength + mg = ig * inv_color + rg * color_strength + mb = ib * inv_color + rb * color_strength + + result_luma = _luma_rec709(rr, rg, rb) + target_luma = result_luma + if original_input is not None and light_strength != 1.0: + orig = np.asarray(original_input, dtype=np.float32).reshape(-1) + or_, og, ob = float(orig[idx]), float(orig[idx + 1]), float(orig[idx + 2]) + original_luma = _luma_rec709(or_, og, ob) + target_luma = result_luma * light_strength + original_luma * inv_light + + mixed_luma = _luma_rec709(mr, mg, mb) + delta = target_luma - mixed_luma + cr, cg, cb = mr + delta, mg + delta, mb + delta + result_hue = _rgb_to_hue(mr, mg, mb) + cr, cg, cb = _set_hue(cr, cg, cb, result_hue) + out[idx] = np.clip(cr, 0.0, 1.0) + out[idx + 1] = np.clip(cg, 0.0, 1.0) + out[idx + 2] = np.clip(cb, 0.0, 1.0) + + return out + + +def transform_image_with_merged_lut(image: np.ndarray, merged_lut3d: np.ndarray) -> np.ndarray: + """Port of transformImageWithMergedLut() — apply 25³ merged cube to H×W×3 RGB.""" + hw3 = np.asarray(image, dtype=np.float32) + if hw3.ndim == 1: + side = int(round(hw3.size ** (1.0 / 3.0))) + hw3 = hw3.reshape(side, side, 3) if side * side * 3 == hw3.size else hw3.reshape(-1, 3)[:, None, :] + if hw3.ndim == 2: + hw3 = hw3.reshape(-1, 1, 3) + reshaped = reshape_merged_lut_for_apply(merged_lut3d, MERGED_LUT_SIZE) + return apply_3d_lut_vectorized(hw3, reshaped, MERGED_LUT_SIZE) + + +def apply_color_light_strength_vectorized( + image_hw3: np.ndarray, + result_hw3: np.ndarray, + color_strength: float, +) -> np.ndarray: + """Vectorized applyColorLightStrength (lightStrength=1, no originalInput).""" + src = np.asarray(image_hw3, dtype=np.float32).reshape(-1, 3) + res = np.asarray(result_hw3, dtype=np.float32).reshape(-1, 3) + inv = 1.0 - color_strength + mixed = src * inv + res * color_strength + + result_luma = 0.299 * res[:, 0] + 0.587 * res[:, 1] + 0.114 * res[:, 2] + mixed_luma = 0.299 * mixed[:, 0] + 0.587 * mixed[:, 1] + 0.114 * mixed[:, 2] + delta = (result_luma - mixed_luma)[:, np.newaxis] + corrected = mixed + delta + + mx = np.max(corrected, axis=1) + mn = np.min(corrected, axis=1) + d = mx - mn + hue = np.zeros(corrected.shape[0], dtype=np.float32) + nz = d > 1e-8 + r, g, b = corrected[:, 0], corrected[:, 1], corrected[:, 2] + rm = (mx == r) & nz + gm = (mx == g) & nz & ~rm + bm = nz & ~rm & ~gm + hue[rm] = ((g[rm] - b[rm]) / d[rm] + np.where(g[rm] < b[rm], 6.0, 0.0)) / 6.0 + hue[gm] = ((b[gm] - r[gm]) / d[gm] + 2.0) / 6.0 + hue[bm] = ((r[bm] - g[bm]) / d[bm] + 4.0) / 6.0 + hue = hue % 1.0 + + v = mx + s = np.where(mx > 1e-8, d / mx, 0.0) + hi = np.floor(hue * 6.0).astype(np.int32) % 6 + f = hue * 6.0 - np.floor(hue * 6.0) + p = v * (1.0 - s) + q = v * (1.0 - f * s) + t = v * (1.0 - (1.0 - f) * s) + out = np.empty_like(corrected) + for i in range(6): + m = hi == i + if not np.any(m): + continue + if i == 0: + out[m] = np.stack([v[m], t[m], p[m]], axis=1) + elif i == 1: + out[m] = np.stack([q[m], v[m], p[m]], axis=1) + elif i == 2: + out[m] = np.stack([p[m], v[m], t[m]], axis=1) + elif i == 3: + out[m] = np.stack([p[m], q[m], v[m]], axis=1) + elif i == 4: + out[m] = np.stack([t[m], p[m], v[m]], axis=1) + else: + out[m] = np.stack([v[m], p[m], q[m]], axis=1) + return np.clip(out, 0.0, 1.0).reshape(image_hw3.shape) + + +def apply_merged_lut_preview( + rgb: np.ndarray, + merged_lut3d: np.ndarray, + color_strength: float = 1.0, +) -> np.ndarray: + """ + Apply merged LUT to H×W×3 float RGB [0,1], then blend with colorStrength. + + - 0% → unchanged source + - 100% → full merged LUT + - between → color/light strength blend + """ + rgb = np.asarray(rgb, dtype=np.float32) + if color_strength <= 0.0: + return rgb.copy() + lutted = transform_image_with_merged_lut(rgb, merged_lut3d) + if color_strength >= 1.0 - 1e-6: + return lutted + return apply_color_light_strength_vectorized(rgb, lutted, color_strength) + + +def nn_outputs_to_flat(lut1d: np.ndarray, lut3d: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """TF.js outputs [1,3,17] and [1,3,17,17,17] → flat arrays for get_merged_lut.""" + lut1d_flat = np.asarray(lut1d, dtype=np.float32).reshape(-1) + lut3d_arr = np.asarray(lut3d, dtype=np.float32) + if lut3d_arr.ndim == 5: + # NCHW-style from graph: [1, 3, 17, 17, 17] → channel-major flat like TF dataSync + lut3d_flat = lut3d_arr.reshape(-1) + else: + lut3d_flat = lut3d_arr.reshape(-1) + return lut1d_flat, lut3d_flat diff --git a/models/color_match.onnx b/models/color_match.onnx new file mode 100644 index 0000000..b99edfd Binary files /dev/null and b/models/color_match.onnx differ diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..af0c399 --- /dev/null +++ b/nodes.py @@ -0,0 +1,165 @@ +""" +ComfyUI nodes for MAGICMATCH — neural color match via ONNX. +""" + +from __future__ import annotations + +import numpy as np +import torch + +from .magicmatch_core.inference import build_merged_lut +from .magicmatch_core.lut import apply_merged_lut_preview + + +class MagicMatchLUT: + """Cached 25³ merged LUT between Build and Preview nodes.""" + + __slots__ = ("merged_lut",) + + def __init__(self, merged_lut: np.ndarray) -> None: + self.merged_lut = np.asarray(merged_lut, dtype=np.float32).reshape(-1) + + +def _image_batch_to_hwc(image: torch.Tensor) -> np.ndarray: + if image.ndim != 4 or image.shape[-1] != 3: + raise ValueError(f"MAGICMATCH: expected IMAGE [B,H,W,3], got {tuple(image.shape)}") + if image.shape[0] != 1: + raise ValueError( + f"MAGICMATCH: batch size {image.shape[0]} not supported — use batch size 1." + ) + return np.clip(image[0].detach().cpu().numpy(), 0.0, 1.0).astype(np.float32) + + +def _hwc_to_image(hwc: np.ndarray) -> torch.Tensor: + hwc = np.clip(np.asarray(hwc, dtype=np.float32), 0.0, 1.0) + return torch.from_numpy(hwc[np.newaxis, ...]) + + +class MagicMatchBuild: + """ + Analyze source + reference (ONNX). Wire output into MagicMatch Preview. + Re-run only when source or reference images change. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "source": ("IMAGE",), + "reference": ("IMAGE",), + }, + } + + RETURN_TYPES = ("MAGICMATCH_LUT",) + RETURN_NAMES = ("lut",) + FUNCTION = "build" + CATEGORY = "MAGICMATCH" + DESCRIPTION = "Build merged color-match LUT from source + reference (run once per image pair)." + + def build(self, source: torch.Tensor, reference: torch.Tensor) -> tuple[MagicMatchLUT]: + src = _image_batch_to_hwc(source) + ref = _image_batch_to_hwc(reference) + merged = build_merged_lut(src, ref) + return (MagicMatchLUT(merged),) + + +class MagicMatchPreview: + """ + Apply cached LUT with strength slider — fast while tuning. + Connect to Preview Image, adjust strength, re-queue. + Then connect downstream to save/export. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "source": ("IMAGE",), + "lut": ("MAGICMATCH_LUT",), + "strength": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + }, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "preview" + CATEGORY = "MAGICMATCH" + DESCRIPTION = ( + "Preview color match at chosen strength (0=source, 1=full match). " + "Re-queue after moving the slider — only this node re-runs." + ) + + def preview( + self, + source: torch.Tensor, + lut: MagicMatchLUT, + strength: float, + ) -> tuple[torch.Tensor]: + src = _image_batch_to_hwc(source) + strength = float(np.clip(strength, 0.0, 1.0)) + out = apply_merged_lut_preview(src, lut.merged_lut, strength) + return (_hwc_to_image(out),) + + +class MagicMatch: + """Source + reference + strength in one node (re-runs ONNX when strength changes).""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "source": ("IMAGE",), + "reference": ("IMAGE",), + "strength": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + }, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "match" + CATEGORY = "MAGICMATCH" + DESCRIPTION = "Single-node color match (re-runs neural net when strength changes)." + + def match( + self, + source: torch.Tensor, + reference: torch.Tensor, + strength: float, + ) -> tuple[torch.Tensor]: + src = _image_batch_to_hwc(source) + ref = _image_batch_to_hwc(reference) + merged = build_merged_lut(src, ref) + strength = float(np.clip(strength, 0.0, 1.0)) + out = apply_merged_lut_preview(src, merged, strength) + return (_hwc_to_image(out),) + + +NODE_CLASS_MAPPINGS = { + "MagicMatchBuild": MagicMatchBuild, + "MagicMatchPreview": MagicMatchPreview, + "MagicMatch": MagicMatch, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MagicMatchBuild": "MagicMatch Build LUT", + "MagicMatchPreview": "MagicMatch Preview (strength)", + "MagicMatch": "MagicMatch (one-shot)", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c5c0933 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +onnxruntime>=1.18 +numpy>=1.24