hdr_nodes.py

"""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,)