From 311e714d34be1c80714833abc754635aa3c1b59b Mon Sep 17 00:00:00 2001 From: linzhenpl07 <60651557+linzhenpl07@users.noreply.github.com> Date: Sun, 13 Sep 2026 18:17:11 +0800 Subject: [PATCH] [AR-Diffusion] Add session-owned streaming VAE decode (#6533) Signed-off-by: linzhenpl07 Co-authored-by: Claude Sonnet 5 --- .../ar_diffusion_streaming_decode.py | 349 ++++++++++++++++ .../ar_diffusion/test_streaming_decode.py | 361 +++++++++++++++++ .../test_pipeline_lingbot_world.py | 379 +++++++++++++++++- .../models/lingbot_world/pipeline.py | 175 +++++++- .../ar_diffusion/streaming_decode.py | 272 +++++++++++++ 5 files changed, 1520 insertions(+), 16 deletions(-) create mode 100644 examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py create mode 100644 tests/diffusion/ar_diffusion/test_streaming_decode.py create mode 100644 vllm_omni/experimental/ar_diffusion/streaming_decode.py diff --git a/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py new file mode 100644 index 000000000..2aa326f6d --- /dev/null +++ b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py @@ -0,0 +1,349 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stream LingBot-World 2.0 pixels from one stepwise request. + +One ``AsyncOmni.generate()`` drives the whole rollout: each AR block leaves +``post_decode()`` as one video chunk, decoded through that session's own +temporal VAE cache rather than as an isolated clip. That is what this script +shows end to end -- chunk *N + 1* continues chunk *N*, so the streamed timeline +carries the same frame count as an offline whole-clip decode of the same +latents instead of losing each block's opening frame to a restarted decoder. +Only the session's first latent frame expands to a single raw frame; on the +default checkpoint, whose block is three latent frames at a temporal factor of +four, that is 9 frames for the opening chunk and 12 for every later one. + +The decode lives inside the pipeline, keyed by ``request_id``, so nothing here +owns a VAE: this file is a client of the served path, and the same chunks reach +a WebSocket client of ``/v1/realtime/video`` unchanged. + +Examples:: + + python examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py \ + --image scene.png --prompt "a lit hallway" \ + --action-script actions.json --output-dir /tmp/lingbot-stream + +``--action-script`` is a JSON list with one camera action list per AR block, +for example ``[[["w"], ["w"], ["w"]], [["a"], [], []]]`` for a checkpoint that +generates three latent frames per block. The frames per block, the temporal +compression and the number of blocks that fit the condition horizon are read +off the checkpoint rather than assumed, so the script can be repointed with +``--model``. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import os +import time +from collections.abc import Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + # Only for type hints -- the real imports stay deferred inside run() so + # --help and the offline validation helpers don't require a CUDA-enabled + # vLLM install. + import numpy as np + +_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" +# Mirrors _MAX_RAW_FRAMES in vllm_omni/diffusion/models/lingbot_world/pipeline.py: +# the image-condition horizon is the pipeline's constant, not the checkpoint's. +_MAX_RAW_FRAMES = 117 + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Stream decoded LingBot-World 2.0 video chunks from a single stepwise request.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.") + parser.add_argument("--image", required=True, help="Initial RGB image.") + parser.add_argument("--prompt", required=True, help="Scene prompt for the rollout.") + parser.add_argument( + "--action-script", + required=True, + help=( + "JSON file holding one camera action list per AR block. Both the per-block frame count " + "and the number of blocks that fit are read off the checkpoint, not assumed." + ), + ) + parser.add_argument("--output-dir", required=True, help="Directory for decoded frames and metadata.") + parser.add_argument("--request-id", default="lingbot-world-stream", help="Request id; also the session id.") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--gpu-memory-fraction", type=float, default=0.6) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--enforce-eager", action="store_true") + return parser.parse_args(argv) + + +def _load_action_script(path: Path) -> list[list[list[str]]]: + """Parse the per-block camera action script, without assuming its geometry. + + The shape is the one ``sampling_params.extra_args["camera_action_script"]`` + takes: one entry per AR block, each entry one per-frame action list. How + many frames a block holds is a property of the checkpoint, so it is checked + in :func:`_validate_against_checkpoint` once the configs have been read. + """ + try: + value = json.loads(path.read_text()) + except json.JSONDecodeError as exc: + raise ValueError(f"--action-script is not valid JSON: {exc.msg}.") from None + if not isinstance(value, list) or not value: + raise ValueError("--action-script must be a non-empty JSON list of per-block action lists.") + script: list[list[list[str]]] = [] + for block_index, block in enumerate(value): + if not isinstance(block, list) or not block: + raise ValueError(f"--action-script block {block_index} must be a non-empty list of frame action lists.") + frames: list[list[str]] = [] + for frame in block: + if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame): + raise ValueError(f"--action-script block {block_index} frames must contain only action strings.") + frames.append(list(frame)) + script.append(frames) + return script + + +def _model_config_json(model: str, relative_path: str) -> dict[str, Any]: + """Read one config file out of a local checkpoint or a Hub repo.""" + if os.path.isdir(model): + return json.loads((Path(model) / relative_path).read_text()) + from huggingface_hub import hf_hub_download + + return json.loads(Path(hf_hub_download(repo_id=model, filename=relative_path)).read_text()) + + +def _checkpoint_geometry(model: str) -> tuple[int, int]: + """Read ``(frames_per_block, temporal_compression)`` off the checkpoint. + + The pipeline takes both from the model -- + ``transformer.config.num_frames_per_block`` and + ``vae.config.scale_factor_temporal`` -- so hardcoding 3 and 4 here would + silently build a wrong action script and an inconsistent ``num_frames`` the + moment ``--model`` points at a checkpoint with different geometry. + """ + transformer_config = _model_config_json(model, "transformer/config.json") + vae_config = _model_config_json(model, "vae/config.json") + if "num_frames_per_block" not in transformer_config: + raise ValueError( + f"{model}'s transformer/config.json declares no num_frames_per_block, " + "so this checkpoint does not generate AR blocks and cannot be streamed this way." + ) + frames_per_block = int(transformer_config["num_frames_per_block"]) + temporal_compression = int(vae_config.get("scale_factor_temporal", 4)) + if frames_per_block <= 0 or temporal_compression <= 0: + raise ValueError( + "the checkpoint reports a non-positive block geometry " + f"(num_frames_per_block={frames_per_block}, scale_factor_temporal={temporal_compression})." + ) + return frames_per_block, temporal_compression + + +def _validate_against_checkpoint( + script: list[list[list[str]]], *, frames_per_block: int, temporal_compression: int +) -> int: + """Check the script against the checkpoint and return the frame count to request.""" + for block_index, block in enumerate(script): + if len(block) != frames_per_block: + raise ValueError( + f"--action-script block {block_index} holds {len(block)} frame action lists, but this " + f"checkpoint generates {frames_per_block} latent frames per AR block." + ) + max_chunks = ((_MAX_RAW_FRAMES - 1) // temporal_compression + 1) // frames_per_block + if len(script) > max_chunks: + raise ValueError( + f"--action-script holds {len(script)} blocks, but this checkpoint's " + f"{_MAX_RAW_FRAMES}-frame image-condition horizon fits at most {max_chunks}." + ) + return (len(script) * frames_per_block - 1) * temporal_compression + 1 + + +def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]: + image = Path(args.image).expanduser().resolve() + action_script = Path(args.action_script).expanduser().resolve() + output_dir = Path(args.output_dir).expanduser().resolve() + if not image.is_file(): + raise ValueError("--image must point to an existing file.") + if not action_script.is_file(): + raise ValueError("--action-script must point to an existing JSON file.") + if not args.prompt.strip(): + raise ValueError("--prompt must contain non-whitespace text.") + if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16: + raise ValueError("--height and --width must be positive multiples of 16.") + if args.tensor_parallel_size <= 0: + raise ValueError("--tensor-parallel-size must be positive.") + if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1: + raise ValueError("--gpu-memory-fraction must be in (0, 1].") + if not args.request_id.strip(): + raise ValueError("--request-id must contain non-whitespace text.") + return image, action_script, output_dir + + +def _as_frames(value: Any) -> np.ndarray: + """Normalize one streamed chunk to ``[T, H, W, 3]`` uint8. + + ``output_type="np"`` reaches the client as a float array in ``[0, 1]``, + with or without the leading batch axis depending on how the formatter + unwrapped it; a nested list is the ``"pil"`` shape. Anything else means + the chunk is not pixels -- most likely the request still asked for + latents -- and is worth failing on rather than saving. + """ + import numpy as np + import PIL.Image + import torch + + if isinstance(value, list): + if len(value) == 1 and isinstance(value[0], list): + value = value[0] + if value and all(isinstance(frame, PIL.Image.Image) for frame in value): + return np.stack([np.asarray(frame.convert("RGB")) for frame in value]) + if isinstance(value, torch.Tensor): + value = value.detach().float().cpu().numpy() + if not isinstance(value, np.ndarray): + raise RuntimeError(f"Expected decoded pixels from each AR block, got {type(value).__name__}.") + if value.ndim == 5 and value.shape[0] == 1: + value = value[0] + if value.ndim != 4 or value.shape[-1] != 3: + raise RuntimeError(f"Expected one [T, H, W, 3] pixel chunk per AR block, got shape {tuple(value.shape)}.") + if value.dtype != np.uint8: + value = (np.clip(value, 0.0, 1.0) * 255.0).round().astype(np.uint8) + return value + + +def _chunk_metadata(output: Any) -> dict[str, Any]: + multimodal = getattr(output, "multimodal_output", None) or {} + metadata = multimodal.get("metadata") if isinstance(multimodal, dict) else None + ar_diffusion = metadata.get("ar_diffusion") if isinstance(metadata, dict) else None + if not isinstance(ar_diffusion, dict): + raise RuntimeError("Streamed chunk is missing its ar_diffusion metadata envelope.") + return dict(ar_diffusion) + + +async def run(argv: Sequence[str] | None = None) -> Path: + args = parse_args(argv) + image, action_script_path, output_dir = _validate_args(args) + camera_action_script = _load_action_script(action_script_path) + output_dir.mkdir(parents=True, exist_ok=True) + + # Imports stay below pure input validation so --help and helper tests do not + # require a CUDA-enabled vLLM installation. + import numpy as np + + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + frames_per_block, temporal_compression = _checkpoint_geometry(args.model) + num_chunks = len(camera_action_script) + num_frames = _validate_against_checkpoint( + camera_action_script, + frames_per_block=frames_per_block, + temporal_compression=temporal_compression, + ) + + engine = AsyncOmni( + model=args.model, + engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine", + enforce_eager=args.enforce_eager, + tensor_parallel_size=args.tensor_parallel_size, + max_num_seqs=1, + # One request, many chunks: the stepwise contract from #6844. Streaming + # output is what turns each post_decode() into a delivered chunk. + step_execution=True, + diffusion_streaming_output=True, + model_config={ + "ar_diffusion_height": args.height, + "ar_diffusion_width": args.width, + "ar_diffusion_kv_config": { + "gpu_memory_fraction": args.gpu_memory_fraction, + "warmup_cudagraph": True, + }, + }, + ) + sampling = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_frames=num_frames, + num_inference_steps=4, + max_sequence_length=512, + seed=args.seed, + # Pixels, not latents: this is what routes the chunk through the + # session's streaming VAE decode instead of handing back a tensor. + output_type="np", + extra_args={"flow_shift": 5.0, "camera_action_script": camera_action_script}, + ) + prompt = {"prompt": args.prompt, "multi_modal_data": {"image": str(image)}} + + measurements: list[dict[str, Any]] = [] + t_start = time.perf_counter() + first_frame_at: float | None = None + chunk_started = t_start + try: + async for output in engine.generate(prompt, sampling, request_id=args.request_id): + images = getattr(output, "images", None) + if not images: + continue + if len(images) != 1: + raise RuntimeError("Expected exactly one decoded chunk per AR block.") + frames = _as_frames(images[0]) + now = time.perf_counter() + if first_frame_at is None: + first_frame_at = now - t_start + chunk_index = len(measurements) + metadata = _chunk_metadata(output) + + np.save(output_dir / f"chunk_{chunk_index:03d}.npy", frames) + (output_dir / f"chunk_{chunk_index:03d}.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n" + ) + measurements.append( + { + "chunk_index": chunk_index, + "chunk_latency_seconds": now - chunk_started, + "frames": int(frames.shape[0]), + "frame_shape": list(frames.shape), + "metadata": metadata, + } + ) + chunk_started = now + print(json.dumps(measurements[-1], sort_keys=True), flush=True) + finally: + engine.shutdown() + + delivered_frames = sum(int(chunk["frames"]) for chunk in measurements) + summary_path = output_dir / "summary.json" + summary_path.write_text( + json.dumps( + { + "request_id": args.request_id, + "chunks_delivered": len(measurements), + "chunks_expected": num_chunks, + "frames_delivered": delivered_frames, + "frames_expected": num_frames, + # The point of session-owned decode: a streamed rollout carries + # the same frame timeline as an offline decode of the same + # latents. Per-chunk decode would deliver fewer, because every + # block would re-expand its own opening frame. + "matches_offline_frame_timeline": delivered_frames == num_frames, + "time_to_first_frame_seconds": first_frame_at, + "total_seconds": time.perf_counter() - t_start, + "chunks": measurements, + }, + indent=2, + sort_keys=True, + ) + + "\n" + ) + return summary_path + + +def main(argv: Sequence[str] | None = None) -> Path: + return asyncio.run(run(argv)) + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/ar_diffusion/test_streaming_decode.py b/tests/diffusion/ar_diffusion/test_streaming_decode.py new file mode 100644 index 000000000..52889bfbd --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_streaming_decode.py @@ -0,0 +1,361 @@ +# 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) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index a58eee935..86d8e70b8 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -170,10 +170,49 @@ class _RecordingTransformer(nn.Module): return {name for name, _ in self.loaded_weights} +class _StubCausalDecoder: + """Stand-in for the Wan causal decoder's cache protocol. + + Reproduces only what streaming depends on: each call walks ``feat_idx`` + across one cache slot per causal convolution, and the session's opening + frame (``first_chunk``) expands to a single raw frame while every later + frame expands by the temporal factor. + """ + + NUM_CONVS = 3 + + def __init__(self) -> None: + self.first_chunk_flags: list[bool] = [] + # Identity of the cache each call wrote through, so an interleaved run + # can assert that a session never advanced another session's context. + self.cache_ids: list[int] = [] + + def __call__(self, x: torch.Tensor, *, feat_cache, feat_idx, first_chunk: bool = False): + self.first_chunk_flags.append(bool(first_chunk)) + self.cache_ids.append(id(feat_cache)) + for _ in range(self.NUM_CONVS): + index = feat_idx[0] + feat_cache[index] = x.detach().clone() + feat_idx[0] = index + 1 + num_frames = 1 if first_chunk else 4 + return x.new_zeros(x.shape[0], 3, num_frames, x.shape[-2] * 8, x.shape[-1] * 8) + + class _StubVAE(_FakePretrained): dtype = torch.float32 - def __init__(self): + def __init__(self, *, streaming: bool = True): + if streaming: + self.decoder = _StubCausalDecoder() + # Recorded, not discarded: on the streaming branch this is the only + # place the tensor handed to the decoder can be observed, and it is + # what carries the checkpoint's latent rescale. + self.post_quant_inputs: list[torch.Tensor] = [] + self.post_quant_conv = self._record_post_quant + self._cached_conv_counts = {"decoder": _StubCausalDecoder.NUM_CONVS, "encoder": 2} + # Module-owned cache the shared VAE keeps for whole-clip decode; + # streaming must never write through it. + self._feat_map = ["module-owned"] self.config = SimpleNamespace( z_dim=16, scale_factor_temporal=4, @@ -185,6 +224,10 @@ class _StubVAE(_FakePretrained): self.decode_inputs: list[torch.Tensor] = [] self.on_decode = None + def _record_post_quant(self, latent: torch.Tensor) -> torch.Tensor: + self.post_quant_inputs.append(latent.detach().clone()) + return latent + def encode(self, video: torch.Tensor): self.encode_inputs.append(video.detach().clone()) latent_frames = (video.shape[2] - 1) // 4 + 1 @@ -455,7 +498,41 @@ def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None: assert spec.sink_frames == 3 assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)] assert spec.cross_attention_lengths == {"text": 512} - assert spec.model_owned_state_bytes_per_session == 9_600 + # The image condition (9,600 B) plus the streaming decoder's own per-session + # temporal cache. The cache dwarfs the condition at any real resolution, and + # the runner subtracts this field from the KV budget before sizing its pools, + # so leaving the cache out of it over-commits device memory. + assert spec.model_owned_state_bytes_per_session == 9_600 + 2_421_248 + + +def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None: + """The declaration scales with output area, and is the condition alone when + the VAE cannot stream.""" + module = _load_pipeline_module() + + def spec_pipeline(): + pipeline = _pipeline(module) + # The spec reads the constructed attention's head count; the head count + # plays no part in model-owned bytes, so any value will do here. + pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1) + return pipeline + + small = spec_pipeline().ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session + + wide = spec_pipeline() + wide._ar_width = wide._ar_width * 2 + widened = wide.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session + # Twice the area: the decoder cache doubles, so the total grows by far more + # than the image condition alone could account for. + assert widened - small > 2_000_000 + + # A VAE without the causal-cache contract cannot stream, so there is no + # decoder cache to declare and the figure falls back to the condition. + bare = spec_pipeline() + for attribute in ("decoder", "post_quant_conv", "_cached_conv_counts"): + if hasattr(bare.vae, attribute): + delattr(bare.vae, attribute) + assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 9_600 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None: @@ -1852,6 +1929,27 @@ def _run_stepwise(pipeline, state): return outputs +def _stepwise_chunks(pipeline, state, ar_state): + """Yield one chunk at a time, binding AR state around each invocation. + + The runner binds runner-owned KV for the duration of one stepwise + invocation and releases it again, so a generator that suspends at each + chunk boundary is what lets two requests be driven the way the scheduler + drives them at session_capacity > 1: A-chunk0, B-chunk0, A-chunk1, ... + """ + with pipeline.bind_ar_diffusion_state(state.request_id, ar_state): + pipeline.prepare_encode(state) + while not state.request_denoise_completed: + output = None + with pipeline.bind_ar_diffusion_state(state.request_id, ar_state): + noise = pipeline.denoise_step(None, states=[state]) + pipeline.step_scheduler(state, noise) + if state.chunk_denoise_completed: + output = pipeline.post_decode(state) + if output is not None: + yield output + + def test_pipeline_declares_step_execution_support() -> None: module = _load_pipeline_module() pipeline = _pipeline(module) @@ -2039,12 +2137,11 @@ def test_stepwise_emits_latents_without_touching_the_vae() -> None: assert pipeline.vae.decode_inputs == [] -def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None: - """A non-latent request must stream pixels under the primary "video" key.""" +def _capture_video_processor(monkeypatch) -> list[tuple[tuple[int, ...], str]]: + """Record what reaches postprocess_video, and name each chunk's output.""" import diffusers.video_processor as video_processor_module - module = _load_pipeline_module() - processed = [] + processed: list[tuple[tuple[int, ...], str]] = [] class VideoProcessor: def __init__(self, *, vae_scale_factor): @@ -2055,9 +2152,21 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non return f"frames-{len(processed)}" monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor) + return processed + + +def _streaming_pipeline(module): pipeline = _pipeline(module, transformer=_RecordingTransformer()) pipeline._ar_height = 16 pipeline._ar_width = 16 + return pipeline + + +def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None: + """A non-latent request must stream pixels under the primary "video" key.""" + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) state = _stepwise_state(num_frames=21) state.sampling.output_type = "np" fake = _FakeARState(state.request_id) @@ -2065,14 +2174,268 @@ def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> Non with pipeline.bind_ar_diffusion_state(state.request_id, fake): outputs = _run_stepwise(pipeline, state) - # One decode per AR block, each seeing only that block's three latent frames. - assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2 assert [output_type for _, output_type in processed] == ["np", "np"] assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}] # Identity metadata survives the switch to pixel output. assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1] +def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> None: + """Chunk N + 1 must continue chunk N instead of restarting the decoder. + + Only the session's opening latent frame collapses to a single raw frame, + so the streamed timeline carries the same frame count an offline decode of + the same latents would: 9 frames, then 12. + """ + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + _run_stepwise(pipeline, state) + + assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)] + # One opening frame for the session, not one per chunk. + assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5 + # The whole-clip decode path is not involved, and the module-owned cache + # the shared VAE keeps for it is untouched. + assert pipeline.vae.decode_inputs == [] + assert pipeline.vae._feat_map == ["module-owned"] + + +def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None: + """Decoder state is keyed by request id and released with the AR session.""" + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(request_id="req-stream", num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + _run_stepwise(pipeline, state) + + decode_state = pipeline._streaming_decode_states["req-stream"] + assert decode_state.session_id == "req-stream" + assert decode_state.frames_decoded == 6 + assert decode_state.nbytes() > 0 + + pipeline.close_ar_diffusion_session("req-stream") + assert "req-stream" not in pipeline._streaming_decode_states + # Release drops the cache itself, not just the pipeline's reference to it. + assert decode_state.nbytes() == 0 + assert decode_state.frames_decoded == 0 + + +def test_streaming_decode_state_is_dropped_on_session_reset(monkeypatch) -> None: + """A reset session must not resume the temporal context it just abandoned.""" + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(request_id="req-reset", num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + _run_stepwise(pipeline, state) + + pipeline.reset_ar_diffusion_session("req-reset") + assert pipeline._streaming_decode_states == {} + # Releasing a session that never decoded is a no-op, not an error. + pipeline.reset_ar_diffusion_session("req-reset") + + +def test_streaming_decode_keeps_interleaved_sessions_isolated(monkeypatch) -> None: + """Two sessions ticking alternately must not share one temporal cache. + + Running one rollout to completion before the other starts would only show + that the dict is keyed by request id. The failure worth catching is chunk N + of one session advancing the other's cache mid-rollout, which needs the + interleaving the scheduler actually produces. + """ + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + states = [] + for request_id in ("req-a", "req-b"): + state = _stepwise_state(request_id=request_id, num_frames=21) + state.sampling.output_type = "np" + states.append(state) + runs = [_stepwise_chunks(pipeline, state, _FakeARState(state.request_id)) for state in states] + + chunks: list[tuple[str, int]] = [] + while runs: + for run in list(runs): + try: + output = next(run) + except StopIteration: + runs.remove(run) + continue + metadata = output.output["metadata"]["ar_diffusion"] + chunks.append((metadata["session_id"], metadata["chunk_index"])) + + # A-chunk0, B-chunk0, A-chunk1, B-chunk1: the sessions really are interleaved. + assert chunks == [("req-a", 0), ("req-b", 0), ("req-a", 1), ("req-b", 1)] + + caches = pipeline._streaming_decode_states + assert set(caches) == {"req-a", "req-b"} + assert caches["req-a"].feat_map is not caches["req-b"].feat_map + # Every frame of a session was decoded through that session's own cache, + # three frames per chunk, and each session opened exactly one rollout. + cache_a, cache_b = (id(caches["req-a"].feat_map), id(caches["req-b"].feat_map)) + assert pipeline.vae.decoder.cache_ids == [cache_a] * 3 + [cache_b] * 3 + [cache_a] * 3 + [cache_b] * 3 + assert pipeline.vae.decoder.first_chunk_flags == ( + [True, False, False] + [True, False, False] + [False] * 3 + [False] * 3 + ) + assert [shape for shape, _ in processed] == [ + (1, 3, 9, 16, 16), + (1, 3, 9, 16, 16), + (1, 3, 12, 16, 16), + (1, 3, 12, 16, 16), + ] + + +def test_streaming_decode_failure_drops_the_half_advanced_cache(monkeypatch) -> None: + """A failed chunk must not leave a cache a later chunk would continue from.""" + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(request_id="req-boom", num_frames=21) + state.sampling.output_type = "np" + + def explode(*args, **kwargs): + raise RuntimeError("decoder blew up") + + pipeline.vae.decoder = explode + + with pytest.raises(RuntimeError, match="decoder blew up"): + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + _run_stepwise(pipeline, state) + + assert pipeline._streaming_decode_states == {} + + +def _enable_vae_tiling(pipeline, *, tile_sample_min: int) -> None: + """Configure the stub VAE the way registry.py configures the real one.""" + pipeline.vae.use_tiling = True + pipeline.vae.spatial_compression_ratio = 8 + pipeline.vae.tile_sample_min_height = tile_sample_min + pipeline.vae.tile_sample_min_width = tile_sample_min + + +def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None: + """The streaming branch must invert the checkpoint's latent statistics. + + Nothing downstream can catch a miss here: the pixels stay finite and + correctly shaped, only wrongly scaled. Request mode and the fallback path + observe the rescale through ``vae.decode``'s recorded input, which streaming + never reaches, so this asserts on what ``post_quant_conv`` received. + """ + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + while not state.chunk_denoise_completed: + pipeline.step_scheduler(state, pipeline.denoise_step(None, states=[state])) + model_space = state.latents.clone() + pipeline.post_decode(state) + + shape = (1, -1, 1, 1, 1) + latent_mean = torch.as_tensor(pipeline.vae.config.latents_mean, dtype=model_space.dtype).view(*shape) + latent_std = torch.as_tensor(pipeline.vae.config.latents_std, dtype=model_space.dtype).view(*shape) + expected = (model_space * latent_std + latent_mean).to(dtype=pipeline.vae.dtype) + + recorded = torch.cat(pipeline.vae.post_quant_inputs, dim=2) + assert recorded.shape == expected.shape, "one recorded frame per latent frame in the chunk" + torch.testing.assert_close(recorded, expected, rtol=0, atol=0) + # And the rescale is not a no-op on this fixture, so the assertion bites. + assert not torch.equal(expected, model_space.to(dtype=pipeline.vae.dtype)) + + +def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None: + """A streamed rollout must report decode time, not zero. + + The profiler wraps ``vae.decode``, which streaming never calls, so without + its own stage the one stage this path changes would be invisible -- and a + perf regression in the streaming decoder would be too. + """ + module = _load_pipeline_module() + _capture_video_processor(monkeypatch) + pipeline = _pipeline( + module, + transformer=_RecordingTransformer(), + od_config=_od_config( + enable_diffusion_pipeline_profiler=True, + ), + ) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + outputs = _run_stepwise(pipeline, state) + + stage = f"{type(pipeline).__name__}._streaming_decode_chunk" + durations = pipeline.stage_durations + assert stage in durations, f"streamed decode is unprofiled; stages seen: {sorted(durations)}" + assert durations[stage] > 0.0 + # The stage reaches the consumer on every streamed chunk, not just at the end. + assert all(stage in (output.stage_durations or {}) for output in outputs) + + +def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> None: + """Tiled decode drives its own cache, so a tiling shape keeps per-chunk decode.""" + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + # A latent of 2x2 exceeds a one-latent-cell tile, so vae.decode would tile. + _enable_vae_tiling(pipeline, tile_sample_min=8) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + outputs = _run_stepwise(pipeline, state) + + # One whole-clip decode per AR block, each seeing only that block's frames. + assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2 + assert pipeline.vae.decoder.first_chunk_flags == [] + assert pipeline._streaming_decode_states == {} + # Nine frames per chunk instead of 9 then 12: the frame loss being removed. + assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 9, 16, 16)] + assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}] + + +def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monkeypatch) -> None: + """``use_tiling`` alone does not tile, so such a shape must still stream. + + ``AutoencoderKLWan._decode`` tiles only when the latent exceeds a tile. A + gate on the flag alone would give up streaming -- and a frame per block -- + on a configuration that would never have tiled. + """ + module = _load_pipeline_module() + processed = _capture_video_processor(monkeypatch) + pipeline = _streaming_pipeline(module) + # A 2x2 latent is within a tile this size, so vae.decode would not tile. + _enable_vae_tiling(pipeline, tile_sample_min=256) + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + _run_stepwise(pipeline, state) + + assert pipeline.vae.decode_inputs == [] + assert pipeline.vae.decoder.first_chunk_flags == [True] + [False] * 5 + assert [shape for shape, _ in processed] == [(1, 3, 9, 16, 16), (1, 3, 12, 16, 16)] + + def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None: module = _load_pipeline_module() resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry( diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index d2dfde3d4..6e32844fc 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -18,6 +18,7 @@ import torch.nn.functional as F from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel +from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -62,11 +63,26 @@ from vllm_omni.experimental.ar_diffusion.capability import ( ARDiffusionKVBranchSpec, ARDiffusionKVCacheSpec, ) +from vllm_omni.experimental.ar_diffusion.streaming_decode import ( + StreamingDecodeState, + WanStreamingDecoder, +) from vllm_omni.experimental.ar_diffusion.tick_protocol import ( ARDiffusionChunkMetadata, ARDiffusionTickRequest, ) +logger = init_logger(__name__) + +# Resident streaming-decoder cache per session, as bytes per output pixel in +# fp32: measured once on the reference checkpoint at 37832 KiB for a 64x64 +# output. Every cached tensor is [B, C, <= CACHE_T, H/8, W/8] per causal +# convolution, so the total scales with height * width and never with session +# length; tests/diffusion/ar_diffusion/test_streaming_decode.py pins that +# scaling, and at 832x480 in bf16 this yields the 1801 MiB the decoder was +# measured to hold. +_STREAMING_DECODE_BYTES_PER_PIXEL_FP32 = 37832 * 1024 / (64 * 64) + if TYPE_CHECKING: from diffusers.video_processor import VideoProcessor from tqdm.std import tqdm as TqdmProgressBar @@ -552,10 +568,20 @@ class LingBotWorldCausalDMDPipeline( self._ar_width = int(model_config.get("ar_diffusion_width", 832)) self._ar_diffusion_kv_state: ARDiffusionKVState | None = None self._ar_sessions: dict[str, _LingBotARSessionState] = {} + # One temporal decoder cache per stepwise session, keyed the way the + # runner keys AR sessions (session_id == request_id) so both are + # released together rather than through two independent lifecycles. + self._streaming_decode_states: dict[str, StreamingDecodeState] = {} + self._cached_streaming_decoder: WanStreamingDecoder | None = None + self._streaming_decode_unsupported = False self.setup_diffusion_pipeline_profiler( profiler_targets=[ "vae.encode", "vae.decode", + # Streaming decode never reaches vae.decode, so it needs its + # own stage or a streamed rollout reports no decode time at + # all -- the one stage this path changes. + "_streaming_decode_chunk", "_generate_block", "text_encoder.forward", "tokenizer.forward", @@ -608,7 +634,32 @@ class LingBotWorldCausalDMDPipeline( _MAX_SEQUENCE_LENGTH, ), ), - model_owned_state_bytes_per_session=condition_bytes_per_session, + model_owned_state_bytes_per_session=( + condition_bytes_per_session + self._streaming_decode_bytes_per_session() + ), + ) + + def _streaming_decode_bytes_per_session(self) -> int: + """Resident streaming-decoder bytes one session can hold, for admission. + + ``model_owned_state_bytes_per_session`` is what the runner subtracts + from the KV budget before sizing the block pools, and this cache is + three orders of magnitude larger than the image condition that used to + be the only entry: leaving it out lets the profiler report a + configuration as fitting and then OOM partway through the second + session's rollout. + + Declared for every session rather than only for pixel-output ones: the + field is documented as a conservative upper bound, it is fixed at load + time, and whether a request asks for pixels is not known then. + """ + decoder = self._streaming_decoder() + if decoder is None: + return 0 + return decoder.declared_state_bytes( + height=self._ar_height, + width=self._ar_width, + dtype=self.vae.dtype, ) @contextmanager @@ -629,9 +680,11 @@ class LingBotWorldCausalDMDPipeline( def reset_ar_diffusion_session(self, session_id: str) -> None: self._ar_sessions.pop(session_id, None) + self._release_streaming_decode_state(session_id) def close_ar_diffusion_session(self, session_id: str) -> None: self._ar_sessions.pop(session_id, None) + self._release_streaming_decode_state(session_id) def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs: if req.num_reqs != 1 or len(req.prompts) != 1: @@ -1279,15 +1332,90 @@ class LingBotWorldCausalDMDPipeline( stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) + def _vae_would_tile(self, latents: torch.Tensor) -> bool: + """Whether ``vae.decode`` would take its tiled path for this shape. + + Tiled decode splits a frame spatially and drives the module's own + ``feat_cache`` across those tiles, so a per-session cache threaded + through it is a separate design and is not attempted here. The test + mirrors ``AutoencoderKLWan._decode``: the flag alone does not tile, the + latent also has to exceed a tile, so a sub-threshold resolution with + tiling enabled still streams. A VAE that sets the flag but reports no + tile size is treated as tiling, because the safe answer when the + threshold is unknown is the one that keeps ``vae.decode`` in charge. + """ + if not bool(getattr(self.vae, "use_tiling", False)): + return False + compression = int(getattr(self.vae, "spatial_compression_ratio", None) or self.vae_scale_factor_spatial) + tile_latent_min_height = int(getattr(self.vae, "tile_sample_min_height", None) or 0) // compression + tile_latent_min_width = int(getattr(self.vae, "tile_sample_min_width", None) or 0) // compression + return bool(latents.shape[-2] > tile_latent_min_height or latents.shape[-1] > tile_latent_min_width) + + def _streaming_decoder(self) -> WanStreamingDecoder | None: + """The session-owned streaming decoder, or ``None`` on a VAE that cannot stream. + + Only the decoder's own capability is decided here, and the decision is + cached including the negative. Whether a particular *shape* would have + been tiled is a per-call question and lives in :meth:`_vae_would_tile`. + """ + if self._streaming_decode_unsupported: + return None + if self._cached_streaming_decoder is not None: + return self._cached_streaming_decoder + try: + decoder = WanStreamingDecoder( + self.vae, + bytes_per_pixel_fp32=_STREAMING_DECODE_BYTES_PER_PIXEL_FP32, + ) + # Materialize the cache geometry now: a decoder that cannot report + # it would otherwise fail on the first chunk instead. + _ = decoder.num_causal_convs + except TypeError as exc: + logger.warning( + "LingBot streaming decode is unavailable (%s); AR blocks will be decoded independently, " + "which drops each block's opening frame and restarts temporal context at every boundary.", + exc, + ) + self._streaming_decode_unsupported = True + return None + self._cached_streaming_decoder = decoder + return decoder + + def _streaming_decode_state(self, decoder: WanStreamingDecoder, session_id: str) -> StreamingDecodeState: + state = self._streaming_decode_states.get(session_id) + if state is None: + state = decoder.new_decode_state(session_id) + self._streaming_decode_states[session_id] = state + return state + + def _release_streaming_decode_state(self, session_id: str) -> None: + """Drop one session's decoder cache. Safe for a session that never decoded. + + The state is released directly rather than through + ``WanStreamingDecoder.release``, which is the same one-line call: the + decoder's method exists for a caller holding only the + ``SupportsStreamingDecode`` protocol, and this pipeline holds the state. + """ + state = self._streaming_decode_states.pop(session_id, None) + if state is not None: + state.release() + def _decode_chunk_to_pixels( - self, latents: torch.Tensor, *, output_type: str + self, latents: torch.Tensor, *, output_type: str, session_id: str | None = None ) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]: """Decode one AR block so streaming consumers receive pixels, not latents. - Blocks are decoded independently, which keeps the shared VAE stateless - and lets a rollout be served without a per-session temporal cache. That - is the same shape Helios streams today; cross-block ``feat_cache`` - continuity is a separate decision and is not made here. + With a ``session_id`` the block is decoded through that session's own + temporal cache, so chunk *N + 1* continues chunk *N*: the causal decoder + expands the session's opening latent frame to one raw frame and every + later one to the full temporal factor, which is the frame timeline an + offline whole-clip decode of the same latents produces. The cache lives + on the session rather than on the shared VAE module, so concurrent + sessions cannot overwrite one another's temporal context. + + Without a ``session_id`` -- or on a VAE that cannot stream -- blocks are + decoded independently: every block re-expands its own opening frame, so + the rollout loses frames and continuity at each block boundary. The registered post-process hook returns a ``{"payload", "metadata"}`` envelope untouched, so the pixel conversion it would do for a bare @@ -1295,9 +1423,34 @@ class LingBotWorldCausalDMDPipeline( """ latent_mean, latent_std = self._vae_latent_stats(latents) vae_latents = (latents * latent_std + latent_mean).to(dtype=self.vae.dtype) - video = self.vae.decode(vae_latents, return_dict=False)[0] + decoder = None + if session_id is not None and not self._vae_would_tile(vae_latents): + decoder = self._streaming_decoder() + if decoder is None: + video = self.vae.decode(vae_latents, return_dict=False)[0] + else: + video = self._streaming_decode_chunk(decoder, vae_latents, cast(str, session_id)) return self._video_processor().postprocess_video(video, output_type=output_type) + def _streaming_decode_chunk( + self, decoder: WanStreamingDecoder, vae_latents: torch.Tensor, session_id: str + ) -> torch.Tensor: + """Decode one chunk through ``session_id``'s temporal cache. + + A separate method so the profiler has a stage to wrap: streaming never + calls ``vae.decode``, and without this a profiled rollout would report + no decode time for the one stage this path changes. + """ + state = self._streaming_decode_state(decoder, session_id) + try: + return decoder.decode_chunk(vae_latents, state) + except Exception: + # A half-advanced cache cannot be resumed: the session's next chunk + # would continue from a temporal context that was never completed, + # so drop it here rather than at close. + self._release_streaming_decode_state(session_id) + raise + def _video_processor(self) -> VideoProcessor: processor = getattr(self, "_cached_video_processor", None) if processor is None: @@ -1566,7 +1719,13 @@ class LingBotWorldCausalDMDPipeline( if inputs.output_type == "latent": payload: dict[str, Any] = {"latents": latents} else: - payload = {"video": self._decode_chunk_to_pixels(latents, output_type=inputs.output_type)} + payload = { + "video": self._decode_chunk_to_pixels( + latents, + output_type=inputs.output_type, + session_id=state.request_id, + ) + } output = { "payload": payload, "metadata": { diff --git a/vllm_omni/experimental/ar_diffusion/streaming_decode.py b/vllm_omni/experimental/ar_diffusion/streaming_decode.py new file mode 100644 index 000000000..e4316c83c --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/streaming_decode.py @@ -0,0 +1,272 @@ +# 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)