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

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

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: