# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Session-owned streaming VAE decode for realtime AR-Diffusion. A causal video decoder already decodes one latent frame at a time, carrying a bounded temporal cache between frames. What it does not do today is carry that cache *between calls*: every decode call clears it, so a chunk can only be decoded as part of a whole clip. This module moves the cache from the decoder module to the session, which is the entire difference between "decode a clip" and "stream a session": * the cache becomes :class:`StreamingDecodeState`, threaded in and out by the caller, so two concurrent sessions cannot overwrite each other's temporal context -- a failure that produces a plausible but wrong video while every timing metric stays green; * ``first_chunk`` becomes a property of the session's first frame rather than of the call, which is what makes chunk ``N + 1`` continue chunk ``N``. Resident state is bounded by construction: each causal convolution retains at most ``CACHE_T`` temporal slices, so the cache is a function of resolution and never of session length. It is not free -- see :meth:`StreamingDecodeState.nbytes`, which is the quantity session admission has to account for. Measured on the reference checkpoint at 832x480 in bf16, on one RTX PRO 6000 with torch 2.11: chunking is neutral -- a session decoded in four chunks is bit-identical to the same session decoded in one call -- and resident decoder state holds at 1801 MiB across every chunk. What a session gains is structural rather than a speedup: a chunk becomes decodable on its own, so pixels leave on the block that produced them instead of after the last one, and the session's opening frame is expanded once rather than once per block. Note what "neutral" is measured against. ``AutoencoderKLWan._decode`` applies ``post_quant_conv`` to the whole latent at once, while the tiled path in this repo -- and this decoder -- apply it per frame. It is a 1x1x1 convolution, so the two are mathematically identical, but they can select different kernels; on bf16 CUDA that seeds a 0.004 difference which the causal convolutions downstream amplify to 0.047. Comparing against ``_decode`` therefore measures where ``post_quant_conv`` runs, not anything about streaming, and the result moves with the torch version. The control that isolates chunking is the same session decoded in a single call through this same path. The consumer is the stepwise execution path: one ``generate()`` is one session, ``post_decode`` emits one chunk per AR block, and the pipeline holds one state per ``request_id``, releasing it when the runner closes or resets that AR session -- so decoder state and KV state share a single lifetime. Only ``torch`` is imported here so the contract can be exercised without a device, a checkpoint, or the distributed VAE stack. """ from __future__ import annotations from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from typing import Any, Protocol, runtime_checkable import torch @dataclass class StreamingDecodeState: """One session's temporal decoder cache. ``feat_map`` holds one entry per causal convolution. Entries are tensors, ``None`` before that convolution has run, or a non-tensor sentinel the decoder uses for its first upsample step, so byte accounting must skip anything that is not a tensor. The cache has **no recompute source**: it is the compressed tail of every frame decoded so far and cannot be rebuilt from the prompt or the latents. Releasing it ends the stream; there is no partial eviction. """ session_id: str feat_map: list[Any] conv_idx: list[int] = field(default_factory=lambda: [0]) frames_decoded: int = 0 chunks_decoded: int = 0 def nbytes(self) -> int: """Resident bytes of the cache, skipping non-tensor entries.""" return sum(t.numel() * t.element_size() for t in self.feat_map if torch.is_tensor(t)) def nbytes_by_device(self) -> dict[str, int]: """Resident bytes grouped by device, for admission accounting.""" totals: dict[str, int] = {} for tensor in self.feat_map: if torch.is_tensor(tensor): key = str(tensor.device) totals[key] = totals.get(key, 0) + tensor.numel() * tensor.element_size() return totals @property def started(self) -> bool: return self.frames_decoded > 0 def release(self) -> None: """Drop the cache and return the session to its pre-stream state.""" self.feat_map = [None] * len(self.feat_map) self.conv_idx = [0] self.frames_decoded = 0 self.chunks_decoded = 0 @runtime_checkable class SupportsStreamingDecode(Protocol): """A decoder that can emit a session's chunks incrementally. The lifecycle mirrors the AR session's: ``new_decode_state`` at session creation, ``decode_chunk`` per committed chunk, and ``release`` on reset or close, so decoder state and KV state are released together rather than through two independent lifecycles. """ def new_decode_state(self, session_id: str) -> StreamingDecodeState: """Create empty decoder state for a new session.""" ... def decode_chunk(self, latent: torch.Tensor, state: StreamingDecodeState) -> torch.Tensor: """Decode one committed chunk, advancing ``state`` in place.""" ... def declared_state_bytes(self, *, height: int, width: int, dtype: torch.dtype) -> int: """Upper bound on resident decoder bytes per session at this shape.""" ... def release(self, state: StreamingDecodeState) -> None: """Drop ``state``'s cache and return it to its pre-stream state. Declared here, not just implemented on :class:`WanStreamingDecoder`, so a caller programming against this protocol -- a session runtime releasing decoder state alongside KV state on close or failure -- can rely on it being present without depending on the concrete class. """ ... class WanStreamingDecoder: """Streaming decode over a Wan-family causal autoencoder. Wraps any module exposing ``decoder``, ``post_quant_conv`` and the cached causal-convolution counts. It never touches the module's own ``_feat_map``, so a module shared by several sessions stays stateless between calls. """ def __init__(self, vae: Any, *, bytes_per_pixel_fp32: float | None = None) -> None: for attribute in ("decoder", "post_quant_conv"): if not hasattr(vae, attribute): raise TypeError(f"Streaming decode requires a Wan-family autoencoder exposing {attribute!r}.") self._vae = vae # A one-off measured constant for this checkpoint; see declared_state_bytes. # Optional at construction because most callers never need admission # accounting, but a constructor argument (rather than requiring an # external attribute assignment after the fact) means the value can't # be silently unset for a caller that does. self._bytes_per_pixel_fp32 = bytes_per_pixel_fp32 def _execution_context(self) -> AbstractContextManager[None]: """The autoencoder's own execution context, when it exposes one. Every public compute entry point on ``OmniAutoencoderKLWan`` -- ``encode``, ``decode``, ``decode_with_chunks``, ``tile_exec`` -- runs inside ``_execution_context()``, which is where autocast is entered for fp16/bf16 parameters (and on NPU is the only place ``torch.npu.amp.autocast`` is entered at all). This decoder drives ``post_quant_conv`` and ``decoder`` directly, so it has to enter the same context or streamed chunks compute under a different numeric and dispatch context from the whole-clip decode they are compared against. A plain ``diffusers`` autoencoder has no such hook, and then there is nothing to enter. """ context = getattr(self._vae, "_execution_context", None) if callable(context): return context() return nullcontext() @property def num_causal_convs(self) -> int: counts = getattr(self._vae, "_cached_conv_counts", None) if not isinstance(counts, dict) or "decoder" not in counts: raise TypeError("The autoencoder does not expose cached decoder convolution counts.") return int(counts["decoder"]) def new_decode_state(self, session_id: str) -> StreamingDecodeState: if not isinstance(session_id, str) or not session_id.strip(): raise ValueError("session_id must be a non-empty string.") return StreamingDecodeState(session_id=session_id, feat_map=[None] * self.num_causal_convs) def decode_chunk(self, latent: torch.Tensor, state: StreamingDecodeState) -> torch.Tensor: """Decode ``latent`` as the continuation of ``state``'s session. ``latent`` is ``[1, C, T, H, W]`` in latent space, already rescaled by the pipeline's latent statistics. Frames are decoded one at a time, the same loop the non-streaming path runs, with the cache carried across calls instead of cleared. The returned tensor is ``[1, 3, T', H', W']`` where ``T'`` is smaller for a session's first chunk: the opening latent frame expands to a single raw frame and every later one to the full temporal factor. """ if latent.ndim != 5: raise ValueError(f"latent must be [B, C, T, H, W]; got shape {tuple(latent.shape)}.") if latent.shape[0] != 1: # One state is one session's temporal context, so a batch would # thread several samples through a single cache. ``decode()`` # handles a batch by decoding each sample with its own cleared # cache (``use_slicing``); there is no equivalent here, and the # streaming contract is one session per state, so refuse rather # than silently mix samples. raise ValueError( "streaming decode takes one sample per session state; " f"got batch size {latent.shape[0]}. Decode each sample through its own state." ) if len(state.feat_map) != self.num_causal_convs: raise ValueError( "Decoder state does not belong to this decoder: " f"{len(state.feat_map)} cache slots against {self.num_causal_convs} causal convolutions." ) num_frames = int(latent.shape[2]) if num_frames == 0: raise ValueError("latent must carry at least one frame.") decoded_frames = [] with self._execution_context(): for index in range(num_frames): state.conv_idx[0] = 0 frame = self._vae.post_quant_conv(latent[:, :, index : index + 1]) # first_chunk marks the session's opening frame, not the call's: # that is what makes chunk N + 1 continue chunk N rather than # restart the causal expansion. decoded_frames.append( self._vae.decoder( frame, feat_cache=state.feat_map, feat_idx=state.conv_idx, first_chunk=(state.frames_decoded == 0), ) ) state.frames_decoded += 1 out = torch.cat(decoded_frames, dim=2) patch_size = getattr(getattr(self._vae, "config", None), "patch_size", None) if patch_size is not None: from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify out = unpatchify(out, patch_size=patch_size) out = torch.clamp(out, min=-1.0, max=1.0) state.chunks_decoded += 1 return out def release(self, state: StreamingDecodeState) -> None: state.release() def declared_state_bytes(self, *, height: int, width: int, dtype: torch.dtype) -> int: """Resident decoder bytes per session at this output shape. Every cached tensor is ``[B, C_l, <= CACHE_T, H_l, W_l]`` with ``H_l`` and ``W_l`` fixed fractions of the output, so the total scales exactly with ``height * width``. Callers that need the constant should measure it once for their checkpoint and pass it as ``bytes_per_pixel_fp32`` to the constructor; this helper scales that measured constant so admission does not have to re-measure per resolution. """ constant = self._bytes_per_pixel_fp32 if constant is None: raise RuntimeError( "declared_state_bytes needs a measured bytes-per-pixel constant for this checkpoint; " "pass bytes_per_pixel_fp32= to WanStreamingDecoder() from a one-off measurement." ) element_ratio = torch.empty((), dtype=dtype).element_size() / 4 return int(constant * height * width * element_ratio)