diff --git a/README.md b/README.md index a22c10f..2456182 100644 --- a/README.md +++ b/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 diff --git a/magicmatch_core/gpu/__init__.py b/magicmatch_core/gpu/__init__.py new file mode 100644 index 0000000..dbb1191 --- /dev/null +++ b/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", +] diff --git a/magicmatch_core/gpu/color_torch.py b/magicmatch_core/gpu/color_torch.py new file mode 100644 index 0000000..a2a33fa --- /dev/null +++ b/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) diff --git a/magicmatch_core/gpu/develop_torch.py b/magicmatch_core/gpu/develop_torch.py new file mode 100644 index 0000000..d54886c --- /dev/null +++ b/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) diff --git a/magicmatch_core/gpu/device.py b/magicmatch_core/gpu/device.py new file mode 100644 index 0000000..4e87bb0 --- /dev/null +++ b/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() diff --git a/magicmatch_core/gpu/lut_torch.py b/magicmatch_core/gpu/lut_torch.py new file mode 100644 index 0000000..ffb7679 --- /dev/null +++ b/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) diff --git a/magicmatch_core/gpu/pipeline.py b/magicmatch_core/gpu/pipeline.py new file mode 100644 index 0000000..2b5c241 --- /dev/null +++ b/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) diff --git a/magicmatch_core/gpu/reference_torch.py b/magicmatch_core/gpu/reference_torch.py new file mode 100644 index 0000000..39ae403 --- /dev/null +++ b/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) diff --git a/magicmatch_core/inference_cuda.py b/magicmatch_core/inference_cuda.py index f07f340..5af7b8d 100644 --- a/magicmatch_core/inference_cuda.py +++ b/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]: diff --git a/magicmatch_core/probe_parity/base_adjustments_cuda.py b/magicmatch_core/probe_parity/base_adjustments_cuda.py index 69dfd46..33d23dd 100644 --- a/magicmatch_core/probe_parity/base_adjustments_cuda.py +++ b/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) diff --git a/magicmatch_core/probe_parity/scene_extract_cuda.py b/magicmatch_core/probe_parity/scene_extract_cuda.py index 8d0dbbe..57aa99c 100644 --- a/magicmatch_core/probe_parity/scene_extract_cuda.py +++ b/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) diff --git a/nodes_cuda.py b/nodes_cuda.py index de1d007..bb0ea18 100644 --- a/nodes_cuda.py +++ b/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 = { diff --git a/scripts/bench_cuda_vs_cpu.py b/scripts/bench_cuda_vs_cpu.py index d2f476e..fbe9914 100644 --- a/scripts/bench_cuda_vs_cpu.py +++ b/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)