249 lines
8.8 KiB
Python
249 lines
8.8 KiB
Python
"""
|
|
ComfyUI nodes — full GPU pipeline (all hot path on CUDA/Torch).
|
|
|
|
CPU nodes in nodes.py are unchanged. Category: MAGICMATCH/GPU Full (experimental).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from .magicmatch_core.gpu.device import gpu_pipeline_available, hwc_numpy_to_torch
|
|
from .magicmatch_core.gpu.pipeline_full_gpu import (
|
|
apply_probe_export_full_gpu,
|
|
build_merged_lut_full_gpu,
|
|
color_match_one_shot_full_gpu,
|
|
)
|
|
from .magicmatch_core.probe_parity.reference import prepare_worker_bitmap_source
|
|
from .nodes import (
|
|
MagicMatchLUT,
|
|
MagicMatchPreview,
|
|
_hwc_to_image,
|
|
_image_batch_to_hwc,
|
|
_normalize_lut_encoding,
|
|
_normalize_render_mode,
|
|
_render_options,
|
|
)
|
|
from .magicmatch_core.apply import RENDER_PROBE_EXPORT, apply_merged_lut_output
|
|
|
|
|
|
def _image_batch_to_feed_tensor(image: torch.Tensor) -> torch.Tensor:
|
|
"""Keep GPU tensor when already on CUDA; JPEG normalize once on CPU for parity."""
|
|
if image.ndim != 4 or image.shape[-1] != 3:
|
|
raise ValueError(f"MAGICMATCH: expected IMAGE [B,H,W,3], got {tuple(image.shape)}")
|
|
if image.shape[0] != 1:
|
|
raise ValueError(f"MAGICMATCH: batch size {image.shape[0]} not supported — use batch size 1.")
|
|
t = image[0]
|
|
if gpu_pipeline_available() and t.is_cuda:
|
|
feed_np = prepare_worker_bitmap_source(t.detach().float().clamp(0.0, 1.0).cpu().numpy())
|
|
return hwc_numpy_to_torch(feed_np, t.device)
|
|
feed_np = prepare_worker_bitmap_source(_image_batch_to_hwc(image))
|
|
return hwc_numpy_to_torch(feed_np)
|
|
|
|
|
|
class MagicMatchBuildGPUFull:
|
|
"""Analyze source + reference — full GPU scene extract + develop + CUDA ONNX."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source": ("IMAGE",),
|
|
"reference": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MAGICMATCH_LUT",)
|
|
RETURN_NAMES = ("lut",)
|
|
FUNCTION = "build"
|
|
CATEGORY = "MAGICMATCH/GPU Full (experimental)"
|
|
DESCRIPTION = (
|
|
"Full GPU build: detection buffers, scene extract, develop@1600, CUDA ONNX. "
|
|
"Caches GPU feed tensor for Preview."
|
|
)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, source, reference):
|
|
import hashlib
|
|
|
|
if source is None or reference is None:
|
|
return float("nan")
|
|
s = hashlib.sha1(source.detach().cpu().numpy().tobytes()).hexdigest()[:16]
|
|
r = hashlib.sha1(reference.detach().cpu().numpy().tobytes()).hexdigest()[:16]
|
|
return (s, r)
|
|
|
|
def build(self, source: torch.Tensor, reference: torch.Tensor) -> tuple[MagicMatchLUT]:
|
|
src_np = _image_batch_to_hwc(source)
|
|
ref_np = _image_batch_to_hwc(reference)
|
|
feed_t = _image_batch_to_feed_tensor(source)
|
|
ref_t = reference[0].detach().float().clamp(0.0, 1.0)
|
|
if not (gpu_pipeline_available() and ref_t.is_cuda):
|
|
ref_t = hwc_numpy_to_torch(ref_np)
|
|
state = build_merged_lut_full_gpu(
|
|
src_np,
|
|
ref_np,
|
|
feed_tensor=feed_t,
|
|
reference_tensor=ref_t,
|
|
)
|
|
return (MagicMatchLUT(state.merged_lut, state.base_adjustments, feed_tensor=state.feed_tensor),)
|
|
|
|
|
|
class MagicMatchGPUFull:
|
|
"""One-shot full GPU color match — fused build + apply, no duplicate scene extract."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"source": ("IMAGE",),
|
|
"reference": ("IMAGE",),
|
|
"strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "slider",
|
|
},
|
|
),
|
|
**_render_options(),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "match"
|
|
CATEGORY = "MAGICMATCH/GPU Full (experimental)"
|
|
DESCRIPTION = "Fused full GPU pipeline: scene extract + ONNX + develop/LUT apply on device."
|
|
|
|
@classmethod
|
|
def IS_CHANGED(
|
|
cls,
|
|
source,
|
|
reference,
|
|
strength,
|
|
lut_encoding="srgb_srgb",
|
|
render_mode=RENDER_PROBE_EXPORT,
|
|
profile_stage="current_profile_stages",
|
|
):
|
|
import hashlib
|
|
|
|
if source is None or reference is None:
|
|
return float("nan")
|
|
s = hashlib.sha1(source.detach().cpu().numpy().tobytes()).hexdigest()[:16]
|
|
r = hashlib.sha1(reference.detach().cpu().numpy().tobytes()).hexdigest()[:16]
|
|
return (s, r, float(strength), lut_encoding, render_mode, profile_stage)
|
|
|
|
def match(
|
|
self,
|
|
source: torch.Tensor,
|
|
reference: torch.Tensor,
|
|
strength: float,
|
|
lut_encoding: str = "srgb_srgb",
|
|
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
|
|
|
|
src_np = _image_batch_to_hwc(source)
|
|
ref_np = _image_batch_to_hwc(reference)
|
|
lut_encoding = _normalize_lut_encoding(lut_encoding)
|
|
render_mode = _normalize_render_mode(render_mode)
|
|
profile_stage = normalize_profile_stage(profile_stage)
|
|
src_t = source[0].detach().float().clamp(0.0, 1.0)
|
|
ref_t = reference[0].detach().float().clamp(0.0, 1.0)
|
|
if render_mode == RENDER_PROBE_EXPORT:
|
|
out = color_match_one_shot_full_gpu(
|
|
src_np,
|
|
ref_np,
|
|
strength,
|
|
profile_stage=profile_stage,
|
|
lut_encoding=lut_encoding,
|
|
source_tensor=src_t,
|
|
reference_tensor=ref_t,
|
|
)
|
|
else:
|
|
feed_t = _image_batch_to_feed_tensor(source)
|
|
state = build_merged_lut_full_gpu(
|
|
src_np,
|
|
ref_np,
|
|
feed_tensor=feed_t,
|
|
reference_tensor=ref_t,
|
|
)
|
|
out = apply_merged_lut_output(
|
|
src_np,
|
|
state.merged_lut,
|
|
strength,
|
|
render_mode=render_mode,
|
|
lut_encoding=lut_encoding,
|
|
profile_stage=profile_stage,
|
|
base_adjustments=state.base_adjustments,
|
|
)
|
|
return (_hwc_to_image(out),)
|
|
|
|
|
|
class MagicMatchPreviewGPUFull(MagicMatchPreview):
|
|
"""Full GPU probe_export apply; reuses cached GPU feed tensor from Build GPU Full."""
|
|
|
|
CATEGORY = "MAGICMATCH/GPU Full (experimental)"
|
|
DESCRIPTION = "GPU develop + LUT apply using cached feed tensor from Build GPU Full."
|
|
|
|
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_np = _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)
|
|
feed_t = getattr(lut, "feed_tensor", None)
|
|
if feed_t is None:
|
|
feed_t = _image_batch_to_feed_tensor(source)
|
|
if render_mode == RENDER_PROBE_EXPORT:
|
|
out = apply_probe_export_full_gpu(
|
|
src_np,
|
|
lut.merged_lut,
|
|
strength,
|
|
base_adjustments=lut.base_adjustments,
|
|
profile_stage=profile_stage,
|
|
lut_encoding=lut_encoding,
|
|
feed_tensor=feed_t,
|
|
)
|
|
else:
|
|
out = apply_merged_lut_output(
|
|
src_np,
|
|
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_np, lut.merged_lut)
|
|
return {
|
|
"ui": {"magicmatch_live": [cache]},
|
|
"result": (_hwc_to_image(out),),
|
|
}
|
|
|
|
|
|
GPU_FULL_NODE_CLASS_MAPPINGS = {
|
|
"MagicMatchBuildGPUFull": MagicMatchBuildGPUFull,
|
|
"MagicMatchGPUFull": MagicMatchGPUFull,
|
|
"MagicMatchPreviewGPUFull": MagicMatchPreviewGPUFull,
|
|
}
|
|
|
|
GPU_FULL_NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MagicMatchBuildGPUFull": "MagicMatch Build LUT (GPU Full)",
|
|
"MagicMatchGPUFull": "MagicMatch one-shot (GPU Full)",
|
|
"MagicMatchPreviewGPUFull": "MagicMatch Preview (GPU Full LUT)",
|
|
}
|