[Feature][LingBot World]stepwise execution (#6844)

Signed-off-by: BruceLoveDecimal <liuqihao970610@gmail.com>
Signed-off-by: liuqihao <liuqihao970610@gmail.com>

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -13,6 +14,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
     LingBotCameraControlReducer,
     integrate_lingbot_camera_actions,
     parse_lingbot_camera_action_frames,
+    parse_lingbot_camera_action_script,
 )
 from vllm_omni.experimental.ar_diffusion.session import (
     ARDiffusionSession,
@@ -292,3 +294,8 @@ def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
         assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
 
     asyncio.run(exercise())
+
+
+def test_parse_camera_action_script_rejects_wrong_chunk_width() -> None:
+    with pytest.raises(ValueError, match="exactly 3 per-latent-frame"):
+        parse_lingbot_camera_action_script([[["w"], ["w"]]], frames_per_chunk=3)

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.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
 
 from __future__ import annotations
 
@@ -15,8 +15,10 @@ from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
 from PIL import Image
 from torch import nn
 
+import vllm_omni.diffusion.models.lingbot_world.dmd_block as lingbot_dmd_block
 import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
 from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
+from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
 from vllm_omni.diffusion.models.lingbot_world.actions import (
     integrate_lingbot_camera_actions,
 )
@@ -24,6 +26,8 @@ from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as
 from vllm_omni.diffusion.models.lingbot_world.camera import (
     build_plucker_embedding as _real_build_plucker_embedding,
 )
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.experimental.ar_diffusion.tick_protocol import (
     ARDiffusionControlInput,
     ARDiffusionTickRequest,
@@ -253,6 +257,8 @@ class _SamplingParams:
         if extra_args is not None:
             self.extra_args.update(extra_args)
         self.latents = None
+        self.guidance_scale = None
+        self.guidance_scale_2 = None
 
 
 class _RequestBatch:
@@ -345,7 +351,6 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
         "FlowUniPCMultistepScheduler": _FakeScheduler,
         "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
         "get_local_device": lambda: torch.device("cpu"),
-        "set_forward_context_denoise_step_idx": lambda index: None,
         "prefetch_subfolders": prefetch_subfolders,
         "from_pretrained_with_prefetch": from_pretrained_with_prefetch,
         "load_transformer_config": load_transformer_config,
@@ -358,6 +363,9 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
     }
     for name, value in replacements.items():
         monkeypatch.setattr(lingbot_pipeline, name, value)
+    # The DMD block math lives in its own module; stub the symbols it imports.
+    monkeypatch.setattr(lingbot_dmd_block, "set_forward_context_denoise_step_idx", lambda index: None)
+    monkeypatch.setattr(lingbot_dmd_block, "randn_tensor", _diffusers_randn_tensor)
     monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
 
 
@@ -774,6 +782,7 @@ def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() ->
         return torch.zeros(shape, device=device, dtype=dtype)
 
     module.randn_tensor = randn
+    lingbot_dmd_block.randn_tensor = randn
     result = pipeline(_request())
 
     assert requested_noise_dtypes == [torch.float32] * 4
@@ -1139,7 +1148,7 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
     module = _load_pipeline_module()
     transformer = _RecordingTransformer()
     pipeline = _pipeline(module, transformer=transformer)
-    module.randn_tensor = lambda shape, **kwargs: torch.full(
+    module.randn_tensor = lingbot_dmd_block.randn_tensor = lambda shape, **kwargs: torch.full(
         shape,
         -99.0,
         device=kwargs["device"],
@@ -1745,6 +1754,288 @@ def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
     assert not hasattr(pipeline, "transformer_cache")
 
 
+class _FakeARState:
+    def __init__(self, session_id: str = "req-1") -> None:
+        self.session_id = session_id
+        self.commits: list[str] = []
+
+    def get_kv_caches(self, branch, *, seq_len, commit_current):
+        del branch, seq_len, commit_current
+        return [SimpleNamespace()]
+
+    def commit_paged_context(self, branch):
+        self.commits.append(branch)
+
+    def clear_cross_attention(self):
+        return None
+
+    def is_cross_attention_populated(self, branch, name):
+        del branch, name
+        return True
+
+    def get_cross_attention_kv(self, branch, name):
+        del branch, name
+        zeros = torch.zeros(1, 1, 2, 4)
+        return [{"k": zeros, "v": zeros} for _ in range(2)]
+
+
+def _empty_action_script(num_chunks: int):
+    return tuple(((), (), ()) for _ in range(num_chunks))
+
+
+def _stepwise_state(
+    *,
+    request_id: str = "req-1",
+    num_frames: int = 21,
+    script=None,
+    seed: int = 17,
+):
+    num_chunks = ((num_frames - 1) // 4 + 1) // 3
+    sampling = _SamplingParams(num_frames=num_frames, seed=seed)
+    sampling.extra_args["_lingbot_camera_trajectory"] = None
+    sampling.extra_args["_lingbot_camera_actions"] = None
+    sampling.extra_args["_lingbot_camera_action_script"] = (
+        script if script is not None else _empty_action_script(num_chunks)
+    )
+    return StepRequestState(
+        request_id=request_id,
+        sampling=sampling,
+        prompt=_prompt(),
+    )
+
+
+def _run_stepwise(pipeline, state):
+    outputs = []
+    pipeline.prepare_encode(state)
+    while not state.request_denoise_completed:
+        noise = pipeline.denoise_step(None, states=[state])
+        pipeline.step_scheduler(state, noise)
+        if state.chunk_denoise_completed:
+            outputs.append(pipeline.post_decode(state))
+    return outputs
+
+
+def test_pipeline_declares_step_execution_support() -> None:
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+
+    assert pipeline.supports_step_execution is True
+    assert isinstance(pipeline, SupportsStepExecution)
+    assert supports_step_execution(pipeline) is True
+
+
+def test_preprocess_materializes_camera_action_script_without_action_path() -> None:
+    module = _load_pipeline_module()
+    sampling = _SamplingParams(include_action=False)
+    sampling.extra_args["camera_action_script"] = [[["w"], ["w"], ["w"]], [["a"], [], []]]
+    request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=sampling, request_id="req-1")
+
+    result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+    assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
+    assert result.sampling_params.extra_args["_lingbot_camera_actions"] is None
+    assert result.sampling_params.extra_args["_lingbot_camera_action_script"] == (
+        (("w",), ("w",), ("w",)),
+        (("a",), (), ()),
+    )
+
+
+@pytest.mark.parametrize(
+    "extra_args",
+    [
+        pytest.param({}, id="request_mode"),
+        pytest.param(_tick_extra_args(chunk_index=0), id="tick_mode"),
+    ],
+)
+def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
+    """A script only steers step execution, so request mode must not drop it silently."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module)
+    sampling = _SamplingParams(
+        include_action=False,
+        extra_args={
+            **extra_args,
+            "_lingbot_camera_trajectory": None,
+            "_lingbot_camera_action_script": _empty_action_script(1),
+        },
+    )
+
+    with pytest.raises(ValueError, match="camera_action_script is read only by LingBot step execution"):
+        pipeline(_request(sampling=sampling))
+
+
+def test_stepwise_progress_metadata_and_commit_trace() -> None:
+    module = _load_pipeline_module()
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert pipeline._ar_diffusion_kv_state is None
+    assert len(outputs) == 2
+    assert [output.chunk_index for output in outputs] == [0, 1]
+    assert all(output.total_chunks == 2 for output in outputs)
+    assert outputs[-1].finished is True
+    assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
+    assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
+    assert fake.commits == ["main", "main"]
+    for chunk_index, output in enumerate(outputs):
+        metadata = output.output["metadata"]["ar_diffusion"]
+        assert metadata == {
+            "session_id": "req-1",
+            "request_id": "req-1",
+            "chunk_index": chunk_index,
+            "applied_event_ids": [],
+        }
+        assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
+
+
+def test_stepwise_matches_tick_transformer_trace_and_latents() -> 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",)))
+    fake = _FakeARState("world-1")
+
+    with pipeline.bind_ar_diffusion_state("world-1", fake):
+        tick_latents = []
+        for chunk_index, frames in enumerate(script):
+            tick = ARDiffusionTickRequest(
+                session_id="world-1",
+                request_id="legacy-lingbot-request",
+                chunk_index=chunk_index,
+                controls=(
+                    ARDiffusionControlInput(
+                        track="camera",
+                        schema="lingbot.camera_actions.v1",
+                        data={"mode": "frames", "frames": [list(frame) for frame in frames]},
+                    ),
+                ),
+            )
+            sampling = _SamplingParams(extra_args=tick.to_extra_args(), seed=17)
+            sampling.extra_args["_lingbot_camera_trajectory"] = None
+            sampling.extra_args["_lingbot_camera_actions"] = frames
+            result = pipeline(_request(sampling=sampling))
+            tick_latents.append(result.output["payload"]["latents"].clone())
+    tick_calls = list(transformer.calls)
+    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_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)
+
+    assert [call["update_cache"] for call in transformer.calls] == [call["update_cache"] for call in tick_calls]
+    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)
+
+
+def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None:
+    """A trajectory that slows down between blocks must condition every stepwise
+    chunk exactly as request mode does. Request mode embeds the whole trajectory
+    once; embedding per chunk would re-normalize each block's framewise
+    translations by its own largest step and present a slow block as full-speed
+    motion."""
+    from vllm_omni.diffusion.models.lingbot_world import camera as camera_module
+
+    module = _load_pipeline_module()
+    # The fixture stubs the ray embedding with a pose-blind ramp; parity needs
+    # the real geometry so a scale mismatch actually shows up.
+    monkeypatch.setattr(module, "build_plucker_embedding", camera_module.build_plucker_embedding)
+    transformer = _RecordingTransformer()
+    pipeline = _pipeline(module, transformer=transformer)
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+
+    # Six latent frames -> two 3-frame blocks. The first block moves at full
+    # speed and everything after it at a tenth of that, so a per-chunk
+    # normalization would rescale the second block by 10x.
+    steps = torch.tensor([1.0] * 3 + [0.1] * 18)
+    poses = torch.eye(4).repeat(21, 1, 1)
+    poses[:, 2, 3] = torch.cumsum(steps, dim=0) - steps[0]
+    trajectory = _CameraTrajectory(
+        poses=poses,
+        intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1),
+    )
+
+    sampling = _SamplingParams(num_frames=21)
+    sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
+    pipeline(_request(sampling=sampling))
+    request_cameras = [call["camera_hidden_states"] for call in transformer.calls]
+    transformer.calls.clear()
+
+    state = _stepwise_state(num_frames=21)
+    state.sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
+    state.sampling.extra_args.pop("_lingbot_camera_action_script")
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        _run_stepwise(pipeline, state)
+    stepwise_cameras = [call["camera_hidden_states"] for call in transformer.calls]
+
+    # Two blocks x (four probes + one commit), in the same order on both paths.
+    assert len(request_cameras) == len(stepwise_cameras) == 10
+    for stepwise, request in zip(stepwise_cameras, request_cameras, strict=True):
+        torch.testing.assert_close(stepwise, request)
+
+
+def test_stepwise_emits_latents_without_touching_the_vae() -> None:
+    """Latent mode must stay latent: streaming decode is opt-in per request."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    assert [set(output.output["payload"]) for output in outputs] == [{"latents"}, {"latents"}]
+    assert pipeline.vae.decode_inputs == []
+
+
+def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
+    """A non-latent request must stream pixels under the primary "video" key."""
+    import diffusers.video_processor as video_processor_module
+
+    module = _load_pipeline_module()
+    processed = []
+
+    class VideoProcessor:
+        def __init__(self, *, vae_scale_factor):
+            assert vae_scale_factor == 8
+
+        def postprocess_video(self, video, *, output_type):
+            processed.append((tuple(video.shape), output_type))
+            return f"frames-{len(processed)}"
+
+    monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        outputs = _run_stepwise(pipeline, state)
+
+    # One decode per AR block, each seeing only that block's three latent frames.
+    assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
+    assert [output_type for _, output_type in processed] == ["np", "np"]
+    assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
+    # Identity metadata survives the switch to pixel output.
+    assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
+
+
 def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
     module = _load_pipeline_module()
     resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
@@ -1755,7 +2046,7 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
     assert resolved is module.LingBotWorldCausalDMDPipeline
     assert cache_acceleration_disabled
     assert preprocess_name == "get_lingbot_world_pre_process_func"
-    request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
+    request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=_SamplingParams(), request_id="req-1")
     assert preprocess(request) is request
     assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
     assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"