[Lingbot World] Bound image condition storage and extend temporal RoPE (#6838)

Signed-off-by: wtz2333 <2955110911@qq.com>

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
 

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
 
 

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:

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.")

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)
 

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:

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")
+        )

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:

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)

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
 
 

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,

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"),
             ),