[Frontend] Add midway prompt update for streaming video generation (#4652)

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

examples/online_serving/streaming_video_generation/README.md

@@ -21,13 +21,16 @@ 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. |
+| 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. |
 | 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. |
 | Client to server | `session.stop` | JSON text: `{"type":"session.stop"}` | Requests cancellation of the active session. |
 | Client to server | `session.ping` | JSON text: `{"type":"session.ping"}` | Optional keepalive; refreshes the server stall clock. |
+| Server to client | `session.interaction.queued` | JSON text: `{"type":"session.interaction.queued","request_id":"...","event_id":"xxx"}` | Confirms the server queued the interaction for the active request. |
 | Server to client | `session.done` | JSON text: `{"type":"session.done","request_id":"...","chunks":3,"stopped":false}` | Ends a completed or stopped session. |
 | Server to client | `session.pong` | JSON text: `{"type":"session.pong"}` | Reply to `session.ping`. |
-| Server to client | `error` | JSON text: `{"type":"error","message":"..."}` | Reports invalid input, unsupported formats, generation failures, control-message errors, or stall timeout. |
+| Server to client | `error` | JSON text: `{"type":"error","message":"..."}` | Reports invalid input, unsupported formats, inactive requests, unsupported prompt updates, generation failures, control-message errors, or stall timeout. |
 
 During generation the client normally sends only `session.start` and then receives binary chunks; silence on the client socket is expected. The server closes the session with a stall error only when there is no engine progress and no `session.ping` for about 60 seconds.
 
@@ -59,6 +62,17 @@ python streaming_video_client.py \
 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 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`. Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`:
+
+```bash
+python streaming_video_client.py \
+  --prompt "A serene lakeside sunrise with mist over the water." \
+  --prompt-updates '[
+    {"at": 2.5, "prompt": "A sea turtle glides past the reeds"},
+    {"at": 5.0, "prompt": "Sunlight breaks through the morning mist", "transition_chunks": 2}
+  ]'
+```
+
 ## Run The Gradio Demo
 
 ```bash

examples/online_serving/streaming_video_generation/gradio_demo.py

@@ -14,8 +14,9 @@ from __future__ import annotations
 
 import argparse
 import json
+import math
 from pathlib import Path
-from typing import Any
+from typing import Any, TypedDict
 
 try:
     import gradio as gr
@@ -23,16 +24,49 @@ except ImportError:
     raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 
 DEFAULT_PROMPT = "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement."
+DEFAULT_TRANSITION_CHUNKS = 3
+PROMPT_UPDATE_HEADERS = ["At (s)", "Prompt", "Transition Chunks"]
+PROMPT_UPDATE_VALUES = [
+    [
+        4,
+        (
+            "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. "
+            "The water is murky, swirling violently, carrying debris and marine life into the vortex. "
+            "The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. "
+            "The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. "
+            "Close-up shot emphasizing the turbulent motion and destruction."
+        ),
+        3,
+    ],
+    [
+        11,
+        (
+            "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. "
+            "Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. "
+            "The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. "
+            "Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. "
+            "Close-up shot emphasizing the turbulent motion and destruction."
+        ),
+        3,
+    ],
+]
 VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html")
 VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js")
 
 
+class ScheduledPromptUpdate(TypedDict):
+    at: float
+    prompt: str
+    transition_chunks: int
+
+
 def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None:
     if value is None:
         return
-    if isinstance(value, str) and not value.strip():
-        return
-    payload[key] = value
+    if isinstance(value, str):
+        payload[key] = value.strip()
+    else:
+        payload[key] = value
 
 
 def _optional_int(value: int | float | None) -> int | None:
@@ -43,11 +77,83 @@ def _optional_float(value: int | float | None) -> float | None:
     return None if value is None else float(value)
 
 
+def _normalize_prompt_update_row(row: list[Any] | tuple[Any, ...]) -> tuple[Any, Any, Any]:
+    values = []
+    for v in row[: len(PROMPT_UPDATE_HEADERS)]:
+        if isinstance(v, str):
+            values.append(v.strip() or None)
+        elif isinstance(v, float) and math.isnan(v):
+            values.append(None)
+        else:
+            values.append(v)
+    while len(values) < len(PROMPT_UPDATE_HEADERS):
+        values.append(None)
+    return values[0], values[1], values[2]  # pyright: ignore[reportIndexIssue]
+
+
+def _prompt_update_row_is_empty(row: list[Any] | tuple[Any, ...]) -> bool:
+    at, prompt, transition = _normalize_prompt_update_row(row)
+    return at is None and prompt is None and transition is None
+
+
+def _parse_prompt_update_row(row: list[Any] | tuple[Any, ...], *, index: int) -> ScheduledPromptUpdate:
+    at_val, prompt_val, transition_val = _normalize_prompt_update_row(row)
+    row_label = f"Prompt update row {index + 1}"
+
+    if at_val is None or prompt_val is None:
+        raise gr.Error(f"{row_label} requires both At (s) and Prompt.")
+
+    try:
+        at = float(at_val)
+    except (TypeError, ValueError) as exc:
+        raise gr.Error(f"{row_label}: At (s) must be a number.") from exc
+    if at < 0:
+        raise gr.Error(f"{row_label}: At (s) must be >= 0.")
+
+    prompt = str(prompt_val).strip()
+    if not prompt:
+        raise gr.Error(f"{row_label}: Prompt must be a non-empty string.")
+
+    transition_chunks = DEFAULT_TRANSITION_CHUNKS
+    if transition_val is not None:
+        try:
+            transition_chunks = int(transition_val)
+        except (TypeError, ValueError) as exc:
+            raise gr.Error(f"{row_label}: Transition Chunks must be an integer.") from exc
+        if transition_chunks < 0:
+            raise gr.Error(f"{row_label}: Transition Chunks must be >= 0.")
+
+    return {
+        "at": at,
+        "prompt": prompt,
+        "transition_chunks": transition_chunks,
+    }
+
+
+def _parse_prompt_updates_dataframe(rows: list[list[Any]] | None) -> list[ScheduledPromptUpdate]:
+    if not rows:
+        return []
+
+    # Single row: empty or complete.
+    if len(rows) == 1:
+        row = rows[0]
+        if _prompt_update_row_is_empty(row):
+            return []
+        return [_parse_prompt_update_row(row, index=0)]
+
+    updates: list[ScheduledPromptUpdate] = []
+    for index, row in enumerate(rows):
+        if _prompt_update_row_is_empty(row):
+            raise gr.Error(f"Prompt update row {index + 1} cannot be empty when multiple rows are present.")
+        updates.append(_parse_prompt_update_row(row, index=index))
+    return sorted(updates, key=lambda update: update["at"])
+
+
 def _build_session_start(
     *,
     model: str,
     prompt: str,
-    negative_prompt: str,
+    negative_prompt: str | None,
     width: int | float | None,
     height: int | float | None,
     fps: int | float | None,
@@ -101,6 +207,7 @@ def build_browser_stream_config(
     model: str,
     prompt: str,
     negative_prompt: str,
+    include_negative_prompt: bool,
     width: int | float | None,
     height: int | float | None,
     fps: int | float,
@@ -118,6 +225,7 @@ def build_browser_stream_config(
     flow_shift: int | float | None,
     include_true_cfg_scale: bool,
     true_cfg_scale: int | float | None,
+    prompt_updates_df: list[list[Any]] | None,
 ) -> tuple[str, Any]:
     """Build the browser-side WebSocket request config for the HTML/MSE player."""
     if not model.strip():
@@ -128,7 +236,7 @@ def build_browser_stream_config(
     payload = _build_session_start(
         model=model.strip(),
         prompt=prompt.strip(),
-        negative_prompt=negative_prompt,
+        negative_prompt=negative_prompt if include_negative_prompt else None,
         width=width,
         height=height,
         fps=fps,
@@ -141,15 +249,21 @@ def build_browser_stream_config(
         true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None,
         seed=seed if include_seed else None,
     )
+    prompt_updates = _parse_prompt_updates_dataframe(prompt_updates_df)
     config = {
         "url": f"ws://{host}:{int(port)}/v1/realtime/video",
         "payload": payload,
     }
+    if prompt_updates:
+        config["prompt_updates"] = prompt_updates
     return json.dumps(config, ensure_ascii=False), gr.update(value="Streaming...", interactive=False)
 
 
 def create_demo() -> gr.Blocks:
-    with gr.Blocks(title="Streaming Video Generation") as demo:
+    with gr.Blocks(
+        title="Streaming Video Generation",
+        css="textarea:disabled, input:disabled { color: var(--body-text-color-subdued); -webkit-text-fill-color: var(--body-text-color-subdued); background-color: var(--background-fill-secondary); }",
+    ) as demo:
         gr.Markdown("# Streaming Video Generation")
         gr.Markdown(
             "Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks "
@@ -162,7 +276,6 @@ def create_demo() -> gr.Blocks:
                 port = gr.Number(label="Server Port", value=8000, precision=0)
                 model = gr.Textbox(label="Model", value="BestWishYsh/Helios-Distilled")
                 prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3)
-                negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=2)
 
                 with gr.Row():
                     width = gr.Number(label="Width", value=640, precision=0)
@@ -170,49 +283,66 @@ def create_demo() -> gr.Blocks:
 
                 with gr.Row():
                     fps = gr.Number(label="FPS", value=16, precision=0)
-                    num_frames = gr.Number(label="Num Frames", value=132, precision=0)
+                    num_frames = gr.Number(label="Num Frames", value=429, precision=0)
                     num_inference_steps = gr.Number(label="Num Inference Steps", value=50, precision=0)
 
-                with gr.Row():
-                    with gr.Column():
-                        guidance_scale = gr.Number(
-                            label="Guidance Scale",
-                            value=1.0,
-                            interactive=True,
-                            scale=4,
-                        )
-                        include_guidance_scale = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
-                    with gr.Column():
-                        guidance_scale_2 = gr.Number(
-                            label="Guidance Scale 2",
-                            value=None,
-                            interactive=False,
-                            scale=4,
-                        )
-                        include_guidance_scale_2 = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
-                    with gr.Column():
-                        true_cfg_scale = gr.Number(
-                            label="True CFG Scale",
-                            value=None,
-                            interactive=False,
-                            scale=4,
-                        )
-                        include_true_cfg_scale = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
-                    # I want it to wrap after 3, but it auto wrap after 2. I'll rather fully delegate to auto wrapping.
-                    with gr.Column():
-                        boundary_ratio = gr.Number(
-                            label="Boundary Ratio",
-                            value=None,
-                            interactive=False,
-                            scale=4,
-                        )
-                        include_boundary_ratio = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
-                    with gr.Column():
-                        flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False, scale=4)
-                        include_flow_shift = gr.Checkbox(label="Send", value=False, scale=1, min_width=70)
-                    with gr.Column():
-                        seed = gr.Number(label="Seed", value=42, precision=0, interactive=True, scale=4)
-                        include_seed = gr.Checkbox(label="Send", value=True, scale=1, min_width=70)
+                with gr.Accordion("Advanced Sampling", open=False):
+                    with gr.Row():
+                        with gr.Column():
+                            negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=1, interactive=False)
+                            include_negative_prompt = gr.Checkbox(label="Send", value=False)
+                    with gr.Row():
+                        with gr.Column(scale=0, min_width=150):
+                            guidance_scale = gr.Number(
+                                label="Guidance Scale",
+                                value=1.0,
+                                interactive=True,
+                            )
+                            include_guidance_scale = gr.Checkbox(label="Send", value=True)
+                        with gr.Column(scale=0, min_width=150):
+                            guidance_scale_2 = gr.Number(
+                                label="Guidance Scale 2",
+                                value=None,
+                                interactive=False,
+                            )
+                            include_guidance_scale_2 = gr.Checkbox(label="Send", value=False)
+                        with gr.Column(scale=0, min_width=150):
+                            true_cfg_scale = gr.Number(
+                                label="True CFG Scale",
+                                value=None,
+                                interactive=False,
+                            )
+                            include_true_cfg_scale = gr.Checkbox(label="Send", value=False)
+                        with gr.Column(scale=0, min_width=150):
+                            boundary_ratio = gr.Number(
+                                label="Boundary Ratio",
+                                value=None,
+                                interactive=False,
+                            )
+                            include_boundary_ratio = gr.Checkbox(label="Send", value=False)
+                        with gr.Column(scale=0, min_width=150):
+                            flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False)
+                            include_flow_shift = gr.Checkbox(label="Send", value=False)
+                        with gr.Column(scale=0, min_width=150):
+                            seed = gr.Number(label="Seed", value=42, precision=0, interactive=True)
+                            include_seed = gr.Checkbox(label="Send", value=True)
+
+                with gr.Accordion("Prompt Updates", open=True):
+                    gr.Markdown(
+                        "Schedule midway text `session.interaction` messages after `video.start`. "
+                        "With one row, leave it fully empty or fill both **At (s)** and **Prompt**. "
+                        "With multiple rows, every row must be complete."
+                    )
+                    prompt_updates_df = gr.Dataframe(
+                        headers=PROMPT_UPDATE_HEADERS,
+                        datatype=["number", "str", "number"],
+                        type="array",
+                        row_count=1,
+                        row_limits=(1, None),
+                        column_count=3,
+                        value=PROMPT_UPDATE_VALUES,
+                        interactive=True,
+                    )
 
                 with gr.Row():
                     start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start")
@@ -231,6 +361,7 @@ def create_demo() -> gr.Blocks:
             model,
             prompt,
             negative_prompt,
+            include_negative_prompt,
             width,
             height,
             fps,
@@ -248,8 +379,10 @@ def create_demo() -> gr.Blocks:
             flow_shift,
             include_true_cfg_scale,
             true_cfg_scale,
+            prompt_updates_df,
         ]
         for toggle, field in [
+            (include_negative_prompt, negative_prompt),
             (include_guidance_scale, guidance_scale),
             (include_guidance_scale_2, guidance_scale_2),
             (include_seed, seed),

examples/online_serving/streaming_video_generation/streaming_video_client.py

@@ -13,10 +13,12 @@ Requirements:
 
 import argparse
 import asyncio
+import contextlib
 import io
 import json
 import os
 import time
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
@@ -39,17 +41,69 @@ HELIOS_DISTILLED_EXTRA_PARAMS = {
     "is_amplify_first_chunk": True,
 }
 
+DEFAULT_TRANSITION_CHUNKS = 3
 
-def _json_object(value: str, *, argument_name: str) -> dict[str, Any]:
+
+@dataclass(frozen=True)
+class ScheduledPromptUpdate:
+    """Client-side schedule entry for a midway text ``session.interaction``."""
+
+    at: float
+    prompt: str
+    transition_chunks: int
+
+
+def _parse_extra_params(value: str) -> dict[str, Any]:
     try:
         parsed = json.loads(value)
     except json.JSONDecodeError as exc:
-        raise argparse.ArgumentTypeError(f"{argument_name} must be valid JSON: {exc}") from exc
+        raise argparse.ArgumentTypeError(f"--extra-params must be valid JSON: {exc}") from exc
     if not isinstance(parsed, dict):
-        raise argparse.ArgumentTypeError(f"{argument_name} must be a JSON object")
+        raise argparse.ArgumentTypeError("--extra-params must be a JSON object")
     return parsed
 
 
+def _parse_prompt_updates(value: str) -> list[ScheduledPromptUpdate]:
+    try:
+        parsed = json.loads(value)
+    except json.JSONDecodeError as exc:
+        raise argparse.ArgumentTypeError(f"--prompt-updates must be valid JSON: {exc}") from exc
+    if not isinstance(parsed, list):
+        raise argparse.ArgumentTypeError("--prompt-updates must be a JSON array")
+
+    updates: list[ScheduledPromptUpdate] = []
+    for index, item in enumerate(parsed):
+        try:
+            at = float(item["at"])
+            prompt = item["prompt"]
+        except (KeyError, TypeError, ValueError) as exc:
+            raise argparse.ArgumentTypeError(
+                f"--prompt-updates[{index}] must have valid 'at' and 'prompt' fields"
+            ) from exc
+        if at < 0:
+            raise argparse.ArgumentTypeError(f"--prompt-updates[{index}].at must be >= 0")
+        if not isinstance(prompt, str) or not prompt.strip():
+            raise argparse.ArgumentTypeError(f"--prompt-updates[{index}].prompt must be a non-empty string")
+
+        try:
+            transition_chunks = item.get("transition_chunks", DEFAULT_TRANSITION_CHUNKS)
+            transition_chunks = int(transition_chunks)
+        except (TypeError, ValueError) as exc:
+            raise argparse.ArgumentTypeError(f"--prompt-updates[{index}].transition_chunks must be an integer") from exc
+        if transition_chunks < 0:
+            raise argparse.ArgumentTypeError(f"--prompt-updates[{index}].transition_chunks must be >= 0")
+
+        updates.append(
+            ScheduledPromptUpdate(
+                at=at,
+                prompt=prompt,
+                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
@@ -88,6 +142,31 @@ def build_session_start(args: argparse.Namespace) -> dict[str, Any]:
     return payload
 
 
+async def _run_prompt_update_scheduler(
+    websocket: Any,
+    updates: list[ScheduledPromptUpdate],
+    *,
+    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)
+
+        payload = {
+            "type": "session.interaction",
+            "interaction": {
+                "event_id": f"scheduled-{update.at:.3f}",
+                "event": {"prompt": update.prompt},
+                "transition_chunks": update.transition_chunks,
+            },
+        }
+        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:
@@ -144,11 +223,15 @@ async def stream_video(args: argparse.Namespace) -> None:
     stream_format = "m4s"
     done = False
     ws = None
+    send_lock = asyncio.Lock()
+    prompt_update_task: asyncio.Task[None] | None = None
+    pending_chunk_metadata: dict[str, Any] | None = None
 
     try:
         async with connect(url, max_size=None) as websocket:
             ws = websocket
-            await websocket.send(json.dumps(payload))
+            async with send_lock:
+                await websocket.send(json.dumps(payload))
             print(f"Connected: {url}")
             print(f"Sent session.start: {json.dumps(payload, ensure_ascii=False)}")
             started_at = time.perf_counter()
@@ -159,6 +242,9 @@ async def stream_video(args: argparse.Namespace) -> None:
                 now = time.perf_counter()
 
                 if isinstance(message, bytes):
+                    # Use chunk metadata just received before video chunks
+                    metadata = pending_chunk_metadata or {}
+                    pending_chunk_metadata = None
                     chunks.append(message)
                     chunk_bytes = len(message)
                     chunk_elapsed = now - last_chunk_at
@@ -176,8 +262,11 @@ async def stream_video(args: argparse.Namespace) -> None:
                     total_elapsed = now - started_at
                     print(
                         f"[chunk {len(chunks):03d}] "
+                        f"kind={metadata.get('kind', 'media')} "
                         f"bytes={chunk_bytes} "
-                        f"frames={chunk_frames} "
+                        f"frames={metadata.get('num_frames', chunk_frames)} "
+                        f"generation_chunk={metadata.get('generation_chunk_index')} "
+                        f"active_interactions={metadata.get('active_event_ids', [])} "
                         f"elapsed={chunk_elapsed:.2f}s "
                         f"total_bytes={total_bytes} "
                         f"total_frames={total_frames} "
@@ -190,9 +279,26 @@ 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')}")
+                    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,
+                                args.prompt_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
+                elif msg_type == "session.interaction.queued":
+                    print(f"Interaction queued by server: event_id={msg.get('event_id')}")
                 elif msg_type == "session.done":
                     print(f"Session complete: {json.dumps(msg, ensure_ascii=False)}")
                     done = True
+                    if prompt_update_task is not None:
+                        prompt_update_task.cancel()
                     break
                 elif msg_type == "error":
                     print(f"ERROR: {msg.get('message', msg)}")
@@ -202,12 +308,17 @@ async def stream_video(args: argparse.Namespace) -> None:
         print("\nInterrupted. Asking server to stop the active session...")
         if ws is not None:
             try:
-                await ws.send(json.dumps({"type": "session.stop"}))
+                async with send_lock:
+                    await ws.send(json.dumps({"type": "session.stop"}))
                 await ws.close(code=1000, reason="client interrupted")
             except Exception:
                 pass
         raise
     finally:
+        if prompt_update_task is not None:
+            prompt_update_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await prompt_update_task
         print(
             "Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)"
         )
@@ -232,7 +343,15 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--prompt",
-        default="A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement.",
+        default=(
+            "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. "
+            "The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. "
+            "The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. "
+            "The water is crystal clear, allowing for a view of the sandy ocean floor below. "
+            "The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. "
+            "The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. "
+            "A close-up shot with dynamic movement."
+        ),
         help="Text prompt for video generation",
     )
     parser.add_argument("--output", default="streaming_video_output.mp4", help="Path for the received video bytes")
@@ -241,7 +360,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--size", default=None, help="Video size as WIDTHxHEIGHT; alternative to width/height")
     parser.add_argument("--seconds", default=None, help="Clip duration in seconds")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second")
-    parser.add_argument("--num-frames", type=int, default=99, help="Number of generated frames")
+    parser.add_argument("--num-frames", type=int, default=429, help="Number of generated frames")
 
     parser.add_argument("--negative-prompt", default=None, help="Negative prompt")
     parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps")
@@ -260,10 +379,44 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--extra-params",
-        type=lambda value: _json_object(value, argument_name="--extra-params"),
+        type=_parse_extra_params,
         default=None,
         help="JSON object merged into extra_params; overrides preset keys on conflict.",
     )
+    parser.add_argument(
+        "--prompt-updates",
+        type=_parse_prompt_updates,
+        default=json.dumps(
+            [
+                {
+                    "at": 4,
+                    "prompt": (
+                        "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. "
+                        "The water is murky, swirling violently, carrying debris and marine life into the vortex. "
+                        "The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. "
+                        "The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. "
+                        "Close-up shot emphasizing the turbulent motion and destruction."
+                    ),
+                },
+                {
+                    "at": 11,
+                    "prompt": (
+                        "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. "
+                        "Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. "
+                        "The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. "
+                        "Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. "
+                        "Close-up shot emphasizing the turbulent motion and destruction."
+                    ),
+                    # "transition_chunks": 3,
+                },
+            ]
+        ),
+        help=(
+            "JSON array of scheduled prompt updates. Each object requires "
+            "'at' (seconds after video.start) and 'prompt'. Optional "
+            f"'transition_chunks' defaults to {DEFAULT_TRANSITION_CHUNKS}."
+        ),
+    )
     return parser.parse_args()
 
 

examples/online_serving/streaming_video_generation/video-stream-view.js

@@ -14,6 +14,7 @@
   let totalBytes = 0;
   let done = false;
   let objectUrl = null;
+  let promptUpdateTimers = [];
 
   function el(id) {
     return document.getElementById(id);
@@ -73,7 +74,50 @@
     } catch (_) { }
   }
 
+  function clearPromptUpdateTimers() {
+    for (const timer of promptUpdateTimers) {
+      clearTimeout(timer);
+    }
+    promptUpdateTimers = [];
+  }
+
+  function schedulePromptUpdates(updates) {
+    clearPromptUpdateTimers();
+    if (!Array.isArray(updates) || updates.length === 0) return;
+
+    const videoStartedAt = performance.now();
+    for (const update of updates) {
+      const atSeconds = Number(update.at);
+      if (!Number.isFinite(atSeconds) || atSeconds < 0) continue;
+
+      const delayMs = Math.max(0, atSeconds * 1000 - (performance.now() - videoStartedAt));
+      const transitionChunks = Number.isFinite(Number(update.transition_chunks))
+        ? Number(update.transition_chunks)
+        : 3;
+
+      const timer = setTimeout(() => {
+        if (!ws || ws.readyState !== WebSocket.OPEN) return;
+        const payload = {
+          type: "session.interaction",
+          interaction: {
+            event_id: "scheduled-" + atSeconds.toFixed(3),
+            event: { prompt: update.prompt },
+            transition_chunks: transitionChunks,
+          },
+        };
+        try {
+          ws.send(JSON.stringify(payload));
+          log("Sent session.interaction at t=" + atSeconds.toFixed(2) + "s: " + JSON.stringify(payload));
+        } catch (err) {
+          log("Failed to send session.interaction: " + err.message);
+        }
+      }, delayMs);
+      promptUpdateTimers.push(timer);
+    }
+  }
+
   function stopCurrent() {
+    clearPromptUpdateTimers();
     if (ws && ws.readyState === WebSocket.OPEN) {
       try { ws.send(JSON.stringify({ type: "session.stop" })); } catch (_) { }
       try { ws.close(); } catch (_) { }
@@ -148,6 +192,11 @@
           const msg = JSON.parse(event.data);
           if (msg.type === "video.start") {
             log("Video session started: request_id=" + (msg.request_id || "") + " format=" + (msg.format || ""));
+            schedulePromptUpdates(config.prompt_updates);
+          } else if (msg.type === "video.chunk_metadata") {
+            log("Chunk metadata: " + JSON.stringify(msg));
+          } else if (msg.type === "session.interaction.queued") {
+            log("Interaction queued by server: event_id=" + (msg.event_id || ""));
           } else if (msg.type === "session.done") {
             done = true;
             setStatus("Done");

realtime_video_prompt_interaction_protocol.md

@@ -0,0 +1,138 @@
+# Realtime Video Prompt Interaction Protocol
+
+This document specifies the prompt-update subset of the realtime video
+interaction protocol exposed by:
+
+```text
+WebSocket /v1/realtime/video
+```
+
+A WebSocket connection contains exactly one video-generation request. Clients
+therefore do not send a request or session ID with an interaction. The server
+assigns a `request_id` in `video.start` and includes it in subsequent
+server-to-client messages.
+
+## Interaction Request
+
+After `video.start`, a client may submit a prompt update:
+
+```json
+{
+  "type": "session.interaction",
+  "interaction": {
+    "event_id": "ui-update-42",
+    "event": {
+      "prompt": "A snowy street illuminated by warm lanterns"
+    },
+    "transition_chunks": 3
+  }
+}
+```
+
+`interaction` has the following fields:
+
+- `event_id` (optional string): session-unique idempotency key. The
+  server preserves a client-provided value or generates a UUID when omitted.
+- `event.prompt` (required string): non-empty target prompt.
+  over which the prompt reaches its target.
+- `transition_chunks` (optional non-negative integer): model-generation chunks
+  over which the prompt reaches its target. If omitted, the model's default transition is used.
+
+Note: Although `["interaction"]["event"]` only contains `prompt`, still consider using `OmniTextPrompt` for this part (with total=false) so that other fields will be supported in the future. Currently, `vllm_omni/diffusion/worker/diffusion_model_runner.py` already takes care of enforcing subset format, routing to meaningful input handing routines, and raising `NotImplementedError`. Only need to adapt this design to the input data structure above.
+
+### Transition Semantics
+
+Prompt transitions are target updates and begin at the next eligible model
+generation boundary after the interaction is processed.
+
+- `transition_chunks = 0` is a hard update: the first affected generation
+  chunk uses the target prompt.
+- `transition_chunks = N`, where `N > 0`, reaches the target in the `N`th
+  affected generation chunk. For example, a three-chunk linear transition uses
+  progress values `1/3`, `2/3`, and `1`.
+
+### Idempotency and Replacement
+
+Within a session:
+
+- repeating an `event_id` with the same payload is a retry and must not
+  apply the update twice;
+- repeating an `event_id` with a different payload is rejected with
+  `event_id_conflict`;
+- a newer prompt target supersedes any older pending prompt target;
+- when a new target interrupts an active transition, its transition starts
+  from the prompt state resolved at its effective boundary.
+
+## Queue Acknowledgement
+
+The server quickly acknowledges admission to the interaction queue:
+
+```json
+{
+  "type": "session.interaction.queued",
+  "request_id": "video_stream-...",
+  "event_id": "ui-update-42"
+}
+```
+
+`queued` confirms syntax validation, idempotency handling, and queue admission.
+It does not claim that the model has scheduled or applied the update. A later
+failure is reported as a correlated error:
+
+```json
+{
+  "type": "error",
+  "request_id": "video_stream-...",
+  "event_id": "ui-update-42",
+  "code": "interaction_failed",
+  "message": "The active model does not support prompt updates"
+}
+```
+
+## Video Chunk Metadata
+
+Each emitted binary payload is preceded by a JSON metadata message. The
+metadata message and its binary payload form one ordered pair and must not be
+interleaved with another outbound message.
+
+```json
+{
+  "type": "video.chunk_metadata",
+  "request_id": "video_stream-...",
+  "kind": "media",
+  "transport_chunk_index": 7,
+  "generation_chunk_index": 6,
+  "num_frames": 9,
+  "byte_length": 184320,
+  "started_event_ids": ["ui-update-42"],
+  "active_event_ids": ["ui-update-42"],
+  "completed_event_ids": []
+}
+```
+
+- `transport_chunk_index` identifies the binary WebSocket payload.
+- `generation_chunk_index` identifies the model-generation chunk represented
+  by the payload.
+- `started_event_ids` lists interactions that first affect this
+  generation chunk.
+- `active_event_ids` lists interactions contributing to this generation
+  chunk, including an ongoing transition.
+- `completed_event_ids` lists interactions that reach their target in
+  this generation chunk.
+
+Encoder trailer bytes are emitted as:
+
+```json
+{
+  "type": "video.chunk_metadata",
+  "request_id": "video_stream-...",
+  "kind": "trailer",
+  "transport_chunk_index": 8,
+  "generation_chunk_index": null,
+  "num_frames": 0,
+  "byte_length": 824,
+  "started_event_ids": [],
+  "active_event_ids": [],
+  "completed_event_ids": []
+}
+```

tests/diffusion/test_diffusion_streaming_output.py

@@ -341,7 +341,9 @@ class TestPipelineStreamingOutputToEntrypoint:
                 with client.websocket_connect("/v1/realtime/video") as ws:
                     ws.send_json({"type": "session.start", "prompt": "integration test"})
                     assert ws.receive_json()["type"] == "video.start"
+                    assert ws.receive_json()["type"] == "video.chunk_metadata"
                     assert ws.receive_bytes() == b"pipeline-fmp4-0"
+                    assert ws.receive_json()["type"] == "video.chunk_metadata"
                     assert ws.receive_bytes() == b"pipeline-fmp4-1"
                     done = ws.receive_json()
                     assert done["type"] == "session.done"
@@ -453,7 +455,6 @@ class TestPipelineStreamingOutputToEntrypoint:
         omni.final_output_task = None
         omni.event_resolver = AsyncEventResolver()
         omni._enable_ar_profiler = False
-        omni._is_sleeping = False
         omni.prom_metrics = MagicMock()
         omni.mod_metrics = MagicMock()
         omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params

tests/diffusion/test_prompt_update.py

@@ -0,0 +1,574 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for midway prompt update across runner, batch, pipeline, and entrypoint layers."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import time
+from types import SimpleNamespace
+from typing import Any, cast
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from tests.engine.test_orchestrator import (
+    OrchestratorFixture,
+    _build_harness,
+    _wait_for,
+)
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.input_batch import InputBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
+from vllm_omni.engine.messages import (
+    ErrorMessage,
+    InteractionMessage,
+    ShutdownRequestMessage,
+    StageSubmissionMessage,
+)
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_client import StagePoolClient
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.engine.stage_pool import StagePool
+from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture
+def pipeline() -> HeliosPipeline:
+    pipeline = object.__new__(HeliosPipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.transformer = SimpleNamespace(dtype=torch.float32)  # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.encode_prompt = MagicMock(
+        return_value=(
+            torch.full((1, 4, 2), 2.0),
+            None,
+        )
+    )
+    pipeline._prepare_next_chunk = MagicMock()
+    return pipeline
+
+
+def _make_diffusion_request_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 _make_diffusion_model_runner(*, pipeline, streaming_output: bool = True) -> DiffusionModelRunner:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.pipeline = pipeline
+    runner.state_cache = {}
+    runner.od_config = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
+        model_class_name="HeliosPipeline",
+        streaming_output=streaming_output,
+    )
+    runner._supports_step_mode = lambda: True
+    return runner
+
+
+def _prompt_interaction(
+    prompt: str = "updated",
+    *,
+    event_id: str = "ui-update-1",
+    transition_chunks: int | None = None,
+) -> OmniInteractionPrompt:
+    interaction: dict[str, Any] = {
+        "event_id": event_id,
+        "event": {"prompt": prompt},
+    }
+    if transition_chunks is not None:
+        interaction["transition_chunks"] = transition_chunks
+    return cast(OmniInteractionPrompt, interaction)
+
+
+class TestPromptUpdateExecution:
+    """Runner, InputBatch cache, and pipeline chunk-boundary behavior."""
+
+    def test_runner_interaction_delegates_to_helios_pipeline(self, pipeline: HeliosPipeline) -> None:
+        """Runner encodes the new prompt and queues pending embeds on the request state."""
+        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(transition_chunks=2))
+
+        pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
+        pending = state.extra["pending_prompt_update"]
+        assert pending["event_id"] == "ui-update-1"
+        assert pending["transition_chunks"] == 2
+        assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0))
+
+    def test_runner_prompt_update_rejects_unsupported_pipeline(self) -> None:
+        """Runner rejects prompt updates when the pipeline lacks prompt-update support."""
+
+        class _UnsupportedPipeline:
+            supports_step_execution = True
+
+        runner = _make_diffusion_model_runner(pipeline=_UnsupportedPipeline())
+        with pytest.raises(ValueError, match="not supported"):
+            runner.submit_interaction("req-1", _prompt_interaction())
+
+    def test_runner_prompt_update_rejects_missing_request(self, pipeline) -> None:
+        """Runner rejects prompt updates when no active request state exists."""
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        with pytest.raises(ValueError, match="No active request state"):
+            runner.submit_interaction("missing", _prompt_interaction())
+
+    @pytest.mark.parametrize(
+        "interaction",
+        [
+            {},
+            {"multi_modal_data": {"camera": {"type": "pose"}}},
+            {"event": {"prompt": "updated", "multi_modal_data": {"camera": {"type": "pose"}}}},
+        ],
+    )
+    def test_runner_interaction_rejects_structural_payloads_until_implemented(
+        self,
+        pipeline: HeliosPipeline,
+        interaction: dict[str, Any],
+    ) -> None:
+        """Unsupported interaction dict shapes are preserved to and rejected by the runner."""
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        runner.state_cache["req-1"] = _make_diffusion_request_state()
+
+        with pytest.raises(NotImplementedError, match="Only text-only prompt update interactions"):
+            runner.submit_interaction("req-1", cast(Any, interaction))
+
+    def test_input_batch_refreshes_prompt_embeds_on_version_change(self) -> None:
+        """InputBatch rebuilds prompt_embeds when prompt_update_version changes."""
+        state = StepRequestState(
+            request_id="req-1",
+            sampling=SimpleNamespace(),  # pyright: ignore[reportArgumentType]
+            prompt="hello",
+        )
+        state.prompt_embeds = torch.zeros(1, 2, 3)
+        state.latents = torch.zeros(1, 2)
+        state.timesteps = torch.tensor([1.0])
+        state.extra["prompt_update_version"] = 0
+
+        batch = InputBatch.make_batch([state])
+        assert torch.equal(batch.prompt_embeds, torch.zeros(1, 2, 3))  # pyright: ignore[reportArgumentType]
+
+        state.prompt_embeds = torch.ones(1, 2, 3)
+        state.extra["prompt_update_version"] = 1
+        refreshed = InputBatch.make_batch([state], cached_batch=batch)
+        assert torch.equal(refreshed.prompt_embeds, torch.ones(1, 2, 3))  # pyright: ignore[reportArgumentType]
+
+    def test_helios_prepare_prompt_update_queues_pending_target(self, pipeline: HeliosPipeline) -> None:
+        """HeliosPipeline queues target embeds without mutating current prompt_embeds."""
+        state = _make_diffusion_request_state()
+
+        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+
+        pending = state.extra["pending_prompt_update"]
+        assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0))
+        assert pending["transition_chunks"] == 2
+        assert torch.equal(state.prompt_embeds, torch.zeros(1, 4, 2))  # pyright: ignore[reportArgumentType]
+
+    def test_helios_prepare_prompt_update_rejects_before_initial_generation(self, pipeline: HeliosPipeline) -> None:
+        """Reject prompt updates submitted before initial prompt embeds exist."""
+        state = _make_diffusion_request_state()
+        state.prompt_embeds = None
+
+        with pytest.raises(
+            ValueError,
+            match="prompt_update is not allowed before initial generation has started",
+        ):
+            pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+
+        assert "pending_prompt_update" not in state.extra
+
+    def test_helios_apply_prompt_update_at_chunk_boundary_starts_transition(self, pipeline: HeliosPipeline) -> None:
+        """At chunk boundary, starts transition state and bumps prompt_update_version."""
+        state = _make_diffusion_request_state()
+        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=2)
+
+        pipeline._apply_prompt_update_at_chunk_boundary(state)
+
+        assert state.extra.get("prompt_update_state") is not None
+        assert torch.equal(state.prompt_embeds, torch.ones(1, 4, 2))  # pyright: ignore[reportArgumentType]
+        assert state.extra["prompt_update_version"] == 1
+        assert state.extra["prompt_update_chunk_metadata"] == {
+            "started_event_ids": ["ui-update-1"],
+            "active_event_ids": ["ui-update-1"],
+            "completed_event_ids": [],
+        }
+
+    def test_helios_apply_prompt_update_advances_transition_over_chunks(self, pipeline: HeliosPipeline) -> None:
+        """At chunk boundary, interpolates embeds until the target prompt is reached."""
+        state = _make_diffusion_request_state()
+        pipeline.prepare_prompt_update(state, "new scene", "ui-update-1", transition_chunks=3)
+        pipeline._apply_prompt_update_at_chunk_boundary(state)
+
+        assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0 / 3.0))  # pyright: ignore[reportArgumentType]
+
+        for _ in range(2):
+            pipeline._apply_prompt_update_at_chunk_boundary(state)
+
+        assert torch.allclose(state.prompt_embeds, torch.full((1, 4, 2), 2.0))  # pyright: ignore[reportArgumentType]
+        assert state.extra["prompt_update_version"] == 3
+
+        # Further chunk boundaries after completion must not bump the version.
+        pipeline._apply_prompt_update_at_chunk_boundary(state)
+        assert state.extra["prompt_update_version"] == 3
+
+
+class TestPromptUpdateIntegration:
+    """AsyncOmni prompt update through orchestrator, inline client, and runner."""
+
+    @pytest.mark.asyncio
+    async def test_orchestrator_prompt_update_failure_surfaces_non_fatal_error(
+        self,
+    ) -> None:
+        """Worker-side prompt_update rejection is reported without crashing the orchestrator."""
+
+        class _RejectingStagePool:
+            async def submit_interaction(self, *args: Any, **kwargs: Any) -> None:
+                del args, kwargs
+                raise ValueError("prompt embeds are not ready")
+
+        output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
+        orchestrator = object.__new__(Orchestrator)
+        orchestrator.stage_pools = [_RejectingStagePool()]  # pyright: ignore[reportAttributeAccessIssue]
+        orchestrator.output_async_queue = output_queue  # pyright: ignore[reportAttributeAccessIssue]
+        orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")}
+
+        await orchestrator._handle_interaction(
+            InteractionMessage(
+                request_id="req-1",
+                interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
+            )
+        )
+
+        msg = output_queue.get_nowait()
+        assert msg == ErrorMessage(
+            error="Failed interaction for request req-1: prompt embeds are not ready",
+            fatal=False,
+            request_id="req-1",
+            event_id="ui-update-1",
+            stage_id=0,
+        )
+
+    @pytest.mark.asyncio
+    async def test_runner_prompt_update_failure_surfaces_non_fatal_error(self, pipeline: HeliosPipeline) -> None:
+        """Runner-side prepare_prompt_update rejection is reported through the orchestrator."""
+        pipeline.prepare_prompt_update = MagicMock(  # pyright: ignore[reportAttributeAccessIssue, reportMethodAssignment]
+            side_effect=ValueError("prompt embeds are not ready")
+        )
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        runner.state_cache["req-1"] = _make_diffusion_request_state(request_id="req-1")
+        prompt_update_engine = self._PromptUpdateEngine(self._IncrementalStreamingPipeline(), runner)
+        inline_client = self._make_inline_pipeline_client(prompt_update_engine)
+
+        output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
+        orchestrator = object.__new__(Orchestrator)
+        orchestrator.stage_pools = [  # pyright: ignore[reportAttributeAccessIssue]
+            StagePool(0, cast(StagePoolClient, inline_client))
+        ]
+        orchestrator.output_async_queue = output_queue  # pyright: ignore[reportAttributeAccessIssue]
+        orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")}
+
+        try:
+            await orchestrator._handle_interaction(
+                InteractionMessage(
+                    request_id="req-1",
+                    interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
+                )
+            )
+        finally:
+            inline_client.shutdown()
+
+        pipeline.prepare_prompt_update.assert_called_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+            runner.state_cache["req-1"],
+            "new scene",
+            "ui-update-1",
+            2,
+        )
+        msg = output_queue.get_nowait()
+        assert msg == ErrorMessage(
+            error="Failed interaction for request req-1: prompt embeds are not ready",
+            fatal=False,
+            request_id="req-1",
+            event_id="ui-update-1",
+            stage_id=0,
+        )
+
+    @pytest.mark.asyncio
+    async def test_prompt_update_reaches_runner_from_async_omni(self, pipeline: HeliosPipeline) -> None:
+        """Midway prompt update submitted via AsyncOmni reaches the diffusion runner."""
+        streaming_pipeline = self._IncrementalStreamingPipeline()
+        runner = _make_diffusion_model_runner(pipeline=pipeline)
+        prompt_update_engine = self._PromptUpdateEngine(streaming_pipeline, runner)
+        inline_client = self._make_inline_pipeline_client(prompt_update_engine)
+        fixture = _build_harness([inline_client])
+        omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture))
+
+        generate_task: asyncio.Task[list[OmniRequestOutput]] | None = None
+        try:
+            generate_task = asyncio.create_task(self._collect_generate_outputs(omni))
+
+            await _wait_for(
+                lambda: (
+                    len(runner.state_cache) > 0
+                    and any(state.external_request_id == "req-omni" for state in omni.request_states.values())
+                )
+            )
+            internal_request_id = next(iter(runner.state_cache))
+
+            await omni.submit_interaction_async(
+                "req-omni",
+                interaction=_prompt_interaction("new scene", event_id="ui-update-1", transition_chunks=2),
+            )
+
+            outputs = await generate_task
+
+            pending = runner.state_cache[internal_request_id].extra["pending_prompt_update"]
+            assert pending["transition_chunks"] == 2
+            assert torch.equal(pending["target_prompt_embeds"], torch.full((1, 4, 2), 2.0))
+            pipeline.encode_prompt.assert_called_once()  # pyright: ignore[reportAttributeAccessIssue]
+            assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+        finally:
+            if generate_task is not None and not generate_task.done():
+                generate_task.cancel()
+                await asyncio.gather(generate_task, return_exceptions=True)
+            await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
+
+    @staticmethod
+    async def _collect_generate_outputs(omni: AsyncOmni) -> list[OmniRequestOutput]:
+        outputs: list[OmniRequestOutput] = []
+        async for output in omni.generate(
+            prompt={"prompt": "a cat"},
+            request_id="req-omni",
+            sampling_params_list=[OmniDiffusionSamplingParams()],
+            output_modalities=["image"],
+        ):
+            outputs.append(output)
+        return outputs
+
+    @classmethod
+    def _make_inline_pipeline_client(cls, engine: _PromptUpdateEngine) -> InlineStageDiffusionClient:
+        metadata = StageMetadata(
+            stage_id=0,
+            stage_type="diffusion",
+            engine_output_type="image",
+            is_comprehension=False,
+            requires_multimodal_data=False,
+            engine_input_source=[],
+            final_output=True,
+            final_output_type="image",
+            default_sampling_params=OmniDiffusionSamplingParams(),
+            custom_process_input_func=None,
+            model_stage=None,
+            runtime_cfg=None,
+        )
+        with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+            with patch(
+                "vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine",
+                return_value=engine,
+            ):
+                od_config = MagicMock(spec=OmniDiffusionConfig)
+                od_config.streaming_output = True
+                return InlineStageDiffusionClient(
+                    model="test_model",
+                    od_config=od_config,
+                    metadata=metadata,
+                    batch_size=1,
+                )
+
+    @staticmethod
+    def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni:
+        omni = object.__new__(AsyncOmni)
+        omni.engine = engine  # pyright: ignore[reportAttributeAccessIssue]
+        omni.log_stats = False
+        omni._pause_cond = asyncio.Condition()
+        omni._paused = False
+        omni.request_states = {}
+        omni.final_output_task = None
+        omni.event_resolver = AsyncEventResolver()
+        omni._enable_ar_profiler = False
+        omni.prom_metrics = MagicMock()
+        omni.mod_metrics = MagicMock()
+        omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params  # pyright: ignore[reportAttributeAccessIssue]
+        omni._compute_final_stage_id = lambda output_modalities: 0
+        omni._compute_final_output_stage_ids = lambda output_modalities: [0]
+        omni.default_sampling_params_list = engine.default_sampling_params_list  # pyright: ignore[reportAttributeAccessIssue]
+        omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)  # pyright: ignore[reportAttributeAccessIssue]
+        return omni
+
+    @staticmethod
+    async def _shutdown_pipeline_omni_harness(
+        omni: AsyncOmni,
+        fixture: OrchestratorFixture,
+        inline_client: InlineStageDiffusionClient,
+    ) -> None:
+        if omni.final_output_task is not None:
+            omni.final_output_task.cancel()
+            await asyncio.gather(omni.final_output_task, return_exceptions=True)
+        inline_client.shutdown()
+        fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        await asyncio.to_thread(fixture.thread.join, 5)
+
+    class _IncrementalStreamingPipeline:
+        """Returns one streaming chunk per ``step_outputs`` call."""
+
+        supports_step_execution = True
+
+        def __init__(self) -> None:
+            self.requests: list[Any] = []
+            self._step_index = 0
+
+        def step_outputs(self, request):
+            self.requests.append(request)
+            if self._step_index == 0:
+                self._step_index += 1
+                return [DiffusionOutput(output={"chunk": 0}, finished=False)]
+            return [DiffusionOutput(output={"chunk": 1}, finished=True)]
+
+    class _OrchestratorBridgeEngine:
+        """Minimal AsyncOmni engine facade backed by a live Orchestrator harness."""
+
+        def __init__(self, fixture: OrchestratorFixture) -> None:
+            self._fixture = fixture
+            self.stage_metadata = [
+                StageRuntimeInfo(
+                    stage_type="diffusion",
+                    final_output=True,
+                    final_output_type="image",
+                )
+            ]
+            self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")]
+            self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
+            self.num_stages = 1
+            self.supported_tasks = ("generate",)
+            self._alive = True
+
+        async def add_request_async(
+            self,
+            *,
+            request_id: str,
+            prompt: Any,
+            sampling_params_list: list[Any],
+            final_stage_id: int,
+            **kwargs: Any,
+        ) -> None:
+            self._fixture.request_sync_q.put_nowait(
+                StageSubmissionMessage(
+                    type="add_request",
+                    request_id=request_id,
+                    prompt=prompt,
+                    original_prompt=prompt,
+                    output_prompt_text=None,
+                    sampling_params_list=sampling_params_list,
+                    final_stage_id=final_stage_id,
+                    final_output_stage_ids=kwargs.get("final_output_stage_ids"),
+                    preprocess_ms=0.0,
+                    request_timestamp=kwargs.get("arrival_time", time.time()),
+                    enqueue_ts=time.perf_counter(),
+                )
+            )
+
+        async def submit_interaction_async(
+            self,
+            request_id: str,
+            *,
+            interaction: dict[str, Any],
+        ) -> None:
+            self._fixture.request_sync_q.put_nowait(
+                InteractionMessage(
+                    request_id=request_id,
+                    interaction=cast(Any, interaction),
+                )
+            )
+
+        async def try_get_output_async(self) -> Any | None:
+            try:
+                return self._fixture.output_sync_q.get_nowait()
+            except queue.Empty:
+                return None
+
+        def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo:
+            return self.stage_metadata[stage_id]
+
+        def is_alive(self) -> bool:
+            return self._fixture.thread.is_alive()
+
+        async def abort_async(self, request_ids: list[str]) -> None:
+            del request_ids
+
+    class _PromptUpdateEngine:
+        """DiffusionEngine stand-in that streams chunks and routes prompt_update to a runner."""
+
+        def __init__(
+            self,
+            streaming_pipeline: TestPromptUpdateIntegration._IncrementalStreamingPipeline,
+            runner: DiffusionModelRunner,
+        ) -> None:
+            self.streaming_pipeline = streaming_pipeline
+            self._runner = runner
+            self.executor = SimpleNamespace(
+                register_failure_callback=MagicMock(),
+                check_health=MagicMock(),
+            )
+
+        async def step_streaming(self, request):
+            if request.request_id not in self._runner.state_cache:
+                state = _make_diffusion_request_state(request_id=request.request_id)
+                self._runner.state_cache[request.request_id] = state
+
+            state = self._runner.state_cache[request.request_id]
+
+            outputs = self.streaming_pipeline.step_outputs(request)
+            output = outputs[0]
+            custom_output = cast(dict[str, Any], output.output) or {}
+            yield [
+                OmniRequestOutput.from_diffusion(
+                    request_id=request.request_id,
+                    images=list(custom_output.get("images") or []),
+                    custom_output=custom_output,
+                    finished=output.finished,
+                )
+            ]
+
+            deadline = time.monotonic() + 5.0
+            while "pending_prompt_update" not in state.extra:
+                if time.monotonic() >= deadline:
+                    raise TimeoutError("timed out waiting for prompt update during streaming")
+                await asyncio.sleep(0.01)
+
+            outputs = self.streaming_pipeline.step_outputs(request)
+            output = outputs[0]
+            custom_output = cast(dict[str, Any], output.output) or {}
+            yield [
+                OmniRequestOutput.from_diffusion(
+                    request_id=request.request_id,
+                    images=list(custom_output.get("images") or []),
+                    custom_output=custom_output,
+                    finished=output.finished,
+                )
+            ]
+
+        def collective_rpc(self, method, timeout, args, kwargs, unique_reply_rank):
+            del timeout, kwargs, unique_reply_rank
+            if method == "submit_interaction":
+                request_id, interaction = args
+                self._runner.submit_interaction(request_id, interaction)
+                return None
+            raise NotImplementedError(f"collective_rpc not mocked for {method!r}")
+
+        def abort(self, request_id: str) -> None:
+            del request_id

tests/entrypoints/openai_api/test_serving_video_output_stream.py

@@ -16,6 +16,7 @@ from fastapi import FastAPI, WebSocket
 from pytest_mock import MockerFixture
 from starlette.testclient import TestClient
 
+from tests.helpers.media import generate_synthetic_image
 from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -27,6 +28,12 @@ def _fake_video_frames(num_frames: int = 2) -> list[np.ndarray]:
     return [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(num_frames)]
 
 
+def _receive_video_chunk(ws) -> tuple[dict, bytes]:
+    metadata = ws.receive_json()
+    assert metadata["type"] == "video.chunk_metadata"
+    return metadata, ws.receive_bytes()
+
+
 def _build_test_app(
     *,
     mocker: MockerFixture,
@@ -37,6 +44,7 @@ def _build_test_app(
 ) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]:
     engine_client = mocker.MagicMock()
     engine_client.abort = mocker.AsyncMock()
+    engine_client.submit_interaction_async = mocker.AsyncMock()
     engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
     engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     engine_client.generate = mock_generate
@@ -126,8 +134,13 @@ class TestStreamingVideoOutputWebSocket:
                 assert "format" not in start["config"]
                 assert "request_id" in start
 
-                assert ws.receive_bytes() == b"mp4-chunk-0"
-                assert ws.receive_bytes() == b"mp4-chunk-1"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["transport_chunk_index"] == 0
+                assert metadata["kind"] == "media"
+                assert chunk == b"mp4-chunk-0"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["transport_chunk_index"] == 1
+                assert chunk == b"mp4-chunk-1"
 
                 done = ws.receive_json()
                 assert done["type"] == "session.done"
@@ -138,6 +151,48 @@ class TestStreamingVideoOutputWebSocket:
         engine_client.streaming_encoder_factory.assert_called_once()
         assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s"
 
+    def test_streaming_session_accepts_initial_image_reference(self, mocker: MockerFixture):
+        """session.start forwards image_reference as prompt multi_modal_data.image."""
+        captured_prompts: list[dict[str, object]] = []
+
+        async def mock_generate(*_args, **kwargs):
+            captured_prompts.append(kwargs["prompt"])
+            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", True)],
+            mock_generate=mock_generate,
+        )
+
+        synth_image = generate_synthetic_image(4, 4)
+        image_url = f"data:image/jpeg;base64,{synth_image['base64']}"
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json(
+                    {
+                        "type": "session.start",
+                        "prompt": "animate this image",
+                        "width": 4,
+                        "height": 4,
+                        "image_reference": {"image_url": image_url},
+                    }
+                )
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+                assert ws.receive_json()["type"] == "session.done"
+
+        prompt = captured_prompts[0]
+        image = prompt["multi_modal_data"]["image"]  # type: ignore[index]
+        assert image.size == (4, 4)
+        engine_client.abort.assert_not_called()
+
     def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture):
         async def mock_generate(*_args, **_kwargs):
             if False:
@@ -158,6 +213,33 @@ class TestStreamingVideoOutputWebSocket:
 
         engine_client.abort.assert_not_called()
 
+    def test_first_message_must_be_session_start(self, mocker: MockerFixture):
+        """Control messages sent before session.start are rejected at the handshake.
+
+        This test can practically ensure ``session.interaction`` before ``session.start`` is rejected,
+        but the implementation is more generic to also guard against other control messages.
+        """
+
+        async def mock_generate(*_args, **_kwargs):
+            if False:
+                yield None
+
+        app, _handler, engine_client = _build_test_app(
+            mocker=mocker,
+            streaming_chunks=[],
+            mock_generate=mock_generate,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.interaction", "interaction": {"event": {"prompt": "too early"}}})
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert err["message"] == "Expected session.start, got: session.interaction"
+
+        engine_client.abort.assert_not_called()
+        engine_client.submit_interaction_async.assert_not_called()
+
     def test_streaming_session_emits_final_encoder_delta_before_done(self, mocker: MockerFixture):
         """A non-empty encoder close delta is delivered as the final binary chunk."""
 
@@ -180,8 +262,13 @@ class TestStreamingVideoOutputWebSocket:
             with client.websocket_connect("/v1/realtime/video") as ws:
                 ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"})
                 assert ws.receive_json()["type"] == "video.start"
-                assert ws.receive_bytes() == b"mp4-chunk-0"
-                assert ws.receive_bytes() == b"mp4-trailer"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["kind"] == "media"
+                assert chunk == b"mp4-chunk-0"
+                metadata, chunk = _receive_video_chunk(ws)
+                assert metadata["kind"] == "trailer"
+                assert metadata["generation_chunk_index"] is None
+                assert chunk == b"mp4-trailer"
 
                 done = ws.receive_json()
                 assert done["type"] == "session.done"
@@ -213,7 +300,7 @@ class TestStreamingVideoOutputWebSocket:
             with client.websocket_connect("/v1/realtime/video") as ws:
                 ws.send_json({"type": "session.start", "prompt": "hello"})
                 assert ws.receive_json()["type"] == "video.start"
-                assert ws.receive_bytes() == b"mp4-chunk-0"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
                 err = ws.receive_json()
                 assert err["type"] == "error"
                 assert error_message in err["message"]
@@ -259,6 +346,7 @@ def _build_handler_for_async_tests(
 ) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]:
     engine_client = mocker.MagicMock()
     engine_client.abort = mocker.AsyncMock()
+    engine_client.submit_interaction_async = mocker.AsyncMock()
     engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
     engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     engine_client.generate = mock_generate
@@ -340,8 +428,8 @@ class TestStreamingVideoOutputStallTimeout:
             with client.websocket_connect("/v1/realtime/video") as ws:
                 ws.send_json({"type": "session.start", "prompt": "slow generation"})
                 assert ws.receive_json()["type"] == "video.start"
-                assert ws.receive_bytes() == b"mp4-chunk-0"
-                assert ws.receive_bytes() == b"mp4-chunk-1"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
                 done = ws.receive_json()
                 assert done["type"] == "session.done"
                 assert done["stopped"] is False
@@ -444,3 +532,374 @@ class TestStreamingVideoOutputStallTimeout:
         assert any(m.get("type") == "session.pong" for m in ws.sent_json)
         assert any(m.get("type") == "session.done" for m in ws.sent_json)
         engine_client.abort.assert_not_called()
+
+
+class TestStreamingVideoOutputPromptUpdate:
+    """session.interaction control messages during streaming."""
+
+    def test_prompt_update_accepted_during_streaming(self, mocker: MockerFixture):
+        """A valid interaction during generation is forwarded to the engine and acknowledged."""
+
+        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,
+        )
+
+        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": "ui-update-42",
+                            "event": {"prompt": "new scene"},
+                        },
+                    }
+                )
+                queued = ws.receive_json()
+                assert queued["type"] == "session.interaction.queued"
+                assert queued["request_id"] == request_id
+                assert queued["event_id"] == "ui-update-42"
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                done = ws.receive_json()
+                assert done["type"] == "session.done"
+                assert done["stopped"] is False
+
+        engine_client.submit_interaction_async.assert_awaited_once_with(
+            request_id,
+            interaction={"event_id": "ui-update-42", "event": {"prompt": "new scene"}},
+        )
+
+    def test_prompt_update_rejects_duplicate_event_id(self, mocker: MockerFixture):
+        """Any repeated event_id is rejected at the session API boundary."""
+
+        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,
+        )
+
+        interaction = {
+            "type": "session.interaction",
+            "interaction": {
+                "event_id": "ui-update-42",
+                "event": {"prompt": "new scene"},
+            },
+        }
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial scene"})
+                request_id = ws.receive_json()["request_id"]
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(interaction)
+                queued = ws.receive_json()
+                assert queued["type"] == "session.interaction.queued"
+
+                ws.send_json(interaction)
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert err["code"] == "event_id_conflict"
+                assert err["request_id"] == request_id
+                assert err["event_id"] == "ui-update-42"
+
+                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": "ui-update-42", "event": {"prompt": "new scene"}},
+        )
+
+    def test_prompt_update_forwards_transition_chunks(self, mocker: MockerFixture):
+        """transition_chunks is passed through to the engine client."""
+
+        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,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                request_id = ws.receive_json()["request_id"]
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "fade to sunset"},
+                            "transition_chunks": 5,
+                        },
+                    }
+                )
+                assert ws.receive_json()["type"] == "session.interaction.queued"
+                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": "ui-update-42",
+                "event": {"prompt": "fade to sunset"},
+                "transition_chunks": 5,
+            },
+        )
+
+    def test_prompt_update_rejects_empty_prompt(self, mocker: MockerFixture):
+        """An empty prompt is rejected before calling 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,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json({"type": "session.interaction", "interaction": {"event": {"prompt": "   "}}})
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert "non-empty event.prompt" in err["message"]
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_not_called()
+
+    @pytest.mark.parametrize(
+        ("transition_chunks"),
+        ["not-an-int", 2.9, True, -1],
+    )
+    def test_prompt_update_rejects_invalid_transition_chunks(
+        self,
+        mocker: MockerFixture,
+        transition_chunks: object,
+    ):
+        """Non-integral or negative transition_chunks returns an error."""
+
+        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,
+        )
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event": {"prompt": "new scene"},
+                            "transition_chunks": transition_chunks,
+                        },
+                    }
+                )
+                err = ws.receive_json()
+                assert err["type"] == "error"
+
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-1"
+                assert ws.receive_json()["type"] == "session.done"
+
+        engine_client.submit_interaction_async.assert_not_called()
+
+    def test_prompt_update_engine_failure_returns_error(self, mocker: MockerFixture):
+        """Engine failures during interaction are surfaced to the client."""
+
+        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,
+        )
+        engine_client.submit_interaction_async.side_effect = RuntimeError("engine unavailable")
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/realtime/video") as ws:
+                ws.send_json({"type": "session.start", "prompt": "initial"})
+                assert ws.receive_json()["type"] == "video.start"
+                assert _receive_video_chunk(ws)[1] == b"mp4-chunk-0"
+
+                ws.send_json(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "new scene"},
+                        },
+                    }
+                )
+                err = ws.receive_json()
+                assert err["type"] == "error"
+                assert "Failed to apply interaction" in err["message"]
+                assert err["code"] == "interaction_failed"
+                assert err["event_id"] == "ui-update-42"
+
+                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()
+
+    @pytest.mark.asyncio
+    async def test_prompt_update_accepted_via_async_handler(self, mocker: MockerFixture):
+        """interaction acknowledgement is sent while generation is in progress."""
+
+        release = asyncio.Event()
+
+        async def mock_generate(*_args, **_kwargs):
+            await release.wait()
+            yield OmniRequestOutput.from_diffusion(
+                request_id="req-ws",
+                images=[_fake_video_frames(2)],
+                final_output_type="image",
+                finished=True,
+            )
+
+        handler, engine_client, ws = _build_handler_for_async_tests(
+            mocker,
+            mock_generate=mock_generate,
+        )
+
+        async def _send_interaction_after_delay() -> None:
+            await asyncio.sleep(0.05)
+            ws.enqueue(
+                json.dumps(
+                    {
+                        "type": "session.interaction",
+                        "interaction": {
+                            "event_id": "ui-update-42",
+                            "event": {"prompt": "mid-stream update"},
+                            "transition_chunks": 2,
+                        },
+                    }
+                )
+            )
+
+        prompt_task = asyncio.create_task(_send_interaction_after_delay())
+        session_task = asyncio.create_task(handler.handle_session(ws))  # type: ignore[arg-type]
+        await asyncio.sleep(0.15)
+        release.set()
+        await session_task
+        await prompt_task
+
+        start = next(m for m in ws.sent_json if m.get("type") == "video.start")
+        request_id = start["request_id"]
+        assert any(m.get("type") == "session.interaction.queued" for m in ws.sent_json)
+        engine_client.submit_interaction_async.assert_awaited_once_with(
+            request_id,
+            interaction={
+                "event_id": "ui-update-42",
+                "event": {"prompt": "mid-stream update"},
+                "transition_chunks": 2,
+            },
+        )

tests/entrypoints/test_async_omni_prompt_update.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Entrypoint contract tests for ``AsyncOmni.submit_interaction_async``."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.client_request_state import ClientRequestState
+from vllm_omni.inputs.data import OmniInteractionPrompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_async_omni(*, num_stages: int = 1, stage_type: str = "diffusion") -> AsyncOmni:
+    omni = object.__new__(AsyncOmni)
+    omni.log_stats = False
+    omni.request_states = {
+        "external-abc-uuid-1": ClientRequestState(
+            request_id="external-abc-uuid-1",
+            external_request_id="external-abc",
+            queue=AsyncMock(),
+        ),
+        "external-abc-uuid-2": ClientRequestState(
+            request_id="external-abc-uuid-2",
+            external_request_id="external-abc",
+            queue=AsyncMock(),
+        ),
+        "other-req-uuid": ClientRequestState(
+            request_id="other-req-uuid",
+            external_request_id="other-req",
+            queue=AsyncMock(),
+        ),
+    }
+    omni.engine = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
+        num_stages=num_stages,
+        get_stage_metadata=lambda stage_id: SimpleNamespace(stage_type=stage_type),
+        submit_interaction_async=AsyncMock(),
+    )
+    return omni
+
+
+def _prompt_interaction(prompt: str = "new prompt", transition_chunks: int | None = None) -> OmniInteractionPrompt:
+    interaction: OmniInteractionPrompt = {"event_id": "ui-update-1", "event": {"prompt": prompt}}
+    if transition_chunks is not None:
+        interaction["transition_chunks"] = transition_chunks
+    return interaction
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_maps_external_to_internal_id() -> None:
+    omni = _make_async_omni()
+    with pytest.raises(ValueError, match="exactly one active request"):
+        await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction(transition_chunks=2))
+
+    omni.request_states.pop("external-abc-uuid-2")
+    await omni.submit_interaction_async(
+        "external-abc",
+        interaction=_prompt_interaction(transition_chunks=2),
+    )
+    omni.engine.submit_interaction_async.assert_awaited_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+        "external-abc-uuid-1",
+        interaction=_prompt_interaction(transition_chunks=2),
+    )
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_passes_missing_transition_chunks() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction())
+    omni.engine.submit_interaction_async.assert_awaited_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+        "external-abc-uuid-1",
+        interaction=_prompt_interaction(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_empty_prompt() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    with pytest.raises(ValueError, match="prompt must be non-empty"):
+        await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction(""))
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_negative_transition_chunks() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    with pytest.raises(ValueError, match="transition_chunks must be >= 0"):
+        await omni.submit_interaction_async(
+            "external-abc",
+            interaction=_prompt_interaction(transition_chunks=-1),
+        )
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_inactive_request() -> None:
+    omni = _make_async_omni()
+    omni.request_states.clear()
+    with pytest.raises(ValueError, match="No active request"):
+        await omni.submit_interaction_async("missing", interaction=_prompt_interaction(transition_chunks=2))
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_rejects_non_diffusion() -> None:
+    omni = _make_async_omni(stage_type="llm")
+    omni.request_states = {
+        "req-uuid": ClientRequestState(
+            request_id="req-uuid",
+            external_request_id="req",
+            queue=AsyncMock(),
+        ),
+    }
+    with pytest.raises(ValueError, match="requires a diffusion stage"):
+        await omni.submit_interaction_async("req", interaction=_prompt_interaction(transition_chunks=2))
+
+
+@pytest.mark.asyncio
+async def test_submit_interaction_async_forwards_unsupported_dict_shape_to_engine() -> None:
+    omni = _make_async_omni()
+    omni.request_states.pop("external-abc-uuid-2")
+    payload: OmniInteractionPrompt = {  # pyright: ignore[reportAssignmentType]  # deliberately wrong
+        "event": {"multi_modal_data": {"unsupported_modality": {"type": "pose"}}},
+        "transition_chunks": 2,
+    }
+
+    await omni.submit_interaction_async(
+        "external-abc",
+        interaction=payload,
+    )
+
+    omni.engine.submit_interaction_async.assert_awaited_once_with(  # pyright: ignore[reportAttributeAccessIssue]
+        "external-abc-uuid-1",
+        interaction=payload,
+    )

vllm_omni/diffusion/data.py

@@ -1274,6 +1274,9 @@ class DiffusionOutput:
     finished: bool = True
     chunk_index: int = 0
     total_chunks: int = 1
+    started_event_ids: list[str] = field(default_factory=list)
+    active_event_ids: list[str] = field(default_factory=list)
+    completed_event_ids: list[str] = field(default_factory=list)
 
     # logged duration of stages
     stage_durations: dict[str, float] = field(default_factory=dict)

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -21,7 +21,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata
 from vllm_omni.errors import client_error_metadata
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -179,6 +179,25 @@ class InlineStageDiffusionClient(StageClientBase):
                 task.cancel()
             self._engine.abort(rid)
 
+    async def submit_interaction_async(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+        timeout: float | None = None,
+    ) -> Any:
+        """Apply a midway interaction to an active streaming request."""
+        logger.debug(
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] interaction: %s",
+            self.stage_id,
+            self.replica_id,
+            request_id,
+        )
+        return await self.collective_rpc_async(
+            "submit_interaction",
+            timeout=timeout,
+            args=(request_id, interaction),
+        )
+
     async def collective_rpc_async(
         self,
         method: str,

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.prompt_update import PromptUpdateMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.platforms import current_omni_platform
@@ -155,7 +156,12 @@ def get_helios_pre_process_func(
 
 
 class HeliosPipeline(
-    nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+    nn.Module,
+    CFGParallelMixin,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+    PromptUpdateMixin,
+    SupportsComponentDiscovery,
 ):
     """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni.
 
@@ -894,9 +900,12 @@ class HeliosPipeline(
 
         output = current_latents if extra["output_type"] == "latent" else current_video
         completed_chunk_index = state.chunk_index
+        prompt_update_metadata = state.extra.pop("prompt_update_chunk_metadata", {})
         state.chunk_index += 1
         finished = state.request_denoise_completed
         if not finished:
+            # Apply queued/advancing prompt updates before preparing the next chunk.
+            self._apply_prompt_update_at_chunk_boundary(state)
             self._prepare_next_chunk(state)
         else:
             self._current_timestep = None
@@ -908,6 +917,9 @@ class HeliosPipeline(
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else {},
             chunk_index=completed_chunk_index,
             total_chunks=state.total_chunks,
+            started_event_ids=prompt_update_metadata.get("started_event_ids", []),
+            active_event_ids=prompt_update_metadata.get("active_event_ids", []),
+            completed_event_ids=prompt_update_metadata.get("completed_event_ids", []),
             finished=finished,
         )
 
@@ -1585,12 +1597,13 @@ class HeliosPipeline(
         negative_prompt: str | list[str] | None = None,
         do_classifier_free_guidance: bool = True,
         num_videos_per_prompt: int = 1,
-        max_sequence_length: int = 226,
+        max_sequence_length: int | None = None,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ):
         device = device or self.device
         dtype = dtype or self.text_encoder.dtype
+        max_sequence_length = max_sequence_length or 226
 
         prompt = [prompt] if isinstance(prompt, str) else prompt
         prompt_clean = [self._prompt_clean(p) for p in prompt]

vllm_omni/diffusion/models/interface.py

@@ -13,6 +13,8 @@ from typing import (
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import torch
 
     from vllm_omni.diffusion.data import DiffusionOutput
@@ -58,7 +60,9 @@ class SupportsStepExecution(Protocol):
         """Prepare request-level inputs and return initialized state."""
         ...
 
-    def denoise_step(self, input_batch: InputBatch, **kwargs: Any) -> torch.Tensor | None:
+    def denoise_step(
+        self, input_batch: InputBatch, *, states: Sequence[StepRequestState] | None = None, **kwargs: Any
+    ) -> torch.Tensor | None:
         """Run one denoise forward on the runner-assembled batch."""
         ...
 
@@ -100,3 +104,30 @@ def supports_step_execution(pipeline: object) -> bool:
     """Return whether `pipeline` implements :class:`SupportsStepExecution`."""
 
     return isinstance(pipeline, SupportsStepExecution)
+
+
+@runtime_checkable
+class SupportsPromptUpdate(Protocol):
+    """Optional protocol for pipelines that support midway prompt updates.
+
+    Pipelines typically implement this via
+    :class:`~vllm_omni.diffusion.prompt_update.PromptUpdateMixin`.
+    """
+
+    supports_prompt_update: ClassVar[bool] = True
+
+    def prepare_prompt_update(
+        self,
+        state: StepRequestState,
+        prompt: str,
+        event_id: str,
+        transition_chunks: int | None = None,
+    ) -> None:
+        """Encode and queue a prompt update on request-local state."""
+        ...
+
+
+def supports_prompt_update(pipeline: object) -> bool:
+    """Return whether ``pipeline`` implements :class:`SupportsPromptUpdate`."""
+
+    return isinstance(pipeline, SupportsPromptUpdate)

vllm_omni/diffusion/output_formatter.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 from collections.abc import Mapping
 from dataclasses import dataclass, field
-from typing import TypeGuard
+from typing import TypeGuard, cast
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.io_support import supports_audio_output
@@ -150,6 +150,13 @@ def format_diffusion_outputs(
         model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
         audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
 
+    stream_metadata = _stream_metadata_from_diffusion_output(diffusion_output)
+    if stream_metadata:
+        existing_stream = postprocess_output.metadata.get("stream")
+        merged_stream = dict(existing_stream) if isinstance(existing_stream, Mapping) else {}
+        merged_stream.update(stream_metadata)
+        postprocess_output.metadata["stream"] = merged_stream
+
     return _format_single_prompt_output(
         request=request,
         prompt=request.prompt,
@@ -164,6 +171,27 @@ def format_diffusion_outputs(
     )
 
 
+def _stream_metadata_from_diffusion_output(diffusion_output: DiffusionOutput) -> DiffusionMetadata | None:
+    has_stream_data = (
+        diffusion_output.total_chunks > 1
+        or not diffusion_output.finished
+        or bool(diffusion_output.started_event_ids)
+        or bool(diffusion_output.active_event_ids)
+        or bool(diffusion_output.completed_event_ids)
+    )
+    if not has_stream_data:
+        return None
+    return cast(
+        DiffusionMetadata,
+        {
+            "generation_chunk_index": diffusion_output.chunk_index,
+            "started_event_ids": diffusion_output.started_event_ids,
+            "active_event_ids": diffusion_output.active_event_ids,
+            "completed_event_ids": diffusion_output.completed_event_ids,
+        },
+    )
+
+
 def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]:
     if isinstance(outputs, list):
         return outputs

vllm_omni/diffusion/prompt_update.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Midway prompt updates for chunked diffusion step execution.
+
+``PromptUpdateMixin`` exposes
+- a pipeline interface for the external to queue a prompt update,
+- a pipeline internal method to apply prompt transitions at chunk boundaries.
+
+``prompt_update_versions``: a helper for ``InputBatch`` cache invalidation.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from dataclasses import dataclass
+from threading import Lock
+from typing import ClassVar, TypedDict, cast
+
+import torch
+
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+DEFAULT_TRANSITION_CHUNKS = 3
+
+logger = logging.getLogger(__name__)
+
+
+class PromptUpdateExtra(TypedDict, total=False):
+    """A "protocol" for ``StepRequestState.extra`` keys used by the prompt-update feature."""
+
+    pending_prompt_update: _PendingPromptUpdate
+    prompt_update_state: _PromptUpdateState
+    prompt_update_chunk_metadata: _PromptUpdateChunkMetadata
+    prompt_update_version: int  # Bumped when prompt_embeds change; invalidates InputBatch cache.
+
+    # The lock for storing/updating pending prompt updates.
+    # Saved in state.extra to be naturally cleaned up with the request state.
+    # As PromptUpdateMixin (the pipeline) currently lacks an obvious request-finish hook
+    prompt_update_lock: Lock
+
+
+def prompt_update_versions(states: Sequence[StepRequestState]) -> tuple[int, ...]:
+    """Return per-request prompt-update versions for batch cache comparison.
+
+    To be used by ``InputBatch`` or other external to determine if the cached batch is still valid.
+    """
+    return tuple(int(cast(PromptUpdateExtra, state.extra).get("prompt_update_version", 0)) for state in states)
+
+
+class PromptUpdateMixin:
+    """Mixin for chunked streaming pipelines that support midway prompt updates.
+
+    Implements the behavior expected by :class:`~vllm_omni.diffusion.models.interface.SupportsPromptUpdate`.
+    All per-request transition state lives on ``StepRequestState.extra``; this mixin itself is stateless.
+    """
+
+    supports_prompt_update: ClassVar[bool] = True
+
+    @staticmethod
+    def _prompt_update_lock(extra: PromptUpdateExtra) -> Lock:
+        return extra.setdefault("prompt_update_lock", Lock())
+
+    def prepare_prompt_update(
+        self,
+        state: StepRequestState,
+        prompt: str,
+        event_id: str,
+        transition_chunks: int | None = None,
+    ) -> None:
+        """Encode and queue a prompt update to apply at the next chunk boundary.
+        It does not actually apply the prompt transition.
+
+        This is an extended interface for the pipeline, to be called by the runner.
+        """
+        if not prompt:
+            raise ValueError("prompt must be non-empty")
+        if not event_id:
+            raise ValueError("event_id must be non-empty")
+        if state.prompt_embeds is None:
+            raise ValueError(
+                f"prompt_update is not allowed before initial generation has started (request_id={state.request_id!r})"
+            )
+        duration = DEFAULT_TRANSITION_CHUNKS if transition_chunks is None else transition_chunks
+        if duration < 0:
+            raise ValueError("transition_chunks must be >= 0")
+
+        target_prompt_embeds, _ = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=None,
+            do_classifier_free_guidance=False,
+            num_videos_per_prompt=state.sampling.num_outputs_per_prompt,
+            max_sequence_length=state.sampling.max_sequence_length,
+            device=self.device,
+            dtype=self.transformer.dtype,
+        )
+        extra = cast(PromptUpdateExtra, state.extra)
+        with self._prompt_update_lock(extra):
+            extra["pending_prompt_update"] = {
+                "event_id": event_id,
+                "prompt": prompt,
+                "target_prompt_embeds": target_prompt_embeds,
+                "transition_chunks": duration,
+            }
+
+    def _apply_prompt_update_at_chunk_boundary(self, state: StepRequestState) -> None:
+        """Advance or start prompt interpolation before the next chunk."""
+        extra = cast(PromptUpdateExtra, state.extra)
+        update_state: _PromptUpdateState | None = extra.get("prompt_update_state")
+        with self._prompt_update_lock(extra):
+            pending = extra.pop("pending_prompt_update", None)
+        embeds_changed = False
+        next_chunk_index = state.chunk_index
+        started_event_ids: list[str] = []
+        active_event_ids: list[str] = []
+        completed_event_ids: list[str] = []
+
+        # If current transition is not complete, advance it.
+        # After completion, leave prompt_update_state in place (so a later pending
+        # update can abort/overwrite), but do not keep bumping the version.
+        if update_state is not None:
+            in_transition = (
+                update_state.transition_chunks > 0
+                and update_state.elapsed_transition_chunks < update_state.transition_chunks
+            )
+            if in_transition:
+                update_state.advance_transition()
+                state.prompt_embeds = update_state.blended_prompt_embeds()
+                embeds_changed = True
+                active_event_ids.append(update_state.event_id)
+                if update_state.elapsed_transition_chunks >= update_state.transition_chunks:
+                    state.prompt_embeds = update_state.target_prompt_embeds
+                    update_state.source_prompt_embeds = update_state.target_prompt_embeds
+                    completed_event_ids.append(update_state.event_id)
+                    if update_state.target_prompt is not None:
+                        logger.debug(
+                            "prompt_update transition complete request_id=%s next_chunk_index=%d prompt=%.20s...",
+                            state.request_id,
+                            next_chunk_index,
+                            update_state.target_prompt,
+                        )
+
+        # If a new prompt update is pending, start a new transition.
+        if pending is not None:
+            if state.prompt_embeds is None:
+                raise RuntimeError(
+                    "internal error: trying to apply a pending prompt update but "
+                    f"current prompt_embeds is None (request_id={state.request_id!r})"
+                )
+            source = state.prompt_embeds.detach().clone()
+            target = pending["target_prompt_embeds"]
+            duration = int(pending["transition_chunks"])
+            prompt = str(pending.get("prompt", ""))
+            event_id = pending.get("event_id")
+            update_state = _PromptUpdateState(
+                source_prompt_embeds=source,
+                target_prompt_embeds=target,
+                transition_chunks=duration,
+                event_id=event_id,
+                target_prompt=prompt,
+            )
+            extra["prompt_update_state"] = update_state
+            started_event_ids.append(event_id)
+            active_event_ids.append(event_id)
+            if duration <= 0:
+                # A hard/immediate transition
+                state.prompt_embeds = target
+                completed_event_ids.append(event_id)
+                logger.debug(
+                    "prompt_update sharp transition request_id=%s next_chunk_index=%d prompt=%.20s...",
+                    state.request_id,
+                    next_chunk_index,
+                    prompt,
+                )
+            else:
+                # A transition that really takes some time
+                update_state.advance_transition()
+                state.prompt_embeds = update_state.blended_prompt_embeds()
+                if update_state.elapsed_transition_chunks >= update_state.transition_chunks:
+                    state.prompt_embeds = update_state.target_prompt_embeds
+                    update_state.source_prompt_embeds = update_state.target_prompt_embeds
+                    completed_event_ids.append(event_id)
+                logger.debug(
+                    "prompt_update transition start request_id=%s next_chunk_index=%d prompt=%.20s...",
+                    state.request_id,
+                    next_chunk_index,
+                    prompt,
+                )
+            embeds_changed = True
+
+        # Indicate that prompt embeddings have changed---clear input batch cache.
+        if embeds_changed:
+            new_version = int(extra.get("prompt_update_version", 0)) + 1
+            extra["prompt_update_version"] = new_version
+
+        extra["prompt_update_chunk_metadata"] = {
+            "started_event_ids": started_event_ids,
+            "active_event_ids": active_event_ids,
+            "completed_event_ids": completed_event_ids,
+        }
+
+
+@dataclass
+class _PromptUpdateState:
+    """Per-request prompt interpolation stored in ``state.extra["prompt_update_state"]``."""
+
+    source_prompt_embeds: torch.Tensor
+    target_prompt_embeds: torch.Tensor
+    transition_chunks: int
+    event_id: str
+    elapsed_transition_chunks: int = 0
+    target_prompt: str | None = None
+
+    def blended_prompt_embeds(self) -> torch.Tensor:
+        alpha = self._current_alpha()
+        if alpha >= 1.0:
+            return self.target_prompt_embeds
+        return (1.0 - alpha) * self.source_prompt_embeds + alpha * self.target_prompt_embeds
+
+    def advance_transition(self) -> None:
+        if self.transition_chunks <= 0:
+            self.source_prompt_embeds = self.target_prompt_embeds
+            return
+        self.elapsed_transition_chunks += 1
+        if self.elapsed_transition_chunks >= self.transition_chunks:
+            self.source_prompt_embeds = self.target_prompt_embeds
+            self.elapsed_transition_chunks = self.transition_chunks
+
+    def _current_alpha(self) -> float:
+        if self.transition_chunks <= 0:
+            return 1.0
+        return min(1.0, self.elapsed_transition_chunks / self.transition_chunks)
+
+
+class _PendingPromptUpdate(TypedDict):
+    event_id: str
+    prompt: str
+    target_prompt_embeds: torch.Tensor
+    transition_chunks: int
+
+
+class _PromptUpdateChunkMetadata(TypedDict):
+    started_event_ids: list[str]
+    active_event_ids: list[str]
+    completed_event_ids: list[str]

vllm_omni/diffusion/stage_diffusion_client.py

@@ -34,7 +34,7 @@ from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt, OmniPromptType
 
 logger = init_logger(__name__)
 _MISSING_RPC_RESULT = object()
@@ -395,6 +395,19 @@ class StageDiffusionClient(StageClientBase):
             )
         )
 
+    async def submit_interaction_async(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+        timeout: float | None = None,
+    ) -> Any:
+        """Apply a midway interaction to an active streaming request."""
+        return await self.collective_rpc_async(
+            "submit_interaction",
+            timeout=timeout,
+            args=(request_id, interaction),
+        )
+
     async def collective_rpc_async(
         self,
         method: str,

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -14,7 +14,7 @@ import copy
 import time
 from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from torch.profiler import record_function
@@ -32,7 +32,7 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import supports_step_execution
+from vllm_omni.diffusion.models.interface import SupportsPromptUpdate, supports_prompt_update, supports_step_execution
 from vllm_omni.diffusion.offloader import get_offload_backend
 from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -53,6 +53,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
+if TYPE_CHECKING:
+    from vllm_omni.inputs.data import OmniInteractionPrompt
+
 logger = init_logger(__name__)
 
 
@@ -803,3 +806,48 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 self._update_states_after(states, input_batch, pipeline_interrupted)
 
                 return BatchRunnerOutput.from_list(runner_output_list)
+
+    def submit_interaction(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Route a midway interaction to the matching active stepwise request feature."""
+        assert self.pipeline 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")
+
+        event = interaction.get("event")
+        if isinstance(event, dict) and "prompt" in event and "multi_modal_data" not in event:
+            # Is a prompt update interaction.
+            self._submit_prompt_update_interaction(request_id, interaction)
+            return
+
+        raise NotImplementedError(
+            "Only text-only prompt update interactions with 'event.prompt' and optional "
+            "'transition_chunks' are supported in this release"
+        )
+
+    def _submit_prompt_update_interaction(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Queue a prompt-update interaction for an active stepwise request."""
+        if not supports_prompt_update(self.pipeline):
+            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 prompt_update: {request_id!r}")
+
+        event = cast(dict[str, Any], interaction.get("event"))
+        prompt = event["prompt"]
+        transition_chunks = interaction.get("transition_chunks")
+        event_id = interaction.get("event_id")
+        if not isinstance(event_id, str) or not event_id:
+            raise ValueError("event_id must be non-empty")
+        pipeline = cast(SupportsPromptUpdate, self.pipeline)
+        pipeline.prepare_prompt_update(state, prompt, event_id, transition_chunks)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -50,6 +50,7 @@ from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPref
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput
 from vllm_omni.engine.stage_init_utils import set_death_signal
+from vllm_omni.inputs.data import OmniInteractionPrompt
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.profiler import OmniTorchProfilerWrapper, create_omni_profiler
@@ -511,6 +512,15 @@ class DiffusionWorker:
         # for the frontend server yet.
         return self.lora_manager.add_adapter(lora_request)
 
+    def submit_interaction(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Apply a midway interaction to an active stepwise request."""
+        assert self.model_runner is not None, "Model runner not initialized"
+        self.model_runner.submit_interaction(request_id, interaction)
+
     def list_loras(self) -> list[int]:
         return self.lora_manager.list_adapters()
 

vllm_omni/diffusion/worker/input_batch.py

@@ -17,6 +17,7 @@ from dataclasses import dataclass, field
 import numpy as np
 import torch
 
+from vllm_omni.diffusion.prompt_update import prompt_update_versions
 from vllm_omni.diffusion.worker.utils import StepRequestState
 
 
@@ -535,12 +536,17 @@ def _same_composition(
     cached_batch: InputBatch | None,
     request_ids: list[str],
     idx_mapping_np: np.ndarray,
+    states: Sequence[StepRequestState],
 ) -> bool:
     if cached_batch is None:
         return False
     if cached_batch.request_ids != request_ids:
         return False
-    return np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np)
+    if not np.array_equal(cached_batch.idx_mapping_np, idx_mapping_np):
+        return False
+    # Midway prompt updates (typically for video generation) can change prompt_embeds without changing ids/mapping.
+    # In this case, each request state manages the embedding's "version". Use it to determine if cache is still valid.
+    return cached_batch._prompt_update_versions == prompt_update_versions(states)
 
 
 def _scatter_batch_tensor_by_mapping(
@@ -615,6 +621,10 @@ class InputBatch:
     negative_txt_seq_lens: list[int] | None = None
     states: Sequence[StepRequestState] = field(default_factory=tuple)
 
+    # For midway prompt updates (typically for video generation) that changes embeddings without changing ids,
+    # Keep a snapshot of the current per-request versions of prompt embeddings for later runtime comparison
+    _prompt_update_versions: tuple[int, ...] = ()
+
     def __post_init__(self) -> None:
         if len(self.request_ids) != int(self.idx_mapping.numel()):
             raise ValueError("`request_ids` and `idx_mapping` must have the same length.")
@@ -653,6 +663,7 @@ class InputBatch:
         self.img_shapes = _prepare_img_shapes(states)
         self.txt_seq_lens = _prepare_seq_lens(states, "txt_seq_lens")
         self.negative_txt_seq_lens = _prepare_seq_lens(states, "negative_txt_seq_lens")
+        self._prompt_update_versions = prompt_update_versions(states)
 
     def _repack_dynamic_fields(
         self,
@@ -693,7 +704,7 @@ class InputBatch:
         selected_states, idx_mapping, idx_mapping_np = _select_states(states, idx_mapping)
         request_ids = _prepare_request_ids(selected_states)
 
-        if _same_composition(cached_batch, request_ids, idx_mapping_np):
+        if _same_composition(cached_batch, request_ids, idx_mapping_np, selected_states):
             assert cached_batch is not None
             cached_batch._repack_dynamic_fields(selected_states)
             return cached_batch
@@ -733,6 +744,7 @@ class InputBatch:
                 "negative_txt_seq_lens",
             ),
             states=tuple(selected_states),
+            _prompt_update_versions=prompt_update_versions(selected_states),
         )
 
 

vllm_omni/diffusion/worker/request_batch.py

@@ -48,6 +48,9 @@ def split_diffusion_output_by_request(
             peak_memory_mb=result.peak_memory_mb,
             chunk_index=result.chunk_index,
             total_chunks=result.total_chunks,
+            started_event_ids=result.started_event_ids,
+            active_event_ids=result.active_event_ids,
+            completed_event_ids=result.completed_event_ids,
         )
         for idx in range(req.num_reqs)
     ]

vllm_omni/engine/async_omni_engine.py

@@ -55,6 +55,7 @@ from vllm_omni.engine.messages import (
     CollectiveRPCResultMessage,
     EngineQueueMessage,
     ErrorMessage,
+    InteractionMessage,
     StageSubmissionMessage,
 )
 from vllm_omni.engine.orchestrator import Orchestrator
@@ -71,7 +72,7 @@ from vllm_omni.entrypoints.utils import (
     load_and_resolve_stage_configs,
     parse_stage_overrides,
 )
-from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 
 logger = init_logger(__name__)
@@ -1703,6 +1704,30 @@ class AsyncOmniEngine:
         """Async abort API."""
         self.abort(request_ids)
 
+    def submit_interaction(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Send an interaction control message to the Orchestrator."""
+        if self.request_queue is None:
+            raise RuntimeError("request_queue is not initialized")
+
+        self.request_queue.sync_q.put_nowait(
+            InteractionMessage(
+                request_id=request_id,
+                interaction=interaction,
+            )
+        )
+
+    async def submit_interaction_async(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Async interaction API."""
+        self.submit_interaction(request_id, interaction)
+
     def collective_rpc(
         self,
         method: str,

vllm_omni/engine/messages.py

@@ -6,7 +6,7 @@ import msgspec
 from vllm.inputs import PromptType
 from vllm.v1.engine import EngineCoreRequest
 
-from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.inputs.data import OmniInteractionPrompt, OmniSamplingParams
 from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -44,6 +44,12 @@ class AbortRequestMessage(EngineQueueMessage, kw_only=True):
     request_ids: list[str]
 
 
+class InteractionMessage(EngineQueueMessage, kw_only=True):
+    type: Literal["interaction"] = "interaction"
+    request_id: str
+    interaction: OmniInteractionPrompt
+
+
 class CollectiveRPCRequestMessage(EngineQueueMessage, kw_only=True):
     type: Literal["collective_rpc"] = "collective_rpc"
     rpc_id: str
@@ -78,6 +84,7 @@ class ErrorMessage(EngineQueueMessage, kw_only=True):
     fatal: bool = False
     request_id: str | None = None
     stage_id: int | None = None
+    event_id: str | None = None  # for interactions on diffusion generation requests
 
 
 class OutputMessage(EngineQueueMessage, kw_only=True):

vllm_omni/engine/orchestrator.py

@@ -40,6 +40,7 @@ from vllm_omni.engine.messages import (
     CollectiveRPCResultMessage,
     EngineQueueMessage,
     ErrorMessage,
+    InteractionMessage,
     OutputMessage,
     RegisterRemoteReplicaMessage,
     ShutdownRequestMessage,
@@ -577,6 +578,8 @@ class Orchestrator:
                 self.duplex_control_plane.dispatch(msg)
             elif msg_type == "abort":
                 await self._handle_abort(msg)
+            elif msg_type == "interaction":
+                await self._handle_interaction(msg)
             elif msg_type == "collective_rpc":
                 await self._handle_collective_rpc(msg)
             elif isinstance(msg, RegisterRemoteReplicaMessage):
@@ -768,6 +771,44 @@ class Orchestrator:
         )
         logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
 
+    async def _handle_interaction(self, msg: InteractionMessage) -> None:
+        """Handle a midway interaction for an active streaming diffusion request."""
+        stage_id = 0
+        request_id = msg.request_id
+        event_id = msg.interaction.get("event_id")
+        req_state = self.request_states.get(request_id)
+        if req_state is None:
+            logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id)
+            await self.output_async_queue.put(
+                ErrorMessage(
+                    error=f"No active request for interaction: {request_id}",
+                    fatal=False,
+                    request_id=request_id,
+                    event_id=event_id,
+                    stage_id=stage_id,
+                )
+            )
+            return
+
+        try:
+            await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction)
+        except Exception as exc:
+            logger.info(
+                "[Orchestrator] Failed interaction for req %s: %s",
+                request_id,
+                exc,
+                exc_info=True,
+            )
+            await self.output_async_queue.put(
+                ErrorMessage(
+                    error=f"Failed interaction for request {request_id}: {exc}",
+                    fatal=False,
+                    request_id=request_id,
+                    event_id=event_id,
+                    stage_id=stage_id,
+                )
+            )
+
     async def _abort_request_ids(self, request_ids: list[str]) -> None:
         """Forward abort requests to all stage pools."""
         if not request_ids:

vllm_omni/engine/stage_client.py

@@ -8,7 +8,12 @@ from typing import TYPE_CHECKING, Any, Protocol
 if TYPE_CHECKING:
     from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest
 
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniTokensPrompt
+    from vllm_omni.inputs.data import (
+        OmniDiffusionSamplingParams,
+        OmniInteractionPrompt,
+        OmniPromptType,
+        OmniTokensPrompt,
+    )
     from vllm_omni.outputs import OmniRequestOutput
 
 from vllm_omni.inputs.data import OmniSamplingParams
@@ -101,3 +106,10 @@ class StagePoolDiffusionClient(StagePoolClient, Protocol):
     ) -> None: ...
 
     def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: ...
+
+    async def submit_interaction_async(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+        timeout: float | None = None,
+    ) -> Any: ...

vllm_omni/engine/stage_pool.py

@@ -24,7 +24,7 @@ from vllm_omni.engine.stage_client import (
     StagePoolDiffusionClient,
     StagePoolLLMClient,
 )
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
 from vllm_omni.metrics import (
     count_audio_frames,
     count_image_pixels,
@@ -1048,6 +1048,24 @@ class StagePool:
                 raise
         return replica_id
 
+    async def submit_interaction(
+        self,
+        request_id: str,
+        interaction: OmniInteractionPrompt,
+    ) -> int:
+        """Submit a midway interaction to an active diffusion (typically video generation) request."""
+        replica_id = self.get_bound_replica_id(request_id)
+        if replica_id is None or self.clients[replica_id] is None:
+            replica_id = await self._pick_or_select(request_id)
+
+        client = self._diffusion_client(replica_id)
+        result = await client.submit_interaction_async(request_id, interaction)
+        # Interaction may raise error if underlying check against ODConfig fails.
+        if isinstance(result, dict) and result.get("error"):
+            reason = result.get("reason") or "Unknown interaction RPC error"
+            raise ValueError(str(reason))
+        return replica_id
+
     async def _pick_or_select(
         self,
         request_id: str,

vllm_omni/entrypoints/async_omni.py

@@ -50,7 +50,7 @@ if TYPE_CHECKING:
         DuplexSessionLifecycleMessage,
     )
     from vllm_omni.experimental.fullduplex.request_client import DuplexRequestClient
-    from vllm_omni.inputs.data import OmniPromptType
+    from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType
 
 logger = init_logger(__name__)
 _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001
@@ -1006,6 +1006,46 @@ class AsyncOmni(EngineClient, OmniBase):
         internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id in request_ids]
         await self._abort(internal_ids)
 
+    async def submit_interaction_async(
+        self,
+        request_id: str,
+        *,
+        interaction: OmniInteractionPrompt,
+    ) -> None:
+        """Apply a midway interaction to an active streaming diffusion request.
+
+        ``request_id`` is the external id created by the server-side session,
+        matching the value passed to :meth:`generate`.
+        """
+        event = interaction.get("event")
+        prompt = event.get("prompt") if isinstance(event, dict) else None
+        if isinstance(event, dict) and "prompt" in event and (not isinstance(prompt, str) or not prompt):
+            raise ValueError("prompt must be non-empty")
+        transition_chunks = interaction.get("transition_chunks")
+        if transition_chunks is not None and transition_chunks < 0:
+            raise ValueError("transition_chunks must be >= 0")
+
+        if self.num_stages != 1:
+            raise ValueError("interaction requires single-stage diffusion")
+        stage_meta = self.engine.get_stage_metadata(0)
+        if stage_meta.stage_type != "diffusion":
+            raise ValueError("interaction requires a diffusion stage")
+
+        internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id == request_id]
+        if not internal_ids:
+            raise ValueError(f"No active request for interaction: {request_id!r}")
+        if len(internal_ids) > 1:
+            raise ValueError(
+                f"interaction requires exactly one active request for {request_id!r}, found {len(internal_ids)}"
+            )
+
+        await self.engine.submit_interaction_async(
+            internal_ids[0],
+            interaction=interaction,
+        )
+        if self.log_stats:
+            logger.info("[AsyncOmni] Queued interaction for request %s", request_id)
+
     async def _abort_internal_requests(self, request_id: str | Iterable[str]):
         """Abort request(s) via the Orchestrator given internal request IDs,
         which take the format <external_request_id>-<UUID>.

vllm_omni/entrypoints/openai/serving_video_output_stream.py

@@ -5,13 +5,22 @@
 Protocol:
     Client -> Server:
         {"type": "session.start", "model": "...", "prompt": "...", "format": "m4s", ...}
-        {"type": "session.prompt_update", "prompt": "..."}  # currently unsupported
+        {
+            "type": "session.interaction",
+            "interaction": {
+                "event_id": "...",
+                "event": {"prompt": "..."},
+                "transition_chunks": (optional, integer),
+            },
+        }
         {"type": "session.stop"}
         {"type": "session.ping"}  # optional; refreshes stall clock
 
     Server -> Client:
         {"type": "video.start", "request_id": "...", "config": {...}, "format": "m4s"}
+        {"type": "video.chunk_metadata", ...}
         <binary frame: fragmented MP4 bytes>
+        {"type": "session.interaction.queued", "request_id": "...", "event_id": "..."}
         {"type": "session.done", ...}
         {"type": "session.pong"}  # reply to session.ping
         {"type": "error", "message": "..."}
@@ -21,7 +30,6 @@ from __future__ import annotations
 
 import asyncio
 import copy
-import inspect
 import json
 import time
 import uuid
@@ -29,19 +37,34 @@ from collections.abc import AsyncGenerator
 from http import HTTPStatus
 from typing import Any, cast
 
+import numpy as np
+import torch
 from fastapi import HTTPException, WebSocket, WebSocketDisconnect
+from PIL import Image
 from pydantic import ValidationError
-from vllm.engine.protocol import EngineClient
 from vllm.logger import init_logger
 
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoParams
 from vllm_omni.entrypoints.openai.stage_params import (
     build_stage_sampling_params_list,
     get_default_sampling_params_list,
 )
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
-from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat, create_streaming_video_encoder
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.entrypoints.openai.video_api_utils import (
+    StreamingVideoFormat,
+    create_streaming_video_encoder,
+    decode_input_reference,
+)
+from vllm_omni.inputs.data import (
+    OmniDiffusionSamplingParams,
+    OmniInteractionEvent,
+    OmniInteractionPrompt,
+    OmniTextPrompt,
+)
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.outputs.output_metadata import DiffusionPayloadValue
 
 logger = init_logger(__name__)
 
@@ -70,7 +93,7 @@ class OmniStreamingVideoOutputHandler:
 
     def __init__(
         self,
-        engine_client: EngineClient,
+        engine_client: AsyncOmni,
         model_name: str | None = None,
         stage_configs: list[Any] | None = None,
         stall_timeout: float = _DEFAULT_STALL_TIMEOUT,
@@ -91,6 +114,7 @@ class OmniStreamingVideoOutputHandler:
         control_task: asyncio.Task[None] | None = None
         stopped = False
         chunks_sent = 0
+        interaction_payloads_by_id: dict[str, OmniInteractionPrompt] = {}
 
         try:
             start = await self._receive_start(websocket)
@@ -112,13 +136,22 @@ class OmniStreamingVideoOutputHandler:
             progress.touch()
 
             control_task = asyncio.create_task(
-                self._control_loop(websocket, request_id, stop_event, send_lock, progress)
+                self._control_loop(
+                    websocket,
+                    request_id,
+                    stop_event,
+                    send_lock,
+                    progress,
+                    interaction_payloads_by_id,
+                )
             )
-            async for chunk in self._stream_video_bytes(request, request_id, output_format, progress):
+            async for chunk, metadata in self._stream_video_bytes(request, request_id, output_format, progress):
                 if stop_event.is_set():
                     stopped = True
                     break
                 async with send_lock:
+                    metadata["transport_chunk_index"] = chunks_sent
+                    await websocket.send_json(metadata)
                     await websocket.send_bytes(chunk)
                 chunks_sent += 1
                 progress.touch()
@@ -234,6 +267,7 @@ class OmniStreamingVideoOutputHandler:
         stop_event: asyncio.Event,
         send_lock: asyncio.Lock,
         progress: _SessionProgress,
+        interaction_payloads_by_id: dict[str, OmniInteractionPrompt],
     ) -> None:
         while not stop_event.is_set():
             if progress.stalled_for(self._stall_timeout):
@@ -282,34 +316,151 @@ class OmniStreamingVideoOutputHandler:
                 except Exception:
                     pass
                 continue
-            if msg_type == "session.prompt_update":
-                await self._send_error(
+            if msg_type == "session.interaction":
+                await self._handle_interaction(
                     websocket,
-                    "session.prompt_update is not supported by this backend yet",
+                    msg,
+                    request_id=request_id,
                     send_lock=send_lock,
+                    interaction_payloads_by_id=interaction_payloads_by_id,
                 )
                 continue
             await self._send_error(websocket, f"Unknown message type: {msg_type}", send_lock=send_lock)
 
     async def _abort_request(self, request_id: str) -> None:
         try:
-            abort = getattr(self._engine_client, "abort", None)
-            if callable(abort):
-                result = abort(request_id)
-                if inspect.isawaitable(result):
-                    await result
+            await self._engine_client.abort(request_id)
         except Exception:
             logger.debug("Failed to abort streaming video request %s", request_id, exc_info=True)
 
+    async def _handle_interaction(
+        self,
+        websocket: WebSocket,
+        msg: dict[str, Any],
+        *,
+        request_id: str,
+        send_lock: asyncio.Lock,
+        interaction_payloads_by_id: dict[str, OmniInteractionPrompt],
+    ) -> None:
+        interaction = msg.get("interaction")
+        if not isinstance(interaction, dict):
+            await self._send_error(websocket, "session.interaction requires an interaction object", send_lock=send_lock)
+            return
+
+        event_id = interaction.get("event_id")
+        if event_id is None:
+            event_id = uuid.uuid4().hex
+        elif not isinstance(event_id, str) or not event_id:
+            await self._send_error(
+                websocket,
+                "session.interaction event_id must be a non-empty string",
+                send_lock=send_lock,
+            )
+            return
+
+        event = interaction.get("event")
+        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():
+            await self._send_error(
+                websocket,
+                "session.interaction requires a non-empty event.prompt",
+                send_lock=send_lock,
+            )
+            return
+
+        if "negative_prompt" in event and event["negative_prompt"] is not None:
+            await self._send_error(
+                websocket,
+                "session.interaction negative_prompt is not supported in this release",
+                send_lock=send_lock,
+            )
+            return
+
+        normalized: OmniInteractionPrompt = {
+            "event_id": event_id,
+            "event": cast(OmniInteractionEvent, event),
+        }
+        transition_chunks = interaction.get("transition_chunks")
+        if transition_chunks is not None:
+            try:
+                if (
+                    isinstance(transition_chunks, bool)
+                    or not isinstance(transition_chunks, (int, float))
+                    or (isinstance(transition_chunks, float) and not transition_chunks.is_integer())
+                ):
+                    raise ValueError
+                normalized["transition_chunks"] = int(transition_chunks)
+            except (TypeError, ValueError):
+                await self._send_error(
+                    websocket,
+                    "session.interaction transition_chunks must be an integer",
+                    send_lock=send_lock,
+                )
+                return
+            if normalized["transition_chunks"] < 0:
+                await self._send_error(
+                    websocket,
+                    "session.interaction transition_chunks must be >= 0",
+                    send_lock=send_lock,
+                )
+                return
+
+        if event_id in interaction_payloads_by_id:
+            await self._send_error(
+                websocket,
+                "event_id_conflict",
+                code="event_id_conflict",
+                request_id=request_id,
+                event_id=event_id,
+                send_lock=send_lock,
+            )
+            return
+
+        interaction_payloads_by_id[event_id] = normalized
+
+        try:
+            await self._engine_client.submit_interaction_async(
+                request_id,
+                interaction=normalized,
+            )
+        except Exception:
+            logger.exception("Failed to apply interaction for request %s", request_id, exc_info=True)
+            interaction_payloads_by_id.pop(event_id, None)
+            await self._send_error(
+                websocket,
+                "Failed to apply interaction",
+                code="interaction_failed",
+                request_id=request_id,
+                event_id=event_id,
+                send_lock=send_lock,
+            )
+            return
+
+        try:
+            async with send_lock:
+                await websocket.send_json(
+                    {
+                        "type": "session.interaction.queued",
+                        "request_id": request_id,
+                        "event_id": event_id,
+                    }
+                )
+        except Exception:
+            logger.debug("Failed to send interaction acknowledgement for %s", request_id, exc_info=True)
+
     async def _stream_video_bytes(
         self,
         request: VideoGenerationRequest,
         request_id: str,
         output_format: StreamingVideoFormat,
         progress: _SessionProgress,
-    ) -> AsyncGenerator[bytes, None]:
+    ) -> AsyncGenerator[tuple[bytes, dict[str, Any]], None]:
         """Yield encoded video bytes from diffusion streaming outputs."""
-        prompt, gen_params, vp = self._build_prompt_and_sampling_params(request)
+        prompt, gen_params, vp = await self._build_prompt_and_sampling_params(request)
         video_codec_options = {"preset": "ultrafast", "threads": "0", "tune": "zerolatency"}
         if isinstance(request.extra_params, dict) and "video_codec_options" in request.extra_params:
             video_codec_options = request.extra_params["video_codec_options"]
@@ -323,23 +474,41 @@ class OmniStreamingVideoOutputHandler:
         completed = False
         try:
             async for result in self._iter_generation_outputs(prompt, gen_params, request_id, progress):
-                if error := getattr(result, "error", None):
-                    raise RuntimeError(str(error))
+                if result.error:
+                    raise RuntimeError(str(result.error))
                 videos = self._extract_video_outputs(result)
                 for video in videos:
                     chunk = encoder.encode(video)
                     if chunk:
-                        yield chunk
+                        yield (
+                            chunk,
+                            self._build_chunk_metadata(
+                                request_id=request_id,
+                                kind="media",
+                                byte_length=len(chunk),
+                                num_frames=self._num_video_frames(video),
+                                stream_metadata=self._extract_stream_metadata(result),
+                            ),
+                        )
 
             completed = True
             final_chunk = encoder.close()
             if final_chunk:
-                yield final_chunk
+                yield (
+                    final_chunk,
+                    self._build_chunk_metadata(
+                        request_id=request_id,
+                        kind="trailer",
+                        byte_length=len(final_chunk),
+                        num_frames=0,
+                        stream_metadata={},
+                    ),
+                )
         finally:
             if not completed:
                 encoder.close()
 
-    def _build_prompt_and_sampling_params(
+    async def _build_prompt_and_sampling_params(
         self,
         request: VideoGenerationRequest,
     ) -> tuple[OmniTextPrompt, OmniDiffusionSamplingParams, VideoParams]:
@@ -350,6 +519,10 @@ class OmniStreamingVideoOutputHandler:
 
         gen_params = self._resolve_default_sampling_params()
         vp = request.resolve_video_params()
+        input_image = await self._decode_image_reference(request)
+        if input_image is not None:
+            prompt["multi_modal_data"] = {"image": input_image}
+
         if vp.width is not None and vp.height is not None:
             gen_params.width = vp.width
             gen_params.height = vp.height
@@ -386,12 +559,28 @@ class OmniStreamingVideoOutputHandler:
         self._apply_lora(request.lora, gen_params)
         return prompt, gen_params, vp
 
+    @staticmethod
+    async def _decode_image_reference(request: VideoGenerationRequest) -> Image.Image | None:
+        try:
+            media_data = await decode_input_reference(
+                request.image_reference,
+                None,
+                None,
+            )
+        except InvalidInputReferenceError as exc:
+            raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc
+        if media_data is not None and not isinstance(media_data, Image.Image):
+            raise HTTPException(
+                HTTPStatus.BAD_REQUEST.value,
+                detail="image_reference must resolve to an image.",
+            )
+        return media_data
+
     def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams:
-        default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None)
-        if default_sampling_params_list:
-            for params in default_sampling_params_list:
-                if isinstance(params, OmniDiffusionSamplingParams):
-                    return copy.deepcopy(params)
+        default_sampling_params_list = self._engine_client.default_sampling_params_list
+        for params in default_sampling_params_list:
+            if isinstance(params, OmniDiffusionSamplingParams):
+                return copy.deepcopy(params)
         return OmniDiffusionSamplingParams()
 
     @staticmethod
@@ -417,8 +606,8 @@ class OmniStreamingVideoOutputHandler:
         gen_params: OmniDiffusionSamplingParams,
         request_id: str,
         progress: _SessionProgress,
-    ) -> AsyncGenerator[Any, None]:
-        stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None)
+    ) -> AsyncGenerator[OmniRequestOutput, None]:
+        stage_configs = self._stage_configs or self._engine_client.stage_configs
         if not stage_configs:
             raise HTTPException(
                 status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
@@ -433,15 +622,14 @@ class OmniStreamingVideoOutputHandler:
                     detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.",
                 )
 
-        engine_client = cast(Any, self._engine_client)
         sampling_params_list = build_stage_sampling_params_list(
             list(stage_configs),
-            get_default_sampling_params_list(engine_client),
+            get_default_sampling_params_list(self._engine_client),
             diffusion_params=gen_params,
             replace_diffusion_params=True,
         )
 
-        async for output in engine_client.generate(
+        async for output in self._engine_client.generate(
             prompt=prompt,
             request_id=request_id,
             sampling_params_list=sampling_params_list,
@@ -450,44 +638,49 @@ class OmniStreamingVideoOutputHandler:
             yield output
 
     @staticmethod
-    def _normalize_video_outputs(videos: Any) -> list[Any]:
+    def _normalize_video_outputs(videos: DiffusionPayloadValue | None) -> list[DiffusionPayloadValue]:
         if videos is None:
             return []
-        if hasattr(videos, "ndim") and videos.ndim == 5:
+        if isinstance(videos, (torch.Tensor, np.ndarray)) and videos.ndim == 5:
             return [videos[i] for i in range(videos.shape[0])]
         if isinstance(videos, list):
             if not videos:
                 return []
             first = videos[0]
-            if hasattr(first, "ndim") and first.ndim == 5:
-                flattened: list[Any] = []
+            if isinstance(first, (torch.Tensor, np.ndarray)) and first.ndim == 5:
+                flattened: list[DiffusionPayloadValue] = []
                 for item in videos:
-                    if hasattr(item, "ndim") and item.ndim == 5:
+                    if isinstance(item, (torch.Tensor, np.ndarray)) and item.ndim == 5:
                         flattened.extend([item[i] for i in range(item.shape[0])])
                     else:
                         flattened.append(item)
                 return flattened
             if isinstance(first, list):
                 return videos
-            if hasattr(first, "ndim") and first.ndim == 3:
+            if isinstance(first, (torch.Tensor, np.ndarray)) and first.ndim == 3:
+                return [videos]
+            if isinstance(first, Image.Image):
                 return [videos]
             return videos
         return [videos]
 
-    def _extract_video_outputs(self, result: Any) -> list[Any]:
-        videos = None
-        if hasattr(result, "images") and result.images:
-            videos = result.images
-        elif hasattr(result, "request_output"):
+    def _extract_video_outputs(self, result: OmniRequestOutput) -> list[DiffusionPayloadValue]:
+        videos: DiffusionPayloadValue | None = None
+        if result.images:
+            # Video frames may be list[list[Image.Image]] despite images: list[Image.Image].
+            videos = result.images  # type: ignore[assignment]
+        elif result.request_output is not None:
             request_output = result.request_output
             if isinstance(request_output, dict) and request_output.get("images"):
                 videos = request_output["images"]
-            elif not isinstance(request_output, dict):
-                if hasattr(request_output, "images") and request_output.images:
-                    videos = request_output.images
-                elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
-                    videos = request_output.multimodal_output.get("video")
-        if videos is None and hasattr(result, "multimodal_output") and result.multimodal_output:
+            elif isinstance(request_output, OmniRequestOutput):
+                if request_output.images:
+                    # Same nested-frames mismatch as result.images above.
+                    videos = request_output.images  # type: ignore[assignment]
+                else:
+                    request_multimodal_output = request_output.multimodal_output
+                    videos = request_multimodal_output.get("video")
+        if videos is None:
             videos = result.multimodal_output.get("video")
 
         normalized = self._normalize_video_outputs(videos)
@@ -498,18 +691,64 @@ class OmniStreamingVideoOutputHandler:
             )
         return normalized
 
+    @staticmethod
+    def _extract_stream_metadata(result: OmniRequestOutput) -> dict[str, Any]:
+        multimodal_output = result.multimodal_output
+        if not isinstance(multimodal_output, dict):
+            return {}
+        metadata = multimodal_output.get("metadata")
+        if not isinstance(metadata, dict):
+            return {}
+        stream = metadata.get("stream")
+        return dict(stream) if isinstance(stream, dict) else {}
+
+    @staticmethod
+    def _num_video_frames(video: DiffusionPayloadValue) -> int:
+        if isinstance(video, (torch.Tensor, np.ndarray)) and len(video.shape) > 0:
+            return int(video.shape[0])
+        if isinstance(video, list):
+            return len(video)
+        return 0
+
+    @staticmethod
+    def _build_chunk_metadata(
+        *,
+        request_id: str,
+        kind: str,
+        byte_length: int,
+        num_frames: int,
+        stream_metadata: dict[str, Any],
+    ) -> dict[str, Any]:
+        return {
+            "type": "video.chunk_metadata",
+            "request_id": request_id,
+            "kind": kind,
+            "num_frames": num_frames,
+            "byte_length": byte_length,
+            "generation_chunk_index": (stream_metadata.get("generation_chunk_index") if kind == "media" else None),
+            "started_event_ids": stream_metadata.get("started_event_ids", []),
+            "active_event_ids": stream_metadata.get("active_event_ids", []),
+            "completed_event_ids": stream_metadata.get("completed_event_ids", []),
+        }
+
     @staticmethod
     async def _send_error(
         websocket: WebSocket,
         message: str,
-        code: int | None = None,
+        code: int | str | None = None,
         *,
+        request_id: str | None = None,
+        event_id: str | None = None,
         send_lock: asyncio.Lock | None = None,
     ) -> None:
         try:
             payload: dict[str, Any] = {"type": "error", "message": message}
             if code is not None:
                 payload["code"] = code
+            if request_id is not None:
+                payload["request_id"] = request_id
+            if event_id is not None:
+                payload["event_id"] = event_id
             if send_lock is None:
                 await websocket.send_json(payload)
             else:

vllm_omni/inputs/data.py

@@ -43,6 +43,20 @@ class OmniTextPrompt(TextPrompt):
     model_intermediate_buffer: NotRequired[dict[str, Any]]
 
 
+class OmniInteractionEvent(OmniTextPrompt, total=False):
+    """Prompt-event payload for mid-stream interactions."""
+
+    pass
+
+
+class OmniInteractionPrompt(TypedDict, total=False):
+    """Mid-stream interaction payload."""
+
+    event_id: NotRequired[str]
+    event: NotRequired[OmniInteractionEvent]
+    transition_chunks: NotRequired[int]
+
+
 class OmniTokensPrompt(TokensPrompt):
     """Tokens prompt with optional embeddings and stage payloads.