From 03611260456d6bd2bf523ca58f4707e64e644a74 Mon Sep 17 00:00:00 2001 From: poxular Date: Tue, 2 Jun 2026 01:54:41 +0200 Subject: [PATCH] Restore 1ccd29d CUDA one-shot path without breaking node imports. 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 --- magicmatch_core/gpu/pipeline_v1.py | 91 ++++++++++++++++++++++++++++++ magicmatch_core/inference_cuda.py | 6 +- nodes_cuda.py | 91 +++++++++++++++++++++++------- 3 files changed, 165 insertions(+), 23 deletions(-) create mode 100644 magicmatch_core/gpu/pipeline_v1.py diff --git a/magicmatch_core/gpu/pipeline_v1.py b/magicmatch_core/gpu/pipeline_v1.py new file mode 100644 index 0000000..2cd7e6d --- /dev/null +++ b/magicmatch_core/gpu/pipeline_v1.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_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) diff --git a/magicmatch_core/inference_cuda.py b/magicmatch_core/inference_cuda.py index dbf1a8d..7dbc27b 100644 --- a/magicmatch_core/inference_cuda.py +++ b/magicmatch_core/inference_cuda.py @@ -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]: diff --git a/nodes_cuda.py b/nodes_cuda.py index c2c716e..45db138 100644 --- a/nodes_cuda.py +++ b/nodes_cuda.py @@ -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 = {