Add Torch GPU pipeline for CUDA ComfyUI nodes.

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

README.md

@@ -81,7 +81,10 @@ pip install -r MAGICMATCH/requirements-cuda.txt
 ```
 
 Requires NVIDIA driver + CUDA libs compatible with your `onnxruntime-gpu` wheel.
-On Windows or machines without CUDA, CUDA nodes fall back to CPU EP automatically.
+ComfyUI already provides PyTorch with CUDA — the GPU pipeline uses Torch for develop,
+detection buffers, and LUT apply (ONNX uses CUDA EP).
+
+On machines without CUDA, CUDA nodes fall back to CPU EP + CPU Torch automatically.
 
 ### ComfyUI nodes
 
@@ -93,9 +96,13 @@ After restart, look under **MAGICMATCH/CUDA (experimental)**:
 | `MagicMatchPreviewCUDA` | MagicMatch Preview (CUDA LUT) |
 | `MagicMatchCUDA` | MagicMatch one-shot (CUDA) |
 
-Preview/apply still uses the CPU NumPy develop stack; only ONNX `sess.run` moves to GPU.
+**GPU pipeline (CUDA nodes):** CUDA ONNX + GPU detection buffers + GPU develop@1600 + GPU full-res apply.
 
-Set `MAGICMATCH_CUDA_NODES=0` before starting ComfyUI to hide CUDA nodes (CPU-only menu).
+**CPU default nodes** under **MAGICMATCH** are unchanged.
+
+Set `MAGICMATCH_CUDA_NODES=0` before starting ComfyUI to hide CUDA nodes.
+
+Use `python3.12` on RunPod if `python` is not on PATH.
 
 ### Benchmark / parity
 

magicmatch_core/gpu/__init__.py

@@ -0,0 +1,11 @@
+"""Torch/CUDA pixel pipeline (develop, detection buffers, LUT apply). CPU parity path unchanged."""
+
+from .device import gpu_pipeline_available, get_torch_device
+from .pipeline import apply_probe_export_gpu, build_merged_lut_probe_style_gpu
+
+__all__ = [
+    "apply_probe_export_gpu",
+    "build_merged_lut_probe_style_gpu",
+    "get_torch_device",
+    "gpu_pipeline_available",
+]

magicmatch_core/gpu/color_torch.py

@@ -0,0 +1,89 @@
+"""Torch color-space helpers matching polarr_color_space.py."""
+
+from __future__ import annotations
+
+import torch
+
+from ..polarr_color_space import (
+    GAMMA_SRGB,
+    MAT_LUT_DEC_SRGB,
+    MAT_LUT_ENC_SRGB,
+    MAT_PROPHOTO_TO_SRGB,
+    MAT_SRGB_TO_PROPHOTO,
+    MAT_SRGB_TO_XYZ,
+    MAT_XYZ_TO_PROPHOTO,
+)
+
+
+def _const(device: torch.device, arr) -> torch.Tensor:
+    return torch.tensor(arr, dtype=torch.float32, device=device)
+
+
+def srgb_to_linear(t: torch.Tensor) -> torch.Tensor:
+    t = torch.clamp(t, 0.0, 1.0)
+    return torch.where(t <= 0.04045, t / 12.92, ((t + 0.055) / 1.055) ** 2.4)
+
+
+def linear_to_srgb(t: torch.Tensor) -> torch.Tensor:
+    t = torch.clamp(t, 0.0, None)
+    return torch.clamp(
+        torch.where(t <= 0.0031308, t * 12.92, 1.055 * torch.pow(t, 1.0 / 2.4) - 0.055),
+        0.0,
+        1.0,
+    )
+
+
+def srgb_to_prophoto(t: torch.Tensor) -> torch.Tensor:
+    shape = t.shape
+    flat = t.reshape(-1, 3)
+    lin = srgb_to_linear(flat)
+    m = _const(flat.device, MAT_SRGB_TO_PROPHOTO)
+    return (lin.double() @ m.T).float().reshape(shape)
+
+
+def prophoto_to_srgb(t: torch.Tensor) -> torch.Tensor:
+    shape = t.shape
+    flat = t.reshape(-1, 3)
+    m = _const(flat.device, MAT_PROPHOTO_TO_SRGB)
+    xyz = (flat.double() @ m.T).float()
+    return linear_to_srgb(xyz).reshape(shape)
+
+
+def lut_primaries_encode(t: torch.Tensor, primaries: int) -> torch.Tensor:
+    if primaries != 0:
+        raise NotImplementedError("GPU LUT path supports srgb primaries only")
+    shape = t.shape
+    flat = t.reshape(-1, 3)
+    m = _const(flat.device, MAT_LUT_ENC_SRGB)
+    return (flat.double() @ m.T).float().reshape(shape)
+
+
+def lut_primaries_decode(t: torch.Tensor, primaries: int) -> torch.Tensor:
+    if primaries != 0:
+        raise NotImplementedError("GPU LUT path supports srgb primaries only")
+    shape = t.shape
+    flat = t.reshape(-1, 3)
+    m = _const(flat.device, MAT_LUT_DEC_SRGB)
+    return (flat.double() @ m.T).float().reshape(shape)
+
+
+def lut_gamma_encode(t: torch.Tensor, gamma: int) -> torch.Tensor:
+    if gamma == GAMMA_SRGB:
+        return linear_to_srgb(t)
+    return t
+
+
+def lut_gamma_decode(t: torch.Tensor, gamma: int) -> torch.Tensor:
+    if gamma == GAMMA_SRGB:
+        return srgb_to_linear(t)
+    return t
+
+
+def srgb_to_prophoto_via_xyz(t: torch.Tensor) -> torch.Tensor:
+    """bitmap.frag path: sRGB -> linear -> XYZ -> ProPhoto."""
+    flat = t.reshape(-1, 3)
+    lin = srgb_to_linear(flat)
+    m1 = _const(flat.device, MAT_SRGB_TO_XYZ)
+    m2 = _const(flat.device, MAT_XYZ_TO_PROPHOTO)
+    xyz = (lin.double() @ m1.T).float()
+    return (xyz.double() @ m2.T).float().reshape(t.shape)

magicmatch_core/gpu/develop_torch.py

@@ -0,0 +1,273 @@
+"""GPU develop stack — port of probe_parity/develop.py render_srgb_develop."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from ..polarr_color_space import ENCODING_PRESETS
+from ..probe_parity.wb import DEFAULT_AS_SHOT_TEMP, DEFAULT_AS_SHOT_TINT, build_wb_matrix
+from .color_torch import linear_to_srgb, prophoto_to_srgb, srgb_to_linear, srgb_to_prophoto
+from .device import get_torch_device, hwc_numpy_to_torch, hwc_torch_to_numpy
+from .lut_torch import apply_polarr_rgb_lut_prophoto_torch
+
+PROPHOTO_LUMA = (0.242655, 0.755158, 0.002187)
+T_A, T_B, T_C, T_D, T_E = 1.2, 0.0, 0.96, 0.22, 0.02
+
+
+def _luma(rgb: torch.Tensor) -> torch.Tensor:
+    w = rgb.new_tensor(PROPHOTO_LUMA)
+    return (rgb * w).sum(dim=-1)
+
+
+def _smoothstep(edge0, edge1, x: torch.Tensor) -> torch.Tensor:
+    t = torch.clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0)
+    return t * t * (3.0 - 2.0 * t)
+
+
+def _rgb_to_hsv(rgb: torch.Tensor) -> torch.Tensor:
+    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
+    mx = torch.maximum(torch.maximum(r, g), b)
+    mn = torch.minimum(torch.minimum(r, g), b)
+    d = mx - mn
+    h = torch.zeros_like(mx)
+    safe_d = torch.where(d > 1e-8, d, torch.ones_like(d))
+    h = torch.where((mx == r) & (d > 1e-8), ((g - b) / safe_d + torch.where(g < b, 6.0, 0.0)) / 6.0, h)
+    h = torch.where((mx == g) & (d > 1e-8), ((b - r) / safe_d + 2.0) / 6.0, h)
+    h = torch.where((mx == b) & (d > 1e-8), ((r - g) / safe_d + 4.0) / 6.0, h)
+    h = h % 1.0
+    s = torch.where(mx <= 0, torch.zeros_like(mx), d / torch.maximum(mx, torch.full_like(mx, 1e-8)))
+    return torch.stack([h, s, mx], dim=-1)
+
+
+def _hsv_to_rgb(hsv: torch.Tensor) -> torch.Tensor:
+    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
+    i = (torch.floor(h * 6.0).to(torch.int64) % 6).to(torch.int64)
+    f = h * 6.0 - torch.floor(h * 6.0)
+    p = v * (1.0 - s)
+    q = v * (1.0 - f * s)
+    t = v * (1.0 - (1.0 - f) * s)
+    out = torch.zeros_like(hsv)
+    for idx, (c0, c1, c2) in enumerate([(v, t, p), (q, v, p), (p, v, t), (p, q, v), (t, p, v), (v, p, q)]):
+        m = i == idx
+        out[..., 0] = torch.where(m, c0, out[..., 0])
+        out[..., 1] = torch.where(m, c1, out[..., 1])
+        out[..., 2] = torch.where(m, c2, out[..., 2])
+    return out
+
+
+def _set_hue(rgb: torch.Tensor, hue: torch.Tensor) -> torch.Tensor:
+    hsv = _rgb_to_hsv(rgb)
+    hsv[..., 0] = hue
+    return _hsv_to_rgb(hsv)
+
+
+def _rgb_to_hue(rgb: torch.Tensor) -> torch.Tensor:
+    return _rgb_to_hsv(rgb)[..., 0]
+
+
+def _tonemap_inv(rgb: torch.Tensor) -> torch.Tensor:
+    out = torch.empty_like(rgb)
+    for ch in range(3):
+        x = rgb[..., ch]
+        a = x * T_C - T_A
+        b_coef = x * T_D - T_B
+        c = x * T_E
+        disc = torch.clamp(b_coef * b_coef - 4.0 * a * c, min=0.0)
+        out[..., ch] = (-b_coef - torch.sqrt(disc)) / (2.0 * a + 1e-12)
+    return out
+
+
+def _tonemap(rgb: torch.Tensor) -> torch.Tensor:
+    x = torch.clamp(rgb, 0.0, 32.0)
+    return (x * (T_A * x + T_B)) / (x * (T_C * x + T_D) + T_E)
+
+
+def _apply_wb(rgb: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
+    return torch.clamp(rgb @ matrix.T, min=0.0)
+
+
+def _apply_shadows_highlights(
+    rgb: torch.Tensor, shadows: float, highlights: float, luma_map: torch.Tensor
+) -> torch.Tensor:
+    color_luma = _luma(rgb)
+    mixed = torch.clamp(color_luma + luma_map, min=0.0) * 0.5
+    luma = _smoothstep(0.0, 1.6, torch.sqrt(mixed))
+    specular = torch.clamp(1.0 - luma * luma, min=0.0)
+    hi_mask = _smoothstep(0.05, 1.0, luma_map) * 1.2
+    if highlights > 0.0:
+        hi_mask = hi_mask * torch.clamp(1.0 - color_luma, min=0.0) * 1.1
+    sh_mask = 1.0 - _smoothstep(0.0, 0.5, luma_map)
+    exposure_amount = shadows * sh_mask + highlights * hi_mask * specular
+    return rgb * (2.0 ** (exposure_amount * 2.0)).unsqueeze(-1)
+
+
+def _apply_saturation(rgb: torch.Tensor, saturation: float) -> torch.Tensor:
+    luma = _luma(rgb).unsqueeze(-1)
+    mix_amount = min(-saturation, 1.0)
+    inv = 1.0 - mix_amount
+    return torch.clamp(rgb * inv + luma * mix_amount, min=0.0)
+
+
+def _contract_blacks_whites(rgb: torch.Tensor, whites: float, blacks: float) -> torch.Tensor:
+    blacks = max(blacks, 0.0)
+    whites = min(whites, 0.0)
+    luma = _luma(rgb)
+    x = torch.clamp(torch.sqrt(luma), 0.0, 1.0)
+    whites_mask = _smoothstep(0.05, 1.0, x)
+    specular = torch.clamp(1.0 - luma * luma * luma * luma, min=0.0)
+    whites_mask = whites_mask * specular
+    blacks_mask = _smoothstep(0.8, -0.4, x)
+    mult = (2.0 ** (whites * whites_mask + blacks * blacks_mask)).unsqueeze(-1)
+    return rgb * mult
+
+
+def _falloff(x: torch.Tensor, m: float) -> torch.Tensor:
+    xm = x * m
+    xm = 1.0 / ((xm + 1.0) ** 2)
+    mm = 1.0 / ((m + 1.0) ** 2)
+    return (xm - mm) / (1.0 - mm)
+
+
+def _expand_whites_blacks(rgb: torch.Tensor, whites: float, blacks: float) -> torch.Tensor:
+    blacks = min(blacks, 0.0)
+    whites = max(whites, 0.0)
+    x = torch.clamp(_luma(rgb), 0.0, 1.0)
+    whites_mask = _falloff(x, -0.7)
+    whites_mask_low = _smoothstep(0.0, 1.0, torch.sqrt(x))
+    blacks_mask = _falloff(1.0 - x, -0.9)
+    mask = whites_mask * blacks_mask
+    blacks_v = blacks * 0.33 * mask
+    whites_v = 1.0 + whites * 0.33 * mask * whites_mask_low
+    out = (rgb * (1.0 - blacks_v.unsqueeze(-1)) + blacks_v.unsqueeze(-1)) * whites_v.unsqueeze(-1)
+    return torch.clamp(out, 0.0, 1.0)
+
+
+def _contrast_curve(contrast: float, width: int, device: torch.device) -> torch.Tensor:
+    if contrast < 0:
+        pts = [(0, 0), (8, 21), (16, 34), (32, 54), (64, 88), (96, 115), (128, 135), (160, 152), (192, 174), (224, 206), (240, 228), (255, 255)]
+    else:
+        pts = [(0, 0), (32, 8), (64, 35), (96, 71), (128, 117), (160, 171), (192, 212), (224, 241), (255, 255)]
+    intensity = abs(contrast)
+    xs = torch.tensor([p[0] / 255.0 for p in pts], dtype=torch.float64, device=device)
+    ys = torch.tensor(
+        [p[1] / 255.0 + intensity * (p[1] / 255.0 - p[0] / 255.0) * (1.0 - p[0] / 255.0) for p in pts],
+        dtype=torch.float64,
+        device=device,
+    )
+    xs = torch.clamp(xs, 0, 1)
+    ys = torch.clamp(ys, 0, 1)
+    samples = torch.linspace(0, 1, width, device=device, dtype=torch.float32)
+    return torch.from_numpy(
+        __import__("numpy").interp(samples.cpu().numpy(), xs.cpu().numpy(), ys.cpu().numpy()).astype("float32")
+    ).to(device)
+
+
+def _apply_curve(rgb: torch.Tensor, curve: torch.Tensor) -> torch.Tensor:
+    rgb = torch.clamp(rgb, 0.0, 1.0)
+    size = curve.shape[0]
+    out = torch.empty_like(rgb)
+    for ch in range(3):
+        values = rgb[..., ch]
+        idx = torch.clamp(torch.floor(values * (size - 1)).to(torch.int64), 0, size - 2)
+        p0x = idx.float() / (size - 1)
+        p1x = (idx + 1).float() / (size - 1)
+        t = (values - p0x) / torch.clamp(p1x - p0x, min=1e-12)
+        out[..., ch] = torch.clamp(curve[idx] + t * (curve[idx + 1] - curve[idx]), 0.0, 1.0)
+    return out
+
+
+@torch.inference_mode()
+def render_srgb_develop_torch(
+    srgb_hwc: torch.Tensor,
+    adjustments: dict,
+    *,
+    merged_lut=None,
+    lut_strength: float = 1.0,
+    lut_encoding: str = "srgb_srgb",
+    input_tone_map_inversed: bool = False,
+    as_shot_temp: float = DEFAULT_AS_SHOT_TEMP,
+    as_shot_tint: float = DEFAULT_AS_SHOT_TINT,
+) -> torch.Tensor:
+    """Develop + optional LUT on GPU; input/output H×W×3 float [0,1] on device."""
+    shape = srgb_hwc.shape
+    device = srgb_hwc.device
+    if input_tone_map_inversed:
+        rgb = torch.clamp(srgb_hwc, 0.0, 1.0).reshape(-1, 3)
+    else:
+        rgb = srgb_to_prophoto(srgb_hwc).reshape(-1, 3)
+        hue0 = _rgb_to_hue(rgb)
+        rgb = _tonemap_inv(rgb)
+        rgb = _set_hue(rgb, hue0)
+        rgb = torch.clamp(rgb, min=0.0)
+
+    gamma = linear_to_srgb(rgb)
+    luma_map = _luma(gamma)
+
+    wb_np = build_wb_matrix(
+        adjustments.get("temperature", as_shot_temp),
+        adjustments.get("tint", as_shot_tint),
+        as_shot_temp=as_shot_temp,
+        as_shot_tint=as_shot_tint,
+    )
+    wb = torch.tensor(wb_np, dtype=torch.float32, device=device)
+    rgb = _apply_wb(rgb, wb)
+
+    exposure_factor = float(2.0 ** adjustments.get("exposure", 0.0))
+    shadows = float(adjustments.get("shadows", 0.0))
+    highlights = float(adjustments.get("highlights", 0.0))
+    whites = float(adjustments.get("whites", 0.0))
+    blacks = float(adjustments.get("blacks", 0.0))
+    saturation = float(adjustments.get("saturation", 0.0))
+    contrast = float(adjustments.get("contrast", 0.0))
+
+    if exposure_factor != 1.0:
+        rgb = rgb * exposure_factor
+    if shadows != 0.0 or highlights != 0.0:
+        rgb = _apply_shadows_highlights(rgb, shadows, highlights, luma_map)
+    if whites < 0.0 or blacks > 0.0:
+        rgb = _contract_blacks_whites(rgb, whites, blacks)
+
+    hue = _rgb_to_hue(rgb)
+    rgb = torch.clamp(rgb, 0.0, 1.0)
+    rgb = _set_hue(rgb, hue)
+    if saturation != 0.0:
+        rgb = _apply_saturation(rgb, saturation)
+
+    hue = _rgb_to_hue(rgb)
+    rgb = _tonemap(rgb)
+    rgb = _set_hue(rgb, hue)
+    rgb = linear_to_srgb(rgb)
+    hue = _rgb_to_hue(rgb)
+    if whites > 0.0 or blacks < 0.0:
+        rgb = _expand_whites_blacks(rgb, whites, blacks)
+    if contrast != 0.0:
+        rgb = _apply_curve(rgb, _contrast_curve(contrast, 256, device))
+    rgb = _set_hue(rgb, hue)
+    rgb = srgb_to_linear(rgb)
+
+    if merged_lut is not None and lut_strength > 0.0:
+        gamma_id, primaries_id = ENCODING_PRESETS.get(lut_encoding, ENCODING_PRESETS["srgb_srgb"])
+        rgb = apply_polarr_rgb_lut_prophoto_torch(
+            rgb.reshape(shape),
+            merged_lut,
+            lut_strength,
+            rgb_gamma=gamma_id,
+            rgb_primaries=primaries_id,
+        ).reshape(-1, 3)
+
+    return prophoto_to_srgb(rgb.reshape(shape))
+
+
+def render_srgb_develop_numpy(
+    srgb_hwc,
+    adjustments: dict,
+    *,
+    device: torch.device | None = None,
+    **kwargs,
+):
+    """NumPy in/out wrapper for pipeline integration."""
+    device = device or get_torch_device()
+    t = hwc_numpy_to_torch(srgb_hwc, device)
+    out = render_srgb_develop_torch(t, adjustments, **kwargs)
+    return hwc_torch_to_numpy(out)

magicmatch_core/gpu/device.py

@@ -0,0 +1,29 @@
+"""Torch device selection for MAGICMATCH GPU pipeline."""
+
+from __future__ import annotations
+
+import torch
+
+
+def get_torch_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def gpu_pipeline_available() -> bool:
+    return torch.cuda.is_available()
+
+
+def hwc_numpy_to_torch(hwc, device: torch.device | None = None) -> torch.Tensor:
+    import numpy as np
+
+    device = device or get_torch_device()
+    t = torch.from_numpy(np.asarray(hwc, dtype=np.float32)).to(device)
+    if t.ndim == 3:
+        return t
+    raise ValueError(f"expected H×W×3, got {tuple(t.shape)}")
+
+
+def hwc_torch_to_numpy(t: torch.Tensor):
+    return t.detach().float().clamp(0.0, 1.0).cpu().numpy()

magicmatch_core/gpu/lut_torch.py

@@ -0,0 +1,77 @@
+"""GPU Polarr 25³ RGB LUT apply."""
+
+from __future__ import annotations
+
+import torch
+
+from .color_torch import lut_gamma_decode, lut_gamma_encode, lut_primaries_decode, lut_primaries_encode
+
+LUT_SIZE = 25
+
+
+def _merged_lut_to_texture(merged_lut_rgb, device: torch.device) -> torch.Tensor:
+    import numpy as np
+
+    tex = np.asarray(merged_lut_rgb, dtype=np.float32).reshape(-1)
+    expected = LUT_SIZE * LUT_SIZE * LUT_SIZE * 3
+    if tex.size != expected:
+        raise ValueError(f"merged LUT must be {expected} floats, got {tex.size}")
+    return torch.from_numpy(tex.reshape(LUT_SIZE, LUT_SIZE * LUT_SIZE, 3)).to(device)
+
+
+def _sample_bilinear_hw3(tex: torch.Tensor, u: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    height, width = tex.shape[:2]
+    x = torch.clamp(u, 0.0, 1.0) * (width - 1)
+    y = torch.clamp(v, 0.0, 1.0) * (height - 1)
+    x0 = torch.floor(x).to(torch.int64)
+    y0 = torch.floor(y).to(torch.int64)
+    x1 = torch.minimum(x0 + 1, torch.tensor(width - 1, device=tex.device))
+    y1 = torch.minimum(y0 + 1, torch.tensor(height - 1, device=tex.device))
+    fx = (x - x0.float()).unsqueeze(-1)
+    fy = (y - y0.float()).unsqueeze(-1)
+    c00 = tex[y0, x0]
+    c01 = tex[y0, x1]
+    c10 = tex[y1, x0]
+    c11 = tex[y1, x1]
+    c0 = c00 * (1.0 - fx) + c01 * fx
+    c1 = c10 * (1.0 - fx) + c11 * fx
+    return c0 * (1.0 - fy) + c1 * fy
+
+
+@torch.inference_mode()
+def apply_polarr_rgb_lut_prophoto_torch(
+    prophoto_hwc: torch.Tensor,
+    merged_lut_rgb,
+    strength: float = 1.0,
+    *,
+    rgb_gamma: int = 1,
+    rgb_primaries: int = 0,
+) -> torch.Tensor:
+    device = prophoto_hwc.device
+    shape = prophoto_hwc.shape
+    orig = prophoto_hwc.reshape(-1, 3)
+    tex = _merged_lut_to_texture(merged_lut_rgb, device)
+
+    tmp = lut_primaries_encode(prophoto_hwc, rgb_primaries)
+    tmp = torch.clamp(tmp, 0.0, 1.0)
+    tmp = lut_gamma_encode(tmp, rgb_gamma).reshape(-1, 3)
+
+    size_index = float(LUT_SIZE - 1)
+    r, g, b = tmp[:, 0], tmp[:, 1], tmp[:, 2]
+    tex_x_base = (b * size_index + 0.5) / LUT_SIZE
+    tex_y = (r * size_index + 0.5) / LUT_SIZE
+    z = g * size_index
+    z_floor = torch.floor(z).to(torch.int64)
+    z_next = torch.minimum(z_floor + 1, torch.tensor(int(size_index), device=device))
+    zf = (z_floor.float() + tex_x_base) / LUT_SIZE
+    zc = (z_next.float() + tex_x_base) / LUT_SIZE
+
+    col1 = _sample_bilinear_hw3(tex, zf, tex_y)
+    col2 = _sample_bilinear_hw3(tex, zc, tex_y)
+    fract_z = (z - z_floor.float()).unsqueeze(-1)
+    mapped = col1 * (1.0 - fract_z) + col2 * fract_z
+
+    mapped = lut_gamma_decode(mapped.reshape(shape), rgb_gamma).reshape(-1, 3)
+    mapped = lut_primaries_decode(mapped.reshape(shape), rgb_primaries).reshape(-1, 3)
+    s = float(torch.clamp(torch.tensor(strength), 0.0, 1.0))
+    return torch.clamp(orig * (1.0 - s) + mapped * s, 0.0, 1.0).reshape(shape)

magicmatch_core/gpu/pipeline.py

@@ -0,0 +1,91 @@
+"""Full GPU-accelerated probe build + apply (ONNX CUDA + Torch develop/buffers)."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ..inference_cuda import run_inference_from_images_cuda
+from ..lut import get_merged_lut, nn_outputs_to_flat
+from ..probe_parity.base_adjustments_cuda import estimate_base_adjustments_cuda
+from ..probe_parity.profile_stage import ProfileStageFlags, normalize_profile_stage, profile_stage_flags
+from ..probe_parity.reference import NET_LONG_EDGE, fit_long_edge, prepare_net_reference, prepare_worker_bitmap_source
+from ..probe_parity.wb import DEFAULT_AS_SHOT_TEMP, DEFAULT_AS_SHOT_TINT
+from .develop_torch import render_srgb_develop_torch
+from .device import get_torch_device, gpu_pipeline_available, hwc_numpy_to_torch, hwc_torch_to_numpy
+
+
+def _render_for_net_gpu(source_hwc: np.ndarray, adjustments: dict, device) -> np.ndarray:
+    long = fit_long_edge(source_hwc, NET_LONG_EDGE)
+    t = hwc_numpy_to_torch(long, device)
+    out = render_srgb_develop_torch(t, adjustments)
+    return hwc_torch_to_numpy(out)
+
+
+def build_merged_lut_probe_style_gpu(
+    source_hwc: np.ndarray,
+    reference_hwc: np.ndarray,
+) -> tuple[np.ndarray, dict]:
+    """
+    Probe-parity build with GPU develop + CUDA ONNX.
+    Scene extract still uses CPU stats on GPU-generated detection buffers.
+    """
+    device = get_torch_device()
+    source_hwc = prepare_worker_bitmap_source(source_hwc)
+    base = estimate_base_adjustments_cuda(
+        source_hwc,
+        worker_feed_prepared=True,
+        use_gpu_detection_buffers=gpu_pipeline_available(),
+    )
+    net_source = _render_for_net_gpu(source_hwc, base, device)
+    net_ref = prepare_net_reference(reference_hwc)
+    lut3d, lut1d = run_inference_from_images_cuda(net_source, net_ref)
+    merged = get_merged_lut(*nn_outputs_to_flat(lut1d, lut3d))
+    return merged, base
+
+
+def apply_probe_export_gpu(
+    source_hwc: np.ndarray,
+    merged_lut: np.ndarray,
+    strength: float,
+    *,
+    base_adjustments: dict | None = None,
+    profile_stage: str = "current_profile_stages",
+    lut_encoding: str = "srgb_srgb",
+) -> np.ndarray:
+    """Full-res develop + LUT on GPU when CUDA is available."""
+    strength = float(np.clip(strength, 0.0, 1.0))
+    if strength <= 0.0:
+        return np.asarray(source_hwc, dtype=np.float32).copy()
+
+    if not gpu_pipeline_available():
+        from ..probe_parity.pipeline import apply_probe_export
+
+        return apply_probe_export(
+            source_hwc,
+            merged_lut,
+            strength,
+            base_adjustments=base_adjustments,
+            profile_stage=profile_stage,
+            lut_encoding=lut_encoding,
+        )
+
+    device = get_torch_device()
+    source_hwc = prepare_worker_bitmap_source(source_hwc)
+    base = base_adjustments
+    if base is None:
+        base = estimate_base_adjustments_cuda(source_hwc, use_gpu_detection_buffers=True)
+    stage_name = normalize_profile_stage(profile_stage)
+    stage: ProfileStageFlags = profile_stage_flags(stage_name)
+    lut_adjustments = {**base, "userLutStrength": strength}
+    t = hwc_numpy_to_torch(source_hwc, device)
+    out = render_srgb_develop_torch(
+        t,
+        lut_adjustments,
+        merged_lut=merged_lut,
+        lut_strength=strength,
+        lut_encoding=lut_encoding,
+        force_color_look=stage["forceColorLookTableWithUserLut"],
+        as_shot_temp=DEFAULT_AS_SHOT_TEMP,
+        as_shot_tint=DEFAULT_AS_SHOT_TINT,
+    )
+    return hwc_torch_to_numpy(out)

magicmatch_core/gpu/reference_torch.py

@@ -0,0 +1,145 @@
+"""GPU detection buffer chain — port of probe_parity/reference.py hot paths."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+from ..probe_parity.reference import fit_to_size
+from .color_torch import linear_to_srgb, srgb_to_prophoto_via_xyz
+from .develop_torch import _rgb_to_hue, _set_hue, _tonemap, _tonemap_inv
+from .device import get_torch_device, hwc_torch_to_numpy
+
+DETECTION_LARGE_EDGE = 2000
+DETECTION_SMALL_EDGE = 300
+
+
+def _to_bchw(t: torch.Tensor) -> torch.Tensor:
+    return t.permute(2, 0, 1).unsqueeze(0)
+
+
+def _to_hwc(t: torch.Tensor) -> torch.Tensor:
+    return t.squeeze(0).permute(1, 2, 0)
+
+
+def _build_mipmap_chain_bchw(img: torch.Tensor) -> list[torch.Tensor]:
+    mips = [img]
+    while mips[-1].shape[-2] > 1 or mips[-1].shape[-1] > 1:
+        h, w = mips[-1].shape[-2:]
+        if h == 1 and w == 1:
+            break
+        pad_h = h % 2
+        pad_w = w % 2
+        cur = mips[-1]
+        if pad_h or pad_w:
+            cur = F.pad(cur, (0, pad_w, 0, pad_h), mode="replicate")
+        mips.append(F.avg_pool2d(cur, kernel_size=2, stride=2))
+    return mips
+
+
+def _bilinear_at_bchw(img: torch.Tensor, xs: torch.Tensor, ys: torch.Tensor) -> torch.Tensor:
+    """Sample B×1×H×W or B×3×H×W; xs/ys shape (out_h, out_w) in source pixel coords."""
+    _, _, in_h, in_w = img.shape
+    gx = 2.0 * xs / max(in_w - 1, 1) - 1.0
+    gy = 2.0 * ys / max(in_h - 1, 1) - 1.0
+    grid = torch.stack((gx, gy), dim=-1).unsqueeze(0)
+    return F.grid_sample(img, grid, mode="bilinear", align_corners=True)
+
+
+def bilinear_sample_grid_torch(hwc: torch.Tensor, out_w: int, out_h: int) -> torch.Tensor:
+    in_h, in_w, _ = hwc.shape
+    device = hwc.device
+    xs = (torch.arange(out_w, device=device, dtype=torch.float32) + 0.5) / out_w * in_w - 0.5
+    ys = (torch.arange(out_h, device=device, dtype=torch.float32) + 0.5) / out_h * in_h - 0.5
+    yy, xx = torch.meshgrid(ys, xs, indexing="ij")
+    return _to_hwc(_bilinear_at_bchw(_to_bchw(hwc), xx, yy))
+
+
+def mipmap_linear_sample_grid_torch(hwc: torch.Tensor, out_w: int, out_h: int) -> torch.Tensor:
+    in_h, in_w, _ = hwc.shape
+    if out_w >= in_w and out_h >= in_h:
+        return bilinear_sample_grid_torch(hwc, out_w, out_h)
+
+    mips_bchw = _build_mipmap_chain_bchw(_to_bchw(hwc))
+    device = hwc.device
+    xs = (torch.arange(out_w, device=device, dtype=torch.float32) + 0.5) / out_w * in_w - 0.5
+    ys = (torch.arange(out_h, device=device, dtype=torch.float32) + 0.5) / out_h * in_h - 0.5
+    xx, yy = torch.meshgrid(xs, ys, indexing="xy")
+
+    level = float(math.log2(max(in_w / out_w, in_h / out_h)))
+    l0 = int(math.floor(level))
+    l1 = min(l0 + 1, len(mips_bchw) - 1)
+    frac = level - l0
+    scale0 = 2.0**l0
+    scale1 = 2.0**l1
+    s0 = _to_hwc(_bilinear_at_bchw(mips_bchw[l0], xx / scale0, yy / scale0))
+    s1 = _to_hwc(_bilinear_at_bchw(mips_bchw[l1], xx / scale1, yy / scale1))
+    return torch.clamp(s0 * (1.0 - frac) + s1 * frac, 0.0, 1.0)
+
+
+@torch.inference_mode()
+def bitmap_shader_import_torch(hwc: torch.Tensor, scale: int = 2) -> torch.Tensor:
+    hwc = torch.clamp(hwc, 0.0, 1.0)
+    h, w, _ = hwc.shape
+    out_h, out_w = max(1, h // scale), max(1, w // scale)
+    src_y = torch.arange(out_h, device=hwc.device, dtype=torch.int64) * scale
+    src_x = torch.arange(out_w, device=hwc.device, dtype=torch.int64) * scale
+    yy, xx = torch.meshgrid(src_y, src_x, indexing="ij")
+    rgb = hwc[yy, xx]
+    pro = srgb_to_prophoto_via_xyz(rgb.reshape(-1, 3))
+    hue = _rgb_to_hue(pro)
+    pro = _tonemap_inv(pro)
+    pro = _set_hue(pro, hue)
+    return linear_to_srgb(torch.clamp(pro, min=0.0)).reshape(out_h, out_w, 3)
+
+
+@torch.inference_mode()
+def transform_export_srgb_torch(tone_map_inversed_hwc: torch.Tensor, *, linear_gain: float = 1.0) -> torch.Tensor:
+    from ..polarr_color_space import MAT_PROPHOTO_TO_SRGB
+    from .color_torch import srgb_to_linear as stl
+
+    hwc = torch.clamp(tone_map_inversed_hwc, 0.0, 1.0)
+    rgb = stl(hwc.reshape(-1, 3).double()).float()
+    rgb = _tonemap(rgb).double()
+    if linear_gain != 1.0:
+        rgb = rgb * linear_gain
+    m = torch.tensor(MAT_PROPHOTO_TO_SRGB, device=hwc.device, dtype=torch.float64)
+    out = linear_to_srgb(torch.clamp((rgb @ m.T).float(), 0.0, 1.0))
+    return torch.clamp(out.reshape(hwc.shape), 0.0, 1.0)
+
+
+@torch.inference_mode()
+def render_detection_export_torch(
+    worker_feed_hwc: torch.Tensor,
+    edge: int,
+    *,
+    half_res_hwc: torch.Tensor | None = None,
+):
+    feed = torch.clamp(worker_feed_hwc, 0.0, 1.0)
+    logical_h, logical_w, _ = feed.shape
+    half = half_res_hwc if half_res_hwc is not None else bitmap_shader_import_torch(feed, scale=2)
+    nw, nh = fit_to_size(logical_w, logical_h, (edge, edge))
+    upsampled = mipmap_linear_sample_grid_torch(half, nw, nh)
+    gain = 1.008 if edge >= 2000 else 1.0
+    return transform_export_srgb_torch(upsampled, linear_gain=gain)
+
+
+def render_detection_inputs_torch(
+    hwc: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GPU renderDetectionInputs; returns small/large H×W×3 on device."""
+    half = bitmap_shader_import_torch(hwc, scale=2)
+    small = render_detection_export_torch(hwc, DETECTION_SMALL_EDGE, half_res_hwc=half)
+    large = render_detection_export_torch(hwc, DETECTION_LARGE_EDGE, half_res_hwc=half)
+    return small, large
+
+
+def render_detection_inputs_numpy(hwc, device=None):
+    from .device import hwc_numpy_to_torch
+
+    device = device or get_torch_device()
+    t = hwc_numpy_to_torch(hwc, device)
+    small, large = render_detection_inputs_torch(t)
+    return hwc_torch_to_numpy(small), hwc_torch_to_numpy(large)

magicmatch_core/inference_cuda.py

@@ -63,9 +63,9 @@ def build_merged_lut_with_base_cuda(
     source_hwc: np.ndarray,
     reference_hwc: np.ndarray,
 ) -> tuple[np.ndarray, dict]:
-    from .probe_parity.pipeline_cuda import build_merged_lut_probe_style_cuda
+    from .gpu.pipeline import build_merged_lut_probe_style_gpu
 
-    return build_merged_lut_probe_style_cuda(source_hwc, reference_hwc)
+    return build_merged_lut_probe_style_gpu(source_hwc, reference_hwc)
 
 
 def color_match_session_info() -> dict[str, str | bool]:

magicmatch_core/probe_parity/base_adjustments_cuda.py

@@ -12,6 +12,11 @@ def estimate_base_adjustments_cuda(
     source_hwc: np.ndarray,
     *,
     worker_feed_prepared: bool = False,
+    use_gpu_detection_buffers: bool = False,
 ) -> dict[str, float]:
-    scene = extract_scene_info_bitmap_cuda(source_hwc, worker_feed_prepared=worker_feed_prepared)
+    scene = extract_scene_info_bitmap_cuda(
+        source_hwc,
+        worker_feed_prepared=worker_feed_prepared,
+        use_gpu_detection_buffers=use_gpu_detection_buffers,
+    )
     return get_original_color_match_base_adjustments(scene)

magicmatch_core/probe_parity/scene_extract_cuda.py

@@ -34,6 +34,7 @@ def extract_scene_info_bitmap_cuda(
     source_path: Path | str | None = None,
     use_probe_browser: bool | None = None,
     worker_feed_prepared: bool = False,
+    use_gpu_detection_buffers: bool = False,
 ) -> SceneInfo:
     """Bitmap scene extract; face detect/parse ONNX runs on CUDA when available."""
     if use_probe_browser is None:
@@ -44,7 +45,21 @@ def extract_scene_info_bitmap_cuda(
         path = Path(source_path) if source_path is not None else None
         return extract_scene_via_probe_browser(path or source_hwc)
 
-    small, large = render_detection_inputs(source_hwc, worker_feed_prepared=worker_feed_prepared)
+    if use_gpu_detection_buffers:
+        from ..gpu.device import gpu_pipeline_available
+        from ..gpu.reference_torch import render_detection_inputs_numpy
+
+        if gpu_pipeline_available():
+            feed = source_hwc
+            if not worker_feed_prepared:
+                from .reference import prepare_worker_bitmap_source
+
+                feed = prepare_worker_bitmap_source(source_hwc)
+            small, large = render_detection_inputs_numpy(feed)
+        else:
+            small, large = render_detection_inputs(source_hwc, worker_feed_prepared=worker_feed_prepared)
+    else:
+        small, large = render_detection_inputs(source_hwc, worker_feed_prepared=worker_feed_prepared)
     stats_noface = get_luminance_statistics(small, large, [])
     adjusted_gamma = get_adjusted_gamma(stats_noface)
 

nodes_cuda.py

@@ -1,14 +1,14 @@
 """
-Experimental ComfyUI nodes — CUDA ONNX inference path.
+Experimental ComfyUI nodes — full GPU pipeline (CUDA ONNX + Torch develop/buffers/LUT).
 
-CPU nodes in nodes.py are unchanged. Use these for GPU testing on H100/Linux;
-Preview/apply still uses the CPU NumPy develop stack.
+CPU nodes in nodes.py are unchanged. Use these on H100/Linux for speed testing.
 """
 
 from __future__ import annotations
 
 import torch
 
+from .magicmatch_core.gpu.pipeline import apply_probe_export_gpu
 from .magicmatch_core.inference_cuda import build_merged_lut_with_base_cuda
 from .nodes import (
     MagicMatchLUT,
@@ -19,11 +19,11 @@ from .nodes import (
     _normalize_render_mode,
     _render_options,
 )
-from .magicmatch_core.apply import apply_merged_lut_output
+from .magicmatch_core.apply import RENDER_PROBE_EXPORT, apply_merged_lut_output
 
 
 class MagicMatchBuildCUDA:
-    """Analyze source + reference (CUDA ONNX when available). Wire into MagicMatch Preview."""
+    """Analyze source + reference (GPU ONNX + GPU detection/develop prep)."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -39,8 +39,8 @@ class MagicMatchBuildCUDA:
     FUNCTION = "build"
     CATEGORY = "MAGICMATCH/CUDA (experimental)"
     DESCRIPTION = (
-        "Build merged LUT via CUDA ONNX (color_match + face models). "
-        "Falls back to CPU EP if CUDA unavailable. CPU Build node unchanged."
+        "Build merged LUT: CUDA ONNX + GPU detection buffers + GPU develop@1600. "
+        "CPU Build node unchanged."
     )
 
     @classmethod
@@ -61,7 +61,7 @@ class MagicMatchBuildCUDA:
 
 
 class MagicMatchCUDA:
-    """One-shot CUDA color match (re-runs ONNX when strength changes)."""
+    """One-shot GPU color match (CUDA ONNX + GPU develop/LUT apply)."""
 
     @classmethod
     def INPUT_TYPES(cls):
@@ -87,7 +87,7 @@ class MagicMatchCUDA:
     RETURN_NAMES = ("image",)
     FUNCTION = "match"
     CATEGORY = "MAGICMATCH/CUDA (experimental)"
-    DESCRIPTION = "One-shot color match with CUDA ONNX build; apply path is CPU NumPy."
+    DESCRIPTION = "One-shot: GPU build + GPU probe_export apply when CUDA available."
 
     @classmethod
     def IS_CHANGED(
@@ -96,7 +96,7 @@ class MagicMatchCUDA:
         reference,
         strength,
         lut_encoding="srgb_srgb",
-        render_mode="probe_export",
+        render_mode=RENDER_PROBE_EXPORT,
         profile_stage="current_profile_stages",
     ):
         import hashlib
@@ -113,7 +113,7 @@ class MagicMatchCUDA:
         reference: torch.Tensor,
         strength: float,
         lut_encoding: str = "srgb_srgb",
-        render_mode: str = "probe_export",
+        render_mode: str = RENDER_PROBE_EXPORT,
         profile_stage: str = "current_profile_stages",
     ) -> tuple[torch.Tensor]:
         from .magicmatch_core.probe_parity.profile_stage import normalize_profile_stage
@@ -124,26 +124,74 @@ class MagicMatchCUDA:
         render_mode = _normalize_render_mode(render_mode)
         profile_stage = normalize_profile_stage(profile_stage)
         merged, base = build_merged_lut_with_base_cuda(src, ref)
-        out = apply_merged_lut_output(
-            src,
-            merged,
-            strength,
-            render_mode=render_mode,
-            lut_encoding=lut_encoding,
-            profile_stage=profile_stage,
-            base_adjustments=base,
-        )
+        if render_mode == RENDER_PROBE_EXPORT:
+            out = apply_probe_export_gpu(
+                src,
+                merged,
+                strength,
+                base_adjustments=base,
+                profile_stage=profile_stage,
+                lut_encoding=lut_encoding,
+            )
+        else:
+            out = apply_merged_lut_output(
+                src,
+                merged,
+                strength,
+                render_mode=render_mode,
+                lut_encoding=lut_encoding,
+                profile_stage=profile_stage,
+                base_adjustments=base,
+            )
         return (_hwc_to_image(out),)
 
 
 class MagicMatchPreviewCUDA(MagicMatchPreview):
-    """Same as MagicMatch Preview — apply is CPU NumPy; included for workflow clarity."""
+    """GPU probe_export apply for LUTs built with MagicMatch Build CUDA."""
 
     CATEGORY = "MAGICMATCH/CUDA (experimental)"
-    DESCRIPTION = (
-        "Full-res apply for LUTs built with MagicMatch Build CUDA. "
-        "Identical to MagicMatch Preview (CPU develop stack)."
-    )
+    DESCRIPTION = "Full-res GPU develop + LUT apply (probe_export). Use with Build CUDA."
+
+    def preview(
+        self,
+        source: torch.Tensor,
+        lut: MagicMatchLUT,
+        strength: float,
+        lut_encoding: str = "srgb_srgb",
+        render_mode: str = RENDER_PROBE_EXPORT,
+        profile_stage: str = "current_profile_stages",
+    ) -> dict:
+        from .magicmatch_core.live_cache import pack_live_cache
+        from .magicmatch_core.probe_parity.profile_stage import normalize_profile_stage
+
+        src = _image_batch_to_hwc(source)
+        lut_encoding = _normalize_lut_encoding(lut_encoding)
+        render_mode = _normalize_render_mode(render_mode)
+        profile_stage = normalize_profile_stage(profile_stage)
+        if render_mode == RENDER_PROBE_EXPORT:
+            out = apply_probe_export_gpu(
+                src,
+                lut.merged_lut,
+                strength,
+                base_adjustments=lut.base_adjustments,
+                profile_stage=profile_stage,
+                lut_encoding=lut_encoding,
+            )
+        else:
+            out = apply_merged_lut_output(
+                src,
+                lut.merged_lut,
+                strength,
+                render_mode=render_mode,
+                lut_encoding=lut_encoding,
+                profile_stage=profile_stage,
+                base_adjustments=lut.base_adjustments,
+            )
+        cache = pack_live_cache(src, lut.merged_lut)
+        return {
+            "ui": {"magicmatch_live": [cache]},
+            "result": (_hwc_to_image(out),),
+        }
 
 
 CUDA_NODE_CLASS_MAPPINGS = {

scripts/bench_cuda_vs_cpu.py

@@ -29,6 +29,7 @@ from magicmatch_core.inference_cuda import (  # noqa: E402
     run_inference_from_images_cuda,
 )
 from magicmatch_core.onnx_providers import cuda_available, get_available_providers  # noqa: E402
+from magicmatch_core.gpu.device import gpu_pipeline_available  # noqa: E402
 from magicmatch_core.probe_parity.face_detection import detect_faces  # noqa: E402
 from magicmatch_core.probe_parity.face_detection_cuda import (  # noqa: E402
     detect_faces_cuda,
@@ -157,6 +158,15 @@ def main() -> int:
         lambda: build_merged_lut_with_base_cuda(src, ref),
         repeats=args.repeats,
     )
+    if gpu_pipeline_available():
+        from magicmatch_core.gpu.pipeline import build_merged_lut_probe_style_gpu
+
+        _, t_gpu_build = _timed(
+            "GPU full pipeline build_merged_lut_probe_style_gpu",
+            lambda: build_merged_lut_probe_style_gpu(src, ref),
+            repeats=args.repeats,
+        )
+        report["gpu_full_build_ms"] = round(t_gpu_build * 1000, 2)
     merged_cpu, _ = build_cpu
     merged_cuda, _ = build_cuda
     hash_cpu = lut_hash(merged_cpu)