434 lines
17 KiB
Python
434 lines
17 KiB
Python
|
# 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
|