Automated PR - 2026-04-23

packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py

@@ -1,23 +1,34 @@
+import enum
 import logging
 import math
 from collections.abc import Generator, Iterator
 from fractions import Fraction
 from io import BytesIO
+from pathlib import Path
 
 import av
 import numpy as np
+import OpenImageIO
 import torch
 from einops import rearrange
 from PIL import Image
 from torch._prims_common import DeviceLikeType
 from tqdm import tqdm
 
+from ltx_core.hdr import LogC3
 from ltx_core.types import Audio, VideoPixelShape
 from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
 
 logger = logging.getLogger(__name__)
 
 
+class ResizeMode(enum.Enum):
+    """How to fit a conditioning video to the target resolution."""
+
+    CENTER_CROP = "center_crop"
+    REFLECT_PAD = "reflect_pad"
+
+
 def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
     """
     Resize image preserving aspect ratio (filling target long side).
@@ -79,6 +90,16 @@ def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dt
     return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
 
 
+def to_vae_range(x: torch.Tensor) -> torch.Tensor:
+    """Map [0, 1] to [-1, 1] (VAE input convention)."""
+    return torch.clamp(x, 0.0, 1.0) * 2.0 - 1.0
+
+
+def from_vae_range(z: torch.Tensor) -> torch.Tensor:
+    """Map [-1, 1] (VAE output convention) to [0, 1]."""
+    return torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
+
+
 def load_image_and_preprocess(
     image_path: str,
     height: int,
@@ -124,6 +145,108 @@ def video_preprocess(
     return result
 
 
+def align_resolution(
+    width: int,
+    height: int,
+    resize_mode: ResizeMode,
+    divisor: int = 64,
+) -> tuple[int, int, int, int]:
+    """Compute aligned generation dimensions and crop-back size.
+    Args:
+        width: Source video width (need not be aligned).
+        height: Source video height (need not be aligned).
+        resize_mode: CENTER_CROP rounds down; REFLECT_PAD rounds up.
+        divisor: Alignment divisor (default 64 for two-stage pipelines).
+    Returns:
+        ``(gen_width, gen_height, crop_width, crop_height)`` where
+        ``gen_*`` are multiples of *divisor* and ``crop_*`` are the
+        original dimensions to trim back to after decoding.  When no
+        cropping is needed ``crop_*`` equals ``gen_*``.
+    """
+    if resize_mode is ResizeMode.REFLECT_PAD:
+        gen_w = ((width + divisor - 1) // divisor) * divisor
+        gen_h = ((height + divisor - 1) // divisor) * divisor
+    else:
+        gen_w = (width // divisor) * divisor
+        gen_h = (height // divisor) * divisor
+
+    crop_w = width if gen_w != width else gen_w
+    crop_h = height if gen_h != height else gen_h
+    return gen_w, gen_h, crop_w, crop_h
+
+
+def resize_and_reflect_pad(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """Resize tensor to fit within target, then reflect-pad to exact dimensions.
+    Unlike resize_and_center_crop which stretches and crops, this preserves the
+    original aspect ratio and pads the shorter dimension with reflected pixels.
+    When the target is already >= the source in both dimensions, interpolation
+    is skipped entirely to preserve original pixels.
+    Args:
+        tensor: Input with shape (H, W, C) or (F, H, W, C)
+        height: Target height
+        width: Target width
+    Returns:
+        Tensor with shape (1, C, 1, height, width) for 3D or (1, C, F, height, width) for 4D
+    """
+    if tensor.ndim == 3:
+        tensor = rearrange(tensor, "h w c -> 1 c h w")
+    elif tensor.ndim == 4:
+        tensor = rearrange(tensor, "f h w c -> f c h w")
+    else:
+        raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")
+
+    _, _, src_h, src_w = tensor.shape
+
+    if height >= src_h and width >= src_w:
+        new_h, new_w = src_h, src_w
+    else:
+        scale = min(height / src_h, width / src_w)
+        new_h = round(src_h * scale)
+        new_w = round(src_w * scale)
+        tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)
+
+    pad_bottom = height - new_h
+    pad_right = width - new_w
+    if pad_bottom > 0 or pad_right > 0:
+        pad_mode = "reflect" if pad_bottom < new_h and pad_right < new_w else "replicate"
+        tensor = torch.nn.functional.pad(tensor, (0, pad_right, 0, pad_bottom), mode=pad_mode)
+
+    tensor = rearrange(tensor, "f c h w -> 1 c f h w")
+    return tensor
+
+
+def load_video_conditioning_hdr(
+    video_path: str,
+    height: int,
+    width: int,
+    frame_cap: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    hdr_transform: str = "logc3",
+    resize_mode: ResizeMode = ResizeMode.CENTER_CROP,
+) -> Iterator[torch.Tensor]:
+    """Load a video and yield preprocessed frames for HDR IC-LoRA conditioning.
+    Decodes through the standard path and applies the LDR compression that
+    matches training. Callers are responsible for providing Rec.709 SDR
+    input — the HDR IC-LoRA was trained on that color space.
+    Args:
+        hdr_transform: LDR-compression name (currently only ``logc3``).
+        resize_mode: How to fit the video to the target resolution.
+    Yields:
+        Per-frame tensors of shape ``(1, C, 1, height, width)``.
+    """
+    if hdr_transform != "logc3":
+        raise ValueError(f"Unsupported HDR transform: {hdr_transform}")
+
+    resize_fn = resize_and_reflect_pad if resize_mode is ResizeMode.REFLECT_PAD else resize_and_center_crop
+
+    for f in decode_video_by_frame(path=video_path, frame_cap=frame_cap, device=device):
+        frame = resize_fn(f.to(torch.float32), height, width)
+        ldr = (frame / 255.0).clamp(0.0, 1.0)
+        compressed = LogC3().compress_ldr(ldr)
+        yield to_vae_range(compressed).to(device=device, dtype=dtype)
+
+
 def decode_image(image_path: str) -> np.ndarray:
     image = Image.open(image_path)
     np_array = np.array(image)[..., :3]
@@ -481,3 +604,85 @@ def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array:
     with BytesIO(video_bytes) as video_file:
         image_array = decode_single_frame(video_file)
     return image_array
+
+
+def save_exr_tensor(tensor: torch.Tensor, file_path: str | Path, half: bool = False) -> None:
+    """Save a single tensor frame as EXR with linear sRGB colorspace metadata.
+    Args:
+        tensor: ``[H, W, C]`` or ``[C, H, W]`` float tensor.
+        file_path: Output path (e.g. ``frame_0000.exr``).
+        half: Force float16 output with ZIP compression.
+    """
+    if tensor.dim() == 3 and tensor.shape[0] == 3:
+        tensor = tensor.permute(1, 2, 0)
+    use_half = half or tensor.dtype in (torch.float16, torch.half)
+    img_np = np.ascontiguousarray(tensor.cpu().numpy().astype(np.float32))
+    file_path = str(file_path)
+
+    h, w = img_np.shape[:2]
+    fmt = OpenImageIO.HALF if use_half else OpenImageIO.FLOAT
+    spec = OpenImageIO.ImageSpec(w, h, 3, fmt)
+    spec.channelnames = ("R", "G", "B")
+    spec.attribute("compression", "zip")
+    spec.attribute("chromaticities", "float[8]", (0.64, 0.33, 0.30, 0.60, 0.15, 0.06, 0.3127, 0.3290))
+    spec.attribute("colorSpace", "sRGB")
+
+    out = OpenImageIO.ImageOutput.create(file_path)
+    if out is None:
+        raise RuntimeError(
+            f"Failed to create EXR writer for '{file_path}'. Ensure OpenImageIO is built with OpenEXR support."
+        )
+    try:
+        if not out.open(file_path, spec):
+            raise RuntimeError(f"Failed to open EXR file '{file_path}': {out.geterror()}")
+        if not out.write_image(img_np):
+            raise RuntimeError(f"Failed to write EXR image '{file_path}': {out.geterror()}")
+    finally:
+        out.close()
+
+
+def _linear_to_srgb(x: np.ndarray) -> np.ndarray:
+    """Linear -> sRGB OETF per IEC 61966-2-1. Input assumed in [0, 1]."""
+    x = np.clip(x, 0.0, 1.0)
+    return np.where(x <= 0.0031308, x * 12.92, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)
+
+
+def encode_exr_sequence_to_mp4(exr_dir: Path, output_mp4: Path, frame_rate: float) -> None:
+    """Convert a linear EXR frame sequence to sRGB and encode to H.264 .mp4 via PyAV.
+    Exposure is fixed at EV=0 (no gain). Each EXR frame is clamped to [0, 1],
+    passed through the sRGB OETF, quantised to 8-bit BGR, and fed to a libx264
+    stream (crf 18, yuv420p). ``frame_rate`` is the original source video's
+    frame rate so playback matches the input timing.
+    """
+    import os  # noqa: PLC0415
+
+    os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+    import cv2  # noqa: PLC0415
+
+    exr_files = sorted(exr_dir.glob("frame_*.exr"))
+    if not exr_files:
+        raise FileNotFoundError(f"No EXR frames found in {exr_dir}")
+
+    container = av.open(str(output_mp4), mode="w")
+    stream = container.add_stream("libx264", rate=Fraction(frame_rate).limit_denominator(1000))
+    stream.pix_fmt = "yuv420p"
+    stream.options = {"crf": "18", "movflags": "+faststart"}
+
+    try:
+        for i, exr_path in enumerate(exr_files):
+            hdr = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED).astype(np.float32)
+            sdr = _linear_to_srgb(np.maximum(hdr, 0.0))
+            bgr8 = (sdr * 255.0 + 0.5).astype(np.uint8)
+
+            if i == 0:
+                stream.height = bgr8.shape[0]
+                stream.width = bgr8.shape[1]
+
+            frame = av.VideoFrame.from_ndarray(bgr8, format="bgr24")
+            for packet in stream.encode(frame):
+                container.mux(packet)
+
+        for packet in stream.encode():
+            container.mux(packet)
+    finally:
+        container.close()