[Core][Diffusion] Camera interaction for diffusion streaming generation (LingBot World 2 as example) (#7198)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py

@@ -12,6 +12,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
     LINGBOT_CAMERA_ACTION_SCHEMA,
     LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
     LingBotCameraControlReducer,
+    camera_trajectory_from_absolute_pose,
     integrate_lingbot_camera_actions,
     parse_lingbot_camera_action_frames,
     parse_lingbot_camera_action_script,
@@ -212,6 +213,11 @@ def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
         torch.tensor([500.0, 500.0, 208.0, 120.0]),
     )
 
+    with pytest.raises(ValueError, match="resolution must be positive"):
+        camera_trajectory_from_absolute_pose(trajectory.poses[:1], width=0, height=480)
+    with pytest.raises(ValueError, match=r"\[frames, 4, 4\]"):
+        camera_trajectory_from_absolute_pose(torch.zeros(4, 4), width=832, height=480)
+
 
 @pytest.mark.parametrize(
     ("data", "message"),

tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py

@@ -422,7 +422,17 @@ def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_transla
     torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
 
 
-def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
+@pytest.mark.parametrize(
+    ("translation_scale", "expected_origins"),
+    [
+        (None, [[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]]),
+        (0.05, [[0.0, 0.0, 0.0], [60.0, 80.0, 0.0], [120.0, 160.0, 0.0]]),
+    ],
+)
+def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm(
+    translation_scale: float | None,
+    expected_origins: list[list[float]],
+) -> None:
     poses = torch.eye(4).repeat(3, 1, 1)
     poses[:, :3, 3] = torch.tensor(
         [
@@ -440,10 +450,37 @@ def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm()
         target_width=1,
         device=torch.device("cpu"),
         dtype=torch.float32,
+        translation_scale=translation_scale,
     )
 
-    expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
-    torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
+    torch.testing.assert_close(result[:, :3, 0, 0], torch.tensor(expected_origins))
+
+    from vllm_omni.diffusion.models.lingbot_world.actions import LINGBOT_CONTROLLER_TRANSLATION_UNIT
+
+    if translation_scale != LINGBOT_CONTROLLER_TRANSLATION_UNIT:
+        return
+    # Fixed controller unit keeps relative speed; max-norm would collapse both to 1.
+    slow = torch.eye(4).repeat(2, 1, 1)
+    slow[1, 2, 3] = 0.005
+    fast = torch.eye(4).repeat(2, 1, 1)
+    fast[1, 2, 3] = 0.5
+    kwargs = dict(
+        height=1,
+        width=1,
+        target_height=1,
+        target_width=1,
+        device=torch.device("cpu"),
+        dtype=torch.float32,
+        translation_scale=LINGBOT_CONTROLLER_TRANSLATION_UNIT,
+    )
+    slow_origin = build_plucker_embedding(_trajectory(slow), **kwargs)[1, :3, 0, 0]
+    fast_origin = build_plucker_embedding(_trajectory(fast), **kwargs)[1, :3, 0, 0]
+    torch.testing.assert_close(slow_origin, torch.tensor([0.0, 0.0, 0.1]))
+    torch.testing.assert_close(fast_origin, torch.tensor([0.0, 0.0, 10.0]))
+    wasd = torch.eye(4).repeat(2, 1, 1)
+    wasd[1, 2, 3] = LINGBOT_CONTROLLER_TRANSLATION_UNIT
+    wasd_origin = build_plucker_embedding(_trajectory(wasd), **kwargs)[1, :3, 0, 0]
+    torch.testing.assert_close(wasd_origin, torch.tensor([0.0, 0.0, 1.0]))
 
 
 def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:

tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py

@@ -18,6 +18,7 @@ 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.interaction.modality_handlers.camera import CameraSession
 from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
 from vllm_omni.diffusion.models.lingbot_world.actions import (
     integrate_lingbot_camera_actions,
@@ -411,8 +412,10 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
             intrinsics=value.intrinsics[:num_frames],
         )
 
-    def build_plucker_embedding(value, *, height, width, target_height, target_width, device, dtype):
-        del target_height, target_width
+    def build_plucker_embedding(
+        value, *, height, width, target_height, target_width, device, dtype, translation_scale=None
+    ):
+        del target_height, target_width, translation_scale
         frames = value.poses.shape[0]
         data = torch.arange(frames * 6 * height * width, device=device, dtype=torch.float32)
         return data.reshape(frames, 6, height, width).to(dtype=dtype)
@@ -1783,6 +1786,7 @@ def test_first_typed_yaw_action_uses_pre_action_identity_anchor(
         target_width=16,
         device=torch.device("cpu"),
         dtype=torch.float32,
+        translation_scale=None,  # None for camera_actions / action_script input (using max-norm)
     )[1:]
 
     assert not torch.equal(action_embedding, neutral_embedding)
@@ -2015,6 +2019,7 @@ def test_stepwise_failed_chunk_releases_pending_encoder_on_close(monkeypatch: py
     transformer.raise_on_call = 1
     with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
         pipeline.prepare_encode(state)
+        pipeline.prepare_next_chunk(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"):
@@ -2037,6 +2042,7 @@ def test_stateful_condition_rejects_tiled_encoder(monkeypatch: pytest.MonkeyPatc
             state = _stepwise_state()
             with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
                 pipeline.prepare_encode(state)
+                pipeline.prepare_next_chunk(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())
 
@@ -2168,11 +2174,14 @@ def _stepwise_state(
 def _run_stepwise(pipeline, state):
     outputs = []
     pipeline.prepare_encode(state)
+    pipeline.prepare_next_chunk(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))
+            if not state.request_denoise_completed:
+                pipeline.prepare_next_chunk(state)
     return outputs
 
 
@@ -2186,6 +2195,7 @@ def _stepwise_chunks(pipeline, state, ar_state):
     """
     with pipeline.bind_ar_diffusion_state(state.request_id, ar_state):
         pipeline.prepare_encode(state)
+        pipeline.prepare_next_chunk(state)
     while not state.request_denoise_completed:
         output = None
         with pipeline.bind_ar_diffusion_state(state.request_id, ar_state):
@@ -2193,6 +2203,8 @@ def _stepwise_chunks(pipeline, state, ar_state):
             pipeline.step_scheduler(state, noise)
             if state.chunk_denoise_completed:
                 output = pipeline.post_decode(state)
+                if not state.request_denoise_completed:
+                    pipeline.prepare_next_chunk(state)
         if output is not None:
             yield output
 
@@ -2463,6 +2475,64 @@ def test_stepwise_chunks_continue_one_sessions_temporal_decode(monkeypatch) -> N
     assert pipeline.vae._feat_map == ["module-owned"]
 
 
+def test_prepare_next_chunk_rejects_live_camera_with_scripted_paths() -> None:
+    """Scripted/cached camera must not silently ignore mid-generation camera events."""
+    module = _load_pipeline_module()
+    pipeline = _pipeline(module, transformer=_RecordingTransformer())
+    pipeline._ar_height = 16
+    pipeline._ar_width = 16
+
+    # Scenario 1, input pre-scripted camera trajectory via action_script
+    state = _stepwise_state(num_frames=21)
+    with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
+        pipeline.prepare_encode(state)
+        assert state.extra.get("camera_action_script") is not None
+        state.interaction_sessions["camera"] = CameraSession(has_received_input=True)
+        with pytest.raises(ValueError, match="camera_action_script"):
+            pipeline.prepare_next_chunk(state)
+
+    # Scenario 2, input pre-scripted camera trajectory via embedding cache
+    poses = torch.eye(4).repeat(21, 1, 1)
+    trajectory = _CameraTrajectory(
+        poses=poses,
+        intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1),
+    )
+    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)):  # pyright: ignore[reportArgumentType]
+        pipeline.prepare_encode(state)
+        assert state.extra.get("camera_embedding_cache") is not None
+        state.interaction_sessions["camera"] = CameraSession(has_received_input=True)
+        with pytest.raises(ValueError, match="camera_embedding_cache"):
+            pipeline.prepare_next_chunk(state)
+
+
+def test_peek_chunk_media_matches_streaming_decoder_frame_counts(monkeypatch) -> None:
+    """Camera timelines must track the streaming decoder's 9-then-12 media counts."""
+    module = _load_pipeline_module()
+    _capture_video_processor(monkeypatch)
+    pipeline = _streaming_pipeline(module)
+    state = _stepwise_state(num_frames=21)
+    state.sampling.output_type = "np"
+    state.sampling.fps = 16.0
+    fake = _FakeARState(state.request_id)
+
+    with pipeline.bind_ar_diffusion_state(state.request_id, fake):
+        pipeline.prepare_encode(state)
+        first = pipeline.peek_chunk_media(state)
+        assert (first.num_media_frames, first.fps, first.num_latent_frames) == (9, 16.0, 3)
+
+        pipeline.prepare_next_chunk(state)
+        while not state.chunk_denoise_completed:
+            noise = pipeline.denoise_step(None, states=[state])
+            pipeline.step_scheduler(state, noise)
+        pipeline.post_decode(state)
+
+        second = pipeline.peek_chunk_media(state)
+        assert (second.num_media_frames, second.fps, second.num_latent_frames) == (12, 16.0, 3)
+
+
 def test_streaming_decode_state_is_owned_by_the_session(monkeypatch) -> None:
     """Decoder state is keyed by request id and released with the AR session."""
     module = _load_pipeline_module()
@@ -2599,6 +2669,7 @@ def test_streaming_decode_receives_rescaled_latents(monkeypatch) -> None:
 
     with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
         pipeline.prepare_encode(state)
+        pipeline.prepare_next_chunk(state)
         while not state.chunk_denoise_completed:
             pipeline.step_scheduler(state, pipeline.denoise_step(None, states=[state]))
         model_space = state.latents.clone()