diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md index 426ee5afd..2f86744b1 100644 --- a/docs/design/feature/realtime_ar_diffusion.md +++ b/docs/design/feature/realtime_ar_diffusion.md @@ -143,8 +143,20 @@ fraction to select additional resident sessions up to the model-declared cap. All model-owned reservations are deducted before the paged self-attention pool is allocated. -LingBot reports its persistent full-horizon image-condition tensor. DreamZero -reports the measured 603 MiB per-session Wan VAE causal-convolution state. +LingBot advances a session-owned causal Wan VAE encoder for each condition +block: the opening pixel frame is followed by four zero pixel frames per later +latent frame. It retains one current condition block, with separate committed +and in-flight encoder histories so a failed block cannot overwrite committed +context. The state reservation includes both encoder histories, derived from +the encoder's convolution input grids, and the streaming decoder cache. These +caches are released through the same close/reset hooks. Condition encoding +currently requires the unpatched, non-tiled Wan encoder path; the existing +decoder's tiling fallback remains independent. Real-weight GPU parity with a +full-horizon encode still requires validation because the pointwise posterior +projection now runs per block. + +DreamZero reports the measured 603 MiB per-session Wan VAE causal-convolution +state. ## Concurrency and routing limits diff --git a/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py index 2aa326f6d..6c87d4e63 100644 --- a/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py +++ b/examples/offline_inference/diffusion/ar_diffusion_streaming_decode.py @@ -24,10 +24,9 @@ Examples:: ``--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``. +generates three latent frames per block. The frames per block and temporal +compression are read off the checkpoint rather than assumed, so the script +can be repointed with ``--model``. """ from __future__ import annotations @@ -49,9 +48,6 @@ if TYPE_CHECKING: 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: @@ -67,8 +63,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "--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." + "JSON file holding one camera action list per AR block. " + "The per-block frame count is read off the checkpoint, not assumed." ), ) parser.add_argument("--output-dir", required=True, help="Directory for decoded frames and metadata.") @@ -154,12 +150,6 @@ def _validate_against_checkpoint( 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 diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py index 4bbea2f4f..0eb2027b1 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2.py @@ -32,7 +32,6 @@ import numpy as np _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" _MAX_SEQUENCE_LENGTH = 512 _MAX_PIXEL_AREA = 480 * 832 -_MAX_RAW_FRAMES = 117 _MAX_SOURCE_CAMERA_FRAMES = 4096 _TEMPORAL_COMPRESSION = 4 _LATENT_FRAMES_PER_BLOCK = 3 @@ -191,8 +190,6 @@ def build_request( raise ValueError("--height * --width must not exceed 480 * 832 pixels.") if args.num_frames <= 0: raise ValueError("--num-frames must be a positive integer.") - if args.num_frames > _MAX_RAW_FRAMES: - raise ValueError("--num-frames must not exceed 117 raw frames.") if args.num_frames > paths.camera_frames: raise ValueError("Camera arrays must contain at least --num-frames frames.") if (args.num_frames - 1) % _TEMPORAL_COMPRESSION: diff --git a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py index d98fe06f9..f0f1043ab 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session. This example exercises the internal session API. It is not a public HTTP or WebSocket protocol. Each JSONL event advances the world by one three-latent -frame AR block and writes the latent plus its identity metadata. The current -scoped implementation supports at most ten ticks per generation epoch because -its image condition is bounded to 117 pixel frames; reset or create a session -to start a new world. +frame AR block and writes the latent plus its identity metadata. Session-owned +causal VAE encoder state produces each condition block without retaining or +repeating a full image-condition horizon. """ from __future__ import annotations @@ -24,8 +23,6 @@ from typing import Any _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" _CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1" _FRAMES_PER_BLOCK = 3 -# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames. -_MAX_REALTIME_TICKS = 10 def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]: @@ -41,7 +38,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument( "--events", required=True, - help="JSONL file with one prompt/action event per AR block; at most 10 events.", + help="JSONL file with one prompt/action event per AR block.", ) parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.") parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.") @@ -84,12 +81,6 @@ def _load_events(path: Path) -> list[dict[str, Any]]: events.append({"event_id": event_id, "prompt": prompt, "frames": frames}) if not events: raise ValueError("events file must contain at least one event.") - if len(events) > _MAX_REALTIME_TICKS: - raise ValueError( - "events file must contain at most " - f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime " - "image-condition horizon is 117 pixel frames." - ) event_ids = [event["event_id"] for event in events] if event_ids != sorted(set(event_ids)): raise ValueError("event_id values must be unique and strictly increasing.") diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md index c38e24fd2..ab8c498fa 100644 --- a/recipes/Robbyant/LingBot-World-2.0.md +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -48,9 +48,11 @@ python examples/offline_inference/diffusion/lingbot_world_v2.py \ ``` The current causal-fast checkpoint uses four DMD sampling steps. Raw frame -counts must be `9 + 12k`, up to 117 frames. Camera files may contain a -longer bounded source trajectory (the official example contains 269 frames); -the runtime consumes only the prefix needed by the request. +counts must be `9 + 12k` for a non-negative integer `k`; for example, 129 frames +produce eleven AR blocks. Both camera files must contain the same number of +source frames, at least as many as requested and at most 4096. The runtime +consumes only the prefix needed by the request. Longer rollouts also require +sufficient device memory. ## Realtime in-process generation (deprecated) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py index 121364c07..9cb67a142 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py @@ -8,6 +8,7 @@ import math import subprocess import sys from pathlib import Path +from typing import TYPE_CHECKING import numpy as np import pytest @@ -20,7 +21,10 @@ _CAMERA_MODULE = importlib.util.module_from_spec(_SPEC) sys.modules[_SPEC.name] = _CAMERA_MODULE _SPEC.loader.exec_module(_CAMERA_MODULE) -CameraTrajectory = _CAMERA_MODULE.CameraTrajectory +if TYPE_CHECKING: + from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory +else: + CameraTrajectory = _CAMERA_MODULE.CameraTrajectory build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory @@ -138,21 +142,20 @@ def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_pa load_camera_trajectory(_trusted(tmp_path)) -def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None: - _write_trajectory( - tmp_path, - _identity_poses(269), - np.arange(269 * 4, dtype=np.float32).reshape(269, 4), - ) +@pytest.mark.parametrize("num_frames", [129, 269, 4096]) +def test_load_camera_trajectory_preserves_all_bounded_source_frames(tmp_path: Path, num_frames: int) -> None: + poses = _identity_poses(num_frames) + poses[:, 0, 3] = np.arange(num_frames) + intrinsics = np.arange(num_frames * 4, dtype=np.float32).reshape(num_frames, 4) + _write_trajectory(tmp_path, poses, intrinsics) trajectory = load_camera_trajectory(_trusted(tmp_path)) - assert trajectory.poses.shape == (117, 4, 4) - assert trajectory.intrinsics.shape == (117, 4) - np.testing.assert_array_equal( - trajectory.intrinsics.numpy(), - np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117], - ) + assert trajectory.poses.shape == (num_frames, 4, 4) + assert trajectory.intrinsics.shape == (num_frames, 4) + assert trajectory.poses.dtype == trajectory.intrinsics.dtype == torch.float32 + np.testing.assert_array_equal(trajectory.poses.numpy(), poses.astype(np.float32)) + np.testing.assert_array_equal(trajectory.intrinsics.numpy(), intrinsics) def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None: diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py index b4bf7805f..96c0b37f1 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py @@ -35,11 +35,13 @@ def _tiny_model( num_layers: int = 2, num_frames_per_block: int = 1, sliding_window_num_frames: int = 3, + attention_head_dim: int = 2, + rope_max_seq_len: int = 16, ): return module.CausalLingBotWorldTransformer3DModel( patch_size=(1, 2, 2), num_attention_heads=2, - attention_head_dim=2, + attention_head_dim=attention_head_dim, in_channels=36, out_channels=2, text_dim=6, @@ -48,7 +50,7 @@ def _tiny_model( num_layers=num_layers, cross_attn_norm=True, eps=1e-6, - rope_max_seq_len=16, + rope_max_seq_len=rope_max_seq_len, sink_size=1, num_frames_per_block=num_frames_per_block, sliding_window_num_frames=sliding_window_num_frames, @@ -121,7 +123,8 @@ def _assert_self_cache_unchanged(cache, snapshot) -> None: @pytest.mark.cpu -def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None: +@pytest.mark.parametrize("num_chunks", [4, 18]) +def test_tiny_transformer_runs_contiguous_chunks_with_explicit_cache_commit_and_camera_path(num_chunks) -> None: torch.manual_seed(7) module = attention_tests._load_module() model = _tiny_model(module).eval() @@ -135,7 +138,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera assert model.patch_embedding.in_channels == 36 assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2 - for start_frame in range(4): + for start_frame in range(num_chunks): hidden_states = torch.randn(1, 36, 1, 4, 4) timestep = torch.tensor([float(start_frame + 1)]) snapshot = _self_cache_snapshot(cache) @@ -169,7 +172,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera expected_token_end = (start_frame + 1) * 2 * 2 assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention) - assert len(outputs) == 4 + assert len(outputs) == num_chunks assert all(len(outputs) == 1 for outputs in cross_key_outputs) assert len(text_projection_outputs) == 1 @@ -991,3 +994,42 @@ def test_sp4_direct_matches_sp1_bf16(tmp_path): @hardware_test(res={"cuda": "L4"}, num_cards=4) def test_tp2_sp2_paged_matches_sp1_bf16(tmp_path): _run(tmp_path, 2, 2, "paged", torch.bfloat16, 1) + + +@pytest.mark.cpu +@pytest.mark.parametrize("start_frame", [0, 1021, 1023, 1024, 21600]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_temporal_rope_matches_absolute_positions_without_growing_cache(start_frame, dtype): + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1, attention_head_dim=128, rope_max_seq_len=1024).to(dtype) + buffers = {name: (value.data_ptr(), value.shape) for name, value in model.named_buffers()} + cosine, sine = model._rotary_embedding( + frames=3, height=2, width=3, start_frame=start_frame, dtype=dtype, device=torch.device("cpu") + ) + # Independent complex-valued reference uses absolute (time, height, width) + # coordinates, including blocks crossing the original precomputed table. + grids = torch.meshgrid( + torch.arange(start_frame, start_frame + 3, dtype=torch.float64), + torch.arange(2, dtype=torch.float64), + torch.arange(3, dtype=torch.float64), + indexing="ij", + ) + reference = [] + for positions, dim in zip(grids, (44, 42, 42), strict=True): + frequencies = 1.0 / 10000 ** (torch.arange(0, dim, 2, dtype=torch.float64) / dim) + phase = positions.flatten()[:, None] * frequencies[None, :] + reference.append(torch.polar(torch.ones_like(phase), phase)) + expected = torch.cat(reference, dim=-1) + torch.testing.assert_close(cosine, expected.real.float().to(dtype), rtol=0, atol=0) + torch.testing.assert_close(sine, expected.imag.float().to(dtype), rtol=0, atol=0) + assert {name: (value.data_ptr(), value.shape) for name, value in model.named_buffers()} == buffers + + +@pytest.mark.cpu +def test_temporal_rope_extension_keeps_spatial_limit(): + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1) + with pytest.raises(ValueError, match="Spatial RoPE"): + model._rotary_embedding( + frames=3, height=17, width=1, start_frame=16, dtype=torch.float32, device=torch.device("cpu") + ) 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 86d8e70b8..fb9e93a5d 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -198,10 +198,43 @@ class _StubCausalDecoder: return x.new_zeros(x.shape[0], 3, num_frames, x.shape[-2] * 8, x.shape[-1] * 8) +class _StubCausalEncoder(nn.Module): + """Minimal Wan cache layout; the real tiny VAE below checks numerical parity.""" + + def __init__(self): + super().__init__() + self.conv_in = nn.Conv3d(3, 16, 1) + self.down_blocks = nn.ModuleList([lingbot_pipeline.WanResample(16, mode="downsample2d") for _ in range(3)]) + self.mid_block = SimpleNamespace(resnets=[]) + self.conv_out = nn.Conv3d(16, 32, 1) + self.inputs: list[torch.Tensor] = [] + + def forward(self, video, *, feat_cache, feat_idx): + self.inputs.append(video.detach().clone()) + first = feat_cache[0] is None + tail = video[:, :, -2:].clone() + if tail.shape[2] == 1 and feat_cache[0] is not None: + tail = torch.cat((feat_cache[0][:, :, -1:], tail), dim=2) + feat_cache[0] = tail + height, width = video.shape[-2] // 8, video.shape[-1] // 8 + previous_frames = 0 if feat_cache[1] is None else feat_cache[1].shape[2] + feat_cache[1] = video.new_zeros(1, 16, min(2, previous_frames + 1), height, width) + feat_idx[0] += 2 + moments = video.new_zeros(1, 32, 1, height, width) + if first: + moments[:, :16] = 2.0 + return moments + + class _StubVAE(_FakePretrained): dtype = torch.float32 def __init__(self, *, streaming: bool = True): + self.encoder = _StubCausalEncoder() + self.quant_conv = nn.Identity() + self._cached_conv_counts = {"encoder": 2} + self._enc_feat_map = ["module-owned"] + self._enc_conv_idx = [71] if streaming: self.decoder = _StubCausalDecoder() # Recorded, not discarded: on the streaming branch this is the only @@ -209,7 +242,7 @@ class _StubVAE(_FakePretrained): # 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} + self._cached_conv_counts["decoder"] = _StubCausalDecoder.NUM_CONVS # Module-owned cache the shared VAE keeps for whole-clip decode; # streaming must never write through it. self._feat_map = ["module-owned"] @@ -498,16 +531,14 @@ 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} - # 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 + # One condition block, committed + pending encoder caches, and decoder state. + # The stub encoder retains 2 RGB frames at 16x16 and 2 feature frames at 2x2. + encoder_bytes = (2 * 3 * 16 * 16 + 2 * 16 * 2 * 2) * 4 + assert spec.model_owned_state_bytes_per_session == 960 + 2 * encoder_bytes + 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.""" + """Admission includes encoder state even when the decoder cannot stream.""" module = _load_pipeline_module() def spec_pipeline(): @@ -526,13 +557,13 @@ def test_session_admission_accounts_for_the_streaming_decoder_cache() -> None: # 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. + # Removing decoder support leaves the condition and both encoder histories. bare = spec_pipeline() - for attribute in ("decoder", "post_quant_conv", "_cached_conv_counts"): + for attribute in ("decoder", "post_quant_conv"): if hasattr(bare.vae, attribute): delattr(bare.vae, attribute) - assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 9_600 + bare.vae._cached_conv_counts.pop("decoder") + assert bare.ar_diffusion_kv_cache_spec().model_owned_state_bytes_per_session == 960 + 2 * 6_656 def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None: @@ -1200,20 +1231,19 @@ def test_request_validation_rejects_unsupported_contracts(request_batch, message _pipeline(module)._parse_request(request_batch) -def test_resource_limits_accept_exact_documented_boundaries() -> None: +def test_request_validation_accepts_frames_beyond_previous_limit() -> None: module = _load_pipeline_module() - sampling = _SamplingParams(height=480, width=832, num_frames=117, max_sequence_length=512) + sampling = _SamplingParams(height=480, width=832, num_frames=129, max_sequence_length=512) parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) - assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 117, 512) + assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 129, 512) @pytest.mark.parametrize( ("sampling", "message"), [ (_SamplingParams(height=480, width=848), "pixel area|480.*832"), - (_SamplingParams(num_frames=129), "num_frames.*117"), (_SamplingParams(max_sequence_length=511), "max_sequence_length.*512"), (_SamplingParams(max_sequence_length=513), "max_sequence_length.*512"), (_SamplingParams(max_sequence_length=512.0), "max_sequence_length.*512"), @@ -1550,9 +1580,9 @@ def test_multi_chunk_generation_uses_one_request_local_cache_and_decodes_accumul assert cache_ref() is None -def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room"): +def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room", session_id: str = "world-1"): return ARDiffusionTickRequest( - session_id="world-1", + session_id=session_id, request_id="legacy-lingbot-request", chunk_index=chunk_index, applied_event_ids=(chunk_index,), @@ -1567,6 +1597,16 @@ def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room") ).to_extra_args() +def test_realtime_rejects_pixel_output_before_condition_encoding() -> None: + pipeline = _pipeline(_load_pipeline_module()) + pipeline._ar_diffusion_kv_state = object() + sampling = _SamplingParams(output_type="np", extra_args=_tick_extra_args(chunk_index=0)) + with pytest.raises(ValueError, match="require output_type='latent'"): + pipeline(_request(sampling=sampling)) + assert pipeline._ar_sessions == {} + assert pipeline.vae.encoder.inputs == [] + + def test_typed_ticks_generate_one_global_block_and_return_standard_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1769,7 +1809,7 @@ def test_typed_tick_rejects_non_contiguous_chunk_index( pipeline(_request(sampling=sampling)) -def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon( +def test_typed_ticks_keep_bounded_encoder_state_beyond_ten_chunks( monkeypatch: pytest.MonkeyPatch, ) -> None: module = _load_pipeline_module() @@ -1777,22 +1817,229 @@ def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon( pipeline._ar_height = 16 pipeline._ar_width = 16 pipeline._ar_diffusion_kv_state = object() - pipeline._ar_sessions["world-1"] = module._LingBotARSessionState( - next_chunk_index=10, - ) + generated = [] monkeypatch.setattr( pipeline, - "encode_prompt", - lambda *args, **kwargs: pytest.fail("the realtime horizon must be validated before prompt encoding"), + "_ar_text_caches", + lambda *args, **kwargs: [SimpleNamespace()], ) - sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=10)) - with pytest.raises( - ValueError, - match=r"at most 10 ticks.*chunk_index 0 through 9.*117 pixel frames", - ): + def generate_block(**kwargs): + generated.append( + { + "condition": kwargs["condition"].clone(), + "start_frame": kwargs["start_frame"], + } + ) + return torch.zeros_like(kwargs["condition"][:, :16]) + + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + + for chunk_index in range(11): + sampling = _SamplingParams( + extra_args=_tick_extra_args(chunk_index=chunk_index), + ) pipeline(_request(sampling=sampling)) + assert pipeline.vae.encode_inputs == [] + assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == [1] + [4] * 32 + assert all(torch.count_nonzero(video) == 0 for video in pipeline.vae.encoder.inputs[1:]) + encoder_cache = pipeline._ar_sessions["world-1"].encoder_cache + assert sum(t.numel() * t.element_size() for t in encoder_cache) == pipeline._condition_encoder_cache_bytes() + assert pipeline.vae._enc_feat_map == ["module-owned"] + assert pipeline.vae._enc_conv_idx == [71] + assert pipeline._ar_sessions["world-1"].next_chunk_index == 11 + assert [item["start_frame"] for item in generated] == list(range(0, 33, 3)) + assert torch.count_nonzero(generated[0]["condition"][:, :4]) > 0 + torch.testing.assert_close( + generated[1]["condition"][:, :4], + torch.zeros_like(generated[1]["condition"][:, :4]), + ) + for item in generated[2:]: + torch.testing.assert_close(item["condition"], generated[1]["condition"]) + + +def _tiny_condition_pipeline(monkeypatch: pytest.MonkeyPatch): + from diffusers import AutoencoderKLWan + + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import retrieve_latents + + pipeline = _pipeline(_load_pipeline_module()) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(7) + pipeline.vae = AutoencoderKLWan( + base_dim=4, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ).eval() + monkeypatch.setattr(lingbot_pipeline, "retrieve_latents", retrieve_latents) + monkeypatch.setattr(pipeline, "_ar_text_caches", lambda *args, **kwargs: [SimpleNamespace()]) + return pipeline + + +@pytest.mark.parametrize("mode", ["realtime", "stepwise"]) +def test_condition_continues_real_causal_vae_beyond_initial_horizon( + monkeypatch: pytest.MonkeyPatch, + mode: str, +) -> None: + pipeline = _tiny_condition_pipeline(monkeypatch) + blocks = 12 + inputs = pipeline._parse_request(_request(sampling=_SamplingParams(num_frames=(blocks * 3 - 1) * 4 + 1))) + with torch.inference_mode(): + expected = pipeline._prepare_condition(inputs, dtype=torch.float32) + module_cache = pipeline.vae._enc_feat_map + seen = [] + + def check(condition): + index = len(seen) + torch.testing.assert_close(condition, expected[:, :, index * 3 : (index + 1) * 3], rtol=1e-6, atol=1e-6) + seen.append(condition.clone()) + + if mode == "realtime": + pipeline._ar_diffusion_kv_state = object() + + def generate_block(**kwargs): + check(kwargs["condition"]) + return torch.zeros_like(kwargs["condition"][:, :16]) + + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + session_id = "world-1" + outputs = ( + pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=index)))) + for index in range(blocks) + ) + else: + + def probe_step(**kwargs): + if kwargs["step_index"] == 0: + check(kwargs["condition"]) + return torch.ones_like(kwargs["current_latents"]) + + monkeypatch.setattr(pipeline._dmd_blocks, "probe_step", probe_step) + monkeypatch.setattr(pipeline._dmd_blocks, "commit_block_kv", lambda **kwargs: None) + state = _stepwise_state(num_frames=(blocks * 3 - 1) * 4 + 1) + session_id = state.request_id + outputs = _stepwise_chunks(pipeline, state, _FakeARState(session_id)) + + cache_sizes = [] + for _ in outputs: + state = pipeline._ar_sessions.get(session_id) + if state is not None: + for cache in (state.encoder_cache, state.pending_encoder_cache): + if cache is not None: + size = sum(t.numel() * t.element_size() for t in cache if t is not None) + assert size == pipeline._condition_encoder_cache_bytes() + assert all(t.grad_fn is None for t in cache if t is not None) + cache_sizes.append(size) + assert len(seen) == blocks and len(set(cache_sizes)) == 1 + assert pipeline.vae._enc_feat_map is module_cache and all(t is None for t in module_cache) + pipeline.close_ar_diffusion_session(session_id) + assert pipeline._ar_sessions == {} + + +def test_condition_encoder_interleaving_retry_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _tiny_condition_pipeline(monkeypatch) + pipeline._ar_diffusion_kv_state = object() + prompts = { + "a": _prompt(images=Image.new("RGB", (16, 16), (255, 0, 0))), + "b": _prompt(images=Image.new("RGB", (16, 16), (0, 255, 0))), + } + expected = {} + with torch.inference_mode(): + for name, prompt in prompts.items(): + inputs = pipeline._parse_request(_request(prompt=prompt, sampling=_SamplingParams(num_frames=45))) + expected[name] = pipeline._prepare_condition(inputs, dtype=torch.float32) + assert not torch.equal(expected["a"], expected["b"]) + active = None + failure = None + original_encoder = pipeline.vae.encoder.forward + + def encoder(*args, **kwargs): + output = original_encoder(*args, **kwargs) + if failure == "encoder": + raise RuntimeError("injected encoder failure") + return output + + def generate_block(**kwargs): + index = kwargs["start_frame"] + torch.testing.assert_close(kwargs["condition"], expected[active][:, :, index : index + 3], rtol=1e-6, atol=1e-6) + if failure == "dit": + raise RuntimeError("injected DiT failure") + return torch.zeros_like(kwargs["condition"][:, :16]) + + monkeypatch.setattr(pipeline.vae.encoder, "forward", encoder) + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + + def run(name, index): + nonlocal active + active = name + sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=index, session_id=name)) + return pipeline(_request(prompt=prompts[name], sampling=sampling)) + + run("a", 0) + for index, kind in [(1, "encoder"), (2, "dit")]: + state = pipeline._ar_sessions["a"] + committed = state.encoder_cache + saved = [t.clone() if t is not None else None for t in committed] + failure = kind + with pytest.raises(RuntimeError, match="injected"): + run("a", index) + assert state.next_chunk_index == index and state.encoder_cache is committed + assert state.pending_encoder_cache is None + for before, after in zip(saved, committed, strict=True): + if before is not None: + torch.testing.assert_close(before, after, rtol=0, atol=0) + failure = None + run("b", index - 1) + run("a", index) + + old_a = pipeline._ar_sessions["a"] + pipeline.reset_ar_diffusion_session("a") + assert old_a.encoder_cache is None and old_a.pending_encoder_cache is None + assert "a" not in pipeline._ar_sessions and "b" in pipeline._ar_sessions + run("a", 0) + old_b = pipeline._ar_sessions["b"] + pipeline.close_ar_diffusion_session("b") + assert old_b.encoder_cache is None and old_b.pending_encoder_cache is None + assert "a" in pipeline._ar_sessions and "b" not in pipeline._ar_sessions + pipeline.close_ar_diffusion_session("a") + assert pipeline._ar_sessions == {} + + +def test_stepwise_failed_chunk_releases_pending_encoder_on_close(monkeypatch: pytest.MonkeyPatch) -> None: + pipeline = _pipeline(_load_pipeline_module()) + state = _stepwise_state() + transformer = pipeline.transformer + transformer.raise_on_call = 1 + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + session = pipeline._ar_sessions[state.request_id] + assert session.encoder_cache is None and session.pending_encoder_cache is not None + with pytest.raises(RuntimeError, match="forced transformer failure"): + pipeline.denoise_step(None, states=[state]) + assert session.next_chunk_index == 0 and session.encoder_cache is None + pipeline.close_ar_diffusion_session(state.request_id) + assert session.encoder_cache is None and session.pending_encoder_cache is None + assert state.request_id not in pipeline._ar_sessions + + +@pytest.mark.parametrize("mode", ["realtime", "stepwise"]) +def test_stateful_condition_rejects_tiled_encoder(monkeypatch: pytest.MonkeyPatch, mode: str) -> None: + pipeline = _pipeline(_load_pipeline_module()) + _enable_vae_tiling(pipeline, tile_sample_min=8) + with pytest.raises(ValueError, match="does not support tiled VAE encoding"): + if mode == "realtime": + pipeline._ar_diffusion_kv_state = object() + pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=0)))) + else: + state = _stepwise_state() + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + pipeline.prepare_encode(state) + assert pipeline.vae.encoder.inputs == [] + assert all(s.encoder_cache is None and s.pending_encoder_cache is None for s in pipeline._ar_sessions.values()) + def test_request_cache_is_released_before_vae_decode() -> None: module = _load_pipeline_module() @@ -1820,25 +2067,25 @@ def test_request_cache_is_released_before_vae_decode() -> None: assert result.output.shape[2] == 21 -def test_117_frame_request_generates_ten_complete_latent_blocks() -> None: +def test_129_frame_request_generates_eleven_complete_latent_blocks() -> None: module = _load_pipeline_module() transformer = _RecordingTransformer() pipeline = _pipeline(module, transformer=transformer) trajectory = _CameraTrajectory( - poses=torch.eye(4).repeat(117, 1, 1), - intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(117, 1), + poses=torch.eye(4).repeat(129, 1, 1), + intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(129, 1), ) sampling = _SamplingParams( - num_frames=117, + num_frames=129, output_type="latent", extra_args={"_lingbot_camera_trajectory": trajectory}, ) result = pipeline(_request(sampling=sampling)) - assert result.output.shape == (1, 16, 30, 2, 2) - assert len(transformer.calls) == 50 - assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 30, 3)) + assert result.output.shape == (1, 16, 33, 2, 2) + assert len(transformer.calls) == 55 + assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 33, 3)) def test_request_cache_becomes_unreachable_after_transformer_error() -> None: @@ -2030,13 +2277,14 @@ def test_stepwise_progress_metadata_and_commit_trace() -> None: assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2) -def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: +@pytest.mark.parametrize("num_chunks", [2, 12]) +def test_stepwise_matches_tick_transformer_trace_and_latents(num_chunks) -> None: module = _load_pipeline_module() transformer = _RecordingTransformer() pipeline = _pipeline(module, transformer=transformer) pipeline._ar_height = 16 pipeline._ar_width = 16 - script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",))) + script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",))) * (num_chunks // 2) fake = _FakeARState("world-1") with pipeline.bind_ar_diffusion_state("world-1", fake): @@ -2063,7 +2311,7 @@ def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: transformer.calls.clear() pipeline.close_ar_diffusion_session("world-1") - stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17) + stepwise_state = _stepwise_state(request_id="req-1", num_frames=12 * num_chunks - 3, script=script, seed=17) stepwise_fake = _FakeARState(stepwise_state.request_id) with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake): stepwise_outputs = _run_stepwise(pipeline, stepwise_state) @@ -2072,6 +2320,15 @@ def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls] for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True): torch.testing.assert_close(output.output["payload"]["latents"], tick_latent) + for step_call, tick_call in zip(transformer.calls, tick_calls, strict=True): + torch.testing.assert_close(step_call["hidden_states"], tick_call["hidden_states"]) + torch.testing.assert_close(step_call["camera_hidden_states"], tick_call["camera_hidden_states"]) + assert stepwise_state.extra["condition"].shape[2] == 3 + assert pipeline.vae.encode_inputs == [] + assert [video.shape[2] for video in pipeline.vae.encoder.inputs] == ([1] + [4] * (num_chunks * 3 - 1)) * 2 + assert stepwise_fake.commits == ["main"] * num_chunks + assert len(stepwise_outputs) == num_chunks and stepwise_outputs[-1].finished + assert pipeline.vae.decode_inputs == [] def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None: @@ -2387,6 +2644,7 @@ def test_streaming_decode_is_reported_to_the_profiler(monkeypatch) -> None: durations = pipeline.stage_durations assert stage in durations, f"streamed decode is unprofiled; stages seen: {sorted(durations)}" assert durations[stage] > 0.0 + assert durations[f"{type(pipeline).__name__}._prepare_condition_chunk"] > 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) @@ -2398,11 +2656,10 @@ def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> Non 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) + outputs = [ + pipeline._decode_chunk_to_pixels(torch.zeros(1, 16, 3, 2, 2), output_type="np", session_id="tiled") + for _ in range(2) + ] # 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 @@ -2410,7 +2667,7 @@ def test_streaming_decode_falls_back_when_the_vae_would_tile(monkeypatch) -> Non 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"}] + assert outputs == ["frames-1", "frames-2"] def test_streaming_decode_survives_tiling_enabled_below_the_tile_threshold(monkeypatch) -> None: diff --git a/tests/examples/offline_inference/test_lingbot_world_v2.py b/tests/examples/offline_inference/test_lingbot_world_v2.py index 41d2e73b7..d83288853 100644 --- a/tests/examples/offline_inference/test_lingbot_world_v2.py +++ b/tests/examples/offline_inference/test_lingbot_world_v2.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import json from pathlib import Path import pytest +from examples.offline_inference.diffusion import ar_diffusion_streaming_decode as streaming_example from examples.offline_inference.diffusion import lingbot_world_v2 as example pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -22,3 +24,25 @@ def test_offline_ulysses_argument(tmp_path): args.ulysses_degree = 0 with pytest.raises(ValueError, match="--ulysses-degree"): example.build_omni_kwargs(args, paths) + + +@pytest.mark.parametrize(("num_blocks", "num_frames"), [(10, 117), (11, 129)]) +def test_streaming_action_script_accepts_long_rollouts(tmp_path, num_blocks, num_frames): + action_script = tmp_path / "actions.json" + action_script.write_text(json.dumps([[["w"], [], ["a"]] for _ in range(num_blocks)])) + + script = streaming_example._load_action_script(action_script) + + assert len(script) == num_blocks + assert ( + streaming_example._validate_against_checkpoint(script, frames_per_block=3, temporal_compression=4) == num_frames + ) + + +def test_streaming_action_script_rejects_wrong_frame_count_in_eleventh_block(tmp_path): + action_script = tmp_path / "actions.json" + action_script.write_text(json.dumps([[[], [], []] for _ in range(10)] + [[[], []]])) + script = streaming_example._load_action_script(action_script) + + with pytest.raises(ValueError, match="block 10 holds 2 frame action lists"): + streaming_example._validate_against_checkpoint(script, frames_per_block=3, temporal_compression=4) diff --git a/vllm_omni/diffusion/models/lingbot_world/camera.py b/vllm_omni/diffusion/models/lingbot_world/camera.py index 0e78e53e6..504d850c7 100644 --- a/vllm_omni/diffusion/models/lingbot_world/camera.py +++ b/vllm_omni/diffusion/models/lingbot_world/camera.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Camera trajectory loading and geometry for LingBot World v2 conditioning.""" from __future__ import annotations @@ -19,7 +19,6 @@ import torch _REFERENCE_HEIGHT = 480 _REFERENCE_WIDTH = 832 -_MAX_ACTION_FRAMES = 117 _MAX_SOURCE_ACTION_FRAMES = 4096 _MAX_NPY_HEADER_BYTES = 10_000 _DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW @@ -207,9 +206,7 @@ def _load_bounded_npy( array = np.load(snapshot_file, allow_pickle=False) except (EOFError, TypeError, ValueError): raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None - # Official LingBot trajectories may be longer than one request. Only - # materialize the bounded prefix that any supported request can consume. - result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32) + result: np.ndarray = np.asarray(array, dtype=np.float32) return result diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index 6e32844fc..ccaa677e0 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -7,7 +7,7 @@ from __future__ import annotations import math import os from collections.abc import Callable, Iterable, Iterator, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, ClassVar, cast @@ -15,6 +15,7 @@ import numpy as np import PIL.Image import torch import torch.nn.functional as F +from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T, WanAttentionBlock, WanResample, WanResidualBlock from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel @@ -97,7 +98,6 @@ LINGBOT_DMD_TIMESTEPS = (1000, 750, 500, 250) _CAMERA_SPATIAL_FOLD = 8 _MAX_PIXEL_AREA = 480 * 832 _MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096 -_MAX_RAW_FRAMES = 117 _MAX_SEQUENCE_LENGTH = 512 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" @@ -150,12 +150,13 @@ class _LingBotRequestInputs: @dataclass class _LingBotARSessionState: - """Small model-owned state; attention tensors remain runner-owned.""" + """Model-owned state; attention tensors remain runner-owned.""" next_chunk_index: int = 0 prompt: str | None = None generator_state: torch.Tensor | None = None - image_condition: torch.Tensor | None = None + encoder_cache: list[torch.Tensor | None] | None = None + pending_encoder_cache: list[torch.Tensor | None] | None = None camera_tail: CameraTrajectory | None = None camera_pitch: float = 0.0 @@ -577,6 +578,7 @@ class LingBotWorldCausalDMDPipeline( self.setup_diffusion_pipeline_profiler( profiler_targets=[ "vae.encode", + "_prepare_condition_chunk", "vae.decode", # Streaming decode never reaches vae.decode, so it needs its # own stage or a streamed rollout reports no decode time at @@ -609,11 +611,11 @@ class LingBotWorldCausalDMDPipeline( recent_window_frames = total_window_frames - sink_frames if recent_window_frames <= 0: raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.") - horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 + condition_latent_frames = int(self.transformer.config.num_frames_per_block) condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels) condition_bytes_per_session = ( condition_channels - * horizon_latent_frames + * condition_latent_frames * latent_height * latent_width * torch.empty((), dtype=self.transformer.dtype).element_size() @@ -635,10 +637,34 @@ class LingBotWorldCausalDMDPipeline( ), ), model_owned_state_bytes_per_session=( - condition_bytes_per_session + self._streaming_decode_bytes_per_session() + condition_bytes_per_session + + 2 * self._condition_encoder_cache_bytes() # Committed and in-flight histories. + + self._streaming_decode_bytes_per_session() ), ) + def _condition_encoder_cache_bytes(self) -> int: + """Bound one Wan encoder history from its per-convolution input grids.""" + if getattr(self.vae.config, "patch_size", None) is not None: + raise ValueError("Stateful LingBot conditioning requires an unpatched Wan VAE.") + encoder = self.vae.encoder + height, width = self._ar_height, self._ar_width + elements = CACHE_T * encoder.conv_in.in_channels * height * width + for layer in encoder.down_blocks: + if isinstance(layer, WanResidualBlock): + elements += CACHE_T * (layer.conv1.in_channels + layer.conv2.in_channels) * height * width + elif isinstance(layer, WanResample) and layer.mode in ("downsample2d", "downsample3d"): + height, width = height // 2, width // 2 + if layer.mode == "downsample3d": + # Temporal downsampling retains one frame after spatial downsampling. + elements += layer.time_conv.in_channels * height * width + elif not isinstance(layer, WanAttentionBlock): + raise ValueError("Stateful LingBot conditioning requires the Wan encoder cache layout.") + for layer in encoder.mid_block.resnets: + elements += CACHE_T * (layer.conv1.in_channels + layer.conv2.in_channels) * height * width + elements += CACHE_T * encoder.conv_out.in_channels * height * width + return elements * torch.empty((), dtype=self.vae.dtype).element_size() + def _streaming_decode_bytes_per_session(self) -> int: """Resident streaming-decoder bytes one session can hold, for admission. @@ -678,12 +704,18 @@ class LingBotWorldCausalDMDPipeline( finally: self._ar_diffusion_kv_state = None + def _release_condition_encoder_state(self, session_id: str) -> None: + state = self._ar_sessions.pop(session_id, None) + if state is not None: + state.encoder_cache = None + state.pending_encoder_cache = None + def reset_ar_diffusion_session(self, session_id: str) -> None: - self._ar_sessions.pop(session_id, None) + self._release_condition_encoder_state(session_id) 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_condition_encoder_state(session_id) self._release_streaming_decode_state(session_id) def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs: @@ -810,8 +842,6 @@ class LingBotWorldCausalDMDPipeline( num_frames = getattr(sampling, "num_frames", None) if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: raise ValueError(f"num_frames must be a positive integer, got {num_frames!r}.") - if num_frames > _MAX_RAW_FRAMES: - raise ValueError(f"num_frames must not exceed {_MAX_RAW_FRAMES}.") temporal_factor = self.vae_scale_factor_temporal if (num_frames - 1) % temporal_factor: raise ValueError( @@ -914,16 +944,16 @@ class LingBotWorldCausalDMDPipeline( raise RuntimeError( f"vae.encode returned an incompatible image latent shape: got {tuple(latent_condition.shape)}." ) + return self._condition_from_latents(latent_condition, first_frame=True, dtype=dtype) + + def _condition_from_latents( + self, latent_condition: torch.Tensor, *, first_frame: bool, dtype: torch.dtype + ) -> torch.Tensor: latent_mean, latent_std = self._vae_latent_stats(latent_condition) latent_condition = (latent_condition - latent_mean) / latent_std - temporal_mask = latent_condition.new_zeros( - 1, - self.vae_scale_factor_temporal, - inputs.num_latent_frames, - latent_condition.shape[-2], - latent_condition.shape[-1], - ) - temporal_mask[:, :, 0] = 1 + temporal_mask = latent_condition.new_zeros(1, self.vae_scale_factor_temporal, *latent_condition.shape[2:]) + if first_frame: + temporal_mask[:, :, 0] = 1 condition = torch.cat((temporal_mask, latent_condition), dim=1).to(dtype=dtype) if condition.shape[1] != 20: raise RuntimeError( @@ -932,6 +962,78 @@ class LingBotWorldCausalDMDPipeline( ) return condition + @torch.no_grad() + def _prepare_condition_chunk( + self, + inputs: _LingBotRequestInputs, + *, + start_frame: int, + encoder_cache: list[torch.Tensor | None] | None, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: + """Continue the causal encoder without modifying its committed history.""" + if getattr(self.vae.config, "patch_size", None) is not None or self.vae_scale_factor_temporal != 4: + raise ValueError("Stateful LingBot conditioning requires an unpatched Wan VAE with temporal factor 4.") + if getattr(self.vae, "use_tiling", False) and ( + inputs.height > self.vae.tile_sample_min_height or inputs.width > self.vae.tile_sample_min_width + ): + raise ValueError("Stateful LingBot conditioning does not support tiled VAE encoding.") + if (encoder_cache is None) != (start_frame == 0): + raise ValueError("LingBot condition encoder history must match the current temporal position.") + + count = int(self.vae._cached_conv_counts["encoder"]) + # Wan replaces cache entries with new tensors; a separate list keeps the + # committed history intact until this AR block succeeds. + pending_cache = [None] * count if encoder_cache is None else list(encoder_cache) + if len(pending_cache) != count: + raise ValueError("LingBot condition encoder cache has an incompatible number of slots.") + block_frames = int(self.transformer.config.num_frames_per_block) + zeros = torch.zeros( + 1, + 3, + self.vae_scale_factor_temporal, + inputs.height, + inputs.width, + device=self.device, + dtype=self.vae.dtype, + ) + image = ( + self._prepare_image_tensor(inputs.image, height=inputs.height, width=inputs.width) + .unsqueeze(2) + .to(dtype=self.vae.dtype) + if start_frame == 0 + else None + ) + context = getattr(self.vae, "_execution_context", None) + encoded_frames = [] + with context() if callable(context) else nullcontext(): + for index in range(block_frames): + # Match Wan's encode loop: one opening pixel frame, then four + # fresh zero pixel frames per subsequent latent frame. + pixel_frames = image if start_frame == 0 and index == 0 else zeros + encoded_frames.append(self.vae.encoder(pixel_frames, feat_cache=pending_cache, feat_idx=[0])) + # The posterior projection is pointwise in time. GPU bitwise parity + # with the whole-clip projection still requires the real-weight gate. + moments = self.vae.quant_conv(torch.cat(encoded_frames, dim=2)) + latent_condition = moments.chunk(2, dim=1)[0] + expected_shape = ( + 1, + self.transformer.config.out_channels, + block_frames, + inputs.height // self.vae_scale_factor_spatial, + inputs.width // self.vae_scale_factor_spatial, + ) + if latent_condition.shape != expected_shape: + raise RuntimeError( + f"Wan condition encoder returned {tuple(latent_condition.shape)}, expected {expected_shape}." + ) + condition = self._condition_from_latents( + latent_condition, + first_frame=start_frame == 0, + dtype=dtype, + ) + return condition, pending_cache + def _prepare_camera( self, inputs: _LingBotRequestInputs, @@ -1150,6 +1252,11 @@ class LingBotWorldCausalDMDPipeline( raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.") if tick is None and self._ar_diffusion_kv_state is not None: raise ValueError("LingBot ARDiffusionEngine requests must carry ar_diffusion_tick.") + if tick is not None and inputs.output_type != "latent": + raise ValueError( + "LingBot realtime ticks currently require output_type='latent'; " + "stateful streaming VAE decode is a separate integration step." + ) session_state: _LingBotARSessionState | None = None if tick is not None: if tick.prompt is not None and " ".join(tick.prompt.split()) != " ".join(inputs.prompt.split()): @@ -1166,16 +1273,6 @@ class LingBotWorldCausalDMDPipeline( "fixed cache geometry " f"{self._ar_height}x{self._ar_width}." ) - horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 - max_realtime_ticks = horizon_latent_frames // block_frames - if tick.chunk_index >= max_realtime_ticks: - raise ValueError( - "LingBot realtime generation currently supports at most " - f"{max_realtime_ticks} ticks per generation epoch " - f"(chunk_index 0 through {max_realtime_ticks - 1}) because " - f"the image-condition horizon is {_MAX_RAW_FRAMES} pixel " - "frames; reset or create a session to start a new world." - ) session_state = self._ar_sessions.setdefault( tick.session_id, _LingBotARSessionState(), @@ -1200,26 +1297,12 @@ class LingBotWorldCausalDMDPipeline( condition = self._prepare_condition(inputs, dtype=dtype) else: assert session_state is not None - if session_state.image_condition is None: - horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 - session_state.image_condition = self._prepare_condition( - replace( - inputs, - num_frames=_MAX_RAW_FRAMES, - num_latent_frames=horizon_latent_frames, - ), - dtype=dtype, - ) - block_frames = int(self.transformer.config.num_frames_per_block) - condition_start = tick.chunk_index * block_frames - condition_stop = condition_start + block_frames - if condition_stop > session_state.image_condition.shape[2]: - raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") - condition = session_state.image_condition[ - :, - :, - condition_start:condition_stop, - ] + condition, pending_encoder_cache = self._prepare_condition_chunk( + inputs, + start_frame=tick.chunk_index * block_frames, + encoder_cache=session_state.encoder_cache, + dtype=dtype, + ) camera_pitch: float | None = None if inputs.camera_actions is not None: assert session_state is not None @@ -1299,6 +1382,7 @@ class LingBotWorldCausalDMDPipeline( assert session_state is not None session_state.prompt = inputs.prompt session_state.generator_state = inputs.generator.get_state() + session_state.encoder_cache = pending_encoder_cache session_state.camera_tail = camera_tail if camera_pitch is not None: session_state.camera_pitch = camera_pitch @@ -1307,11 +1391,6 @@ class LingBotWorldCausalDMDPipeline( # Phase 4: either expose model-space latents or invert the checkpoint's # latent normalization and decode to pixel-space video. if tick is not None: - if inputs.output_type != "latent": - raise ValueError( - "LingBot realtime ticks currently require output_type='latent'; " - "stateful streaming VAE decode is a separate integration step." - ) output = { "payload": {"latents": generated_latents}, "metadata": {"ar_diffusion": ARDiffusionChunkMetadata.from_tick(tick).to_dict()}, @@ -1460,12 +1539,6 @@ class LingBotWorldCausalDMDPipeline( self._cached_video_processor = processor return processor - def _horizon_latent_frames(self) -> int: - return (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 - - def _max_realtime_chunks(self) -> int: - return self._horizon_latent_frames() // int(self.transformer.config.num_frames_per_block) - def _require_bound_ar_state(self) -> None: if self._ar_diffusion_kv_state is None: raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.") @@ -1492,13 +1565,6 @@ class LingBotWorldCausalDMDPipeline( ) block_frames = int(self.transformer.config.num_frames_per_block) total_chunks = inputs.num_latent_frames // block_frames - max_chunks = self._max_realtime_chunks() - if total_chunks > max_chunks: - raise ValueError( - "LingBot step execution currently supports at most " - f"{max_chunks} chunks because the image-condition horizon is " - f"{_MAX_RAW_FRAMES} pixel frames." - ) if inputs.camera_action_script is not None and len(inputs.camera_action_script) != total_chunks: raise ValueError( "camera_action_script must contain one action list per generated chunk; " @@ -1510,14 +1576,6 @@ class LingBotWorldCausalDMDPipeline( max_sequence_length=inputs.max_sequence_length, dtype=dtype, ) - image_condition = self._prepare_condition( - replace( - inputs, - num_frames=_MAX_RAW_FRAMES, - num_latent_frames=self._horizon_latent_frames(), - ), - dtype=dtype, - ) camera_trajectory_cache = None camera_embedding_cache = None if inputs.camera_trajectory is not None: @@ -1546,7 +1604,6 @@ class LingBotWorldCausalDMDPipeline( state.chunk_num_steps = len(LINGBOT_DMD_TIMESTEPS) state.extra = { "inputs": inputs, - "image_condition": image_condition, "schedule": _build_shifted_flow_schedule(flow_shift=inputs.flow_shift), "dtype": dtype, "block_frames": block_frames, @@ -1557,7 +1614,12 @@ class LingBotWorldCausalDMDPipeline( "camera_embedding_cache": camera_embedding_cache, } self._ar_text_caches(prompt_embeds, invalidate=False) - self._prepare_next_chunk(state) + self._ar_sessions[state.request_id] = _LingBotARSessionState() + try: + self._prepare_next_chunk(state) + except Exception: + self._release_condition_encoder_state(state.request_id) + raise return state def _prepare_next_chunk(self, state: StepRequestState) -> None: @@ -1566,10 +1628,15 @@ class LingBotWorldCausalDMDPipeline( block_frames = int(extra["block_frames"]) start_frame = state.chunk_index * block_frames stop_frame = start_frame + block_frames - image_condition = extra["image_condition"] - if stop_frame > image_condition.shape[2]: - raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") - condition = image_condition[:, :, start_frame:stop_frame] + session_state = self._ar_sessions[state.request_id] + if session_state.next_chunk_index != state.chunk_index: + raise ValueError("LingBot condition chunks must be contiguous within the session.") + condition, pending_encoder_cache = self._prepare_condition_chunk( + inputs, + start_frame=start_frame, + encoder_cache=session_state.encoder_cache, + dtype=extra["dtype"], + ) previous = extra.get("camera_tail") if extra.get("camera_action_script") is not None: chunk_actions = extra["camera_action_script"][state.chunk_index] @@ -1639,6 +1706,7 @@ class LingBotWorldCausalDMDPipeline( ) state.step_in_chunk = 0 state.step_index = 0 + session_state.pending_encoder_cache = pending_encoder_cache def denoise_step( self, @@ -1702,6 +1770,9 @@ class LingBotWorldCausalDMDPipeline( latents = state.latents if latents is None: raise RuntimeError("LingBot step execution requires latents before post_decode.") + session_state = self._ar_sessions[state.request_id] + if session_state.pending_encoder_cache is None: + raise RuntimeError("LingBot condition encoder has no prepared chunk to commit.") self._dmd_blocks.commit_block_kv( latents=latents, condition=extra["condition"], @@ -1737,10 +1808,15 @@ class LingBotWorldCausalDMDPipeline( ).to_dict() }, } + session_state.encoder_cache = session_state.pending_encoder_cache + session_state.pending_encoder_cache = None + session_state.next_chunk_index += 1 state.chunk_index += 1 finished = state.request_denoise_completed if not finished: self._prepare_next_chunk(state) + else: + self._release_condition_encoder_state(state.request_id) return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py index 32931db0c..c7033cee8 100644 --- a/vllm_omni/diffusion/models/lingbot_world/transformer.py +++ b/vllm_omni/diffusion/models/lingbot_world/transformer.py @@ -749,7 +749,7 @@ def _sinusoidal_embedding(dim: int, timestep: torch.Tensor) -> torch.Tensor: return torch.cat((phase.cos(), phase.sin()), dim=1) -def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]: +def _rope_axis(max_seq_len: int, dim: int, *, start: int = 0) -> tuple[torch.Tensor, torch.Tensor]: if dim == 0: empty = torch.empty(max_seq_len, 0, dtype=torch.float32) return empty, empty.clone() @@ -759,7 +759,7 @@ def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]: 10000, torch.arange(0, dim, 2, dtype=torch.float64) / dim, ) - phase = torch.outer(torch.arange(max_seq_len, dtype=torch.float64), frequencies) + phase = torch.outer(torch.arange(start, start + max_seq_len, dtype=torch.float64), frequencies) return phase.cos().float(), phase.sin().float() @@ -1028,23 +1028,29 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): dtype: torch.dtype, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: - if start_frame + frames > self.config.rope_max_seq_len: - raise ValueError("Temporal RoPE positions exceed rope_max_seq_len.") if height > self.config.rope_max_seq_len or width > self.config.rope_max_seq_len: raise ValueError("Spatial RoPE positions exceed rope_max_seq_len.") + if start_frame + frames <= self._rope_temporal_cosine.shape[0]: + temporal_cosine = self._rope_temporal_cosine[start_frame : start_frame + frames] + temporal_sine = self._rope_temporal_sine[start_frame : start_frame + frames] + else: + # Keep absolute positions and checkpoint frequencies, computing only + # this block instead of growing a table for the whole session. + temporal_dim = self._rope_temporal_cosine.shape[1] * 2 + temporal_cosine, temporal_sine = _rope_axis(frames, temporal_dim, start=start_frame) + temporal_cosine = temporal_cosine.to(self._rope_temporal_cosine) + temporal_sine = temporal_sine.to(self._rope_temporal_sine) def expand_axis(table: torch.Tensor, axis: str) -> torch.Tensor: if axis == "temporal": - return ( - table[start_frame : start_frame + frames].view(frames, 1, 1, -1).expand(frames, height, width, -1) - ) + return table.view(frames, 1, 1, -1).expand(frames, height, width, -1) if axis == "height": return table[:height].view(1, height, 1, -1).expand(frames, height, width, -1) return table[:width].view(1, 1, width, -1).expand(frames, height, width, -1) cosine = torch.cat( ( - expand_axis(self._rope_temporal_cosine, "temporal"), + expand_axis(temporal_cosine, "temporal"), expand_axis(self._rope_height_cosine, "height"), expand_axis(self._rope_width_cosine, "width"), ), @@ -1052,7 +1058,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): ) sine = torch.cat( ( - expand_axis(self._rope_temporal_sine, "temporal"), + expand_axis(temporal_sine, "temporal"), expand_axis(self._rope_height_sine, "height"), expand_axis(self._rope_width_sine, "width"), ),