246 lines
8.6 KiB
Python
246 lines
8.6 KiB
Python
"""
|
|
ComfyUI nodes — fast GPU pipeline (speed over strict parity).
|
|
|
|
CPU nodes in nodes.py are unchanged. Category: MAGICMATCH/GPU Fast (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_fast import (
|
|
apply_probe_export_fast_gpu,
|
|
build_merged_lut_fast_gpu,
|
|
color_match_one_shot_fast_gpu,
|
|
)
|
|
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:
|
|
"""Direct GPU upload from Comfy IMAGE tensor — no JPEG q98 normalize."""
|
|
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].detach().float().clamp(0.0, 1.0)
|
|
if gpu_pipeline_available() and t.is_cuda:
|
|
return t
|
|
return hwc_numpy_to_torch(_image_batch_to_hwc(image))
|
|
|
|
|
|
class MagicMatchBuildGPUFast:
|
|
"""Analyze source + reference — fast 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 Fast (experimental)"
|
|
DESCRIPTION = (
|
|
"Fast GPU build: no JPEG normalize, bilinear ref, GPU gray-world WB. "
|
|
"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_fast_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 MagicMatchGPUFast:
|
|
"""One-shot fast GPU color match — fused build + apply."""
|
|
|
|
@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 Fast (experimental)"
|
|
DESCRIPTION = "Fast fused GPU pipeline: approximate scene extract + ONNX + develop/LUT apply."
|
|
|
|
@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_fast_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_fast_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 MagicMatchPreviewGPUFast(MagicMatchPreview):
|
|
"""Fast GPU probe_export apply; reuses cached GPU feed tensor from Build GPU Fast."""
|
|
|
|
CATEGORY = "MAGICMATCH/GPU Fast (experimental)"
|
|
DESCRIPTION = "Fast GPU develop + LUT apply using cached feed tensor from Build GPU Fast."
|
|
|
|
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_fast_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_FAST_NODE_CLASS_MAPPINGS = {
|
|
"MagicMatchBuildGPUFast": MagicMatchBuildGPUFast,
|
|
"MagicMatchGPUFast": MagicMatchGPUFast,
|
|
"MagicMatchPreviewGPUFast": MagicMatchPreviewGPUFast,
|
|
}
|
|
|
|
GPU_FAST_NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MagicMatchBuildGPUFast": "MagicMatch Build LUT (GPU Fast)",
|
|
"MagicMatchGPUFast": "MagicMatch one-shot (GPU Fast)",
|
|
"MagicMatchPreviewGPUFast": "MagicMatch Preview (GPU Fast LUT)",
|
|
}
|