vllm-omni/vllm_omni/diffusion/utils/chunked_video.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Turn committed VAE chunks into MP4 bytes without materializing the video.

The producer side is the model's business: a Wan VAE walks a causal
frame-by-frame loop, a MiniMax-H3 VAE walks overlapping clips, and Wan S2V
decodes one clip per autoregressive iteration. What happens to a chunk once it
is final is not: every one of them quantizes to uint8, moves to the host once,
and lands in a bounded encoder. That consumer lives here so a model only has to
publish finished chunks -- :class:`SupportsChunkedVAEDecode` for producers that
can be driven, :class:`ChunkedVideoMP4Session` for producers that drive
themselves.
"""

from __future__ import annotations

from typing import Any

import numpy as np
import torch

from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder


def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray:
    """Quantize a ``BCTHW`` chunk to ``BTHWC`` uint8, transferring once.

    ``value_range`` is the interval the producer publishes, which differs per
    checkpoint (see ``SupportsChunkedVAEDecode.chunk_value_range``). Quantizing
    on the accelerator means one transfer moves the final bytes rather than
    float frames.

    That transfer is synchronous: ``.cpu()`` blocks inside the producer's chunk
    callback, so what overlaps the remaining decode today is the CPU H.264
    encode of earlier chunks, not this copy. Splitting device preparation,
    device-to-host copy, and encode into three independently scheduled stages --
    a pool of reusable pinned host slots with non-blocking copies on a dedicated
    stream, leases released by the encoder -- fits behind this same callback
    contract, but is not implemented here.
    """
    low, high = float(value_range[0]), float(value_range[1])
    if high <= low:
        raise ValueError(f"value_range must be increasing, got {value_range!r}")
    scale = 255.0 / (high - low)
    frames = chunk.clamp(low, high).sub(low).mul(scale).round().to(torch.uint8)
    return frames.permute(0, 2, 3, 4, 1).cpu().numpy()


class ChunkedVideoMP4Session:
    """Encode committed video chunks into one progressive MP4 per batch entry.

    Push finished ``BCTHW`` chunks as the producer commits them; each batch
    entry gets its own bounded encoder, so host transfer and H.264 encoding
    overlap whatever the producer is still decoding.

    ``audio_waveforms`` holds one waveform per batch entry (``None`` for a
    silent entry), so a caller with several outputs per prompt repeats a
    request's waveform across its entries. ``batch_frames`` coalesces transfers
    for producers that publish finer than a transfer is worth; ``crop`` trims
    the decoder's padding to the requested output size.
    """

    def __init__(
        self,
        *,
        value_range: tuple[float, float],
        fps: float,
        audio_waveforms: list[np.ndarray | None] | None = None,
        audio_sample_rate: int | None = None,
        batch_frames: int = 1,
        max_pending: int = 2,
        video_codec_options: dict[str, str] | None = None,
        crop: tuple[int, int] | None = None,
    ) -> None:
        if batch_frames <= 0:
            raise ValueError("batch_frames must be positive")
        self._value_range = value_range
        self._fps = fps
        self._audio_waveforms = audio_waveforms
        self._audio_sample_rate = audio_sample_rate
        self._batch_frames = batch_frames
        self._max_pending = max_pending
        self._video_codec_options = video_codec_options
        self._crop = crop
        self._encoders: list[ChunkedMP4Encoder] = []
        self._pending: list[torch.Tensor] = []
        self._pending_frames = 0

    def push(self, chunk: torch.Tensor) -> None:
        """Queue one committed ``BCTHW`` chunk."""
        if self._crop is not None:
            height, width = self._crop
            chunk = chunk[..., :height, :width]
        self._pending.append(chunk)
        self._pending_frames += int(chunk.shape[2])
        if self._pending_frames >= self._batch_frames:
            self._flush()

    def finish(self) -> list[bytes]:
        """Flush what is pending and return one MP4 per batch entry."""
        self._flush()
        return [encoder.finish() for encoder in self._encoders]

    def abort(self) -> None:
        for encoder in self._encoders:
            encoder.abort()

    def _flush(self) -> None:
        if not self._pending:
            return
        frames = chunk_to_uint8_frames(torch.cat(self._pending, dim=2), self._value_range)
        if not self._encoders:
            self._encoders = [
                ChunkedMP4Encoder(
                    width=frames.shape[3],
                    height=frames.shape[2],
                    fps=self._fps,
                    audio_waveform=self._waveform_for(index, frames.shape[0]),
                    audio_sample_rate=self._audio_sample_rate,
                    max_pending=self._max_pending,
                    video_codec_options=self._video_codec_options,
                )
                for index in range(frames.shape[0])
            ]
        for index, encoder in enumerate(self._encoders):
            encoder.push(np.ascontiguousarray(frames[index]))
        self._pending.clear()
        self._pending_frames = 0

    def _waveform_for(self, index: int, batch_size: int) -> np.ndarray | None:
        if self._audio_waveforms is None:
            return None
        if len(self._audio_waveforms) != batch_size:
            raise ValueError(
                f"expected one audio waveform per batch entry, got "
                f"{len(self._audio_waveforms)} for {batch_size} entries"
            )
        return self._audio_waveforms[index]


def decode_to_mp4(vae: Any, z: torch.Tensor, **session_kwargs: Any) -> list[bytes]:
    """Decode ``z`` straight into one progressive MP4 per batch entry.

    Drives a VAE that declares :class:`SupportsChunkedVAEDecode`, so host
    transfer and encoding overlap the remaining decode and the full video is
    never materialized. Ranks that own no decode output receive no chunks and
    get an empty list, matching the empty tensor the full-decode path returns
    there.
    """
    if not supports_chunked_vae_decode(vae):
        raise TypeError(f"{type(vae).__name__} does not expose the chunked VAE decode capability")

    session = ChunkedVideoMP4Session(value_range=vae.chunk_value_range, **session_kwargs)
    try:
        vae.decode_with_chunks(z, on_chunk=session.push)
        return session.finish()
    except BaseException:
        session.abort()
        raise