vllm-omni/vllm_omni/diffusion/models/minimax_h3/vae_temporal.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Memory-optimized forks of the checkpoint's temporal VAE loops.

Two method replacements installed on the checkpoint's ``AutoencoderKLLegacy``
at adapter construction:

* ``encode_temporal`` pads misaligned frame counts by concatenating repeated
  last frames onto the *whole* video before the chunk loop (~4GB copy for a
  15s clip). The fork pads inside the tail chunk instead: the chunks see
  bit-identical frames while the copy shrinks to a single clip.

* ``_decode_temporal_streaming`` accumulates the whole decoded video in a
  floating-point buffer that is denormalized, clamped, and quantized to
  uint8 only after the final chunk lands. The fork fuses the adapter's
  in-place revert and the output quantizer into ``write_part`` and
  accumulates directly into a uint8 buffer -- the same per-element op order
  (``(x - mean) / std -> clamp(0, 1) -> *255 -> round -> uint8``), so the
  delivered video is bit-identical while the resident decode buffer shrinks
  4x.

Set ``VLLM_OMNI_VAE_LEGACY_TEMPORAL=1`` to keep the checkpoint's methods.
"""

from __future__ import annotations

import os
from types import MethodType

import torch
from vllm.logger import init_logger

logger = init_logger(__name__)

_LEGACY_TEMPORAL_ENV = "VLLM_OMNI_VAE_LEGACY_TEMPORAL"


def _legacy_temporal_enabled() -> bool:
    value = os.environ.get(_LEGACY_TEMPORAL_ENV, "0")
    return value.strip().lower() not in ("", "0", "false", "off")


def _encode_temporal_tail_pad(self, x):
    """Drop-in fork of the checkpoint's ``encode_temporal``.

    Frame handling is identical to the original; only the padding site
    moves from a whole-video ``torch.cat`` into the tail chunk.
    """
    offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0

    orig_frames = int(x.shape[2])
    pad_size = (offset_frame - orig_frames) % self.clip_length
    padded_frames = orig_frames + pad_size
    num_chunks = (padded_frames - offset_frame) // self.clip_length

    z_list = []
    for i in range(num_chunks):
        start_idx = i * self.clip_length + offset_frame
        end_idx = (i + 1) * self.clip_length + offset_frame
        clip_x = x[:, :, start_idx : min(end_idx, orig_frames), :, :]
        if end_idx > orig_frames:
            # The original pads the whole video up front by repeating the
            # last frame; repeating it inside the tail chunk feeds the
            # identical frames to the identical encodes.
            clip_x = torch.cat(
                [clip_x, clip_x[:, :, -1:].repeat(1, 1, end_idx - orig_frames, 1, 1)],
                dim=2,
            )

        if self.isolated_key_frame:
            key_frame = clip_x[:, :, :1, :, :]
            z_key = self._adaptive_encode(key_frame)

            if clip_x.shape[2] > 1:
                video_frames = clip_x[:, :, 1:, :, :]
                z_video = self._adaptive_encode(video_frames)
                z = torch.cat([z_key, z_video], dim=2)
            else:
                z = z_key
        else:
            z = self._adaptive_encode(clip_x)

        z_list.append(z)

    z = torch.cat(z_list, dim=2)
    if self.token_drop > 0:
        z = z[:, :, : -self.token_drop]

    if self.isolated_first_frame:
        input_first_frame = x[:, :, :1, :, :]
        z_first_frame = self._adaptive_encode(input_first_frame)

        if self.frame_pre_padding == 0:
            z = torch.cat([z_first_frame, z], dim=2)
        else:
            z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2)

    if self.isolated_last_frame:
        # The original indexes into the padded video: pad frames replicate
        # the last real frame, out-of-range indexes produce an empty slice
        # (a no-op), and negative indexes wrap. Reproduce that mapping
        # against the unpadded tensor instead of materializing the pad.
        last_frame_idx = padded_frames - self.frame_drop + offset_frame
        if last_frame_idx < 0:
            last_frame_idx += padded_frames
        if 0 <= last_frame_idx < padded_frames:
            src_idx = min(last_frame_idx, orig_frames - 1)
            input_last_frame = x[:, :, src_idx : src_idx + 1, :, :]
            z_last_frame = self._adaptive_encode(input_last_frame)
            z = torch.cat([z, z_last_frame], dim=2)

    return z


def _decode_temporal_streaming_uint8(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype):
    """Drop-in fork of the checkpoint's ``_decode_temporal_streaming``.

    The chunk loop, overlap blending, frame plan, and assertions are
    identical to the original. ``write_part`` additionally runs the
    denormalize -> clamp -> quantize chain in place (the same op order the
    pipeline applies after decode) and accumulates into a uint8 buffer.
    """
    total_frames, pad_frames, output_frames = self._decode_temporal_output_frame_plan(
        z, z_head, z_tail, num_chunks, pad_tokens
    )
    if output_frames <= 0:
        raise ValueError(
            f"decode_temporal streaming planned non-positive output_frames={output_frames} "
            f"total_frames={total_frames} pad_frames={pad_frames}"
        )

    chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
    split_count = int(self.token_drop > 0) + 1
    dec = None
    dec_overlap = None
    norm = None
    write_pos = 0
    logical_frames = 0
    dropped_frames = 0
    decoded_count = 0

    def write_part(part):
        nonlocal dec, dec_overlap, norm, write_pos, logical_frames, dropped_frames
        part_frames = int(part.shape[2])
        if part_frames <= 0:
            return
        logical_frames += part_frames
        if dec is None:
            out_shape = list(part.shape)
            out_shape[2] = output_frames
            dec = torch.empty(out_shape, dtype=torch.uint8, device=part.device)
            transform_rev = self.processor.transform_rev
            mean = torch.as_tensor(transform_rev.mean, dtype=part.dtype, device=part.device).view(1, 3, 1, 1, 1)
            std = torch.as_tensor(transform_rev.std, dtype=part.dtype, device=part.device).view(1, 3, 1, 1, 1)
            norm = (mean, std)
        # revert_tensor ((x - mean) / std, clamp) followed by the output
        # quantizer (*255, round, uint8), in place on the part: the blended
        # raw values are consumed before this point, and every write into
        # ``dec`` goes through here, so the final buffer matches the legacy
        # post-decode chain element for element. Denormalization and clamping
        # run in the decoded dtype exactly like the legacy adapter; the
        # quantizer then runs in FP32 because the legacy adapter returned
        # ``frames.float()`` before the pipeline multiplied and rounded —
        # rounding fp16/bf16 directly flips borderline pixels by one.
        mean_t, std_t = norm
        part.sub_(mean_t).div_(std_t).clamp_(0.0, 1.0)
        part = part.float()
        part.mul_(255.0).round_()

        remaining = int(dec.shape[2]) - write_pos
        copy_frames = min(part_frames, max(0, remaining))
        if copy_frames > 0:
            dec[:, :, write_pos : write_pos + copy_frames, :, :].copy_(part[:, :, :copy_frames, :, :])
            write_pos += copy_frames
        dropped_frames += part_frames - copy_frames

    for i in range(num_chunks):
        t_start_idx = i * self.tokens_chunk_size
        t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
        clip_z = z[:, :, t_start_idx:t_end_idx, :, :]

        if i == 0 and z_head is not None:
            clip_z = torch.cat([z_head, clip_z], dim=2)

        if i == num_chunks - 1 and z_tail is not None:
            clip_z = torch.cat([clip_z, z_tail], dim=2)

        clip_dec = self._adaptive_decode(clip_z)
        decoded_count += 1
        if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
            clip_dec = clip_dec.to(temporal_cat_dtype)
        if clip_dec.device != z.device:
            clip_dec = clip_dec.to(z.device)

        dec_tail = None
        if i == 0 and z_head is not None:
            write_part(clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :])
            clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :]

        if i == num_chunks - 1 and z_tail is not None:
            dec_tail = clip_dec[:, :, -1:, :, :]
            clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :]

        for j in range(split_count):
            f_start_idx = j * chunk_dec
            f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
            clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
            clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]

            if j == 0:
                if dec_overlap is not None:
                    clip_dec_chunk = self.blend(dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3)
                    dec_overlap = None
                write_part(clip_dec_chunk)
            else:
                # Break the view's reference to the full decoded clip so earlier
                # temporal chunks can be released before the final output exists.
                dec_overlap = clip_dec_chunk.contiguous()

        if i == num_chunks - 1:
            if dec_overlap is not None:
                write_part(dec_overlap)
                dec_overlap = None
            if dec_tail is not None:
                write_part(dec_tail)

        del clip_dec, clip_z

    if dec is None:
        raise RuntimeError("decode_temporal streaming produced no output tensor")
    if logical_frames != total_frames or dropped_frames != pad_frames or write_pos != output_frames:
        raise RuntimeError(
            "decode_temporal streaming frame plan mismatch: "
            f"logical_frames={logical_frames} total_frames={total_frames} "
            f"dropped_frames={dropped_frames} pad_frames={pad_frames} "
            f"write_pos={write_pos} output_frames={output_frames}"
        )

    return dec


def _processor_denorm_available(model) -> bool:
    transform_rev = getattr(getattr(model, "processor", None), "transform_rev", None)
    mean = getattr(transform_rev, "mean", None)
    std = getattr(transform_rev, "std", None)
    return mean is not None and std is not None and len(mean) == 3 and len(std) == 3


def install_temporal_stream_patches(model) -> None:
    """Install the memory-optimized temporal forks on the checkpoint model.

    Each fork is installed only when the checkpoint contract it mirrors is
    discoverable; otherwise the checkpoint's own method keeps running (the
    adapter's post-decode revert already covers that path).
    """
    if _legacy_temporal_enabled():
        logger.info("MiniMax-H3 VAE temporal stream patches disabled by %s", _LEGACY_TEMPORAL_ENV)
        return
    if callable(getattr(model, "encode_temporal", None)) and hasattr(model, "clip_length"):
        model.encode_temporal = MethodType(_encode_temporal_tail_pad, model)
        logger.info("MiniMax-H3 VAE encode_temporal tail-padding patch installed")
    else:
        logger.warning("MiniMax-H3 VAE encode_temporal contract not found; keeping checkpoint method")
    if callable(getattr(model, "_decode_temporal_streaming", None)) and _processor_denorm_available(model):
        model._decode_temporal_streaming = MethodType(_decode_temporal_streaming_uint8, model)
        logger.info("MiniMax-H3 VAE streaming uint8 decode patch installed")
    else:
        logger.warning("MiniMax-H3 VAE streaming decode contract not found; keeping checkpoint method")


__all__ = ["install_temporal_stream_patches"]