507 lines
18 KiB
Python
507 lines
18 KiB
Python
"""ComfyUI HDR nodes: ingest, decode postprocess, and delivery.
|
||
|
||
Stage-oriented layout (shared ACEScct working space):
|
||
|
||
- Ingest: ``LTXVLoadEXRSequence`` (SDR→ACEScct uses Comfy Convert Image Color Space)
|
||
- Runtime: ``LTXVVAEForceFloat32``
|
||
- Delivery: ``LTXVHDRDecodePostprocess`` (HLG video uses Comfy Create Video / Save Video)
|
||
|
||
Color math: ``hdr_color``. File I/O helpers: ``hdr_io``.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
from torch import Tensor
|
||
|
||
from .hdr_color import HDRColorSpace, Primaries, compress_acescct, images_to_acescct
|
||
from .hdr_io import (
|
||
_list_exr_files,
|
||
_read_exr,
|
||
_resolve_path,
|
||
_save_exr_frame,
|
||
_silent_audio,
|
||
_trim_to_8k1,
|
||
)
|
||
from .nodes_registry import comfy_node
|
||
|
||
logger = logging.getLogger("LTXVideo.hdr_nodes")
|
||
|
||
_EXR_COLOR_SPACE_CHOICES = ["linear", "acescct", "srgb_linear", "acescg"]
|
||
_COLOR_SPACE_CHOICES = [c.value for c in HDRColorSpace]
|
||
|
||
|
||
class LogC3:
|
||
"""ARRI LogC3 (EI 800) HDR compression.
|
||
|
||
Maps linear [0, inf) -> [0, 1] via the camera log curve, then scales to
|
||
[-1, 1] for VAE input.
|
||
"""
|
||
|
||
A = 5.555556
|
||
B = 0.052272
|
||
C = 0.247190
|
||
D = 0.385537
|
||
E = 5.367655
|
||
F = 0.092809
|
||
CUT = 0.010591
|
||
|
||
def compress(self, hdr: Tensor) -> Tensor:
|
||
x = torch.clamp(hdr, min=0.0)
|
||
log_part = self.C * torch.log10(self.A * x + self.B) + self.D
|
||
lin_part = self.E * x + self.F
|
||
logc = torch.where(x >= self.CUT, log_part, lin_part)
|
||
logc = torch.clamp(logc, 0.0, 1.0)
|
||
return logc * 2.0 - 1.0
|
||
|
||
def decompress(self, z: Tensor) -> Tensor:
|
||
logc = torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
|
||
cut_log = self.E * self.CUT + self.F
|
||
lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
|
||
lin_from_lin = (logc - self.F) / self.E
|
||
return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Transform registry — singleton (stateless)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_LOGC3 = LogC3()
|
||
|
||
|
||
def _hdr_decompress(decoded_01: Tensor, transfer: str = "acescct") -> Tensor:
|
||
"""Decompress VAE-decoded ``[0,1]`` image to linear HDR ``[0, inf)``.
|
||
|
||
ComfyUI's VAE decode returns images in ``[0, 1]`` via ``(raw + 1) / 2``.
|
||
|
||
* ``acescct`` — ACEScct working codes (SDR→HDR IC-LoRA + native HDR).
|
||
* ``logc3`` — legacy LogC3 SDR→HDR LoRA only.
|
||
"""
|
||
transfer = (transfer or "acescct").lower()
|
||
if transfer == "acescct":
|
||
from .hdr_color import decompress_acescct
|
||
|
||
# Decoded IMAGE is already ACEScct working-space [0, 1].
|
||
return decompress_acescct(decoded_01.float())
|
||
if transfer != "logc3":
|
||
raise ValueError(
|
||
f"Unsupported HDR transfer {transfer!r}; use 'logc3' or 'acescct'."
|
||
)
|
||
raw = decoded_01.float() * 2.0 - 1.0
|
||
return _LOGC3.decompress(raw)
|
||
|
||
|
||
def _linear_to_srgb(x: Tensor) -> Tensor:
|
||
"""Convert linear [0, 1] to sRGB [0, 1]."""
|
||
return torch.where(
|
||
x <= 0.0031308,
|
||
12.92 * x,
|
||
1.055 * torch.pow(x.clamp(min=0.0031308), 1.0 / 2.4) - 0.055,
|
||
).clamp(0.0, 1.0)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ComfyUI Nodes
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _exr_payload_for_color_space(
|
||
image: Tensor,
|
||
hdr_linear: Tensor,
|
||
transfer: str,
|
||
exr_color_space: str,
|
||
) -> tuple[Tensor, Primaries | None, str | None]:
|
||
"""Build EXR RGB frames and optional primaries / colorSpace tags.
|
||
|
||
``linear`` keeps the legacy 2.3 behaviour (raw decompressed buffer, no tags).
|
||
Other values match pipelines ``--hdr`` EXR authoring.
|
||
"""
|
||
space = (exr_color_space or "linear").lower()
|
||
transfer = (transfer or "acescct").lower()
|
||
rgb_in = image[..., :3].float()
|
||
hdr = hdr_linear[..., :3].float()
|
||
|
||
if space == "linear":
|
||
return hdr, None, None
|
||
|
||
if space == "acescct":
|
||
if transfer == "acescct":
|
||
frames = rgb_in.clamp(0.0, 1.0)
|
||
else:
|
||
# LogC3 linear is treated as Rec.709 scene-linear for ACEScct authoring.
|
||
frames = compress_acescct(
|
||
Primaries.REC709.to(Primaries.ACESCG, hdr).clamp(min=0.0)
|
||
)
|
||
return frames, Primaries.ACESCG, "ACEScct"
|
||
|
||
if space == "acescg":
|
||
if transfer == "acescct":
|
||
frames = hdr # decompress_acescct → ACEScg linear
|
||
else:
|
||
frames = Primaries.REC709.to(Primaries.ACESCG, hdr).clamp(min=0.0)
|
||
return frames, Primaries.ACESCG, "ACEScg"
|
||
|
||
if space == "srgb_linear":
|
||
if transfer == "acescct":
|
||
frames = Primaries.ACESCG.to(Primaries.REC709, hdr).clamp(min=0.0)
|
||
else:
|
||
frames = hdr.clamp(min=0.0)
|
||
return frames, Primaries.REC709, "sRGB"
|
||
|
||
raise ValueError(
|
||
f"Unsupported exr_color_space {exr_color_space!r}; "
|
||
f"use one of {_EXR_COLOR_SPACE_CHOICES}."
|
||
)
|
||
|
||
|
||
@comfy_node(name="LTXVHDRDecodePostprocess")
|
||
class LTXVHDRDecodePostprocess:
|
||
"""Decompress HDR from VAE output, tonemap for preview, optionally save EXR.
|
||
|
||
Place after VAE Decode. Recovers linear HDR from the VAE ``[0,1]`` image and
|
||
tonemaps to SDR for display. When ``save_exr`` is enabled, writes an EXR
|
||
sequence (legacy ``linear``, or tagged ``acescct`` / ``srgb_linear`` /
|
||
``acescg``).
|
||
|
||
Outputs:
|
||
tonemapped: SDR preview [0, 1] after Reinhard tonemap + sRGB gamma.
|
||
hdr_linear: Raw linear HDR values [0, inf) for further processing.
|
||
With ``transfer=acescct`` this is ACEScg scene-linear; with
|
||
``logc3`` it is LogC3/Rec.709-style scene-linear.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
},
|
||
"optional": {
|
||
"transfer": (
|
||
["acescct", "logc3"],
|
||
{
|
||
"default": "acescct",
|
||
"tooltip": (
|
||
"Working-space curve of the VAE decode. "
|
||
"Use acescct for SDR→HDR IC-LoRA (after Convert "
|
||
"Image Color Space sRGB→HDR ACEScct) and for native HDR EXR; "
|
||
"logc3 only for legacy LogC3 HDR LoRA checkpoints."
|
||
),
|
||
},
|
||
),
|
||
"exposure": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.1,
|
||
"display": "slider",
|
||
"tooltip": (
|
||
"Exposure in stops (EV). 0 = no change, "
|
||
"+1 = 2x brighter, -1 = half brightness. "
|
||
"Affects the tonemapped preview only."
|
||
),
|
||
},
|
||
),
|
||
"save_exr": (
|
||
"BOOLEAN",
|
||
{
|
||
"default": False,
|
||
"tooltip": "Save HDR frames as an EXR sequence.",
|
||
},
|
||
),
|
||
"exr_color_space": (
|
||
_EXR_COLOR_SPACE_CHOICES,
|
||
{
|
||
"default": "acescg",
|
||
"tooltip": (
|
||
"EXR contents when save_exr is on. "
|
||
"Default 'acescg' is tagged ACEScg scene-linear. "
|
||
"'acescct' is ACEScct log; 'srgb_linear' is tagged "
|
||
"Rec.709 linear; 'linear' is legacy raw decompress "
|
||
"(untagged). All paths use OpenImageIO."
|
||
),
|
||
},
|
||
),
|
||
"output_dir": (
|
||
"STRING",
|
||
{
|
||
"default": "output/hdr_exr",
|
||
"tooltip": (
|
||
"Directory for EXR frames (relative to ComfyUI "
|
||
"output directory, or absolute path)."
|
||
),
|
||
},
|
||
),
|
||
"filename_prefix": (
|
||
"STRING",
|
||
{"default": "frame"},
|
||
),
|
||
"half_precision": (
|
||
"BOOLEAN",
|
||
{
|
||
"default": True,
|
||
"tooltip": (
|
||
"Save EXR as float16 (half). Smaller files, "
|
||
"negligible quality loss for most workflows."
|
||
),
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE")
|
||
RETURN_NAMES = ("tonemapped", "hdr_linear")
|
||
OUTPUT_NODE = True
|
||
FUNCTION = "postprocess"
|
||
CATEGORY = "Lightricks/HDR"
|
||
DESCRIPTION = (
|
||
"Decompresses VAE-decoded HDR (ACEScct or legacy LogC3), tonemaps for "
|
||
"preview, and optionally writes EXR (linear / ACEScct / sRGB linear / "
|
||
"ACEScg)."
|
||
)
|
||
|
||
def postprocess(
|
||
self,
|
||
image: torch.Tensor,
|
||
transfer: str = "acescct",
|
||
exposure: float = 0.0,
|
||
save_exr: bool = False,
|
||
exr_color_space: str = "acescg",
|
||
output_dir: str = "output/hdr_exr",
|
||
filename_prefix: str = "frame",
|
||
half_precision: bool = True,
|
||
) -> tuple:
|
||
|
||
hdr = _hdr_decompress(image, transfer=transfer)
|
||
hdr = torch.clamp(hdr, min=0.0, max=1e4)
|
||
|
||
# Reinhard tonemap with exposure (stops → linear multiplier)
|
||
exposure_mult = 2.0**exposure
|
||
hdr_exposed = hdr * exposure_mult
|
||
tonemapped_linear = (hdr_exposed / (1.0 + hdr_exposed)).clamp(0.0, 1.0)
|
||
tonemapped = _linear_to_srgb(tonemapped_linear)
|
||
|
||
if save_exr:
|
||
frames, primaries, tag = _exr_payload_for_color_space(
|
||
image, hdr, transfer, exr_color_space
|
||
)
|
||
self._save_exr_frames_tagged(
|
||
frames,
|
||
output_dir,
|
||
filename_prefix,
|
||
half_precision,
|
||
primaries=primaries,
|
||
color_space_tag=tag,
|
||
)
|
||
|
||
return (tonemapped, hdr)
|
||
|
||
@staticmethod
|
||
def _resolve_exr_dir(output_dir: str) -> str:
|
||
import folder_paths
|
||
|
||
if not os.path.isabs(output_dir):
|
||
output_dir = os.path.join(folder_paths.get_output_directory(), output_dir)
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
return output_dir
|
||
|
||
@classmethod
|
||
def _save_exr_frames_tagged(
|
||
cls,
|
||
frames: torch.Tensor,
|
||
output_dir: str,
|
||
filename_prefix: str,
|
||
half_precision: bool,
|
||
*,
|
||
primaries: Primaries | None,
|
||
color_space_tag: str | None,
|
||
) -> None:
|
||
output_dir = cls._resolve_exr_dir(output_dir)
|
||
for i, frame in enumerate(frames):
|
||
_save_exr_frame(
|
||
frame,
|
||
Path(output_dir) / f"{filename_prefix}_{i:05d}.exr",
|
||
half_precision=half_precision,
|
||
primaries=primaries,
|
||
color_space_tag=color_space_tag,
|
||
)
|
||
logger.info(
|
||
"Saved %d EXR frame(s) to %s (%s%s)",
|
||
frames.shape[0],
|
||
output_dir,
|
||
"float16" if half_precision else "float32",
|
||
f", {color_space_tag}" if color_space_tag else ", untagged linear",
|
||
)
|
||
|
||
|
||
@comfy_node(name="LTXVLoadEXRSequence")
|
||
class LTXVLoadEXRSequence:
|
||
"""Load an EXR still or folder and convert to ACEScct ``[0,1]`` IMAGE for the VAE.
|
||
|
||
Declares the source colour space the same way as pipelines ``--hdr``:
|
||
``srgb_linear`` / ``acescg`` (compress to ACEScct) or ``acescct`` (passthrough).
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"path": (
|
||
"STRING",
|
||
{
|
||
"default": "",
|
||
"tooltip": (
|
||
"Path to a .exr still or a directory of *.exr frames. "
|
||
"I2V: single still is fine. V2V: use a folder with 8k+1 frames. "
|
||
"Relative paths are resolved under ComfyUI's input folder."
|
||
),
|
||
},
|
||
),
|
||
"color_space": (_COLOR_SPACE_CHOICES, {"default": "srgb_linear"}),
|
||
"frame_rate": (
|
||
"FLOAT",
|
||
{
|
||
"default": 24.0,
|
||
"min": 1.0,
|
||
"max": 120.0,
|
||
"step": 0.001,
|
||
"tooltip": "Required for EXR folders (no container fps).",
|
||
},
|
||
),
|
||
},
|
||
"optional": {
|
||
"frame_start": ("INT", {"default": 0, "min": 0, "max": 100000}),
|
||
"frame_cap": (
|
||
"INT",
|
||
{
|
||
"default": 0,
|
||
"min": 0,
|
||
"max": 100000,
|
||
"tooltip": "Max frames to load (0 = all). Still trimmed to 8k+1.",
|
||
},
|
||
),
|
||
"trim_to_8k1": (
|
||
"BOOLEAN",
|
||
{
|
||
"default": True,
|
||
"tooltip": "Trim to longest 8k+1 prefix required by the video VAE.",
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "FLOAT", "INT", "INT", "INT", "AUDIO")
|
||
RETURN_NAMES = ("images", "frame_rate", "width", "height", "frame_count", "audio")
|
||
FUNCTION = "load"
|
||
CATEGORY = "Lightricks/HDR"
|
||
DESCRIPTION = (
|
||
"Native HDR input: load EXR → ACEScct [0,1] frames for IC-LoRA / VAE encode. "
|
||
"Also emits silent AUDIO matching duration so V2V graphs that freeze audio still wire. "
|
||
"Native EXR path only (not SDR→HDR IC-LoRA ingest). "
|
||
"Downstream: set LTXVPreprocess img_compression=0 — Python --hdr EXR conditioning "
|
||
"skips JPEG/CRF; compressing ACEScct codes darkens HDR highlights."
|
||
)
|
||
|
||
def load(
|
||
self,
|
||
path: str,
|
||
color_space: str,
|
||
frame_rate: float,
|
||
frame_start: int = 0,
|
||
frame_cap: int = 0,
|
||
trim_to_8k1: bool = True,
|
||
):
|
||
root = _resolve_path(path)
|
||
files = _list_exr_files(root)
|
||
files = files[frame_start:]
|
||
if frame_cap > 0:
|
||
files = files[:frame_cap]
|
||
if not files:
|
||
raise ValueError("No EXR frames left after frame_start / frame_cap.")
|
||
|
||
cs = HDRColorSpace(color_space)
|
||
frames = [_read_exr(f) for f in files]
|
||
shapes = {tuple(f.shape) for f in frames}
|
||
if len(shapes) != 1:
|
||
raise ValueError(
|
||
"EXR frames must share the same H×W×C; got "
|
||
+ ", ".join(str(s) for s in sorted(shapes))
|
||
)
|
||
images = torch.stack(frames, dim=0)
|
||
images = images_to_acescct(images, cs)
|
||
if trim_to_8k1:
|
||
images = _trim_to_8k1(images)
|
||
|
||
h, w = int(images.shape[1]), int(images.shape[2])
|
||
n = int(images.shape[0])
|
||
if n < 1:
|
||
raise ValueError(f"No EXR frames loaded from {root}")
|
||
audio = _silent_audio(n, float(frame_rate))
|
||
logger.info(
|
||
"Loaded %d EXR frame(s) from %s as %s → ACEScct (%dx%d @ %.3f fps); "
|
||
"silent audio %d samples @ %d Hz",
|
||
n,
|
||
root if root.is_dir() else (root.parent if root.is_file() else root),
|
||
cs.value,
|
||
w,
|
||
h,
|
||
frame_rate,
|
||
int(audio["waveform"].shape[-1]),
|
||
int(audio["sample_rate"]),
|
||
)
|
||
return (images, float(frame_rate), w, h, n, audio)
|
||
|
||
|
||
@comfy_node(name="LTXVVAEForceFloat32")
|
||
class LTXVVAEForceFloat32:
|
||
"""Force the video VAE to encode/decode in float32 (native HDR requirement).
|
||
|
||
This node does **not** clone the VAE into a separate float32 copy for the
|
||
output. It mutates the same object that was passed in (``vae_dtype`` and
|
||
``first_stage_model`` weights). The output socket is that same instance.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"vae": ("VAE",),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("VAE",)
|
||
FUNCTION = "patch"
|
||
CATEGORY = "Lightricks/HDR"
|
||
DESCRIPTION = (
|
||
"Forces float32 on the video VAE (vae_dtype + first_stage_model) for "
|
||
"native HDR. IMPORTANT: this does NOT clone a float32 VAE for the "
|
||
"output — it mutates the input object in place and returns the same "
|
||
"instance. If you also wire that loader's VAE (or this output) into "
|
||
"another branch, those nodes share the mutated model: they may see "
|
||
"float16 or float32 depending on execution order (a race with the "
|
||
"original dtype). Prefer a single chain VAELoader → this node → all "
|
||
"encode/decode / IC-LoRA guide uses."
|
||
)
|
||
|
||
def patch(self, vae):
|
||
# In-place: Comfy Module.to(dtype) mutates weights; a wrapper copy would
|
||
# still share first_stage_model and waste peak memory for no isolation.
|
||
vae.vae_dtype = torch.float32
|
||
model = getattr(vae, "first_stage_model", None)
|
||
if model is not None:
|
||
try:
|
||
model.to(dtype=torch.float32)
|
||
except Exception as exc:
|
||
raise RuntimeError(
|
||
"LTXVVAEForceFloat32 could not cast VAE weights to float32 "
|
||
f"({type(exc).__name__}: {exc}). Native HDR requires float32."
|
||
) from exc
|
||
logger.info("VAE dtype forced to float32 for native HDR")
|
||
return (vae,)
|