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

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

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.

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],
+        },
+    }

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()