diff --git a/docs/serving/streaming_video_output_api.md b/docs/serving/streaming_video_output_api.md index d01489540..b219ee019 100644 --- a/docs/serving/streaming_video_output_api.md +++ b/docs/serving/streaming_video_output_api.md @@ -41,7 +41,7 @@ bytes. | Server to client | `video.start` | Confirms the request ID, format, and accepted configuration | | Server to client | `video.chunk_metadata` | Describes the next binary media frame | | Server to client | Binary frame | Carries one fragmented MP4 chunk | -| Client to server | `session.interaction` | Queues an optional prompt update during generation | +| Client to server | `session.interaction` | Queues an optional interaction (e.g., prompt. camera action) during generation | | Server to client | `session.interaction.queued` | Confirms that an interaction was queued | | Client to server | `session.ping` | Refreshes the stall clock | | Server to client | `session.pong` | Acknowledges a ping | @@ -83,10 +83,39 @@ Models that implement interaction updates can change the active prompt: } ``` +Models that register the camera modality (for example LingBot-World) also +accept structural SE3 payloads on `event.multi_modal_data.camera`. Coordinates +follow Unity semantics (`+X` right, `+Y` up, `+Z` forward; same as LingBot's +native camera frame). WASD key tokens are not admitted by the engine: + +```json +{ + "type": "session.interaction", + "interaction": { + "event_id": "forward-1", + "event": { + "multi_modal_data": { + "camera": { + "mode": "velocity", + "data": { + "translation": [0.0, 0.0, 0.05], + "rotation": [0.0, 0.0, 0.0, 1.0] + } + } + } + } + } +} +``` + +`mode` is `target` (lerp to an absolute pose over `transition_chunks`) or +`velocity` (hold a per-frame SE3 delta until replaced). Prompt and camera may +be combined in one interaction event. + This feature is pipeline-dependent. An enabled endpoint does not imply that -the loaded model supports midway prompt updates. +the loaded model supports midway prompt or camera updates. The default start timeout is 10 seconds and the server reports a stall after about 60 seconds without generation progress or a `session.ping`. See the [streaming video generation example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/streaming_video_generation) -for prompt-update scheduling and browser playback. +for prompt/camera-update scheduling and browser playback. diff --git a/examples/online_serving/streaming_video_generation/README.md b/examples/online_serving/streaming_video_generation/README.md index 8e535ce0e..ba4724c69 100644 --- a/examples/online_serving/streaming_video_generation/README.md +++ b/examples/online_serving/streaming_video_generation/README.md @@ -21,7 +21,7 @@ The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=Tr | Direction | Message | Format | Description | | --- | --- | --- | --- | | Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. Image-conditioned models take a first frame as `image_reference`: `{"image_url":"https://..."}` or a `data:` URL. | -| Client to server | `session.interaction` | JSON text: `{"type":"session.interaction","interaction":{"event_id":"xxx","event":{"prompt":"..."},"transition_chunks":3}}` | Updates the active prompt midway through generation. `event_id` is optional and `transition_chunks` defaults to the model setting. | +| Client to server | `session.interaction` | JSON text: `{"type":"session.interaction","interaction":{"event_id":"xxx","event":{"prompt":"..."},"transition_chunks":3}}` | Updates the active prompt and/or camera midway through generation. Camera at `event.multi_modal_data.camera` is structural SE3 (`mode` + `translation`/`rotation`). `event_id` is optional and `transition_chunks` defaults to the model setting. | | Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. | | Server to client | `video.chunk_metadata` | JSON text: `{"type":"video.chunk_metadata","request_id":"...","kind":"media","transport_chunk_index":0,"generation_chunk_index":0,"num_frames":9,"byte_length":1234,"started_event_ids":[],"active_event_ids":[],"completed_event_ids":[]}` | Precedes each binary frame and describes the immediately following payload. | | Server to client | Video chunk | Binary WebSocket frame | Fragmented MP4 (`m4s`) video bytes. | @@ -59,10 +59,10 @@ python streaming_video_client.py \ --output helios_stream.mp4 ``` -The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`. +The example client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`. The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration. -Schedule midway prompt updates with `--prompt-updates`. None are sent unless you ask for them, and only pipelines that implement midway prompt updates accept them (LingBot-World, for one, does not, and rejects the whole session). Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`: +Schedule midway prompt updates with `--prompt-updates`. None are sent unless you ask for them, and only pipelines that implement midway prompt updates accept them. Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`: ```bash python streaming_video_client.py \ @@ -73,6 +73,21 @@ python streaming_video_client.py \ ]' ``` +Schedule midway camera updates with `--camera-updates` when the served model +registers the camera modality (e.g., LingBot World 2). +It accepts `["camera"]["data"]["translation"/"rotation"]` in the endpoint's required SE3 format. +It also accepts `["actions"]` at the update level and converts keystrokes to the SE3 format. + +```bash +python streaming_video_client.py \ + --model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --image-reference /path/to/first_frame.png \ + --camera-updates '[ + {"at": 1.0, "actions": ["w"]}, + {"at": 3.0, "camera": {"mode": "target", "data": {"translation": [0.0, 0.0, 1.0], "rotation": [0.0, 0.0, 0.0, 1.0]}}, "transition_chunks": 2} + ]' +``` + The Helios demo continues the default prompt into an underwater storm: ```bash @@ -144,6 +159,7 @@ python streaming_video_client.py \ --output lingbot_world_v2_stream.mp4 ``` -`camera_action_script` carries one three-latent-frame WASD action list per generated -chunk. Mid-session `session.interaction` only updates the prompt, so a LingBot rollout -follows the camera script it started with. +`camera_action_script` is an optional request-scoped WASD script (one +three-latent-frame action list per generated chunk). Prefer mid-session +`session.interaction` camera control via `--camera-updates` when you need +live motion; the engine wire format is structural SE3 only. diff --git a/examples/online_serving/streaming_video_generation/camera_utils.py b/examples/online_serving/streaming_video_generation/camera_utils.py new file mode 100644 index 000000000..82f493b37 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/camera_utils.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Client-side helpers that map WASD/IJKL keys to structural SE3 camera payloads. +Convenience for demos. +""" + +from __future__ import annotations + +import math +from typing import Any + +# Unity-frame (+X right, +Y up, +Z forward) step sizes used by the optional +# WASD helper. Magnitudes mirror LingBot's script integrator; the engine never +# sees key tokens — only structural SE3 payloads. +_WASD_STEP = 0.05 +_WASD_PITCH_DEG = 4.0 +_WASD_YAW_DEG = 6.0 +# Match LingBot's per-frame pitch clamp; without absolute pose this only limits +# a single helper call, not cumulative held velocity. +_WASD_PITCH_LIMIT_DEG = 85.0 + + +def _quat_from_axis_angle(axis: tuple[float, float, float], degrees: float) -> list[float]: + """Return a unit quaternion ``[x, y, z, w]`` for a right-handed axis-angle.""" + radians = math.radians(degrees) + half = 0.5 * radians + sine = math.sin(half) + return [axis[0] * sine, axis[1] * sine, axis[2] * sine, math.cos(half)] + + +def wasd_to_camera_payload(actions: list[str], *, mode: str = "velocity") -> dict[str, Any]: + """Map WASD/IJKL key tokens to a structural camera interaction payload. + + This is a client convenience only. Servers reject ``data.actions``. + + The delta is chosen so that at the identity pose, ``T_base @ T_delta`` + matches LingBot's script integrator rotation + ``R_y(yaw) @ R_x(pitch)``. Without the current absolute pose this remains + an approximation once the camera has already yawed/pitched. + """ + keys = {str(action).lower() for action in actions} + dx = _WASD_STEP * (("d" in keys) - ("a" in keys)) + dy = 0.0 + dz = _WASD_STEP * (("w" in keys) - ("s" in keys)) + pitch = _WASD_PITCH_DEG * (("i" in keys) - ("k" in keys)) + yaw = _WASD_YAW_DEG * (("l" in keys) - ("j" in keys)) + pitch = max(-_WASD_PITCH_LIMIT_DEG, min(_WASD_PITCH_LIMIT_DEG, pitch)) + + # R_delta = R_y(yaw) @ R_x(pitch): apply pitch first, then yaw. + qx, qy, qz, qw = 0.0, 0.0, 0.0, 1.0 + if abs(pitch) > 0.0: + px, py, pz, pw = _quat_from_axis_angle((1.0, 0.0, 0.0), pitch) + qx, qy, qz, qw = px, py, pz, pw + if abs(yaw) > 0.0: + yx, yy, yz, yw = _quat_from_axis_angle((0.0, 1.0, 0.0), yaw) + # q = yaw ⊗ pitch + qx, qy, qz, qw = ( + yw * qx + yx * qw + yy * qz - yz * qy, + yw * qy - yx * qz + yy * qw + yz * qx, + yw * qz + yx * qy - yy * qx + yz * qw, + yw * qw - yx * qx - yy * qy - yz * qz, + ) + + return { + "mode": mode, + "data": { + "translation": [dx, dy, dz], + "rotation": [qx, qy, qz, qw], + }, + } diff --git a/examples/online_serving/streaming_video_generation/streaming_video_client.py b/examples/online_serving/streaming_video_generation/streaming_video_client.py index 268891284..d159234e4 100644 --- a/examples/online_serving/streaming_video_generation/streaming_video_client.py +++ b/examples/online_serving/streaming_video_generation/streaming_video_client.py @@ -26,6 +26,7 @@ from pathlib import Path from typing import Any import av +from camera_utils import wasd_to_camera_payload try: from websockets.asyncio.client import connect # pyright: ignore[reportMissingImports] @@ -56,6 +57,15 @@ class ScheduledPromptUpdate: transition_chunks: int +@dataclass(frozen=True) +class ScheduledCameraUpdate: + """Client-side schedule entry for a midway camera ``session.interaction``.""" + + at: float + camera: dict[str, Any] + transition_chunks: int | None + + def _image_reference(value: str) -> dict[str, Any]: """Accept an http(s) URL as-is, or inline a local image as a data URL.""" if value.startswith(("http://", "https://", "data:")): @@ -119,6 +129,61 @@ def _parse_prompt_updates(value: str) -> list[ScheduledPromptUpdate]: return sorted(updates, key=lambda update: update.at) +def _parse_camera_updates(value: str) -> list[ScheduledCameraUpdate]: + try: + parsed = json.loads(value) + except json.JSONDecodeError as exc: + raise argparse.ArgumentTypeError(f"--camera-updates must be valid JSON: {exc}") from exc + if not isinstance(parsed, list): + raise argparse.ArgumentTypeError("--camera-updates must be a JSON array") + + updates: list[ScheduledCameraUpdate] = [] + for index, item in enumerate(parsed): + if not isinstance(item, dict): + raise argparse.ArgumentTypeError(f"--camera-updates[{index}] must be a JSON object") + try: + at = float(item["at"]) + except (KeyError, TypeError, ValueError) as exc: + raise argparse.ArgumentTypeError(f"--camera-updates[{index}] must have a valid 'at' field") from exc + if at < 0: + raise argparse.ArgumentTypeError(f"--camera-updates[{index}].at must be >= 0") + + if "camera" in item: + camera = item["camera"] + if not isinstance(camera, dict): + raise argparse.ArgumentTypeError(f"--camera-updates[{index}].camera must be an object") + elif "actions" in item: + actions = item["actions"] + if not isinstance(actions, list): + raise argparse.ArgumentTypeError(f"--camera-updates[{index}].actions must be a list of WASD/IJKL keys") + mode = item.get("mode", "velocity") + if mode not in ("target", "velocity"): + raise argparse.ArgumentTypeError(f"--camera-updates[{index}].mode must be target or velocity") + camera = wasd_to_camera_payload(actions, mode=mode) + else: + raise argparse.ArgumentTypeError( + f"--camera-updates[{index}] must provide 'camera' (structural SE3) " + "or 'actions' (client-side WASD helper)" + ) + + transition_chunks: int | None + if "transition_chunks" not in item: + transition_chunks = None if camera.get("mode") == "velocity" else DEFAULT_TRANSITION_CHUNKS + else: + try: + transition_chunks = int(item["transition_chunks"]) + except (TypeError, ValueError) as exc: + raise argparse.ArgumentTypeError( + f"--camera-updates[{index}].transition_chunks must be an integer" + ) from exc + if transition_chunks < 0: + raise argparse.ArgumentTypeError(f"--camera-updates[{index}].transition_chunks must be >= 0") + + updates.append(ScheduledCameraUpdate(at=at, camera=camera, transition_chunks=transition_chunks)) + + return sorted(updates, key=lambda update: update.at) + + def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None: if value is not None: payload[key] = value @@ -173,7 +238,7 @@ async def _run_prompt_update_scheduler( payload = { "type": "session.interaction", "interaction": { - "event_id": f"scheduled-{update.at:.3f}", + "event_id": f"scheduled-prompt-{update.at:.3f}", "event": {"prompt": update.prompt}, "transition_chunks": update.transition_chunks, }, @@ -183,6 +248,30 @@ async def _run_prompt_update_scheduler( print(f"Sent session.interaction at t={update.at:.2f}s: {json.dumps(payload, ensure_ascii=False)}") +async def _run_camera_update_scheduler( + websocket: Any, + updates: list[ScheduledCameraUpdate], + *, + video_started_at: float, + send_lock: asyncio.Lock, +) -> None: + for update in updates: + delay = (video_started_at + update.at) - time.perf_counter() + if delay > 0: + await asyncio.sleep(delay) + + interaction: dict[str, Any] = { + "event_id": f"scheduled-camera-{update.at:.3f}", + "event": {"multi_modal_data": {"camera": update.camera}}, + } + if update.transition_chunks is not None: + interaction["transition_chunks"] = update.transition_chunks + payload = {"type": "session.interaction", "interaction": interaction} + async with send_lock: + await websocket.send(json.dumps(payload, ensure_ascii=False)) + print(f"Sent session.interaction at t={update.at:.2f}s: {json.dumps(payload, ensure_ascii=False)}") + + def _count_decoded_frames(video_bytes: bytes, *, stream_format: str) -> int | None: """Return the number of decodable video frames in a streamed byte sequence, if known.""" if not video_bytes: @@ -241,6 +330,7 @@ async def stream_video(args: argparse.Namespace) -> None: ws = None send_lock = asyncio.Lock() prompt_update_task: asyncio.Task[None] | None = None + camera_update_task: asyncio.Task[None] | None = None pending_chunk_metadata: dict[str, Any] | None = None try: @@ -295,8 +385,8 @@ async def stream_video(args: argparse.Namespace) -> None: if msg_type == "video.start": stream_format = msg.get("format") or stream_format print(f"Video session started: request_id={msg.get('request_id')} format={msg.get('format')}") + video_started_at = time.perf_counter() if args.prompt_updates and prompt_update_task is None: - video_started_at = time.perf_counter() prompt_update_task = asyncio.create_task( _run_prompt_update_scheduler( websocket, @@ -305,6 +395,15 @@ async def stream_video(args: argparse.Namespace) -> None: send_lock=send_lock, ) ) + if args.camera_updates and camera_update_task is None: + camera_update_task = asyncio.create_task( + _run_camera_update_scheduler( + websocket, + args.camera_updates, + video_started_at=video_started_at, + send_lock=send_lock, + ) + ) elif msg_type == "video.chunk_metadata": # Save the metadata for the incoming video chunk pending_chunk_metadata = msg @@ -315,6 +414,8 @@ async def stream_video(args: argparse.Namespace) -> None: done = True if prompt_update_task is not None: prompt_update_task.cancel() + if camera_update_task is not None: + camera_update_task.cancel() break elif msg_type == "error": print(f"ERROR: {msg.get('message', msg)}") @@ -335,6 +436,10 @@ async def stream_video(args: argparse.Namespace) -> None: prompt_update_task.cancel() with contextlib.suppress(asyncio.CancelledError): await prompt_update_task + if camera_update_task is not None: + camera_update_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await camera_update_task print( "Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)" ) @@ -405,9 +510,9 @@ def parse_args() -> argparse.Namespace: default=None, help="JSON object merged into extra_params; overrides preset keys on conflict.", ) - # Midway prompt updates are pipeline-specific: a model that does not implement - # them (LingBot-World, for one) rejects the whole session on the first one, so - # nothing is scheduled unless the caller asks for it. + # Midway interactions are pipeline-specific: unsupported modalities reject + # the update (and may leave the session running). Nothing is scheduled unless + # the caller asks for it. parser.add_argument( "--prompt-updates", type=_parse_prompt_updates, @@ -419,6 +524,19 @@ def parse_args() -> argparse.Namespace: f"'transition_chunks' defaults to {DEFAULT_TRANSITION_CHUNKS}." ), ) + parser.add_argument( + "--camera-updates", + type=_parse_camera_updates, + default=[], + help=( + "JSON array of scheduled camera updates for models that register the " + "camera modality (e.g. LingBot-World). Each object needs 'at' and either " + "'camera' (structural SE3: mode + translation/rotation) or " + "'actions' (client-side WASD/IJKL helper converted before send). " + "Optional 'transition_chunks' defaults to None for velocity and " + f"{DEFAULT_TRANSITION_CHUNKS} for target." + ), + ) return parser.parse_args() diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md index ab8c498fa..24a36653d 100644 --- a/recipes/Robbyant/LingBot-World-2.0.md +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -21,13 +21,12 @@ below. - **Realtime stepwise** — see [Streaming video serving](#streaming-video-serving). The suggested path: one `WS /v1/realtime/video` session produces the whole - rollout, one video chunk per AR block. It cannot take realtime camera - interaction yet; that is being added, and this path is the one that will - remain. + rollout, one video chunk per AR block. Mid-session camera control uses + structural SE3 `session.interaction` payloads; optional + `camera_action_script` remains for request-scoped WASD scripts. - **Realtime tick** *(deprecated)* — see [Realtime in-process generation](#realtime-in-process-generation-deprecated). - The only path that takes realtime camera interaction today, so it stays until - the stepwise path can, but it takes no new work and will be removed. + Older one-block-per-`generate()` control plane with JSONL WASD frames; will be removed in the future. - **Offline** — see [Offline generation](#offline-generation). Replays a fixed pose/intrinsics trajectory in one request and writes an MP4. Use it when the camera path is known up front and streaming is not needed. @@ -56,10 +55,10 @@ sufficient device memory. ## Realtime in-process generation (deprecated) -Use [Streaming video serving](#streaming-video-serving) instead unless you need -mid-session camera interaction, which only this path has. It still works and -emits no runtime warning; removal is tracked as B4 of the LingBot World 2.0 -roadmap ([#6672](https://github.com/vllm-project/vllm-omni/issues/6672)). +Prefer [Streaming video serving](#streaming-video-serving) for mid-session +camera control. This tick example still works and emits no runtime warning; +removal is tracked as B4 of the LingBot World 2.0 roadmap +([#6672](https://github.com/vllm-project/vllm-omni/issues/6672)). The realtime example keeps AR-Diffusion KV and model-owned state across requests. Each JSONL line describes the prompt and/or three latent-frame camera @@ -96,8 +95,8 @@ The stepwise path keeps AR-Diffusion paged KV but issues **one** request for the whole rollout: `prepare_encode` runs once, then every AR block is four DMD steps followed by one streamed chunk. Identity metadata uses `session_id = request_id` with contiguous `chunk_index` values from zero. -Camera motion for this landing is request-scoped: one three-frame action list -per chunk, fixed when the rollout starts. +Camera motion can be request-scoped (`camera_action_script` / `action_path`) or +updated mid-session via `session.interaction`. Serve it with the AR-Diffusion deploy config, which selects the AR-Diffusion engine and enables streamed step execution: @@ -115,11 +114,11 @@ keep their request-mode topology when no deploy config is given, and the stepwise serving topology is only ever an explicit choice. Clients then use the generic WebSocket protocol documented in -[`docs/serving/video_stream_api.md`](../../docs/serving/video_stream_api.md): +[`docs/serving/streaming_video_output_api.md`](../../docs/serving/streaming_video_output_api.md): `session.start` begins one rollout and each AR block arrives as a binary video -chunk. The model is image-conditioned, so the first frame is required and rides -on `image_reference` as an `http(s)` or `data:` URL; per-chunk camera actions -ride on `extra_params`: +chunk. The model is image-conditioned, so the first frame is required +at `image_reference` as an `http(s)` or `data:` URL; an optional request-scoped +pre-scripted list of WASD actions can be provided in `extra_params`: ```json {"type": "session.start", "model": "robbyant/lingbot-world-v2-14b-causal-fast-diffusers", @@ -134,7 +133,26 @@ ride on `extra_params`: chunk, and a request generates `((num_frames - 1) // 4 + 1) // 3` chunks: three for `num_frames: 33`, seven for `num_frames: 81`. -The bundled client speaks this protocol: +Mid-session camera control uses structural SE3 on `session.interaction` +(Unity frame: `+X` right, `+Y` up, `+Z` forward). WASD key tokens belong in +clients and are required to be converted before submitting to the service: + +```json +{"type": "session.interaction", + "interaction": { + "event_id": "forward-1", + "event": { + "multi_modal_data": { + "camera": { + "mode": "velocity", + "data": {"translation": [0.0, 0.0, 0.05], "rotation": [0.0, 0.0, 0.0, 1.0]} + } + } + } + }} +``` + +The bundled client helps translate the protocol, converting `--camera-updates` keystrokes to corresponding SE3 matrices: ```bash python examples/online_serving/streaming_video_generation/streaming_video_client.py \ @@ -142,23 +160,21 @@ python examples/online_serving/streaming_video_generation/streaming_video_client --prompt "The camera moves slowly forward through the scene." \ --image-reference /path/to/first_frame.png \ --width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \ - --extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \ + --camera-updates '[{"at": 1.0, "actions": ["w"]}, {"at": 3.0, "actions": ["a"]}]' \ --output lingbot_world_v2_stream.mp4 ``` A served request may instead point at a pose/intrinsics trajectory with `extra_params.action_path`, which is resolved inside the trusted root set by `model_config.lingbot_action_root` or `VLLM_OMNI_LINGBOT_ACTION_ROOT`; a server -started without that root configured accepts only `camera_action_script`. +started without that root configured accepts only `camera_action_script` or +mid-session camera interaction. Requested `width`/`height` must match `ar_diffusion_width`/`ar_diffusion_height` in the deploy config, because the AR cache geometry is fixed at load time. Blocks are decoded independently, so seams between chunks are possible; a session-owned streaming decoder is tracked separately. -Mid-session `session.interaction` for camera control is not wired yet, so a -served rollout follows the script it started with. - The deprecated tick example remains available for the older one-block-per-`generate()` control plane. @@ -199,8 +215,8 @@ tested commit. - Only the 14B causal-fast checkpoint is supported. - The tick control plane is internal; the public transport is the stepwise - `WS /v1/realtime/video` path, which cannot take mid-session camera - interaction yet. + `WS /v1/realtime/video` path with structural SE3 mid-session camera + interaction (WASD remains a client-side or `camera_action_script` convenience). - Stepwise serving requires an explicit `--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml`; there is no default deploy config for this model. diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py index e26e44c2c..c0e8ce6a7 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py +++ b/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"), diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py index 9cb67a142..98da24c7a 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py +++ b/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: diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index 9d5e958c5..0b4c77ebb 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/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() diff --git a/tests/diffusion/streaming_interactive_generation/test_camera_interaction.py b/tests/diffusion/streaming_interactive_generation/test_camera_interaction.py new file mode 100644 index 000000000..2c61aa805 --- /dev/null +++ b/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 diff --git a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py index 7780f0475..d3b46a00d 100644 --- a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py +++ b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py @@ -78,7 +78,7 @@ def _make_diffusion_request_state(*, request_id: str = "req-1", fps: int | None prompt="hello", ) state.prompt_embeds = torch.zeros(1, 4, 2) - state.extra = {"window_num_frames": 8} + state.extra = {"window_num_frames": 8, "num_latent_frames_per_chunk": 3} return state @@ -109,7 +109,12 @@ def _make_diffusion_model_runner( # instances; unit tests of submit_interaction stub this. Stepwise path tests # rebind to the real method. runner._supports_step_mode = lambda: True - runner._interaction_coordinator = None + # Mirror load_model: coordinator is always present once the runner is "loaded". + # Build now so handlers capture the pipeline's current encode_prompt (tests may + # rebind the mock before calling this helper). + runner._interaction_coordinator = InteractionCoordinator.build(pipeline, runner.od_config) + if hasattr(pipeline, "_interaction_coordinator"): + pipeline._interaction_coordinator = runner._interaction_coordinator return runner @@ -155,6 +160,7 @@ class TestPromptUpdateExecution: supports_step_execution = True runner = _make_diffusion_model_runner(pipeline=_UnsupportedPipeline()) + runner.state_cache["req-1"] = _make_diffusion_request_state() with pytest.raises(ValueError, match="not supported"): runner.submit_interaction("req-1", _prompt_interaction()) @@ -168,25 +174,98 @@ class TestPromptUpdateExecution: "interaction", [ {}, - {"multi_modal_data": {"camera": {"type": "pose"}}}, - {"event": {"prompt": "updated", "multi_modal_data": {"camera": {"type": "pose"}}}}, - {"event": {"prompt": "updated", "multi_modal_data": {}}}, - {"event": {"prompt": "updated", "multi_modal_data": None}}, - {"event": {"prompt": "updated", "multi_modal_data": "bad"}}, + {"event_id": "bad", "event": {"multi_modal_data": {}}}, + {"event_id": "bad", "event": {"multi_modal_data": None}}, + {"event_id": "bad", "event": {"prompt": "updated", "multi_modal_data": "bad"}}, + { + "event_id": "cam-only", + "event": { + "multi_modal_data": { + "camera": { + "mode": "velocity", + "data": {"translation": [0.0, 0.0, 0.05], "rotation": [0.0, 0.0, 0.0, 1.0]}, + }, + } + }, + }, + { + "event_id": "cam-and-prompt", + "event": { + "prompt": "updated", + "multi_modal_data": { + "camera": { + "mode": "velocity", + "data": {"translation": [0.0, 0.0, 0.05], "rotation": [0.0, 0.0, 0.0, 1.0]}, + }, + }, + }, + }, ], ) - def test_runner_interaction_rejects_structural_payloads_until_implemented( + def test_runner_rejects_unsupported_interactions( self, pipeline: HeliosPipeline, interaction: dict[str, Any], ) -> None: - """Unsupported interaction dict shapes are preserved to and rejected by the runner.""" + """Helios has no camera handler; malformed / unsupported payloads are rejected.""" runner = _make_diffusion_model_runner(pipeline=pipeline) - runner.state_cache["req-1"] = _make_diffusion_request_state() + state = _make_diffusion_request_state() + runner.state_cache["req-1"] = state - with pytest.raises(NotImplementedError, match="Only text-only prompt update interactions"): + with pytest.raises((ValueError, KeyError, AttributeError, TypeError)): runner.submit_interaction("req-1", cast(Any, interaction)) + # Rejected composite/malformed events must leave interaction queues unchanged. + pipeline.encode_prompt.assert_not_called() # pyright: ignore[reportAttributeAccessIssue] + assert "prompt" not in state.interaction_sessions + assert "camera" not in state.interaction_sessions + + def test_rejected_composite_leaves_existing_prompt_queue_unchanged( + self, + pipeline: HeliosPipeline, + ) -> None: + """Helios must not apply a composite prompt when camera modality is unsupported.""" + runner = _make_diffusion_model_runner(pipeline=pipeline) + state = _make_diffusion_request_state() + runner.state_cache["req-1"] = state + runner.submit_interaction("req-1", _prompt_interaction("already-queued", event_id="prior")) + session = state.interaction_sessions["prompt"] + assert isinstance(session, PromptSession) + prior = session.pending_event + assert prior is not None + assert prior.event_id == "prior" + pipeline.encode_prompt.reset_mock() # pyright: ignore[reportAttributeAccessIssue] + + with pytest.raises(ValueError, match="camera"): + runner.submit_interaction( + "req-1", + cast( + Any, + { + "event_id": "cam-and-prompt", + "event": { + "prompt": "should-not-replace", + "multi_modal_data": { + "camera": { + "mode": "velocity", + "data": { + "translation": [0.0, 0.0, 0.05], + "rotation": [0.0, 0.0, 0.0, 1.0], + }, + }, + }, + }, + }, + ), + ) + + pipeline.encode_prompt.assert_not_called() # pyright: ignore[reportAttributeAccessIssue] + pending = 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 + @pytest.mark.parametrize("model_class_name", ["HeliosPipeline", "HeliosPyramidPipeline"]) def test_coordinator_registers_prompt_for_helios_aliases( self, @@ -367,7 +446,11 @@ class TestPromptUpdateExecution: f"that use the chunk media timeline, got {fps!r} " f"(request_id={state.request_id!r})" ) - return ChunkMediaSpec(num_frames=int(state.extra["window_num_frames"]), fps=float(fps)) + return ChunkMediaSpec( + num_media_frames=int(state.extra["window_num_frames"]), + fps=float(fps), + num_latent_frames=int(state.extra["window_num_frames"]), + ) def prepare_encode(self, state: StepRequestState) -> StepRequestState: state.prompt_embeds = torch.zeros(1, 4, 2) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 28bba6e4d..de4c28a1a 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -780,6 +780,65 @@ class TestRunner: assert runner.pipeline.denoise_calls == 0 assert runner.state_cache == {} + def test_peer_prepare_encode_failure_skips_interaction_session_initialization(self, monkeypatch): + """One rank's failure in ``_prepare_batch_inputs -> prepare_encode`` must be handled before all ranks initialize interaction sessions. + + ``maybe_prepare_initial_session`` calls ``synchronized_monotonic_time()`` + (a broadcast). If this rank's ``prepare_encode`` succeeded while a peer + failed, entering that broadcast while the peer enters the failure + all-reduce would hang. + """ + from unittest.mock import MagicMock + + from vllm_omni.diffusion.interaction.types import ChunkMediaSpec, InteractionChunkMetadata + + class _InteractionStepPipeline(_StepPipeline): + def __init__(self): + super().__init__() + self.prepare_next_chunk_calls = 0 + + def peek_chunk_media(self, state): + del state + return ChunkMediaSpec(num_media_frames=8, fps=16.0, num_latent_frames=8) + + def apply_interaction_at_chunk_boundary(self, state): + del state + + def prepare_next_chunk(self, state): + del state + self.prepare_next_chunk_calls += 1 + + runner = _make_runner() + pipeline = _InteractionStepPipeline() + runner.pipeline = pipeline + coordinator = MagicMock() + coordinator.maybe_prepare_initial_session.return_value = InteractionChunkMetadata( + started_event_ids=[], + active_event_ids=[], + completed_event_ids=[], + ) + runner._interaction_coordinator = coordinator + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + # Peer failed during prepare_encode; this rank's local encode succeeds. + monkeypatch.setattr( + model_runner_module, + "_dit_any_rank_failed", + lambda local_failed: True, + ) + + req = _make_step_request(num_inference_steps=1) + result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req)) + + output = result.get_request_output("req-1") + assert output.finished is True + assert output.result is not None + assert "another DiT rank" in (output.result.error or "") + assert pipeline.prepare_calls == 1 + coordinator.maybe_prepare_initial_session.assert_not_called() + assert pipeline.prepare_next_chunk_calls == 0 + assert pipeline.denoise_calls == 0 + assert runner.state_cache == {} + def test_receives_kv_payload_before_prepare_encode(self, monkeypatch): runner = _make_runner() captured: dict[str, object] = {} diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index 291eada75..d8eccb6e4 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -959,6 +959,76 @@ class TestStreamingVideoOutputPromptUpdate: }, ) + def test_composite_prompt_and_camera_interaction_forwarded(self, mocker: MockerFixture): + """Prompt + camera multi_modal_data is admitted and forwarded to the engine.""" + + async def mock_generate(*_args, **_kwargs): + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=False, + ) + await asyncio.sleep(0.3) + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=True, + ) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)], + mock_generate=mock_generate, + ) + + camera_payload = { + "mode": "target", + "data": {"translation": [0.0, 0.0, 1.0]}, + } + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "initial scene"}) + start = ws.receive_json() + assert start["type"] == "video.start" + request_id = start["request_id"] + + assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0" + + ws.send_json( + { + "type": "session.interaction", + "interaction": { + "event_id": "combo-1", + "event": { + "prompt": "new scene", + "multi_modal_data": {"camera": camera_payload}, + }, + "transition_chunks": 2, + }, + } + ) + queued = ws.receive_json() + assert queued["type"] == "session.interaction.queued" + assert queued["event_id"] == "combo-1" + + assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1" + assert ws.receive_json()["type"] == "session.done" + + engine_client.submit_interaction_async.assert_awaited_once_with( + request_id, + interaction={ + "event_id": "combo-1", + "event": { + "prompt": "new scene", + "multi_modal_data": {"camera": camera_payload}, + }, + "transition_chunks": 2, + }, + ) + class TestStreamingVideoOutputExtraParams: """extra_params handling when building streaming generation inputs.""" diff --git a/vllm_omni/diffusion/interaction/coordinator.py b/vllm_omni/diffusion/interaction/coordinator.py index 5b4986d88..d99f3e47a 100644 --- a/vllm_omni/diffusion/interaction/coordinator.py +++ b/vllm_omni/diffusion/interaction/coordinator.py @@ -11,8 +11,9 @@ from vllm_omni.diffusion.interaction.types import ( InteractionChunkMetadata, InteractionPayload, merge_interaction_metadata, + synchronized_monotonic_time, ) -from vllm_omni.diffusion.models.interface import supports_interaction_apply +from vllm_omni.diffusion.models.interface import SupportsInteractionApply, supports_interaction_apply from vllm_omni.diffusion.worker.utils import StepRequestState if TYPE_CHECKING: @@ -52,7 +53,7 @@ class InteractionCoordinator: @property def needs_chunk_media(self) -> bool: - """Whether any registered handler needs chunk ``num_frames``/``fps``.""" + """Whether any registered handler needs chunk media/latent counts + fps.""" return any(handler.needs_chunk_media for handler in self._handlers.values()) def get_handler(self, modality: str) -> InteractionHandler: @@ -83,31 +84,109 @@ class InteractionCoordinator: transition_chunks=transition_chunks, ) + def enqueue_parts( + self, + state: StepRequestState, + *, + parts: list[tuple[str, InteractionPayload]], + event_id: str, + received_at: float, + transition_chunks: int | None, + ) -> None: + """Ensure every modality is supported and every payload validates, then enqueue. + + Composite events must not partially mutate queues when a later track is + unsupported or carries an invalid payload. + """ + if not parts: + raise ValueError("interaction event requires prompt and/or multi_modal_data") + + # Reconcile rank-local samples so USP/SP shards share one arrival time. + received_at = synchronized_monotonic_time(received_at) + + handlers = [(self.get_handler(modality), payload) for modality, payload in parts] + for handler, payload in handlers: + handler.validate_payload( + state, + event_id=event_id, + payload=payload, + transition_chunks=transition_chunks, + ) + for (modality, payload), (handler, _) in zip(parts, handlers, strict=True): + del handler + self.enqueue( + state, + modality=modality, + event_id=event_id, + received_at=received_at, + payload=payload, + transition_chunks=transition_chunks, + ) + + def maybe_prepare_initial_session( + self, + state: StepRequestState, + pipeline: SupportsInteractionApply, + ) -> InteractionChunkMetadata: + """Materialize non-lazy modality sessions before chunk 0 denoise. + + Handlers with ``lazy_initialize_session=False`` (e.g. camera) must exist + before the first chunk even with no client enqueue. + """ + num_media_frames: int | None = None + fps: float | None = None + num_latent_frames: int | None = None + if any( + (not handler.lazy_initialize_session) and handler.needs_chunk_media for handler in self._handlers.values() + ): + media = pipeline.peek_chunk_media(state) + num_media_frames = media.num_media_frames + fps = media.fps + num_latent_frames = media.num_latent_frames + + boundary_at = synchronized_monotonic_time() + chunk_index = state.chunk_index + metas: list[InteractionChunkMetadata] = [] + for handler in self._handlers_in_apply_order(): + if handler.lazy_initialize_session: + continue + meta = handler.apply_at_chunk_boundary( + state, + boundary_at=boundary_at, + chunk_index=chunk_index, + num_media_frames=num_media_frames, + fps=fps, + num_latent_frames=num_latent_frames, + ) + if meta is not None: + metas.append(meta) + return merge_interaction_metadata(metas) + def apply_at_chunk_boundary( self, state: StepRequestState, *, boundary_at: float, chunk_index: int | None = None, - num_frames: int | None = None, + num_media_frames: int | None = None, fps: float | None = None, + num_latent_frames: int | None = None, ) -> InteractionChunkMetadata: """Fan out chunk-boundary apply to handlers in stable order (prompt first).""" if chunk_index is None: chunk_index = state.chunk_index metas: list[InteractionChunkMetadata] = [] for handler in self._handlers_in_apply_order(): - if handler.modality not in state.interaction_sessions: - # This function is always called at chunk boundary without knowing if an interaction is really enqueued. - # `state.interaction_sessions[this modality]` is lazy-populated after the first interaction is enqueued. - # So, check here and skip a modality handler if this modality's session is not yet created. + if handler.lazy_initialize_session and handler.modality not in state.interaction_sessions: + # Lazy modalities skip apply until the first enqueue creates the session. continue meta = handler.apply_at_chunk_boundary( state, boundary_at=boundary_at, chunk_index=chunk_index, - num_frames=num_frames, + num_media_frames=num_media_frames, fps=fps, + num_latent_frames=num_latent_frames, ) if meta is not None: metas.append(meta) diff --git a/vllm_omni/diffusion/interaction/mixin.py b/vllm_omni/diffusion/interaction/mixin.py index 366e90588..8627ec591 100644 --- a/vllm_omni/diffusion/interaction/mixin.py +++ b/vllm_omni/diffusion/interaction/mixin.py @@ -4,9 +4,9 @@ from __future__ import annotations -import time from typing import TYPE_CHECKING, cast +from vllm_omni.diffusion.interaction.types import synchronized_monotonic_time from vllm_omni.diffusion.models.interface import SupportsInteractionApply from vllm_omni.diffusion.worker.utils import StepRequestState @@ -35,18 +35,21 @@ class InteractionMixin: "onto the pipeline before chunked generation" ) - num_frames: int | None = None + num_media_frames: int | None = None fps: float | None = None + num_latent_frames: int | None = None if self._interaction_coordinator.needs_chunk_media: media = cast(SupportsInteractionApply, self).peek_chunk_media(state) - num_frames = media.num_frames + num_media_frames = media.num_media_frames fps = media.fps + num_latent_frames = media.num_latent_frames merged = self._interaction_coordinator.apply_at_chunk_boundary( state, - boundary_at=time.monotonic(), - num_frames=num_frames, + boundary_at=synchronized_monotonic_time(), + num_media_frames=num_media_frames, fps=fps, + num_latent_frames=num_latent_frames, ) state.interaction_chunk_metadata = merged diff --git a/vllm_omni/diffusion/interaction/modality_handlers/base.py b/vllm_omni/diffusion/interaction/modality_handlers/base.py index 61f548fcd..9f8ccfe4e 100644 --- a/vllm_omni/diffusion/interaction/modality_handlers/base.py +++ b/vllm_omni/diffusion/interaction/modality_handlers/base.py @@ -5,7 +5,9 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import ClassVar, Self +from typing import ClassVar + +from typing_extensions import Self from vllm_omni.diffusion.interaction.types import ( InteractionChunkMetadata, @@ -23,9 +25,14 @@ class InteractionHandler(ABC): """ modality: ClassVar[str] - # When True, chunk-boundary apply needs ``ChunkMediaSpec`` (num_frames/fps) + # When True, chunk-boundary apply needs ``ChunkMediaSpec`` (frame count & fps) # Those information are useful when interaction handler needs interpolation/integration on a frame-by-frame basis needs_chunk_media: ClassVar[bool] = False + # When True, the coordinator skips ``apply_at_chunk_boundary`` until a session + # exists (normally created on the first enqueue). When False, every chunk + # boundary runs apply even with no prior enqueue so the handler can create + # its session and materialize default chunk data (e.g. identity camera hold). + lazy_initialize_session: ClassVar[bool] = True @classmethod def from_pipeline(cls, pipeline: object) -> Self: @@ -37,6 +44,21 @@ class InteractionHandler(ABC): del pipeline return cls() + @abstractmethod + def validate_payload( + self, + state: StepRequestState, + *, + event_id: str, + payload: InteractionPayload, + transition_chunks: int | None, + ) -> None: + """Validate a modality payload without mutating request state. + + Used by ``InteractionCoordinator.enqueue_parts`` so composite events + fail before any track is queued. + """ + @abstractmethod def enqueue( self, @@ -56,7 +78,9 @@ class InteractionHandler(ABC): *, boundary_at: float, chunk_index: int | None = None, # defaults to state.chunk_index when omitted - num_frames: int | None = None, # only set when at least one interaction handler needs it - fps: float | None = None, # only set when at least one interaction handler needs it + # Optional chunk media information. Only present when at least one initialized interaction handler needs it + num_media_frames: int | None = None, + fps: float | None = None, + num_latent_frames: int | None = None, ) -> InteractionChunkMetadata | None: """Advance request-local state and materialize this chunk's effects.""" diff --git a/vllm_omni/diffusion/interaction/modality_handlers/camera.py b/vllm_omni/diffusion/interaction/modality_handlers/camera.py new file mode 100644 index 000000000..df1c9af1d --- /dev/null +++ b/vllm_omni/diffusion/interaction/modality_handlers/camera.py @@ -0,0 +1,490 @@ +# 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", +] diff --git a/vllm_omni/diffusion/interaction/modality_handlers/prompt.py b/vllm_omni/diffusion/interaction/modality_handlers/prompt.py index e87d9029b..818abe173 100644 --- a/vllm_omni/diffusion/interaction/modality_handlers/prompt.py +++ b/vllm_omni/diffusion/interaction/modality_handlers/prompt.py @@ -90,6 +90,7 @@ class PromptInteractionHandler(InteractionHandler): """ modality: ClassVar[str] = "prompt" + lazy_initialize_session: ClassVar[bool] = True def __init__( self, @@ -113,16 +114,14 @@ class PromptInteractionHandler(InteractionHandler): ) @override - def enqueue( + def validate_payload( self, state: StepRequestState, *, event_id: str, - received_at: float, payload: InteractionPayload, transition_chunks: int | None, ) -> None: - """Prompt updates are last-write-win and unbuffered at chunk boundary.""" prompt = payload.get("prompt") if not isinstance(prompt, str) or not prompt: raise ValueError("prompt must be non-empty") @@ -136,6 +135,27 @@ class PromptInteractionHandler(InteractionHandler): if duration < 0: raise ValueError("transition_chunks must be >= 0") + @override + def enqueue( + self, + state: StepRequestState, + *, + event_id: str, + received_at: float, + payload: InteractionPayload, + transition_chunks: int | None, + ) -> None: + """Prompt updates are last-write-win and unbuffered at chunk boundary.""" + self.validate_payload( + state, + event_id=event_id, + payload=payload, + transition_chunks=transition_chunks, + ) + prompt = payload.get("prompt") + assert isinstance(prompt, str) + duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks + target_prompt_embeds, _ = self._encode_prompt( prompt=prompt, negative_prompt=None, @@ -164,12 +184,13 @@ class PromptInteractionHandler(InteractionHandler): *, boundary_at: float, chunk_index: int | None = None, - num_frames: int | None = None, + num_media_frames: int | None = None, fps: float | None = None, + num_latent_frames: int | None = None, ) -> InteractionChunkMetadata | None: """Advance or start prompt interpolation before the next chunk.""" # Prompt lerp is chunk-LWW; media timeline and caller chunk_index unused. - del chunk_index, num_frames, fps + del chunk_index, num_media_frames, fps, num_latent_frames session = state.interaction_sessions.get("prompt") assert isinstance(session, PromptSession) diff --git a/vllm_omni/diffusion/interaction/registry.py b/vllm_omni/diffusion/interaction/registry.py index f86dab902..9e8d3c9af 100644 --- a/vllm_omni/diffusion/interaction/registry.py +++ b/vllm_omni/diffusion/interaction/registry.py @@ -5,12 +5,13 @@ A pipeline only declares the modalities that it supports. And a pipeline can choose which handler for each modality (there can be more than one handler for a modality). Outer key: pipeline architecture name (``od_config.model_class_name``). -Inner key: modality name (e.g. ``\"prompt\"``). +Inner key: modality name (e.g. ``\"camera\"``). """ from __future__ import annotations from vllm_omni.diffusion.interaction.modality_handlers.base import InteractionHandler +from vllm_omni.diffusion.interaction.modality_handlers.camera import SE3DeltaCameraHandler from vllm_omni.diffusion.interaction.modality_handlers.prompt import PromptInteractionHandler STRUCTURED_HANDLER_REGISTRY: dict[str, dict[str, type[InteractionHandler]]] = { @@ -21,4 +22,7 @@ STRUCTURED_HANDLER_REGISTRY: dict[str, dict[str, type[InteractionHandler]]] = { "HeliosPyramidPipeline": { "prompt": PromptInteractionHandler, }, + "LingBotWorldCausalDMDPipeline": { + "camera": SE3DeltaCameraHandler, + }, } diff --git a/vllm_omni/diffusion/interaction/types.py b/vllm_omni/diffusion/interaction/types.py index 904552129..942940102 100644 --- a/vllm_omni/diffusion/interaction/types.py +++ b/vllm_omni/diffusion/interaction/types.py @@ -1,38 +1,75 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared types for midway interaction handlers.""" from __future__ import annotations +import math +import time from collections.abc import Mapping from dataclasses import dataclass, field from threading import Lock from typing import Literal +import torch + # Wire payload for one modality track. Handlers validate keys they need. InteractionPayload = Mapping[str, object] # Shared command semantics across modalities that support them. -# Prompt specialization uses ``"target"`` only. -InteractionMode = Literal["target"] +# Prompt specialization narrows this to ``"target"`` only. +InteractionMode = Literal["target", "velocity"] + + +def synchronized_monotonic_time(stamp: float | None = None) -> float: + """Return a monotonic timestamp shared across DiT ranks. + + ``submit_interaction`` and chunk-boundary apply run as collectives, but + ``time.monotonic()`` is rank-local. Under USP/SP, inter-rank skew that + straddles a ``1/fps`` bucket edge would schedule the same event onto + different frames per shard. Rank 0's stamp is broadcast (callers may pass + a local sample; non-source ranks' values are ignored). + """ + if stamp is None: + stamp = time.monotonic() + if not torch.distributed.is_initialized(): + return stamp + try: + from vllm_omni.diffusion.distributed.parallel_state import get_world_group + + group = get_world_group() + except Exception: + return stamp + if group.world_size <= 1: + return stamp + return group.broadcast_object(stamp, src=0) @dataclass(frozen=True) class ChunkMediaSpec: - """Media extent of one generation chunk, for interaction timeline mapping.""" + """Decoded media extent and latent step count for one generation chunk. - num_frames: int + Attributes: + num_media_frames: Number of frames in this chunk when decoded into playable media. + fps: Frames per second of the media. + num_latent_frames: Number of latent frames in this chunk (internal DiT units). + """ + + num_media_frames: int fps: float + num_latent_frames: int def __post_init__(self) -> None: - if self.num_frames <= 0: - raise ValueError(f"ChunkMediaSpec.num_frames must be > 0, got {self.num_frames}") + if self.num_media_frames <= 0: + raise ValueError(f"ChunkMediaSpec.num_media_frames must be > 0, got {self.num_media_frames}") + if self.num_latent_frames <= 0: + raise ValueError(f"ChunkMediaSpec.num_latent_frames must be > 0, got {self.num_latent_frames}") if self.fps <= 0: raise ValueError(f"ChunkMediaSpec.fps must be > 0, got {self.fps}") @property def duration_s(self) -> float: - return float(self.num_frames) / float(self.fps) + return float(self.num_media_frames) / float(self.fps) @dataclass(kw_only=True) @@ -41,6 +78,7 @@ class InteractionEvent: Example: - ``QueuedPromptEvent`` in ``vllm_omni.diffusion.interaction.modality_handlers.prompt`` + - ``QueuedCameraEvent`` in ``vllm_omni.diffusion.interaction.modality_handlers.camera`` """ event_id: str @@ -60,12 +98,37 @@ class InteractionSession: Example: - ``PromptSession`` in ``vllm_omni.diffusion.interaction.modality_handlers.prompt`` + - ``CameraSession`` in ``vllm_omni.diffusion.interaction.modality_handlers.camera`` """ lock: Lock = field(default_factory=Lock, repr=False) last_boundary_at: float | None = None +def resolve_event_frame_offset( + *, + received_at: float, + previous_boundary_at: float | None, + num_frames: int, + fps: float, +) -> int: + """Map an event arrival time onto ``[0, num_frames - 1]`` for this chunk. + + * Uses ``floor((received_at - previous_boundary_at) * fps)``. + * Events without a prior boundary (or with non-positive fps) map to frame 0. + * Offsets past the represented media window are clamped to the final frame. + """ + num_frames = max(int(num_frames), 1) + if previous_boundary_at is None or fps <= 0: + return 0 + offset = math.floor((received_at - previous_boundary_at) * float(fps)) + if offset < 0: + return 0 + if offset >= num_frames: + return num_frames - 1 + return int(offset) + + @dataclass class InteractionChunkMetadata: """Per-modality (or merged) event-id bookkeeping for one chunk boundary.""" diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 610fb2b8f..ebea371a4 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -535,19 +535,23 @@ class HeliosPipeline( "zero_steps": int(extra.get("zero_steps", 1)), } ) - self.prepare_next_chunk(state) return state def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec: """Expose this chunk's decoded media extent for interaction timelines.""" - num_frames = int(state.extra["window_num_frames"]) + num_media_frames = int(state.extra["window_num_frames"]) + num_latent_frames = int(state.extra["num_latent_frames_per_chunk"]) fps = state.sampling.fps if fps is None or float(fps) <= 0: raise ValueError( "sampling.fps is required and must be > 0 for interaction modalities that use the " f"chunk media timeline, got {fps!r} (request_id={state.request_id!r})" ) - return ChunkMediaSpec(num_frames=num_frames, fps=float(fps)) + return ChunkMediaSpec( + num_media_frames=num_media_frames, + fps=float(fps), + num_latent_frames=num_latent_frames, + ) @override def prepare_next_chunk(self, state: StepRequestState) -> None: diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index 92216d19e..57d7fe7a7 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -151,10 +151,10 @@ class SupportsInteractionApply(Protocol): """Optional protocol for pipelines with unified mid-generation, chunk-boundary hooks.""" def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec: - """Return the media timeline represented by the upcoming/current chunk. + """Return the media timeline and latent frame count for the upcoming/current chunk. Useful when interaction handler needs interpolation/integration on a frame-by-frame basis, - or for backpressure/pacing. + and/or for backpressure/pacing. """ ... diff --git a/vllm_omni/diffusion/models/lingbot_world/actions.py b/vllm_omni/diffusion/models/lingbot_world/actions.py index 4720b041e..057341c84 100644 --- a/vllm_omni/diffusion/models/lingbot_world/actions.py +++ b/vllm_omni/diffusion/models/lingbot_world/actions.py @@ -25,6 +25,7 @@ from vllm_omni.experimental.ar_diffusion.tick_protocol import ( LINGBOT_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1" LINGBOT_CAMERA_TRAJECTORY_SCHEMA = "lingbot.camera_trajectory.v1" +LINGBOT_CONTROLLER_TRANSLATION_UNIT = 0.05 _ACTION_ORDER = ("w", "a", "s", "d", "i", "j", "k", "l") _VALID_ACTIONS = frozenset(_ACTION_ORDER) _REFERENCE_HEIGHT = 480 @@ -282,7 +283,7 @@ def integrate_lingbot_camera_actions( """Convert latent-frame WASD/IJKL controls to cumulative C2W poses. Calibration and motion constants intentionally match SGLang's LingBot - adapter: movement 0.05, pitch 4 degrees, yaw 6 degrees, pitch clamp 85 + adapter: movement 0.05 (LINGBOT_CONTROLLER_TRANSLATION_UNIT), pitch 4 degrees, yaw 6 degrees, pitch clamp 85 degrees, and target-resolution focal lengths of 500 pixels. """ @@ -325,13 +326,31 @@ def integrate_lingbot_camera_actions( right /= right_norm + 1e-6 movement = np.zeros(3, dtype=np.float64) - movement += forward * 0.05 * (("w" in actions) - ("s" in actions)) - movement += right * 0.05 * (("d" in actions) - ("a" in actions)) + movement += forward * LINGBOT_CONTROLLER_TRANSLATION_UNIT * (("w" in actions) - ("s" in actions)) + movement += right * LINGBOT_CONTROLLER_TRANSLATION_UNIT * (("d" in actions) - ("a" in actions)) + # Keep the cumulative controller pose in FP64 between chunks, matching + # SGLang's NumPy integration and avoiding long-session drift. current_pose = np.eye(4, dtype=np.float64) current_pose[:3, :3] = new_rotation current_pose[:3, 3] = translation + movement poses.append(current_pose) + trajectory = camera_trajectory_from_absolute_pose(torch.from_numpy(np.stack(poses)), width=width, height=height) + return trajectory, current_pitch + + +def camera_trajectory_from_absolute_pose( + poses: torch.Tensor, + *, + width: int, + height: int, +) -> CameraTrajectory: + if width <= 0 or height <= 0: + raise ValueError("camera action resolution must be positive.") + if poses.ndim != 3 or poses.shape[-2:] != (4, 4): + raise ValueError(f"absolute camera poses must have shape [frames, 4, 4], got {tuple(poses.shape)}.") + if poses.shape[0] == 0: + raise ValueError("absolute camera poses must not be empty.") # build_plucker_embedding expects intrinsics in the 832x480 reference # coordinate system. These values become SGLang's [500, 500, W/2, H/2] # after that function scales them to the requested resolution. @@ -342,9 +361,7 @@ def integrate_lingbot_camera_actions( _REFERENCE_HEIGHT / 2, ) trajectory = CameraTrajectory( - # Keep the cumulative controller pose in FP64 between chunks, matching - # SGLang's NumPy integration and avoiding long-session drift. - poses=torch.from_numpy(np.stack(poses)), + poses=poses, intrinsics=torch.tensor(reference_intrinsics, dtype=torch.float32).repeat(len(poses), 1), ) - return trajectory, current_pitch + return trajectory diff --git a/vllm_omni/diffusion/models/lingbot_world/camera.py b/vllm_omni/diffusion/models/lingbot_world/camera.py index 504d850c7..45e3add44 100644 --- a/vllm_omni/diffusion/models/lingbot_world/camera.py +++ b/vllm_omni/diffusion/models/lingbot_world/camera.py @@ -389,7 +389,11 @@ def _invert_camera_poses(poses: torch.Tensor) -> torch.Tensor: return inverse -def _prepare_framewise_poses(raw_poses: torch.Tensor) -> torch.Tensor: +def _prepare_framewise_poses( + raw_poses: torch.Tensor, + *, + translation_scale: float | None = None, +) -> torch.Tensor: """Convert raw C2W poses to normalized first-relative framewise deltas.""" # Remove global placement, then convert the sequence to framewise deltas. @@ -402,12 +406,19 @@ def _prepare_framewise_poses(raw_poses: torch.Tensor) -> torch.Tensor: relative_poses[1:], ) - # Only relative direction/magnitude matters to the checkpoint; normalizing - # by the largest step removes the trajectory's arbitrary translation scale. + # Normalizing the control's scale. Different between scripted and realtime input. translations = relative_poses[:, :3, 3] - max_translation_norm = torch.linalg.vector_norm(translations, dim=-1).max() - if max_translation_norm > 0: - relative_poses[:, :3, 3] = translations / max_translation_norm + if translation_scale is None: + # Scripted action_path input: divide by the largest step so an arbitrary recording scale becomes 1. + scale = torch.linalg.vector_norm(translations, dim=-1).max() + else: + # Realtime interaction input: use the controller translation unit (one WASD step = full speed). + # It cannot see future motion or normalize globally. + if translation_scale <= 0: + raise ValueError(f"translation_scale must be > 0, got {translation_scale}") + scale = translation_scale + if float(scale) > 0: + relative_poses[:, :3, 3] = translations / scale return relative_poses @@ -420,6 +431,7 @@ def build_plucker_embedding( target_width: int, device: torch.device, dtype: torch.dtype, + translation_scale: float | None = None, ) -> torch.Tensor: """Build checkpoint-ordered ``(ray origin, ray direction)`` channels.""" @@ -438,7 +450,10 @@ def build_plucker_embedding( # Compute geometry in at least FP32, then cast at the model boundary. compute_dtype = torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype - poses = _prepare_framewise_poses(trajectory.poses.to(device=device, dtype=compute_dtype)) + poses = _prepare_framewise_poses( + trajectory.poses.to(device=device, dtype=compute_dtype), + translation_scale=translation_scale, + ) intrinsics = trajectory.intrinsics.to(device=device, dtype=compute_dtype).clone() intrinsics[:, (0, 2)] *= width / _REFERENCE_WIDTH intrinsics[:, (1, 3)] *= height / _REFERENCE_HEIGHT diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index 8e8d5dea4..8b0f00736 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -26,16 +26,21 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.interaction.mixin import InteractionMixin +from vllm_omni.diffusion.interaction.modality_handlers.camera import CameraSession +from vllm_omni.diffusion.interaction.types import ChunkMediaSpec from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery, SupportsStepExecution from vllm_omni.diffusion.models.lingbot_world.actions import ( LINGBOT_CAMERA_ACTION_SCHEMA, LINGBOT_CAMERA_TRAJECTORY_SCHEMA, + LINGBOT_CONTROLLER_TRANSLATION_UNIT, LingBotCameraActionFrames, LingBotCameraActionScript, as_camera_action_frames, as_camera_action_script, + camera_trajectory_from_absolute_pose, integrate_lingbot_camera_actions, parse_lingbot_camera_action_frames, parse_lingbot_camera_action_script, @@ -385,25 +390,36 @@ def get_lingbot_world_pre_process_func( camera_actions = None else: action_path = extra_args.get("action_path") - if not isinstance(action_path, (str, os.PathLike)) or not str(action_path): - raise ValueError("action_path is required in sampling_params.extra_args.action_path.") - if not configured_action_root: - raise ValueError( - "sampling_params.extra_args.action_path requires a trusted action root configured by " - f"model_config.lingbot_action_root or {_ACTION_ROOT_ENV}." + if action_path is None or action_path == "": + # Stepwise / streaming sessions may omit a request-scoped camera + # script and instead drive motion via mid-generation camera + # interaction (idle hold until the first camera event). + if getattr(od_config, "step_execution", False) or getattr(od_config, "streaming_output", False): + trajectory = None + camera_actions = None + camera_action_script = None + else: + raise ValueError("action_path is required in sampling_params.extra_args.action_path.") + else: + if not isinstance(action_path, (str, os.PathLike)) or not str(action_path): + raise ValueError("action_path is required in sampling_params.extra_args.action_path.") + if not configured_action_root: + raise ValueError( + "sampling_params.extra_args.action_path requires a trusted action root configured by " + f"model_config.lingbot_action_root or {_ACTION_ROOT_ENV}." + ) + action_directory = resolve_trusted_action_directory( + action_path, + configured_action_root, ) - action_directory = resolve_trusted_action_directory( - action_path, - configured_action_root, - ) - try: - trajectory = load_camera_trajectory(action_directory) - except OSError: - raise ValueError( - "Unable to load camera trajectory from action_path; expected poses.npy and intrinsics.npy." - ) from None - camera_actions = None - camera_action_script = None + try: + trajectory = load_camera_trajectory(action_directory) + except OSError: + raise ValueError( + "Unable to load camera trajectory from action_path; expected poses.npy and intrinsics.npy." + ) from None + camera_actions = None + camera_action_script = None updated_prompt = dict(prompt) updated_multi_modal_data = dict(multi_modal_data) @@ -460,6 +476,7 @@ class LingBotWorldCausalDMDPipeline( SupportImageInput, SupportsComponentDiscovery, SupportsStepExecution, + InteractionMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, ): @@ -779,8 +796,10 @@ class LingBotWorldCausalDMDPipeline( camera_sources = ( int(camera_trajectory is not None) + int(camera_actions is not None) + int(camera_action_script is not None) ) - if camera_sources != 1: - raise ValueError("LingBot pre-processing must materialize exactly one camera input.") + if camera_sources > 1: + raise ValueError("LingBot pre-processing must materialize at most one camera input.") + # Zero camera sources are allowed for stepwise camera-interaction sessions; + # request-mode ``forward`` still requires an explicit trajectory or actions. request_flow_shift = ( extra_args["flow_shift"] if "flow_shift" in extra_args else getattr(self.scheduler.config, "shift", 5.0) @@ -1047,7 +1066,7 @@ class LingBotWorldCausalDMDPipeline( *, dtype: torch.dtype, previous: CameraTrajectory | None = None, - latent_aligned: bool = False, + latent_aligned: bool = False, # iff using realtime interaction inputs ) -> tuple[torch.Tensor, CameraTrajectory]: """Convert raw camera frames to a latent-aligned ray tensor.""" @@ -1083,10 +1102,10 @@ class LingBotWorldCausalDMDPipeline( ), ) drop_anchor = True - elif inputs.camera_actions is not None: - # The action integrator returns post-action poses. Keep the first - # action visible to framewise-delta conditioning by prepending its - # known pre-action state (identity for a new realtime session). + elif inputs.camera_actions is not None or latent_aligned: + # Action/interaction integrators return post-action poses. Prepend the + # known pre-action state (identity for a new session) so the first + # pose still contributes a framewise delta. identity = torch.eye( 4, device=trajectory.poses.device, @@ -1100,6 +1119,8 @@ class LingBotWorldCausalDMDPipeline( ), ) drop_anchor = True + # Realtime interaction cannot see future steps; use the controller translation unit to normalize speed. + # Full action_path keeps max-norm. camera_embedding = build_plucker_embedding( embedding_trajectory, height=inputs.height, @@ -1108,6 +1129,7 @@ class LingBotWorldCausalDMDPipeline( target_width=inputs.width, device=self.device, dtype=dtype, + translation_scale=LINGBOT_CONTROLLER_TRANSLATION_UNIT if latent_aligned else None, ) if drop_anchor: camera_embedding = camera_embedding[1:] @@ -1277,6 +1299,11 @@ class LingBotWorldCausalDMDPipeline( "request mode takes one three-frame camera_actions control per " "block, or a camera_trajectory for the request." ) + if inputs.camera_trajectory is None and inputs.camera_actions is None: + raise ValueError( + "LingBot request mode requires camera_trajectory or camera_actions; " + "omit them only for stepwise camera-interaction sessions." + ) tick = ARDiffusionTickRequest.from_extra_args(req.sampling_params.extra_args) if tick is not None and self._ar_diffusion_kv_state is None: raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.") @@ -1645,11 +1672,6 @@ class LingBotWorldCausalDMDPipeline( } self._ar_text_caches(prompt_embeds, invalidate=False) self._ar_sessions[state.request_id] = _LingBotARSessionState() - try: - self._prepare_next_chunk(state) - except Exception: - self._release_condition_encoder_state(state.request_id) - raise return state def _prepare_next_chunk(self, state: StepRequestState) -> None: @@ -1668,7 +1690,14 @@ class LingBotWorldCausalDMDPipeline( dtype=extra["dtype"], ) previous = extra.get("camera_tail") + if extra.get("camera_action_script") is not None: + if ( + (camera_session := state.interaction_sessions.get("camera")) + and isinstance(camera_session, CameraSession) + and camera_session.has_received_input + ): + raise ValueError("Cannot use mid-generation camera interaction together with camera_action_script; ") chunk_actions = extra["camera_action_script"][state.chunk_index] action_trajectory, camera_pitch = integrate_lingbot_camera_actions( chunk_actions, @@ -1678,11 +1707,12 @@ class LingBotWorldCausalDMDPipeline( initial_pitch=float(extra.get("camera_pitch", 0.0)), ) extra["camera_pitch"] = camera_pitch + media_frames = self._chunk_media_frame_count(state, block_frames=block_frames) chunk_inputs = replace( inputs, camera_trajectory=action_trajectory, camera_actions=chunk_actions, - num_frames=(block_frames - 1) * self.vae_scale_factor_temporal + 1, + num_frames=media_frames, num_latent_frames=block_frames, ) camera, camera_tail = self._prepare_camera( @@ -1691,14 +1721,53 @@ class LingBotWorldCausalDMDPipeline( previous=previous, ) elif extra.get("camera_embedding_cache") is not None: - # Same slice request mode takes from its one full-trajectory - # embedding, so both paths condition a block identically. + if ( + (camera_session := state.interaction_sessions.get("camera")) + and isinstance(camera_session, CameraSession) + and camera_session.has_received_input + ): + raise ValueError("Cannot use mid-generation camera interaction together with camera_embedding_cache; ") trajectory = extra["camera_trajectory_cache"] camera = extra["camera_embedding_cache"][:, :, start_frame:stop_frame] camera_tail = CameraTrajectory( poses=trajectory.poses[stop_frame - 1 : stop_frame].clone(), intrinsics=trajectory.intrinsics[stop_frame - 1 : stop_frame].clone(), ) + elif self._interaction_coordinator is not None and self._interaction_coordinator.has_modality("camera"): + # Poses come from the latest boundary apply: ``prepare_encode`` for + # chunk 0, the diffusion runner for later chunks. This method only + # digests them — it does not apply interactions itself. + camera_session = state.interaction_sessions.get("camera") + if not isinstance(camera_session, CameraSession) or camera_session.last_absolute_poses is None: + raise RuntimeError( + "LingBot camera interaction requires apply_interaction_at_chunk_boundary before prepare_next_chunk." + ) + absolute_poses = camera_session.last_absolute_poses + media_frames = self._chunk_media_frame_count(state, block_frames=block_frames) + num_latent_frames = block_frames + if int(absolute_poses.shape[0]) != num_latent_frames: + raise ValueError( + "camera interaction must produce exactly one pose per latent frame; " + f"got {int(absolute_poses.shape[0])}, expected {num_latent_frames}." + ) + action_trajectory = camera_trajectory_from_absolute_pose( + absolute_poses, + width=inputs.width, + height=inputs.height, + ) + chunk_inputs = replace( + inputs, + camera_trajectory=action_trajectory, + camera_actions=None, + num_frames=media_frames, + num_latent_frames=num_latent_frames, + ) + camera, camera_tail = self._prepare_camera( + chunk_inputs, + dtype=extra["dtype"], + previous=previous, + latent_aligned=True, + ) else: raise RuntimeError("LingBot step execution is missing a camera trajectory or action script.") if camera.shape[2:] != condition.shape[2:]: @@ -1842,19 +1911,48 @@ class LingBotWorldCausalDMDPipeline( session_state.pending_encoder_cache = None session_state.next_chunk_index += 1 state.chunk_index += 1 - finished = state.request_denoise_completed - if not finished: - self._prepare_next_chunk(state) - else: - self._release_condition_encoder_state(state.request_id) return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, chunk_index=completed_chunk_index, total_chunks=state.total_chunks, - finished=finished, + finished=state.request_denoise_completed, ) + def _chunk_media_frame_count(self, state: StepRequestState, *, block_frames: int) -> int: + """Pixel frames this chunk will emit when decoded. + + Independent (non-streaming) decode restarts causal expansion every block, + so each chunk is ``(block_frames - 1) * temporal + 1`` frames. + Streaming decode expands only the session's opening latent to one frame and every later latent to the full + temporal factor, so chunk 0 matches the independent count and later chunks are ``block_frames * temporal``. + """ + decode_state = self._streaming_decode_states.get(state.request_id) + if decode_state is not None and decode_state.started: + return block_frames * self.vae_scale_factor_temporal + return (block_frames - 1) * self.vae_scale_factor_temporal + 1 + + def peek_chunk_media(self, state: StepRequestState) -> ChunkMediaSpec: + """Expose this chunk's decoded media extent and latent step count.""" + num_latent_frames = int(state.extra.get("block_frames") or self.transformer.config.num_frames_per_block) + num_media_frames = self._chunk_media_frame_count(state, block_frames=num_latent_frames) + fps = state.sampling.fps + if fps is None or float(fps) <= 0: + # Media-frame camera controls are not wall-clock paced; a unit fps keeps + # resolve_event_frame_offset well-defined when the client omits sampling.fps. + fps = float(num_media_frames) + return ChunkMediaSpec( + num_media_frames=num_media_frames, + fps=fps, + num_latent_frames=num_latent_frames, + ) + + def prepare_next_chunk(self, state: StepRequestState) -> None: + """Prepare the next AR block after chunk-boundary interaction apply.""" + if state.request_denoise_completed: + return + self._prepare_next_chunk(state) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return cast(set[str], loader.load_weights(weights)) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 6a5b81c58..d4d457bdc 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -43,6 +43,7 @@ from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_rank from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.interaction.coordinator import InteractionCoordinator +from vllm_omni.diffusion.interaction.types import InteractionPayload from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( SupportsInteractionApply, @@ -1080,30 +1081,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): error_outputs: list[RunnerOutput] = [] for state in states: if state.request_id in new_request_ids: - # Everything that runs before ``_dit_any_rank_failed`` must be - # inside the try: an exception in ``_initialize_generator`` or + # Everything that requires rank-synchronization must be called + # inside a try, record the exception and handle with `_dit_any_rank_failed`. + # Reason (example): An exception in ``_initialize_generator`` or # ``clear_pipeline_stage_durations`` on one rank would skip the # all-reduce here while every peer proceeds into it, and the # peers then hang on the NCCL collective until timeout. - per_req_exc: BaseException | None = None - try: - self._initialize_generator(state.sampling) - clear_pipeline_stage_durations(pipeline) - # encode - pipeline.prepare_encode(state) - merge_stage_durations( - state, - consume_pipeline_stage_durations(pipeline), - ) - except Exception as exc: - per_req_exc = exc - # Pipelines that do rank-0-only work (e.g. MiniMax H3 - # reference-video prep) must broadcast per-request failures - # internally so downstream collectives stay in step; even so, - # cross-check that every DiT rank agrees so a rank-local error - # (or a future pipeline that omits the guard) does not leave - # the process group half-way through a new request. - if _dit_any_rank_failed(per_req_exc is not None): + def _abort_prep_failure(per_req_exc: BaseException | None) -> None: self.state_cache.pop(state.request_id, None) if per_req_exc is None: per_req_exc = RuntimeError( @@ -1123,6 +1107,41 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): result=DiffusionOutput.from_exception(per_req_exc), ) ) + + per_req_exc: BaseException | None = None + try: + self._initialize_generator(state.sampling) + clear_pipeline_stage_durations(pipeline) + pipeline.prepare_encode(state) + except Exception as exc: + per_req_exc = exc + # Pipelines that do rank-0-only work (e.g. MiniMax H3 + # reference-video prep) must broadcast per-request failures + # internally so downstream collectives stay in step; even so, + # cross-check that every DiT rank agrees so a rank-local error + # (or a future pipeline that omits the guard) does not leave + # the process group half-way through a new request. + if _dit_any_rank_failed(per_req_exc is not None): + _abort_prep_failure(per_req_exc) + continue + # If the pipeline supports interaction, the interaction session initialization also needs to call + # synchronized_monotonic_time(). Wrap in another try-block to not block on prepare_encode failures. + try: + if supports_interaction_apply(pipeline) and state.chunk_index == 0: + pipe = cast(SupportsInteractionApply, pipeline) + assert self._interaction_coordinator is not None, "Model not loaded. Call load_model() first." + state.interaction_chunk_metadata = self._interaction_coordinator.maybe_prepare_initial_session( + state, pipe + ) + pipe.prepare_next_chunk(state) + merge_stage_durations( + state, + consume_pipeline_stage_durations(pipeline), + ) + except Exception as exc: + per_req_exc = exc + if _dit_any_rank_failed(per_req_exc is not None): + _abort_prep_failure(per_req_exc) continue prepared_states.append(state) @@ -1431,47 +1450,33 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): interaction: OmniInteractionPrompt, ) -> None: """Route a midway interaction through the pipeline interaction coordinator.""" - assert self.pipeline is not None, "Model not loaded. Call load_model() first." + assert self.pipeline is not None and self._interaction_coordinator is not None, ( + "Model not loaded. Call load_model() first." + ) if not self.od_config.streaming_output: raise ValueError("submit_interaction requires streaming_output=True") if not self._supports_step_mode(): raise ValueError("submit_interaction requires step execution support") - coordinator = self._interaction_coordinator - if coordinator is None: - coordinator = InteractionCoordinator.build(self.pipeline, self.od_config) - self._interaction_coordinator = coordinator - if hasattr(self.pipeline, "_interaction_coordinator"): - self.pipeline._interaction_coordinator = coordinator - - event = interaction.get("event") - has_mm = isinstance(event, dict) and "multi_modal_data" in event - has_prompt = isinstance(event, dict) and "prompt" in event and event.get("prompt") is not None - - # Prompt-only interactions in this release; multi_modal_data lands with camera support. - if not isinstance(event, dict) or has_mm or not has_prompt: - raise NotImplementedError( - "Only text-only prompt update interactions with 'event.prompt' and optional " - "'transition_chunks' are supported in this release" - ) - if not coordinator.has_modality("prompt"): - raise ValueError(f"prompt_update is not supported by pipeline {self.od_config.model_class_name!r}") - state = self.state_cache.get(request_id) if state is None: raise ValueError(f"No active request state for interaction: {request_id!r}") - event_id = interaction.get("event_id") - if not isinstance(event_id, str) or not event_id: - raise ValueError("event_id must be non-empty") - prompt = event["prompt"] - if not isinstance(prompt, str) or not prompt: - raise ValueError("prompt must be non-empty") - coordinator.enqueue( + event = interaction["event"] + parts: list[tuple[str, InteractionPayload]] = [] + prompt = event.get("prompt") + if prompt is not None: + parts.append(("prompt", {"prompt": prompt})) + multi_modal_data = event.get("multi_modal_data") + if multi_modal_data: + parts.extend( + (str(modality), cast(InteractionPayload, payload)) for modality, payload in multi_modal_data.items() + ) + + self._interaction_coordinator.enqueue_parts( state, - modality="prompt", - event_id=event_id, + parts=parts, + event_id=interaction["event_id"], received_at=time.monotonic(), - payload={"prompt": prompt}, transition_chunks=interaction.get("transition_chunks"), ) diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index f8f4722ad..2647ec11c 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -362,15 +362,41 @@ class OmniStreamingVideoOutputHandler: if not isinstance(event, dict): await self._send_error(websocket, "session.interaction requires event object", send_lock=send_lock) return - prompt = event.get("prompt") - if not isinstance(prompt, str) or not prompt.strip(): + multi_modal_data = event.get("multi_modal_data") + has_prompt = isinstance(prompt, str) and bool(prompt.strip()) + has_mm = isinstance(multi_modal_data, dict) and bool(multi_modal_data) + if not has_prompt and not has_mm: await self._send_error( websocket, - "session.interaction requires a non-empty event.prompt", + "session.interaction requires a non-empty event.prompt and/or non-empty event.multi_modal_data", send_lock=send_lock, ) return + if "prompt" in event and prompt is not None and not has_prompt: + await self._send_error( + websocket, + "session.interaction event.prompt must be a non-empty string when provided", + send_lock=send_lock, + ) + return + if "multi_modal_data" in event and multi_modal_data is not None and not has_mm: + await self._send_error( + websocket, + "session.interaction event.multi_modal_data must be a non-empty object when provided", + send_lock=send_lock, + ) + return + if has_mm: + assert isinstance(multi_modal_data, dict) + for modality, payload in multi_modal_data.items(): + if not isinstance(payload, dict): + await self._send_error( + websocket, + f"session.interaction multi_modal_data[{modality!r}] must be an object", + send_lock=send_lock, + ) + return if "negative_prompt" in event and event["negative_prompt"] is not None: await self._send_error( @@ -380,9 +406,19 @@ class OmniStreamingVideoOutputHandler: ) return + normalized_event: dict[str, Any] = {} + if has_prompt: + normalized_event["prompt"] = prompt + if has_mm: + normalized_event["multi_modal_data"] = multi_modal_data + # Preserve other event fields that OmniInteractionEvent may carry. + for key, value in event.items(): + if key not in normalized_event and key not in ("prompt", "multi_modal_data"): + normalized_event[key] = value + normalized: OmniInteractionPrompt = { "event_id": event_id, - "event": cast(OmniInteractionEvent, event), + "event": cast(OmniInteractionEvent, normalized_event), } transition_chunks = interaction.get("transition_chunks") if transition_chunks is not None: @@ -572,8 +608,11 @@ class OmniStreamingVideoOutputHandler: @staticmethod async def _decode_image_reference(request: VideoGenerationRequest) -> Image.Image | None: try: + image_reference = request.image_reference + if isinstance(image_reference, list): + raise InvalidInputReferenceError("image_reference must be a single image reference.") media_data = await decode_input_reference( - request.image_reference, + image_reference, None, None, )