491 lines
19 KiB
Python
491 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Camera modality interaction handlers.
|
|
|
|
Payload is structural SE(3): ``translation`` (xyz) + unit
|
|
quaternion ``rotation`` (x, y, z, w). WASD-style key tokens belong in
|
|
clients, not in the engine.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Mapping, Sequence
|
|
from dataclasses import dataclass, field
|
|
from typing import ClassVar, cast
|
|
|
|
import torch
|
|
from typing_extensions import Self, override
|
|
|
|
from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler
|
|
from vllm_omni.diffusion.interaction.types import (
|
|
InteractionChunkMetadata,
|
|
InteractionEvent,
|
|
InteractionMode,
|
|
InteractionPayload,
|
|
InteractionSession,
|
|
resolve_event_frame_offset,
|
|
)
|
|
from vllm_omni.diffusion.worker.utils import StepRequestState
|
|
|
|
Vec3 = tuple[float, float, float]
|
|
# Unit quaternion in ``(x, y, z, w)`` order. Identity is ``(0, 0, 0, 1)``.
|
|
Quat = tuple[float, float, float, float]
|
|
|
|
_IDENTITY_QUAT: Quat = (0.0, 0.0, 0.0, 1.0)
|
|
|
|
|
|
@dataclass
|
|
class CameraPose:
|
|
"""Canonical absolute camera pose used by the shared timeline.
|
|
|
|
Coordinates follow Unity semantics: ``+X`` right, ``+Y`` up,
|
|
``+Z`` forward. Session start is the identity transform;
|
|
``target`` poses are relative to that origin.
|
|
"""
|
|
|
|
translation: Vec3 = (0.0, 0.0, 0.0)
|
|
rotation: Quat = _IDENTITY_QUAT
|
|
|
|
@classmethod
|
|
def identity(cls) -> CameraPose:
|
|
return cls()
|
|
|
|
def clone(self) -> CameraPose:
|
|
return CameraPose(translation=self.translation, rotation=self.rotation)
|
|
|
|
def as_matrix(self) -> torch.Tensor:
|
|
"""Return a 4x4 rigid transform from translation + unit quaternion."""
|
|
mat = torch.eye(4, dtype=torch.float64)
|
|
mat[:3, :3] = _quat_to_rotmat(self.rotation)
|
|
mat[0, 3] = float(self.translation[0])
|
|
mat[1, 3] = float(self.translation[1])
|
|
mat[2, 3] = float(self.translation[2])
|
|
return mat
|
|
|
|
@classmethod
|
|
def from_matrix(cls, matrix: torch.Tensor) -> CameraPose:
|
|
if matrix.shape != (4, 4):
|
|
raise ValueError("camera pose matrix must be 4x4")
|
|
rotation = _rotmat_to_quat(matrix[:3, :3].detach().cpu().double())
|
|
translation = (
|
|
float(matrix[0, 3]),
|
|
float(matrix[1, 3]),
|
|
float(matrix[2, 3]),
|
|
)
|
|
return cls(translation=translation, rotation=rotation)
|
|
|
|
|
|
@dataclass(kw_only=True)
|
|
class QueuedCameraEvent(InteractionEvent):
|
|
"""Timestamped camera command waiting for the next chunk boundary."""
|
|
|
|
# Absolute pose for ``target``; per-frame SE3 delta for ``velocity``.
|
|
pose: CameraPose = field(default_factory=CameraPose.identity)
|
|
# Integer frame progress for ``target`` lerps (avoids float reciprocal drift).
|
|
elapsed_transition_frames: int = 0
|
|
|
|
|
|
@dataclass
|
|
class CameraSession(InteractionSession):
|
|
"""Request-local camera timeline under ``state.interaction_sessions['camera']``."""
|
|
|
|
current_pose: CameraPose = field(default_factory=CameraPose.identity)
|
|
pending_events: list[QueuedCameraEvent] = field(default_factory=list)
|
|
# In-flight command; ``elapsed_transition_chunks`` lives on the event.
|
|
active_event: QueuedCameraEvent | None = None
|
|
target_source: CameraPose | None = None
|
|
# Absolute C2W poses ``[num_latent_frames, 4, 4]`` for the most recent chunk.
|
|
last_absolute_poses: torch.Tensor | None = None
|
|
# Since a camera session is not lazily initialized, use this field to mark that a client enqueue is ever received.
|
|
has_received_input: bool = False
|
|
|
|
|
|
class SE3DeltaCameraHandler(InteractionHandler):
|
|
"""Generic camera timeline that samples absolute poses per frame.
|
|
|
|
* ``mode=target``: lerp toward an absolute pose (relative to session start).
|
|
* ``mode=velocity``: hold a structural SE3 delta until replaced.
|
|
|
|
``received_at`` is resolved on the media grid (``num_media_frames``/``fps``), then mapped onto latent steps.
|
|
Later integration and ``last_absolute_poses`` aligns with ``num_latent_frames``.
|
|
"""
|
|
|
|
modality: ClassVar[str] = "camera"
|
|
needs_chunk_media: ClassVar[bool] = True
|
|
# Camera conditioning must exist every chunk; apply creates the session and
|
|
# materializes identity/hold poses before any client enqueue.
|
|
lazy_initialize_session: ClassVar[bool] = False
|
|
default_transition_chunks: ClassVar[int] = 1
|
|
|
|
@classmethod
|
|
@override
|
|
def from_pipeline(cls, pipeline: object) -> Self:
|
|
"""It doesn't need anything from the pipeline (e.g., `encode_prompt` method)"""
|
|
del pipeline
|
|
return cls()
|
|
|
|
@override
|
|
def validate_payload(
|
|
self,
|
|
state: StepRequestState,
|
|
*,
|
|
event_id: str,
|
|
payload: InteractionPayload,
|
|
transition_chunks: int | None,
|
|
) -> None:
|
|
if not event_id:
|
|
raise ValueError("event_id must be non-empty")
|
|
|
|
mode = payload.get("mode", "target")
|
|
if mode not in ("target", "velocity"):
|
|
raise ValueError("camera mode must be 'target' or 'velocity'")
|
|
data = payload.get("data")
|
|
if not isinstance(data, Mapping):
|
|
raise ValueError("camera data must be an object")
|
|
if "actions" in data:
|
|
raise ValueError(
|
|
"camera data.actions (WASD keys) is not accepted by the engine; "
|
|
"send structural translation/rotation instead"
|
|
)
|
|
_parse_pose(data)
|
|
if mode == "target":
|
|
duration = self.default_transition_chunks if transition_chunks is None else int(transition_chunks)
|
|
if duration < 0:
|
|
raise ValueError("transition_chunks must be >= 0")
|
|
del state
|
|
|
|
@override
|
|
def enqueue(
|
|
self,
|
|
state: StepRequestState,
|
|
*,
|
|
event_id: str,
|
|
received_at: float,
|
|
payload: InteractionPayload,
|
|
transition_chunks: int | None,
|
|
) -> None:
|
|
self.validate_payload(
|
|
state,
|
|
event_id=event_id,
|
|
payload=payload,
|
|
transition_chunks=transition_chunks,
|
|
)
|
|
mode = payload.get("mode", "target")
|
|
camera_mode = cast(InteractionMode, mode)
|
|
data = cast(Mapping[str, object], payload.get("data"))
|
|
pose = _parse_pose(data)
|
|
if camera_mode == "target":
|
|
duration = self.default_transition_chunks if transition_chunks is None else int(transition_chunks)
|
|
else:
|
|
duration = 0
|
|
|
|
session = state.interaction_sessions.setdefault("camera", CameraSession())
|
|
assert isinstance(session, CameraSession)
|
|
with session.lock:
|
|
session.has_received_input = True
|
|
session.pending_events.append(
|
|
QueuedCameraEvent(
|
|
event_id=event_id,
|
|
received_at=received_at,
|
|
mode=camera_mode,
|
|
transition_chunks=duration,
|
|
pose=pose,
|
|
)
|
|
)
|
|
|
|
@override
|
|
def apply_at_chunk_boundary(
|
|
self,
|
|
state: StepRequestState,
|
|
*,
|
|
boundary_at: float,
|
|
chunk_index: int | None = None,
|
|
num_media_frames: int | None = None,
|
|
fps: float | None = None,
|
|
num_latent_frames: int | None = None,
|
|
) -> InteractionChunkMetadata | None:
|
|
"""Advance camera state for one chunk and store latent-aligned poses.
|
|
|
|
.. note::
|
|
Hardcoded the following policy of per-frame alignment:
|
|
|
|
- Resolve ``received_at`` on the media grid, then map each media index onto a latent step.
|
|
- Integrate one velocity step per **latent frame**, so holding speed ``0.05`` for a chunk with
|
|
``num_latent_frames=3`` yields total travel ``0.15``.
|
|
- ``last_absolute_poses.shape[0] == num_latent_frames``: interpolation is latent-aligned.
|
|
|
|
If future models need different scale/alignment policies, can refactor interaction registry
|
|
to be able to mark a policy.
|
|
"""
|
|
del chunk_index
|
|
if num_media_frames is None or fps is None:
|
|
raise ValueError("SE3DeltaCameraHandler requires chunk num_media_frames and fps")
|
|
if num_latent_frames is None:
|
|
num_latent_frames = num_media_frames
|
|
if int(num_latent_frames) <= 0:
|
|
raise ValueError(f"num_latent_frames must be > 0, got {num_latent_frames}")
|
|
|
|
session = state.interaction_sessions.setdefault("camera", CameraSession())
|
|
assert isinstance(session, CameraSession)
|
|
with session.lock:
|
|
poses, started, active, completed = self._step_one_chunk(
|
|
session,
|
|
num_media_frames=num_media_frames,
|
|
fps=fps,
|
|
num_latent_frames=num_latent_frames,
|
|
boundary_at=boundary_at,
|
|
)
|
|
if poses:
|
|
session.last_absolute_poses = torch.stack([p.as_matrix() for p in poses], dim=0)
|
|
else:
|
|
session.last_absolute_poses = torch.zeros((0, 4, 4), dtype=torch.float64)
|
|
|
|
return InteractionChunkMetadata(
|
|
started_event_ids=started,
|
|
active_event_ids=active,
|
|
completed_event_ids=completed,
|
|
)
|
|
|
|
def _step_one_chunk(
|
|
self,
|
|
session: CameraSession,
|
|
*,
|
|
num_media_frames: int,
|
|
fps: float,
|
|
num_latent_frames: int,
|
|
boundary_at: float,
|
|
) -> tuple[list[CameraPose], list[str], list[str], list[str]]:
|
|
"""Sample absolute poses for this chunk under target/velocity semantics."""
|
|
num_media_frames = max(int(num_media_frames), 1)
|
|
num_latent_frames = max(int(num_latent_frames), 1)
|
|
pending = list(session.pending_events)
|
|
session.pending_events.clear()
|
|
|
|
by_latent: dict[int, list[QueuedCameraEvent]] = {}
|
|
for event in pending:
|
|
media_frame = resolve_event_frame_offset(
|
|
received_at=event.received_at,
|
|
previous_boundary_at=session.last_boundary_at,
|
|
num_frames=num_media_frames,
|
|
fps=fps,
|
|
)
|
|
latent_idx = _media_frame_to_latent(
|
|
media_frame,
|
|
num_media_frames=num_media_frames,
|
|
num_latent_frames=num_latent_frames,
|
|
)
|
|
by_latent.setdefault(latent_idx, []).append(event)
|
|
for events in by_latent.values():
|
|
events.sort(key=lambda item: item.received_at)
|
|
|
|
started: list[str] = []
|
|
completed: list[str] = []
|
|
seen_completed: set[str] = set()
|
|
poses: list[CameraPose] = []
|
|
|
|
def _mark_completed(event_id: str) -> None:
|
|
if event_id not in seen_completed:
|
|
seen_completed.add(event_id)
|
|
completed.append(event_id)
|
|
|
|
for latent_idx in range(num_latent_frames):
|
|
for event in by_latent.get(latent_idx, []):
|
|
started.append(event.event_id)
|
|
cancelled_id = self._activate_event(session, event)
|
|
if cancelled_id is not None:
|
|
_mark_completed(cancelled_id)
|
|
just_completed = self._step_one_frame(session, num_latent_frames)
|
|
poses.append(session.current_pose.clone())
|
|
if just_completed is not None:
|
|
_mark_completed(just_completed)
|
|
|
|
active: list[str] = []
|
|
if session.active_event is not None:
|
|
active.append(session.active_event.event_id)
|
|
|
|
session.last_boundary_at = boundary_at
|
|
return poses, started, active, completed
|
|
|
|
def _activate_event(self, session: CameraSession, event: QueuedCameraEvent) -> str | None:
|
|
"""Activate ``event``, returning the event_id cancelled by this replacement if any."""
|
|
cancelled_id: str | None = None
|
|
if session.active_event is not None and session.active_event.event_id != event.event_id:
|
|
cancelled_id = session.active_event.event_id
|
|
|
|
event.elapsed_transition_chunks = 0.0
|
|
event.elapsed_transition_frames = 0
|
|
session.active_event = event
|
|
if event.mode == "target":
|
|
session.target_source = session.current_pose.clone()
|
|
else:
|
|
session.target_source = None
|
|
return cancelled_id
|
|
|
|
def _clear_active_target(self, session: CameraSession) -> str:
|
|
"""Snap to the target pose, clear the active target track, and return its event_id."""
|
|
assert session.active_event is not None and session.active_event.mode == "target"
|
|
event = session.active_event
|
|
event_id = event.event_id
|
|
session.current_pose = event.pose.clone()
|
|
session.active_event = None
|
|
session.target_source = None
|
|
return event_id
|
|
|
|
def _step_one_frame(self, session: CameraSession, total_num_frames_this_chunk: int) -> str | None:
|
|
"""Advance the active camera command by one frame."""
|
|
event = session.active_event
|
|
if event is None:
|
|
return None
|
|
|
|
if event.mode == "target":
|
|
duration = event.transition_chunks
|
|
source = session.target_source or session.current_pose
|
|
target = event.pose
|
|
if duration <= 0:
|
|
return self._clear_active_target(session)
|
|
|
|
total_frames = max(1, int(round(duration * float(total_num_frames_this_chunk))))
|
|
event.elapsed_transition_frames += 1
|
|
event.elapsed_transition_chunks = event.elapsed_transition_frames / float(total_num_frames_this_chunk)
|
|
alpha = min(1.0, event.elapsed_transition_frames / float(total_frames))
|
|
session.current_pose = _lerp_pose(source, target, alpha)
|
|
if event.elapsed_transition_frames >= total_frames:
|
|
return self._clear_active_target(session)
|
|
return None
|
|
|
|
# Velocity: apply the held structural SE3 delta once per frame.
|
|
session.current_pose = _compose_pose(session.current_pose, event.pose)
|
|
return None
|
|
|
|
|
|
def _media_frame_to_latent(
|
|
media_frame: int,
|
|
*,
|
|
num_media_frames: int,
|
|
num_latent_frames: int,
|
|
) -> int:
|
|
"""Map a media-timeline frame index onto the latent control grid."""
|
|
if num_latent_frames <= 1 or num_media_frames <= 1:
|
|
return 0
|
|
return min(int(media_frame) * int(num_latent_frames) // int(num_media_frames), num_latent_frames - 1)
|
|
|
|
|
|
def _as_xyz(value: object, *, name: str) -> Vec3:
|
|
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)) or len(value) != 3:
|
|
raise ValueError(f"camera {name} must be a length-3 list/tuple")
|
|
if any(math.isnan(v) or math.isinf(v) for v in value):
|
|
raise ValueError(f"camera {name} must be non-nan and non-inf")
|
|
return (float(value[0]), float(value[1]), float(value[2]))
|
|
|
|
|
|
def _as_quat(value: object, *, name: str) -> Quat:
|
|
if not isinstance(value, Sequence) or isinstance(value, (str, bytes)) or len(value) != 4:
|
|
raise ValueError(f"camera {name} must be a length-4 list/tuple (x, y, z, w)")
|
|
if any(math.isnan(v) or math.isinf(v) for v in value):
|
|
raise ValueError(f"camera {name} must be non-nan and non-inf")
|
|
return _quat_normalize((float(value[0]), float(value[1]), float(value[2]), float(value[3])))
|
|
|
|
|
|
def _parse_pose(data: Mapping[str, object]) -> CameraPose:
|
|
translation = _as_xyz(data.get("translation", (0.0, 0.0, 0.0)), name="translation")
|
|
rotation = _as_quat(data.get("rotation", _IDENTITY_QUAT), name="rotation")
|
|
return CameraPose(translation=translation, rotation=rotation)
|
|
|
|
|
|
def _lerp(a: float, b: float, alpha: float) -> float:
|
|
return a + (b - a) * alpha
|
|
|
|
|
|
def _lerp_pose(source: CameraPose, target: CameraPose, alpha: float) -> CameraPose:
|
|
alpha = min(1.0, max(0.0, alpha))
|
|
return CameraPose(
|
|
translation=cast(
|
|
Vec3,
|
|
tuple(_lerp(s, t, alpha) for s, t in zip(source.translation, target.translation)),
|
|
),
|
|
rotation=_quat_nlerp(source.rotation, target.rotation, alpha),
|
|
)
|
|
|
|
|
|
def _compose_pose(base: CameraPose, delta: CameraPose) -> CameraPose:
|
|
"""Apply ``delta`` after ``base`` in the session frame (T_new = T_base @ T_delta)."""
|
|
return CameraPose.from_matrix(base.as_matrix() @ delta.as_matrix())
|
|
|
|
|
|
def _quat_normalize(q: Quat) -> Quat:
|
|
x, y, z, w = q
|
|
norm = math.sqrt(x * x + y * y + z * z + w * w)
|
|
if norm <= 0.0:
|
|
raise ValueError("camera rotation quaternion must be non-zero")
|
|
return (x / norm, y / norm, z / norm, w / norm)
|
|
|
|
|
|
def _quat_nlerp(a: Quat, b: Quat, alpha: float) -> Quat:
|
|
"""Normalized linear interpolation; flips ``b`` when needed for shortest path."""
|
|
dot = a[0] * b[0] + a[1] * b[1] + a[2] * b[2] + a[3] * b[3]
|
|
if dot < 0.0:
|
|
b = (-b[0], -b[1], -b[2], -b[3])
|
|
return _quat_normalize(
|
|
(
|
|
_lerp(a[0], b[0], alpha),
|
|
_lerp(a[1], b[1], alpha),
|
|
_lerp(a[2], b[2], alpha),
|
|
_lerp(a[3], b[3], alpha),
|
|
)
|
|
)
|
|
|
|
|
|
def _quat_to_rotmat(q: Quat) -> torch.Tensor:
|
|
x, y, z, w = q
|
|
xx, yy, zz = x * x, y * y, z * z
|
|
xy, xz, yz = x * y, x * z, y * z
|
|
wx, wy, wz = w * x, w * y, w * z
|
|
return torch.tensor(
|
|
[
|
|
[1.0 - 2.0 * (yy + zz), 2.0 * (xy - wz), 2.0 * (xz + wy)],
|
|
[2.0 * (xy + wz), 1.0 - 2.0 * (xx + zz), 2.0 * (yz - wx)],
|
|
[2.0 * (xz - wy), 2.0 * (yz + wx), 1.0 - 2.0 * (xx + yy)],
|
|
],
|
|
dtype=torch.float64,
|
|
)
|
|
|
|
|
|
def _rotmat_to_quat(rotation: torch.Tensor) -> Quat:
|
|
"""Convert a 3x3 rotation matrix to a unit quaternion ``(x, y, z, w)``."""
|
|
m = rotation
|
|
trace = float(m[0, 0] + m[1, 1] + m[2, 2])
|
|
if trace > 0.0:
|
|
s = math.sqrt(trace + 1.0) * 2.0
|
|
w = 0.25 * s
|
|
x = (float(m[2, 1]) - float(m[1, 2])) / s
|
|
y = (float(m[0, 2]) - float(m[2, 0])) / s
|
|
z = (float(m[1, 0]) - float(m[0, 1])) / s
|
|
elif float(m[0, 0]) > float(m[1, 1]) and float(m[0, 0]) > float(m[2, 2]):
|
|
s = math.sqrt(1.0 + float(m[0, 0]) - float(m[1, 1]) - float(m[2, 2])) * 2.0
|
|
w = (float(m[2, 1]) - float(m[1, 2])) / s
|
|
x = 0.25 * s
|
|
y = (float(m[0, 1]) + float(m[1, 0])) / s
|
|
z = (float(m[0, 2]) + float(m[2, 0])) / s
|
|
elif float(m[1, 1]) > float(m[2, 2]):
|
|
s = math.sqrt(1.0 + float(m[1, 1]) - float(m[0, 0]) - float(m[2, 2])) * 2.0
|
|
w = (float(m[0, 2]) - float(m[2, 0])) / s
|
|
x = (float(m[0, 1]) + float(m[1, 0])) / s
|
|
y = 0.25 * s
|
|
z = (float(m[1, 2]) + float(m[2, 1])) / s
|
|
else:
|
|
s = math.sqrt(1.0 + float(m[2, 2]) - float(m[0, 0]) - float(m[1, 1])) * 2.0
|
|
w = (float(m[1, 0]) - float(m[0, 1])) / s
|
|
x = (float(m[0, 2]) + float(m[2, 0])) / s
|
|
y = (float(m[1, 2]) + float(m[2, 1])) / s
|
|
z = 0.25 * s
|
|
return _quat_normalize((x, y, z, w))
|
|
|
|
|
|
__all__ = [
|
|
"CameraPose",
|
|
"CameraSession",
|
|
"QueuedCameraEvent",
|
|
"SE3DeltaCameraHandler",
|
|
]
|