From eaf90d8eff7aa2a9753d5f98344fb34f60e658bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:07:43 +0800 Subject: [PATCH] [Frontend] Add midway prompt update for streaming video generation (#4652) Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> --- .../streaming_video_generation/README.md | 16 +- .../streaming_video_generation/gradio_demo.py | 231 +++++-- .../streaming_video_client.py | 171 +++++- .../video-stream-view.js | 49 ++ realtime_video_prompt_interaction_protocol.md | 138 +++++ .../test_diffusion_streaming_output.py | 3 +- tests/diffusion/test_prompt_update.py | 574 ++++++++++++++++++ .../test_serving_video_output_stream.py | 473 ++++++++++++++- .../test_async_omni_prompt_update.py | 140 +++++ vllm_omni/diffusion/data.py | 3 + .../inline_stage_diffusion_client.py | 21 +- .../models/helios/pipeline_helios.py | 17 +- vllm_omni/diffusion/models/interface.py | 33 +- vllm_omni/diffusion/output_formatter.py | 30 +- vllm_omni/diffusion/prompt_update.py | 245 ++++++++ vllm_omni/diffusion/stage_diffusion_client.py | 15 +- .../worker/diffusion_model_runner.py | 52 +- .../diffusion/worker/diffusion_worker.py | 10 + vllm_omni/diffusion/worker/input_batch.py | 16 +- vllm_omni/diffusion/worker/request_batch.py | 3 + vllm_omni/engine/async_omni_engine.py | 27 +- vllm_omni/engine/messages.py | 9 +- vllm_omni/engine/orchestrator.py | 41 ++ vllm_omni/engine/stage_client.py | 14 +- vllm_omni/engine/stage_pool.py | 20 +- vllm_omni/entrypoints/async_omni.py | 42 +- .../openai/serving_video_output_stream.py | 341 +++++++++-- vllm_omni/inputs/data.py | 14 + 28 files changed, 2615 insertions(+), 133 deletions(-) create mode 100755 realtime_video_prompt_interaction_protocol.md create mode 100644 tests/diffusion/test_prompt_update.py create mode 100644 tests/entrypoints/test_async_omni_prompt_update.py create mode 100644 vllm_omni/diffusion/prompt_update.py diff --git a/examples/online_serving/streaming_video_generation/README.md b/examples/online_serving/streaming_video_generation/README.md index 1cfa29dfd..7caa74906 100644 --- a/examples/online_serving/streaming_video_generation/README.md +++ b/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 diff --git a/examples/online_serving/streaming_video_generation/gradio_demo.py b/examples/online_serving/streaming_video_generation/gradio_demo.py index 2f7a301f1..964b70f89 100644 --- a/examples/online_serving/streaming_video_generation/gradio_demo.py +++ b/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), diff --git a/examples/online_serving/streaming_video_generation/streaming_video_client.py b/examples/online_serving/streaming_video_generation/streaming_video_client.py index c77113420..e8dcd8cdd 100644 --- a/examples/online_serving/streaming_video_generation/streaming_video_client.py +++ b/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() diff --git a/examples/online_serving/streaming_video_generation/video-stream-view.js b/examples/online_serving/streaming_video_generation/video-stream-view.js index 1b0a5024b..46206c0d0 100644 --- a/examples/online_serving/streaming_video_generation/video-stream-view.js +++ b/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"); diff --git a/realtime_video_prompt_interaction_protocol.md b/realtime_video_prompt_interaction_protocol.md new file mode 100755 index 000000000..534d76a51 --- /dev/null +++ b/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": [] +} +``` diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index ae5f9eee8..824902ad4 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/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 diff --git a/tests/diffusion/test_prompt_update.py b/tests/diffusion/test_prompt_update.py new file mode 100644 index 000000000..cdee35979 --- /dev/null +++ b/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 diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index 72d293853..ec1727fe8 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/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, + }, + ) diff --git a/tests/entrypoints/test_async_omni_prompt_update.py b/tests/entrypoints/test_async_omni_prompt_update.py new file mode 100644 index 000000000..f3edeeec7 --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a001c665c..9e049b211 100644 --- a/vllm_omni/diffusion/data.py +++ b/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) diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 4e0922ff5..3c4edda26 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/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, diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 90ee0ab3f..cb11eb847 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/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] diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index bd26c0395..1f2c8c226 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/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) diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index b369ed0c3..d5b03443c 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/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 diff --git a/vllm_omni/diffusion/prompt_update.py b/vllm_omni/diffusion/prompt_update.py new file mode 100644 index 000000000..400838035 --- /dev/null +++ b/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] diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 2b13ddc9d..a8733b8d0 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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, diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index ee1b61c21..c54b721f8 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 948406c73..b7450038a 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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() diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py index c0d10198a..d21e73287 100644 --- a/vllm_omni/diffusion/worker/input_batch.py +++ b/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), ) diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index 1f7c2acaa..488d1eec2 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/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) ] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 9deabf6d3..162f3657f 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 384d8a055..80bf42b0e 100644 --- a/vllm_omni/engine/messages.py +++ b/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): diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 9a1d63170..affd0e452 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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: diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 6db090955..74f65d8c2 100644 --- a/vllm_omni/engine/stage_client.py +++ b/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: ... diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index fe4ec5a78..487134db1 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 34f5f0a91..332e5b921 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 -. diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index f713df07c..64327420c 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/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", ...} + {"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: diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 1eb4dc202..6626176b6 100644 --- a/vllm_omni/inputs/data.py +++ b/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.