vllm-omni/tests/diffusion/ar_diffusion/test_streaming_decode.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Streaming VAE decode against the real Wan autoencoder, on CPU.

The autoencoder is built from a config rather than a checkpoint -- weights are
random, only shapes and the temporal cache protocol matter -- so these run
without a device or a download.

The load-bearing test is equivalence: decoding a session chunk by chunk must
produce exactly what decoding the whole clip produces. Everything else about
streaming is worthless if that does not hold.
"""

from __future__ import annotations

from contextlib import contextmanager

import pytest
import torch

from vllm_omni.experimental.ar_diffusion.streaming_decode import (
    StreamingDecodeState,
    SupportsStreamingDecode,
    WanStreamingDecoder,
)

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

# Small enough to decode on CPU in a test, same architecture as the shipped
# checkpoint: base_dim 96 / dim_mult [1,2,4,4] would be far too slow, so the
# widths are scaled down while the causal structure is unchanged.
VAE_CONFIG = {
    "base_dim": 8,
    "z_dim": 4,
    "dim_mult": [1, 2],
    "num_res_blocks": 1,
    "temperal_downsample": [True],
    "attn_scales": [],
    "dropout": 0.0,
}
LATENT_H = LATENT_W = 4


@pytest.fixture(scope="module")
def vae():
    diffusers = pytest.importorskip("diffusers")
    torch.manual_seed(0)
    model = diffusers.AutoencoderKLWan(**VAE_CONFIG).eval()
    for parameter in model.parameters():
        parameter.requires_grad_(False)
    return model


@pytest.fixture(scope="module")
def decoder(vae):
    return WanStreamingDecoder(vae)


def _latent(num_frames: int, *, seed: int = 0) -> torch.Tensor:
    generator = torch.Generator().manual_seed(seed)
    return torch.randn(1, VAE_CONFIG["z_dim"], num_frames, LATENT_H, LATENT_W, generator=generator)


@torch.no_grad()
def _decode_whole(vae, latent: torch.Tensor) -> torch.Tensor:
    """The non-streaming reference: one call, cache cleared at both ends."""
    return vae._decode(latent, return_dict=False)[0]


# --------------------------------------------------------------------------
# Equivalence: the property that makes streaming worth doing at all
# --------------------------------------------------------------------------


def _stream(decoder, latent: torch.Tensor, chunk_sizes, *, session_id: str = "s") -> torch.Tensor:
    state = decoder.new_decode_state(session_id)
    pieces, offset = [], 0
    for size in chunk_sizes:
        pieces.append(decoder.decode_chunk(latent[:, :, offset : offset + size], state))
        offset += size
    return torch.cat(pieces, dim=2)


@torch.no_grad()
@pytest.mark.parametrize("chunk_sizes", [(3, 3, 3), (1, 1, 1, 1, 1), (2, 4, 3)])
def test_chunking_is_neutral(vae, decoder, chunk_sizes) -> None:
    """Splitting a session into chunks must change nothing about its output.

    This is the property streaming actually introduces, so it is the one
    asserted exactly. The control is the same session decoded in a single
    call through the same code path -- not ``_decode`` -- because ``_decode``
    differs in a way that has nothing to do with chunking; see
    ``test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact``.
    """
    total = sum(chunk_sizes)
    latent = _latent(total)
    reference = _stream(decoder, latent, (total,), session_id="ref")
    streamed = _stream(decoder, latent, chunk_sizes)

    assert streamed.shape == reference.shape
    torch.testing.assert_close(streamed, reference, rtol=0, atol=0)


@torch.no_grad()
def test_distance_from_the_one_shot_path_is_a_post_quant_conv_artifact(vae, decoder) -> None:
    """Explain, and bound, the gap against diffusers' ``_decode``.

    ``_decode`` applies ``post_quant_conv`` to the whole latent at once; 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, and on bf16 CUDA that difference is real and is
    amplified by the causal convolutions downstream.

    Asserting equality with ``_decode`` would therefore pin a kernel-selection
    coincidence rather than a property of streaming. Assert instead that the
    gap is entirely explained by where ``post_quant_conv`` runs.
    """
    latent = _latent(6)
    per_frame_pq = torch.cat([vae.post_quant_conv(latent[:, :, k : k + 1]) for k in range(latent.shape[2])], dim=2)
    whole_pq = vae.post_quant_conv(latent)

    streamed = _stream(decoder, latent, (3, 3))
    one_shot = _decode_whole(vae, latent)

    # Same inputs to the decoder => same output, whatever the backend does.
    if torch.equal(whole_pq, per_frame_pq):
        torch.testing.assert_close(streamed, one_shot, rtol=0, atol=0)
    else:  # pragma: no cover - backend dependent
        pytest.skip("backend batches post_quant_conv differently; the gap is that, not chunking")


@torch.no_grad()
def test_a_restarted_stream_does_not_match_a_continued_one(decoder) -> None:
    """Negative control: the equivalence test above must be discriminating.

    Clearing state between chunks -- what the non-streaming path does on every
    call -- has to change the pixels, otherwise the temporal cache carries
    nothing and the assertion above holds for a write-only cache too.

    Restarting also changes the frame count, because a fresh state expands its
    opening latent frame to one raw frame. Comparing totals would therefore
    pass on the frame count alone and never look at a pixel, so this compares
    the tail both runs agree on: the frames from the chunk's later latents,
    which are the full temporal factor wide on either side.
    """
    latent = _latent(6)
    state = decoder.new_decode_state("continued")
    decoder.decode_chunk(latent[:, :, :3], state)
    continued = decoder.decode_chunk(latent[:, :, 3:], state)
    restarted = decoder.decode_chunk(latent[:, :, 3:], decoder.new_decode_state("restarted"))

    tail = min(continued.shape[2], restarted.shape[2]) - 1
    assert tail > 0, "the comparison needs at least one fully expanded frame on both sides"
    continued_tail, restarted_tail = continued[:, :, -tail:], restarted[:, :, -tail:]
    assert continued_tail.shape == restarted_tail.shape, "the tails must be comparable for this to mean anything"
    assert not torch.equal(continued_tail, restarted_tail), (
        "restarting the decoder produced the same pixels as continuing it, so the carried cache is not being read"
    )


# --------------------------------------------------------------------------
# Cross-session isolation: the failure that timing metrics cannot see
# --------------------------------------------------------------------------


@torch.no_grad()
@pytest.mark.parametrize("seed_a, seed_b", [(1, 2), (11, 12)])
def test_interleaved_sessions_keep_independent_temporal_context(vae, decoder, seed_a: int, seed_b: int) -> None:
    """Two sessions ticking alternately must each match their solo decode.

    A session that reads another's temporal cache produces a perfectly
    plausible but wrong video, and every latency metric stays green, so this
    has to be asserted against a recorded solo run rather than inferred. The
    control is a solo stream rather than ``_decode``, so it differs from the
    run under test only in the interleaving.
    """
    latent_a = _latent(6, seed=seed_a)
    latent_b = _latent(6, seed=seed_b)
    solo_a = _stream(decoder, latent_a, (6,), session_id="solo-a")
    solo_b = _stream(decoder, latent_b, (6,), session_id="solo-b")
    assert not torch.equal(solo_a, solo_b), "negative control: the two sessions must differ"

    state_a = decoder.new_decode_state("a")
    state_b = decoder.new_decode_state("b")
    out_a, out_b = [], []
    for start in (0, 3):
        window = slice(start, start + 3)
        out_a.append(decoder.decode_chunk(latent_a[:, :, window], state_a))
        out_b.append(decoder.decode_chunk(latent_b[:, :, window], state_b))

    torch.testing.assert_close(torch.cat(out_a, dim=2), solo_a, rtol=0, atol=0)
    torch.testing.assert_close(torch.cat(out_b, dim=2), solo_b, rtol=0, atol=0)


@torch.no_grad()
def test_the_module_cache_is_untouched_so_sessions_can_share_one_decoder(vae, decoder) -> None:
    vae.clear_cache()
    before = [entry for entry in vae._feat_map]
    decoder.decode_chunk(_latent(3), decoder.new_decode_state("s"))
    assert vae._feat_map == before, "streaming decode must not write the module-owned cache"


# --------------------------------------------------------------------------
# Causal geometry and boundedness
# --------------------------------------------------------------------------


@torch.no_grad()
def test_opening_chunk_is_shorter_and_later_chunks_are_full_length(vae, decoder) -> None:
    """(n - 1) * factor + 1 for the opening chunk, n * factor after it."""
    factor = 2 ** sum(VAE_CONFIG["temperal_downsample"])
    state = decoder.new_decode_state("s")
    first = decoder.decode_chunk(_latent(3), state)
    second = decoder.decode_chunk(_latent(3, seed=5), state)

    assert first.shape[2] == (3 - 1) * factor + 1
    assert second.shape[2] == 3 * factor
    assert state.chunks_decoded == 2
    assert state.frames_decoded == 6


@torch.no_grad()
def test_resident_state_does_not_grow_with_session_length(vae, decoder) -> None:
    """Each causal convolution retains at most CACHE_T temporal slices."""
    state = decoder.new_decode_state("s")
    decoder.decode_chunk(_latent(3), state)
    baseline = state.nbytes()
    assert baseline > 0

    for step in range(8):
        decoder.decode_chunk(_latent(3, seed=step + 10), state)
        assert state.nbytes() == baseline

    assert sum(state.nbytes_by_device().values()) == baseline


@torch.no_grad()
def test_release_returns_the_session_to_its_pre_stream_state(vae, decoder) -> None:
    state = decoder.new_decode_state("s")
    decoder.decode_chunk(_latent(3), state)
    assert state.started and state.nbytes() > 0

    decoder.release(state)
    assert state.nbytes() == 0
    assert not state.started
    assert state.chunks_decoded == 0

    # A released session restarts rather than resumes, so its first chunk is
    # short again -- the same semantics reset() has for AR KV.
    factor = 2 ** sum(VAE_CONFIG["temperal_downsample"])
    assert decoder.decode_chunk(_latent(3), state).shape[2] == (3 - 1) * factor + 1


# --------------------------------------------------------------------------
# Contract
# --------------------------------------------------------------------------


def test_the_vaes_execution_context_is_entered_for_every_frame(vae) -> None:
    """Decode must run inside the autoencoder's own execution context.

    Every public compute entry point on ``OmniAutoencoderKLWan`` enters
    ``_execution_context()``, which is where autocast is established for
    fp16/bf16 parameters -- and on NPU the only place it is established at all.
    Driving ``post_quant_conv`` and ``decoder`` directly bypasses it unless this
    decoder enters it itself, which would compute streamed chunks under a
    different numeric context from the whole-clip decode they are compared to.
    """
    entered: list[str] = []

    class _ContextualVAE:
        """The parts of the contract this decoder uses, plus the hook."""

        config = vae.config
        dtype = torch.float32

        def __init__(self) -> None:
            self.decoder = vae.decoder
            self.post_quant_conv = vae.post_quant_conv
            self._cached_conv_counts = vae._cached_conv_counts

        @contextmanager
        def _execution_context(self):
            entered.append("enter")
            try:
                yield
            finally:
                entered.append("exit")

    decoder = WanStreamingDecoder(_ContextualVAE())
    with torch.no_grad():
        decoder.decode_chunk(_latent(3), decoder.new_decode_state("s"))

    # One context for the call, not one per frame, and it closes.
    assert entered == ["enter", "exit"]


def test_a_plain_diffusers_autoencoder_needs_no_execution_context(decoder) -> None:
    """The hook is optional: a VAE without one still decodes."""
    with torch.no_grad():
        out = decoder.decode_chunk(_latent(1), decoder.new_decode_state("s"))
    assert out.shape[2] == 1


def test_wan_decoder_satisfies_the_protocol(decoder) -> None:
    assert isinstance(decoder, SupportsStreamingDecode)


def test_state_from_another_decoder_is_rejected(decoder) -> None:
    foreign = StreamingDecodeState(session_id="s", feat_map=[None])
    with pytest.raises(ValueError, match="does not belong to this decoder"):
        decoder.decode_chunk(_latent(1), foreign)


def test_a_non_wan_module_is_rejected() -> None:
    with pytest.raises(TypeError, match="Wan-family autoencoder"):
        WanStreamingDecoder(object())


def test_empty_and_malformed_latents_are_rejected(decoder) -> None:
    state = decoder.new_decode_state("s")
    with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"):
        decoder.decode_chunk(torch.zeros(1, 4, 4, 4), state)
    with pytest.raises(ValueError, match="at least one frame"):
        decoder.decode_chunk(torch.zeros(1, 4, 0, LATENT_H, LATENT_W), state)


def test_a_batch_is_refused_rather_than_threaded_through_one_cache(decoder) -> None:
    """One state is one session, so a batch has no correct interpretation here.

    ``AutoencoderKLWan.decode`` handles a batch by decoding each sample with
    its own cleared cache (``use_slicing``). This decoder carries exactly one
    temporal context, so a batch would mix samples into it and still return
    correctly shaped pixels.
    """
    state = decoder.new_decode_state("s")
    batched = torch.zeros(2, VAE_CONFIG["z_dim"], 3, LATENT_H, LATENT_W)
    with pytest.raises(ValueError, match="one sample per session state"):
        decoder.decode_chunk(batched, state)
    assert state.frames_decoded == 0, "a refused call must not advance the session"


def test_session_id_must_be_meaningful(decoder) -> None:
    with pytest.raises(ValueError, match="session_id"):
        decoder.new_decode_state("  ")


def test_declared_state_bytes_scales_with_area_and_dtype(vae) -> None:
    # A local decoder, not the module-scoped fixture: the measured constant
    # belongs to this test, and leaking it would change what a reordered run
    # of the default-None case sees. Passing it here is also the constructor
    # argument's intended use.
    # Measured for the shipped checkpoint: 37832 KiB of fp32 cache at 64x64.
    decoder = WanStreamingDecoder(vae, bytes_per_pixel_fp32=37832 * 1024 / (64 * 64))
    fp32 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.float32)
    bf16 = decoder.declared_state_bytes(height=480, width=832, dtype=torch.bfloat16)
    half_area = decoder.declared_state_bytes(height=240, width=832, dtype=torch.float32)
    assert bf16 == fp32 // 2
    assert half_area == fp32 // 2
    assert fp32 / 2**20 == pytest.approx(3602.0, abs=5.0)