[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/streaming_interactive_generation/test_camera_interaction.py

@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for camera interaction handlers and coordinator registration."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator
+from vllm_omni.diffusion.interaction.mixin import InteractionMixin
+from vllm_omni.diffusion.interaction.modality_handlers.camera import CameraSession, SE3DeltaCameraHandler
+from vllm_omni.diffusion.interaction.modality_handlers.prompt import PromptSession
+from vllm_omni.diffusion.interaction.registry import STRUCTURED_HANDLER_REGISTRY
+from vllm_omni.diffusion.interaction.types import ChunkMediaSpec, resolve_event_frame_offset
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_FORWARD_VELOCITY = {
+    "mode": "velocity",
+    "data": {"translation": [0.0, 0.0, 0.05], "rotation": [0.0, 0.0, 0.0, 1.0]},
+}
+
+
+class _FakePromptPipeline(InteractionMixin):
+    """Minimal pipeline stub with prompt-update support."""
+
+    def __init__(self) -> None:
+        self.device = torch.device("cpu")
+        self.transformer = SimpleNamespace(dtype=torch.float32)
+        self.encode_prompt = MagicMock(return_value=(torch.full((1, 4, 2), 2.0), None))
+        self._interaction_coordinator = None
+
+    def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
+        del state
+        return ChunkMediaSpec(num_media_frames=8, fps=16.0, num_latent_frames=8)
+
+
+class _FakeLingBotPipeline(InteractionMixin):
+    """Minimal pipeline stub for LingBot camera registration."""
+
+    def __init__(self) -> None:
+        self._interaction_coordinator = None
+
+    def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec:
+        del state
+        return ChunkMediaSpec(num_media_frames=3, fps=3.0, num_latent_frames=3)
+
+
+def _make_state(*, request_id: str = "req-1") -> StepRequestState:
+    state = StepRequestState(
+        request_id=request_id,
+        sampling=SimpleNamespace(num_outputs_per_prompt=1, max_sequence_length=226),  # pyright: ignore[reportArgumentType]
+        prompt="hello",
+    )
+    state.prompt_embeds = torch.zeros(1, 4, 2)
+    state.extra = {}
+    return state
+
+
+def _boundary_at(previous_boundary_at: float | None, num_frames: int, fps: float) -> float:
+    if previous_boundary_at is None:
+        return 0.0
+    if fps <= 0:
+        return previous_boundary_at
+    return previous_boundary_at + num_frames / fps
+
+
+class TestCoordinatorResolution:
+    def test_helios_has_prompt_but_no_camera(self) -> None:
+        pipeline = _FakePromptPipeline()
+        od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+
+        assert coordinator.has_modality("prompt")
+        assert not coordinator.has_modality("camera")
+        assert "camera" not in STRUCTURED_HANDLER_REGISTRY["HeliosPipeline"]
+
+    def test_lingbot_registers_se3_camera_handler(self) -> None:
+        pipeline = _FakeLingBotPipeline()
+        od_config = SimpleNamespace(model_class_name="LingBotWorldCausalDMDPipeline")
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+
+        assert coordinator.has_modality("camera")
+        handler = coordinator.get_handler("camera")
+        assert isinstance(handler, SE3DeltaCameraHandler)
+        assert handler.lazy_initialize_session is False
+
+    def test_non_lazy_camera_apply_creates_session_without_enqueue(self) -> None:
+        """Camera apply must run before any enqueue to materialize identity poses."""
+        pipeline = _FakeLingBotPipeline()
+        od_config = SimpleNamespace(model_class_name="LingBotWorldCausalDMDPipeline")
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+        state = _make_state()
+
+        assert "camera" not in state.interaction_sessions
+        meta = coordinator.maybe_prepare_initial_session(state, pipeline)
+        session = state.interaction_sessions["camera"]
+        assert isinstance(session, CameraSession)
+        assert session.last_absolute_poses is not None
+        assert session.last_absolute_poses.shape == (3, 4, 4)
+        assert meta.started_event_ids == []
+        assert meta.active_event_ids == []
+
+    def test_lazy_prompt_apply_skips_without_enqueue(self) -> None:
+        pipeline = _FakePromptPipeline()
+        od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+        state = _make_state()
+
+        meta = coordinator.apply_at_chunk_boundary(state, boundary_at=0.0, chunk_index=0)
+        assert "prompt" not in state.interaction_sessions
+        assert meta.started_event_ids == []
+        assert meta.active_event_ids == []
+        assert meta.completed_event_ids == []
+
+        initial = coordinator.maybe_prepare_initial_session(state, pipeline)
+        assert "prompt" not in state.interaction_sessions
+        assert initial.started_event_ids == []
+
+    def test_unsupported_modality_includes_model_context(self) -> None:
+        pipeline = _FakePromptPipeline()
+        od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+
+        with pytest.raises(ValueError, match="HeliosPipeline"):
+            coordinator.enqueue(
+                _make_state(),
+                modality="camera",
+                event_id="cam-1",
+                received_at=0.0,
+                payload=_FORWARD_VELOCITY,
+                transition_chunks=None,
+            )
+
+    def test_enqueue_parts_rejects_unsupported_modality_atomically(self) -> None:
+        """Unsupported tracks in a composite event must not mutate any queues."""
+        pipeline = _FakePromptPipeline()
+        od_config = SimpleNamespace(model_class_name="HeliosPipeline")
+        coordinator = InteractionCoordinator.build(pipeline, od_config)
+        state = _make_state()
+        coordinator.enqueue(
+            state,
+            modality="prompt",
+            event_id="prior",
+            received_at=0.0,
+            payload={"prompt": "already-queued"},
+            transition_chunks=1,
+        )
+        prior_session = state.interaction_sessions["prompt"]
+        assert isinstance(prior_session, PromptSession)
+        prior = prior_session.pending_event
+        assert prior is not None
+        pipeline.encode_prompt.reset_mock()
+
+        with pytest.raises(ValueError, match="camera"):
+            coordinator.enqueue_parts(
+                state,
+                parts=[
+                    ("prompt", {"prompt": "should-not-replace"}),
+                    ("camera", _FORWARD_VELOCITY),
+                ],
+                event_id="composite-bad",
+                received_at=1.0,
+                transition_chunks=2,
+            )
+
+        pipeline.encode_prompt.assert_not_called()
+        pending = prior_session.pending_event
+        assert pending is not None
+        assert pending.event_id == "prior"
+        assert pending.prompt == "already-queued"
+        assert "camera" not in state.interaction_sessions
+
+        # Invalid camera payload must also fail before any track is queued.
+        dual = InteractionCoordinator(
+            {
+                "prompt": coordinator.get_handler("prompt"),
+                "camera": SE3DeltaCameraHandler(),
+            },
+            model_class_name="HeliosPipeline",
+        )
+        pipeline.encode_prompt.reset_mock()
+        with pytest.raises(ValueError, match="actions"):
+            dual.enqueue_parts(
+                state,
+                parts=[
+                    ("prompt", {"prompt": "should-not-replace"}),
+                    ("camera", {"mode": "velocity", "data": {"actions": ["w"]}}),
+                ],
+                event_id="composite-bad-payload",
+                received_at=2.0,
+                transition_chunks=2,
+            )
+        pipeline.encode_prompt.assert_not_called()
+        pending = prior_session.pending_event
+        assert pending is not None
+        assert pending.event_id == "prior"
+        assert pending.prompt == "already-queued"
+        assert "camera" not in state.interaction_sessions
+
+
+class TestCameraHandlers:
+    def test_rejects_wasd_actions_payload(self) -> None:
+        handler = SE3DeltaCameraHandler()
+        state = _make_state()
+        with pytest.raises(ValueError, match="actions"):
+            handler.enqueue(
+                state,
+                event_id="cam-wasd",
+                received_at=0.0,
+                payload={"mode": "velocity", "data": {"actions": ["w"]}},
+                transition_chunks=None,
+            )
+
+    def test_target_transition_progress_and_se3_projection(self) -> None:
+        handler = SE3DeltaCameraHandler()
+        state = _make_state()
+        handler.enqueue(
+            state,
+            event_id="cam-1",
+            received_at=0.0,
+            payload={
+                "mode": "target",
+                "data": {"translation": [0.0, 0.0, 3.0], "rotation": [0.0, 0.0, 0.0, 1.0]},
+            },
+            transition_chunks=2,
+        )
+
+        meta = handler.apply_at_chunk_boundary(
+            state,
+            chunk_index=0,
+            num_media_frames=4,
+            fps=16.0,
+            boundary_at=0.25,
+        )
+        session = state.interaction_sessions["camera"]
+        assert isinstance(session, CameraSession)
+        assert session.last_absolute_poses is not None
+        assert meta is not None
+        assert meta.started_event_ids == ["cam-1"]
+        assert meta.active_event_ids == ["cam-1"]
+        assert meta.completed_event_ids == []
+        assert session.last_absolute_poses.shape == (4, 4, 4)
+
+        meta2 = handler.apply_at_chunk_boundary(
+            state,
+            chunk_index=1,
+            num_media_frames=4,
+            fps=16.0,
+            boundary_at=0.5,
+        )
+        assert meta2 is not None
+        assert meta2.started_event_ids == []
+        assert meta2.active_event_ids == []
+        assert meta2.completed_event_ids == ["cam-1"]
+        assert session.active_event is None
+        assert session.current_pose.translation[2] == pytest.approx(3.0)
+        assert session.last_absolute_poses is not None
+        assert session.last_absolute_poses.shape == (4, 4, 4)
+
+    def test_single_chunk_target_completes_once(self) -> None:
+        handler = SE3DeltaCameraHandler()
+        # Include a non-power-of-two frame count (7) so float reciprocal drift
+        # would previously miss the final snap / completion ack.
+        for num_frames in (8, 7):
+            state = _make_state()
+            handler.enqueue(
+                state,
+                event_id="cam-fast",
+                received_at=0.0,
+                payload={
+                    "mode": "target",
+                    "data": {"translation": [1.0, 0.0, 0.0], "rotation": [0.0, 0.0, 0.0, 1.0]},
+                },
+                transition_chunks=1,
+            )
+            meta = handler.apply_at_chunk_boundary(
+                state,
+                chunk_index=0,
+                num_media_frames=num_frames,
+                fps=16.0,
+                boundary_at=0.5,
+            )
+            assert meta is not None
+            assert meta.started_event_ids == ["cam-fast"]
+            assert meta.active_event_ids == []
+            assert meta.completed_event_ids == ["cam-fast"]
+
+            meta2 = handler.apply_at_chunk_boundary(
+                state,
+                chunk_index=1,
+                num_media_frames=num_frames,
+                fps=16.0,
+                boundary_at=1.0,
+            )
+            assert meta2 is not None
+            assert meta2.started_event_ids == []
+            assert meta2.active_event_ids == []
+            assert meta2.completed_event_ids == []
+            session = state.interaction_sessions["camera"]
+            assert isinstance(session, CameraSession)
+            assert session.current_pose.translation[0] == pytest.approx(1.0)
+
+    def test_velocity_composes_structural_se3_delta(self) -> None:
+        handler = SE3DeltaCameraHandler()
+        state = _make_state()
+        handler.enqueue(
+            state,
+            event_id="vel-1",
+            received_at=0.0,
+            payload=_FORWARD_VELOCITY,
+            transition_chunks=None,
+        )
+
+        meta0 = handler.apply_at_chunk_boundary(state, chunk_index=0, num_media_frames=3, fps=16.0, boundary_at=0.0)
+        assert meta0 is not None
+        assert meta0.started_event_ids == ["vel-1"]
+        assert meta0.active_event_ids == ["vel-1"]
+        session = state.interaction_sessions["camera"]
+        assert isinstance(session, CameraSession)
+        assert session.current_pose.translation[2] == pytest.approx(0.15)
+        assert session.last_absolute_poses is not None
+        assert session.last_absolute_poses.shape == (3, 4, 4)
+        torch.testing.assert_close(
+            session.last_absolute_poses[0, :3, 3],
+            torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
+            atol=1e-6,
+            rtol=0,
+        )
+        torch.testing.assert_close(
+            session.last_absolute_poses[1, :3, 3],
+            torch.tensor([0.0, 0.0, 0.10], dtype=torch.float64),
+            atol=1e-6,
+            rtol=0,
+        )
+
+        meta1 = handler.apply_at_chunk_boundary(state, chunk_index=1, num_media_frames=2, fps=16.0, boundary_at=0.125)
+        assert meta1 is not None
+        assert meta1.active_event_ids == ["vel-1"]
+        assert session.current_pose.translation[2] == pytest.approx(0.25)
+
+    def test_velocity_completed_when_replaced_by_target(self) -> None:
+        handler = SE3DeltaCameraHandler()
+        state = _make_state()
+        handler.enqueue(
+            state,
+            event_id="vel-1",
+            received_at=0.0,
+            payload=_FORWARD_VELOCITY,
+            transition_chunks=None,
+        )
+        meta0 = handler.apply_at_chunk_boundary(
+            state,
+            chunk_index=0,
+            num_media_frames=2,
+            fps=16.0,
+            boundary_at=0.0,
+        )
+        assert meta0 is not None
+        assert meta0.active_event_ids == ["vel-1"]
+
+        handler.enqueue(
+            state,
+            event_id="tgt-1",
+            received_at=0.1,
+            payload={
+                "mode": "target",
+                "data": {"translation": [0.0, 0.0, 0.0], "rotation": [0.0, 0.0, 0.0, 1.0]},
+            },
+            transition_chunks=0,
+        )
+        meta1 = handler.apply_at_chunk_boundary(
+            state,
+            chunk_index=1,
+            num_media_frames=2,
+            fps=16.0,
+            boundary_at=0.125,
+        )
+        assert meta1 is not None
+        assert meta1.started_event_ids == ["tgt-1"]
+        assert "vel-1" in meta1.completed_event_ids
+        assert meta1.active_event_ids == []
+
+    def test_same_chunk_velocity_then_target_by_frame_offset(self) -> None:
+        handler = SE3DeltaCameraHandler()
+        state = _make_state()
+        fps = 10.0
+        boundary = 0.0
+        handler.enqueue(
+            state,
+            event_id="vel",
+            received_at=0.0,
+            payload=_FORWARD_VELOCITY,
+            transition_chunks=None,
+        )
+        handler.enqueue(
+            state,
+            event_id="tgt",
+            received_at=0.5,
+            payload={
+                "mode": "target",
+                "data": {"translation": [0.0, 0.0, 0.0], "rotation": [0.0, 0.0, 0.0, 1.0]},
+            },
+            transition_chunks=0,
+        )
+        state.interaction_sessions["camera"].last_boundary_at = boundary
+        meta = handler.apply_at_chunk_boundary(
+            state,
+            chunk_index=0,
+            num_media_frames=10,
+            fps=fps,
+            boundary_at=_boundary_at(boundary, 10, fps),
+        )
+        assert meta is not None
+        assert meta.started_event_ids == ["vel", "tgt"]
+        assert meta.completed_event_ids == ["vel", "tgt"]
+        assert meta.active_event_ids == []
+        session = state.interaction_sessions["camera"]
+        assert isinstance(session, CameraSession)
+        assert session.active_event is None
+        assert session.current_pose.translation[2] == pytest.approx(0.0)
+
+
+class TestResolveEventFrameOffset:
+    def test_clamps_and_floors(self) -> None:
+        assert resolve_event_frame_offset(received_at=0.0, previous_boundary_at=None, num_frames=8, fps=16.0) == 0
+        assert resolve_event_frame_offset(received_at=0.1, previous_boundary_at=0.0, num_frames=8, fps=16.0) == 1
+        assert resolve_event_frame_offset(received_at=10.0, previous_boundary_at=0.0, num_frames=8, fps=16.0) == 7
+        assert resolve_event_frame_offset(received_at=-1.0, previous_boundary_at=0.0, num_frames=8, fps=16.0) == 0