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

Restore 1ccd29d CUDA one-shot path without breaking node imports.

Browse Source 
Add pipeline_v1 (GPU build + GPU apply) as a standalone module so partial
git checkouts of inference_cuda no longer break GPU Full/Fast imports.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
poxular
2026-06-02 01:54:41 +02:00
co-authored by Cursor
parent aeae298c40
commit 0361126045
3 changed files with 165 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
magicmatch_core/gpu/pipeline_v1.py
+91
View File
@@ -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_v1(
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_v1(
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/inference_cuda.py
+3 -3
View File
@@ -103,10 +103,10 @@ def build_merged_lut_with_base_cuda(
source_hwc: np.ndarray,
reference_hwc: np.ndarray,
) -> tuple[np.ndarray, dict]:
"""CPU parity scene extract + develop; CUDA ONNX only (golden LUT path)."""
from .probe_parity.pipeline_cuda import build_merged_lut_probe_style_cuda
"""GPU develop + CUDA ONNX build (1ccd29d-style; used by CUDA ComfyUI nodes)."""
from .gpu.pipeline_v1 import build_merged_lut_probe_style_gpu_v1
return build_merged_lut_probe_style_cuda(source_hwc, reference_hwc)
return build_merged_lut_probe_style_gpu_v1(source_hwc, reference_hwc)
def color_match_session_info() -> dict[str, str | bool]:
nodes_cuda.py
+71 -20
View File
@@ -1,15 +1,15 @@
"""
Experimental ComfyUI nodes — CUDA ONNX build with CPU parity develop/apply.
Experimental ComfyUI nodes — CUDA one-shot (1ccd29d-style GPU develop + apply).
CPU nodes in nodes.py are unchanged. CUDA one-shot matches golden MagicMatch look;
only color_match.onnx (+ face ONNX) run on GPU when available.
CPU nodes in nodes.py are unchanged. Build uses GPU torch develop + CUDA ONNX;
one-shot apply uses GPU develop/LUT (pipeline_v1), not the CPU parity stack.
"""
from __future__ import annotations
import torch
from .magicmatch_core.gpu.pipeline import apply_probe_export_gpu, color_match_one_shot_gpu
from .magicmatch_core.gpu.pipeline_v1 import apply_probe_export_gpu_v1
from .magicmatch_core.inference_cuda import build_merged_lut_with_base_cuda
from .nodes import (
MagicMatchLUT,
@@ -24,7 +24,7 @@ from .magicmatch_core.apply import RENDER_PROBE_EXPORT, apply_merged_lut_output
class MagicMatchBuildCUDA:
"""Analyze source + reference (GPU ONNX + GPU detection/develop prep)."""
"""Analyze source + reference (GPU develop + CUDA ONNX)."""
@classmethod
def INPUT_TYPES(cls):
@@ -40,8 +40,8 @@ class MagicMatchBuildCUDA:
FUNCTION = "build"
CATEGORY = "MAGICMATCH/CUDA (experimental)"
DESCRIPTION = (
"Build merged LUT via CUDA ONNX (color_match + face models). "
"Scene extract and develop stay on CPU parity stack."
"Build merged LUT: GPU detection/develop prep + CUDA ONNX. "
"Same build path as MagicMatch one-shot (CUDA)."
)
@classmethod
@@ -62,7 +62,7 @@ class MagicMatchBuildCUDA:
class MagicMatchCUDA:
"""One-shot GPU color match (CUDA ONNX + GPU develop/LUT apply)."""
"""One-shot GPU color match (1ccd29d path: GPU build + GPU apply)."""
@classmethod
def INPUT_TYPES(cls):
@@ -88,7 +88,7 @@ class MagicMatchCUDA:
RETURN_NAMES = ("image",)
FUNCTION = "match"
CATEGORY = "MAGICMATCH/CUDA (experimental)"
DESCRIPTION = "One-shot with CUDA ONNX build; apply uses CPU parity develop stack."
DESCRIPTION = "One-shot: GPU build + GPU probe_export apply (1ccd29d-style path)."
@classmethod
def IS_CHANGED(
@@ -125,23 +125,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_v1(
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 CPU apply as MagicMatch Preview — for LUTs from Build CUDA."""
"""GPU probe_export apply for LUTs built with MagicMatch Build CUDA."""
CATEGORY = "MAGICMATCH/CUDA (experimental)"
DESCRIPTION = "Full-res CPU parity apply for LUTs built with MagicMatch Build CUDA."
DESCRIPTION = "GPU develop + LUT apply (1ccd29d-style). 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_v1(
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 = {