nodes/lookuters22-MAGICMATCH
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial release: MAGICMATCH ComfyUI color-match nodes with bundled ONNX model.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
poxular
2026-06-01 16:51:00 +02:00
co-authored by Cursor
commit 4d2d3d4236
10 changed files with 761 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3
View File
@@ -0,0 +1,3 @@
__pycache__/
*.pyc
.pytest_cache/
LICENSE
+21
View File
@@ -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.
README.md
+65
View File
@@ -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.
__init__.py
+7
View File
@@ -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"]
magicmatch_core/__init__.py
+12
View File
@@ -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",
]
magicmatch_core/inference.py
+84
View File
@@ -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))
magicmatch_core/lut.py
+402
View File
@@ -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
models/color_match.onnx
BIN
View File
Binary file not shown.
nodes.py
+165
View File
@@ -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)",
}
requirements.txt
+2
View File
@@ -0,0 +1,2 @@
onnxruntime>=1.18
numpy>=1.24