diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 809e3feef..65ad51057 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -1322,6 +1322,7 @@ steps: - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model agents: queue: "mithril-h100-pool" plugins: diff --git a/examples/online_serving/streaming_video_generation/README.md b/examples/online_serving/streaming_video_generation/README.md new file mode 100644 index 000000000..1cfa29dfd --- /dev/null +++ b/examples/online_serving/streaming_video_generation/README.md @@ -0,0 +1,94 @@ +# Streaming Video Generation + +This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced. +It covers text-only video generation. Image/reference input is intentionally not included for now. + +## Start The Server + +Start a diffusion video model with streaming output enabled: + +```bash +vllm serve BestWishYsh/Helios-Distilled \ + --omni \ + --diffusion-streaming-output \ + --port 8000 +``` + +The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=True` in the default diffusion stage `engine_args`, then loaded by `OmniDiffusionConfig.from_kwargs()`. + +## WebSocket Protocol + +| 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. | +| 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 | 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.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. | + +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. + +## Install Client Dependency + +```bash +pip install websockets +# For the Gradio demo: +pip install 'vllm-omni[demo]' websockets +``` + +## Run The Client + +```bash +python streaming_video_client.py \ + --host 127.0.0.1 \ + --port 8000 \ + --model BestWishYsh/Helios-Distilled \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --width 640 \ + --height 384 \ + --fps 16 \ + --num-frames 99 \ + --guidance-scale 1.0 \ + --seed 42 \ + --output helios_stream.mp4 +``` + +The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`. +The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration. + +## Run The Gradio Demo + +```bash +python gradio_demo.py \ + --host 127.0.0.1 \ + --port 7860 +``` + +The Gradio demo requests fMP4 (`m4s`) chunks and appends them directly in the browser with a Media Source Extensions player. + +## Model Choice + +### Helios + +The example uses `BestWishYsh/Helios-Distilled` model by default. + +To ensure streaming-level generation speed, `pyramid_num_inference_steps_list` is suggested to be as low as `[1, 1, 1]`. Both example clients uses the following Helios-Distilled preset by default: + +```json +{ + "is_enable_stage2": true, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": true +} +``` + +Disable it in the CLI example with `--no-helios-distilled-preset`, or override/extend it with `--extra-params`: + +```bash +python streaming_video_client.py \ + --extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}' +``` diff --git a/examples/online_serving/streaming_video_generation/gradio_demo.py b/examples/online_serving/streaming_video_generation/gradio_demo.py new file mode 100644 index 000000000..2f7a301f1 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/gradio_demo.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gradio demo for `/v1/realtime/video`. + +Usage: + python gradio_demo.py --host 127.0.0.1 --port 7860 + +Start a vLLM-Omni server with streaming output enabled first, for example: + vllm serve BestWishYsh/Helios-Distilled --omni --diffusion-streaming-output --port 8000 +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +try: + import gradio as gr +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." +VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html") +VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js") + + +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 + + +def _optional_int(value: int | float | None) -> int | None: + return None if value is None else int(value) + + +def _optional_float(value: int | float | None) -> float | None: + return None if value is None else float(value) + + +def _build_session_start( + *, + model: str, + prompt: str, + negative_prompt: str, + width: int | float | None, + height: int | float | None, + fps: int | float | None, + num_frames: int | float | None, + num_inference_steps: int | float | None, + guidance_scale: int | float | None, + guidance_scale_2: int | float | None, + boundary_ratio: int | float | None, + flow_shift: int | float | None, + true_cfg_scale: int | float | None, + seed: int | float | None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "type": "session.start", + "model": model, + "prompt": prompt, + "format": "m4s", + } + + _maybe_set(payload, "negative_prompt", negative_prompt) + _maybe_set(payload, "width", _optional_int(width)) + _maybe_set(payload, "height", _optional_int(height)) + _maybe_set(payload, "fps", _optional_int(fps)) + _maybe_set(payload, "num_frames", _optional_int(num_frames)) + _maybe_set(payload, "num_inference_steps", _optional_int(num_inference_steps)) + _maybe_set(payload, "guidance_scale", _optional_float(guidance_scale)) + _maybe_set(payload, "guidance_scale_2", _optional_float(guidance_scale_2)) + _maybe_set(payload, "boundary_ratio", _optional_float(boundary_ratio)) + _maybe_set(payload, "flow_shift", _optional_float(flow_shift)) + _maybe_set(payload, "true_cfg_scale", _optional_float(true_cfg_scale)) + _maybe_set(payload, "seed", _optional_int(seed)) + + if "Helios" in model: + payload["extra_params"] = { + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": True, + } + + return payload + + +def _set_optional_field_enabled(enabled: bool) -> Any: + return gr.update(interactive=enabled) + + +def build_browser_stream_config( + host: str, + port: int | float, + model: str, + prompt: str, + negative_prompt: str, + width: int | float | None, + height: int | float | None, + fps: int | float, + num_frames: int | float | None, + num_inference_steps: int | float | None, + include_guidance_scale: bool, + guidance_scale: int | float | None, + include_guidance_scale_2: bool, + guidance_scale_2: int | float | None, + include_seed: bool, + seed: int | float | None, + include_boundary_ratio: bool, + boundary_ratio: int | float | None, + include_flow_shift: bool, + flow_shift: int | float | None, + include_true_cfg_scale: bool, + true_cfg_scale: int | float | None, +) -> tuple[str, Any]: + """Build the browser-side WebSocket request config for the HTML/MSE player.""" + if not model.strip(): + raise gr.Error("Model is required.") + if not prompt.strip(): + raise gr.Error("Prompt is required.") + + payload = _build_session_start( + model=model.strip(), + prompt=prompt.strip(), + negative_prompt=negative_prompt, + width=width, + height=height, + fps=fps, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale if include_guidance_scale else None, + guidance_scale_2=guidance_scale_2 if include_guidance_scale_2 else None, + boundary_ratio=boundary_ratio if include_boundary_ratio else None, + flow_shift=flow_shift if include_flow_shift else None, + true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None, + seed=seed if include_seed else None, + ) + config = { + "url": f"ws://{host}:{int(port)}/v1/realtime/video", + "payload": payload, + } + 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: + gr.Markdown("# Streaming Video Generation") + gr.Markdown( + "Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks " + "directly to a Media Source Extensions video player." + ) + + with gr.Row(): + with gr.Column(scale=1): + host = gr.Textbox(label="Server Host", value="127.0.0.1") + 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) + height = gr.Number(label="Height", value=384, precision=0) + + with gr.Row(): + fps = gr.Number(label="FPS", value=16, precision=0) + num_frames = gr.Number(label="Num Frames", value=132, 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.Row(): + start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start") + hidden_stream_config = gr.Textbox(value="", visible=False) + + with gr.Column(scale=1): + gr.HTML( + VIDEO_STREAM_VIEW_HTML.read_text(encoding="utf-8"), + label="Streaming Preview", + js_on_load=VIDEO_STREAM_VIEW_JS.read_text(encoding="utf-8"), + ) + + inputs = [ + host, + port, + model, + prompt, + negative_prompt, + width, + height, + fps, + num_frames, + num_inference_steps, + include_guidance_scale, + guidance_scale, + include_guidance_scale_2, + guidance_scale_2, + include_seed, + seed, + include_boundary_ratio, + boundary_ratio, + include_flow_shift, + flow_shift, + include_true_cfg_scale, + true_cfg_scale, + ] + for toggle, field in [ + (include_guidance_scale, guidance_scale), + (include_guidance_scale_2, guidance_scale_2), + (include_seed, seed), + (include_boundary_ratio, boundary_ratio), + (include_flow_shift, flow_shift), + (include_true_cfg_scale, true_cfg_scale), + ]: + toggle.change( + fn=_set_optional_field_enabled, + inputs=toggle, + outputs=field, + ) + start_button.click( + fn=build_browser_stream_config, + inputs=inputs, + outputs=[hidden_stream_config, start_button], + ).then( + fn=lambda config: config, + inputs=hidden_stream_config, + outputs=hidden_stream_config, + js=""" + (config) => { + if (window.vllmStreamingVideoStart) { + window.vllmStreamingVideoStart(config); + } + return config; + } + """, + ) + + return demo + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Streaming video generation Gradio demo") + parser.add_argument("--host", default="127.0.0.1", help="Host/IP for gradio launch") + parser.add_argument("--port", type=int, default=7860, help="Port for gradio launch") + parser.add_argument("--share", action="store_true", help="Create a public Gradio link") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + demo = create_demo() + demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/streaming_video_generation/streaming_video_client.py b/examples/online_serving/streaming_video_generation/streaming_video_client.py new file mode 100644 index 000000000..c77113420 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/streaming_video_client.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Text-only WebSocket client for `/v1/realtime/video`. + +The client prints a line for every binary video chunk it receives and writes +the received bytes to disk when the session finishes or is interrupted. +Image/reference input is intentionally not supported in this example yet. + +Requirements: + pip install websockets +""" + +import argparse +import asyncio +import io +import json +import os +import time +from pathlib import Path +from typing import Any + +import av + +try: + from websockets.asyncio.client import connect # pyright: ignore[reportMissingImports] +except ImportError: + try: + from websockets import connect # pyright: ignore[reportMissingImports] + except ImportError: + print("Please install websockets: pip install websockets") + raise SystemExit(1) + + +HELIOS_DISTILLED_EXTRA_PARAMS = { + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": True, +} + + +def _json_object(value: str, *, argument_name: 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 + if not isinstance(parsed, dict): + raise argparse.ArgumentTypeError(f"{argument_name} must be a JSON object") + return parsed + + +def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None: + if value is not None: + payload[key] = value + + +def build_session_start(args: argparse.Namespace) -> dict[str, Any]: + payload: dict[str, Any] = { + "type": "session.start", + "model": args.model, + "prompt": args.prompt, + } + + _maybe_set(payload, "negative_prompt", args.negative_prompt) + _maybe_set(payload, "width", args.width) + _maybe_set(payload, "height", args.height) + _maybe_set(payload, "size", args.size) + _maybe_set(payload, "seconds", args.seconds) + _maybe_set(payload, "fps", args.fps) + _maybe_set(payload, "num_frames", args.num_frames) + _maybe_set(payload, "num_inference_steps", args.num_inference_steps) + _maybe_set(payload, "guidance_scale", args.guidance_scale) + _maybe_set(payload, "guidance_scale_2", args.guidance_scale_2) + _maybe_set(payload, "boundary_ratio", args.boundary_ratio) + _maybe_set(payload, "flow_shift", args.flow_shift) + _maybe_set(payload, "true_cfg_scale", args.true_cfg_scale) + _maybe_set(payload, "seed", args.seed) + + extra_params: dict[str, Any] = {} + if "Helios-Distilled" in args.model and args.helios_distilled_preset: + extra_params.update(HELIOS_DISTILLED_EXTRA_PARAMS) + if args.extra_params: + extra_params.update(args.extra_params) + if extra_params: + payload["extra_params"] = extra_params + + return payload + + +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: + return 0 + + av_format = "mp4" if stream_format == "m4s" else None + try: + with av.open(io.BytesIO(video_bytes), format=av_format, mode="r") as container: + stream = container.streams.video[0] + return sum(1 for _ in container.decode(stream)) + except Exception: + return None + + +def save_video( + chunks: list[bytes], + output_path: Path, + *, + fps: float, + stream_format: str, +) -> None: + if not chunks: + print("No video chunks received; nothing to save.") + return + + from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes + + streamed_bytes = b"".join(chunks) + playback_bytes = finalize_streaming_video_bytes( + streamed_bytes, + input_format="m4s", + fps=fps, + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = output_path.with_suffix(output_path.suffix + ".tmp") + with tmp_path.open("wb") as f: + f.write(playback_bytes) + os.replace(tmp_path, output_path) + print( + f"Saved {len(chunks)} streamed chunk(s), {len(streamed_bytes)} streamed bytes " + f"-> {len(playback_bytes)} playback bytes at {output_path}" + ) + + +async def stream_video(args: argparse.Namespace) -> None: + payload = build_session_start(args) + output_path = Path(args.output) + url = f"ws://{args.host}:{args.port}/v1/realtime/video" + chunks: list[bytes] = [] + total_bytes = 0 + total_frames = 0 + accumulated_stream = bytearray() + stream_format = "m4s" + done = False + ws = None + + try: + async with connect(url, max_size=None) as websocket: + ws = websocket + 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() + last_chunk_at = started_at + + while True: + message = await websocket.recv() + now = time.perf_counter() + + if isinstance(message, bytes): + chunks.append(message) + chunk_bytes = len(message) + chunk_elapsed = now - last_chunk_at + last_chunk_at = now + + total_bytes += chunk_bytes + accumulated_stream.extend(message) + decoded_total = _count_decoded_frames(bytes(accumulated_stream), stream_format=stream_format) + if decoded_total is not None: + chunk_frames = max(decoded_total - total_frames, 0) + total_frames = decoded_total + else: + chunk_frames = 0 + + total_elapsed = now - started_at + print( + f"[chunk {len(chunks):03d}] " + f"bytes={chunk_bytes} " + f"frames={chunk_frames} " + f"elapsed={chunk_elapsed:.2f}s " + f"total_bytes={total_bytes} " + f"total_frames={total_frames} " + f"total_elapsed={total_elapsed:.2f}s" + ) + continue + + msg = json.loads(message) + msg_type = msg.get("type") + 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')}") + elif msg_type == "session.done": + print(f"Session complete: {json.dumps(msg, ensure_ascii=False)}") + done = True + break + elif msg_type == "error": + print(f"ERROR: {msg.get('message', msg)}") + else: + print(f"Received control message: {msg}") + except (KeyboardInterrupt, asyncio.CancelledError): + print("\nInterrupted. Asking server to stop the active session...") + if ws is not None: + try: + await ws.send(json.dumps({"type": "session.stop"})) + await ws.close(code=1000, reason="client interrupted") + except Exception: + pass + raise + finally: + print( + "Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)" + ) + save_video( + chunks, + output_path, + fps=float(payload.get("fps") or args.fps), + stream_format=stream_format, + ) + if not done: + print("Saved partial output because the stream did not finish normally.") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Streaming video generation WebSocket client") + parser.add_argument("--host", default="127.0.0.1", help="API server host") + parser.add_argument("--port", type=int, default=8000, help="API server port") + parser.add_argument( + "--model", + default="BestWishYsh/Helios-Distilled", + help="Model name served by the API server", + ) + 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.", + help="Text prompt for video generation", + ) + parser.add_argument("--output", default="streaming_video_output.mp4", help="Path for the received video bytes") + parser.add_argument("--width", type=int, default=640, help="Video width") + parser.add_argument("--height", type=int, default=384, help="Video height") + 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("--negative-prompt", default=None, help="Negative prompt") + parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps") + parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale") + parser.add_argument("--guidance-scale-2", type=float, default=None, help="High-noise CFG scale for video models") + parser.add_argument("--boundary-ratio", type=float, default=None, help="Boundary split ratio for video models") + parser.add_argument("--flow-shift", type=float, default=None, help="Scheduler flow_shift for video models") + parser.add_argument("--true-cfg-scale", type=float, default=None, help="True CFG scale for models that support it") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + + parser.add_argument( + "--helios-distilled-preset", + action=argparse.BooleanOptionalAction, + default=True, + help="Send the Helios-Distilled extra_params used by the bundled example.", + ) + parser.add_argument( + "--extra-params", + type=lambda value: _json_object(value, argument_name="--extra-params"), + default=None, + help="JSON object merged into extra_params; overrides preset keys on conflict.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + try: + asyncio.run(stream_video(args)) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/streaming_video_generation/video-stream-view.html b/examples/online_serving/streaming_video_generation/video-stream-view.html new file mode 100644 index 000000000..bbef95c69 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/video-stream-view.html @@ -0,0 +1,11 @@ + +
+ +
+ Ready +
+ 
Ready.
+
diff --git a/examples/online_serving/streaming_video_generation/video-stream-view.js b/examples/online_serving/streaming_video_generation/video-stream-view.js new file mode 100644 index 000000000..1b0a5024b --- /dev/null +++ b/examples/online_serving/streaming_video_generation/video-stream-view.js @@ -0,0 +1,198 @@ +/** + * This file and the corresponding HTML file are needed to display a streamable video view in the Gradio demo. + * Gradio's built-in video player does not support modern M4S streaming format. + */ +(function () { + const playerId = "vllm-streaming-video-player"; + const statusId = "vllm-streaming-video-status"; + const logId = "vllm-streaming-video-log"; + let ws = null; + let mediaSource = null; + let sourceBuffer = null; + let queue = []; + let chunkCount = 0; + let totalBytes = 0; + let done = false; + let objectUrl = null; + + function el(id) { + return document.getElementById(id); + } + + function setStatus(text) { + const node = el(statusId); + if (node) node.textContent = text; + } + + function log(line) { + const node = el(logId); + if (!node) return; + node.textContent += "\n" + line; + node.scrollTop = node.scrollHeight; + } + + function setStartButton(disabled, label) { + const btn = document.querySelector("#streaming-video-start button") || document.getElementById("streaming-video-start"); + if (!btn) return; + btn.disabled = disabled; + btn.textContent = label; + } + + function chooseMime() { + const candidates = [ + 'video/mp4; codecs="avc1.42E01E"', + 'video/mp4; codecs="avc1.4D401F"', + 'video/mp4; codecs="avc1.64001F"', + "video/mp4" + ]; + for (const mime of candidates) { + if (window.MediaSource && MediaSource.isTypeSupported(mime)) return mime; + } + return null; + } + + function pump() { + if (!sourceBuffer || sourceBuffer.updating || queue.length === 0) return; + try { + sourceBuffer.appendBuffer(queue.shift()); + } catch (err) { + done = true; + setStatus("SourceBuffer error"); + setStartButton(false, "Restart"); + log("SourceBuffer append error: " + err.message); + finishStream(); + } + } + + function finishStream() { + if (!mediaSource || mediaSource.readyState !== "open") return; + if (sourceBuffer && sourceBuffer.updating) return; + if (queue.length > 0) return; + try { + mediaSource.endOfStream(); + } catch (_) { } + } + + function stopCurrent() { + if (ws && ws.readyState === WebSocket.OPEN) { + try { ws.send(JSON.stringify({ type: "session.stop" })); } catch (_) { } + try { ws.close(); } catch (_) { } + } + ws = null; + sourceBuffer = null; + mediaSource = null; + queue = []; + if (objectUrl) URL.revokeObjectURL(objectUrl); + objectUrl = null; + } + + window.vllmStreamingVideoStart = function (configJson) { + stopCurrent(); + chunkCount = 0; + totalBytes = 0; + done = false; + const logNode = el(logId); + if (logNode) logNode.textContent = "Starting..."; + setStatus("Starting..."); + + let config; + try { + config = JSON.parse(configJson); + } catch (err) { + setStatus("Invalid request"); + setStartButton(false, "Restart"); + log("Invalid request config: " + err.message); + return configJson; + } + + const mime = chooseMime(); + if (!mime) { + setStatus("MSE unsupported"); + setStartButton(false, "Restart"); + log("This browser does not support MP4 Media Source Extensions playback."); + return configJson; + } + + const video = el(playerId); + if (!video) return configJson; + mediaSource = new MediaSource(); + objectUrl = URL.createObjectURL(mediaSource); + video.src = objectUrl; + + mediaSource.addEventListener("sourceopen", () => { + try { + sourceBuffer = mediaSource.addSourceBuffer(mime); + sourceBuffer.mode = "segments"; + sourceBuffer.addEventListener("updateend", () => { + pump(); + if (done) finishStream(); + }); + } catch (err) { + setStatus("SourceBuffer error"); + setStartButton(false, "Restart"); + log("SourceBuffer error: " + err.message); + return; + } + + ws = new WebSocket(config.url); + ws.binaryType = "arraybuffer"; + ws.onopen = () => { + setStatus("Streaming..."); + setStartButton(true, "Streaming..."); + log("Connected: " + config.url); + ws.send(JSON.stringify(config.payload)); + log("Sent session.start: " + JSON.stringify(config.payload)); + }; + ws.onmessage = (event) => { + if (typeof event.data === "string") { + const msg = JSON.parse(event.data); + if (msg.type === "video.start") { + log("Video session started: request_id=" + (msg.request_id || "") + " format=" + (msg.format || "")); + } else if (msg.type === "session.done") { + done = true; + setStatus("Done"); + setStartButton(false, "Restart"); + log("Session complete: " + JSON.stringify(msg)); + finishStream(); + } else if (msg.type === "error") { + done = true; + setStatus("Error"); + setStartButton(false, "Restart"); + log("ERROR: " + (msg.message || JSON.stringify(msg))); + finishStream(); + } else { + log("Control message: " + JSON.stringify(msg)); + } + return; + } + + const data = event.data; + queue.push(data); + chunkCount += 1; + totalBytes += data.byteLength; + log(`[chunk ${String(chunkCount).padStart(3, "0")}] bytes=${data.byteLength} total_bytes=${totalBytes}`); + pump(); + }; + ws.onerror = () => { + done = true; + setStatus("Error"); + setStartButton(false, "Restart"); + log("ERROR: WebSocket error"); + finishStream(); + }; + ws.onclose = () => { + if (!done) { + done = true; + setStatus("Closed"); + setStartButton(false, "Restart"); + log("WebSocket closed before session.done"); + finishStream(); + } + }; + }, { once: true }); + + return configJson; + }; + + console.log("video-stream-view.html loaded"); +})(); diff --git a/pyproject.toml b/pyproject.toml index 82d7af131..196b544e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dev = [ "pytest-asyncio>=0.21.0", "pytest-cov>=4.0.0", "pytest-mock>=3.10.0", + "websockets>=15.0.1", "datasets>=2.14.0", "mypy==1.11.1", "pre-commit==4.0.1", diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index a8985f372..f940c11d6 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -9,6 +9,7 @@ from unittest.mock import Mock import pytest +from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler @@ -33,6 +34,7 @@ def _make_engine() -> DiffusionEngine: engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) engine._out_queue = {} + engine._out_queue_streaming = {} engine._closed = False engine._shutdown_complete = False engine.abort_queue = queue.Queue() @@ -57,6 +59,23 @@ def test_close_completes_pending_async_waiters() -> None: event_loop.close() +def test_close_completes_pending_streaming_waiters() -> None: + engine = _make_engine() + event_loop = asyncio.new_event_loop() + try: + engine.main_loop = event_loop + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_queue_streaming["pending-stream"] = queue + + engine.close() + + output = queue.get_nowait() + assert output.error == "DiffusionEngine is closed." + assert output.finished is True + finally: + event_loop.close() + + def test_handle_finished_requests_ignores_already_drained_waiter() -> None: class RacingOutQueue(dict): def __contains__(self, key): diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 960dd072b..95a0d9c27 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -9,6 +9,7 @@ import torch import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner pytestmark = [pytest.mark.diffusion] @@ -31,6 +32,46 @@ class _DummyPipeline: return self._output +class _ChunkStepPipeline: + device = torch.device("cpu") + supports_step_execution = True + + def __init__(self, outputs): + self._outputs = outputs + self.prepare_calls = 0 + self.decode_calls = 0 + + def prepare_encode(self, state): + self.prepare_calls += 1 + state.prompt_embeds = torch.zeros(1, 1, 1) + state.latents = torch.zeros(1, 1) + state.timesteps = torch.tensor([1.0, 0.0]) + state.step_index = 0 + state.step_in_chunk = 0 + state.chunk_num_steps = 2 + state.total_chunks = len(self._outputs) + return state + + def denoise_step(self, input_batch, states): + del states + return torch.ones_like(input_batch.latents) + + def step_scheduler(self, state, noise_pred): + state.latents = noise_pred + state.step_index += 1 + state.step_in_chunk += 1 + + def post_decode(self, state): + output = self._outputs[self.decode_calls] + self.decode_calls += 1 + state.chunk_index += 1 + state.step_index = 0 + state.step_in_chunk = 0 + if not state.request_denoise_completed: + state.latents = torch.zeros(1, 1) + return output + + def _make_request(skip_cache_refresh: bool = True): sampling_params = SimpleNamespace( generator=None, @@ -52,11 +93,13 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok")) runner.cache_backend = cache_backend runner.offload_backend = None + runner.state_cache = {} runner.prompt_embed_cache = None runner.od_config = SimpleNamespace( cache_backend=cache_backend_name, enable_cache_dit_summary=enable_cache_dit_summary, parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.kv_transfer_manager = SimpleNamespace( receive_kv_cache=lambda req, target_device=None: None, @@ -66,6 +109,49 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar return runner +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): + """Step streaming returns empty step results until a chunk decode boundary.""" + chunks = [ + DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), + DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), + ] + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.pipeline = _ChunkStepPipeline(chunks) + runner.od_config.streaming_output = True + runner.od_config.step_execution = True + req = _make_request(skip_cache_refresh=True) + req.request_id = "req" + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req)], + scheduled_cached_reqs=SimpleNamespace(request_ids=[]), + ) + + first = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert first.get_request_output("req").result is None + + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]), + ) + second = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert second.get_request_output("req").result == chunks[0] + assert second.get_request_output("req").finished is False + + DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + fourth = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert fourth.get_request_output("req").result == chunks[1] + assert fourth.get_request_output("req").finished is True + + @pytest.mark.core_model @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch): @@ -76,6 +162,9 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa cache_summary_calls = [] monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) monkeypatch.setattr( model_runner_module, "cache_summary", @@ -101,6 +190,9 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp cache_summary_calls = [] monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) monkeypatch.setattr( model_runner_module, "cache_summary", @@ -157,6 +249,7 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): cache_config={}, model_class_name="NextStep11Pipeline", enforce_eager=True, + streaming_output=False, ) monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 503ea0842..280b56420 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio import queue import threading from types import SimpleNamespace @@ -238,6 +239,35 @@ class TestRequestScheduler: assert finished == {req_id} assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED + def test_streaming_output_keeps_request_running_until_final_chunk(self) -> None: + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace()) + req_id = scheduler.add_request(_make_request("stream")) + + sched_output = scheduler.schedule() + chunk = RunnerOutput( + request_id=req_id, + step_index=1, + finished=False, + result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), + ) + finished = scheduler.update_from_output(sched_output, chunk) + + assert finished == set() + assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING + assert scheduler.has_requests() is True + + final_chunk = RunnerOutput( + request_id=req_id, + step_index=2, + finished=True, + result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), + ) + finished = scheduler.update_from_output(sched_output, final_chunk) + + assert finished == {req_id} + assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED + def test_fifo_single_request_scheduling(self) -> None: req_id_a = self.scheduler.add_request(_make_request("a")) req_id_b = self.scheduler.add_request(_make_request("b")) @@ -369,6 +399,7 @@ class TestRequestScheduler: class TestDiffusionEngine: def test_add_req_and_wait_for_response_single_path(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(streaming_output=False) engine.scheduler = RequestScheduler() engine.scheduler.initialize(SimpleNamespace()) engine._rpc_lock = threading.RLock() @@ -391,6 +422,7 @@ class TestDiffusionEngine: scheduler = _StubScheduler(request, runner_output) engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(streaming_output=False) engine.scheduler = scheduler engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) @@ -410,7 +442,7 @@ class TestDiffusionEngine: ) -> None: request = _make_request("init") scheduler = _StubScheduler(request, DiffusionOutput(output=None)) - od_config = SimpleNamespace(model_class_name="mock_model") + od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False) fake_executor_cls = mocker.Mock(return_value=mocker.Mock()) monkeypatch.setattr( @@ -472,7 +504,7 @@ class TestDiffusionEngine: def test_finalize_finished_request_returns_aborted_output(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) - engine.scheduler = RequestScheduler() + engine.scheduler = StepScheduler() engine.scheduler.initialize(SimpleNamespace()) req_id = engine.scheduler.add_request(_make_request("req-finalize")) @@ -483,12 +515,77 @@ class TestDiffusionEngine: assert output.aborted is True assert output.abort_message == "Request req-finalize aborted." + @pytest.mark.asyncio + async def test_streaming_runner_output_notifies_each_chunk(self) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = StepScheduler() + engine.scheduler.initialize(SimpleNamespace()) + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._out_queue_streaming = {} + engine.main_loop = asyncio.get_running_loop() + + req_id = engine.scheduler.add_request(_make_request("stream-engine")) + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_queue_streaming[req_id] = queue + sched_output = engine.scheduler.schedule() + + chunk = RunnerOutput( + request_id=req_id, + step_index=1, + finished=False, + result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), + ) + finished_req_ids = engine.scheduler.update_from_output(sched_output, chunk) + engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, chunk) + + notified_chunk = await asyncio.wait_for(queue.get(), timeout=1) + assert notified_chunk.output == "chunk-0" + assert notified_chunk.finished is False + assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING + + final_chunk = RunnerOutput( + request_id=req_id, + step_index=2, + finished=True, + result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), + ) + finished_req_ids = engine.scheduler.update_from_output(sched_output, final_chunk) + engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, final_chunk) + + notified_final = await asyncio.wait_for(queue.get(), timeout=1) + assert notified_final.output == "chunk-1" + assert notified_final.finished is True + assert engine.scheduler.get_request_state(req_id) is None + + @pytest.mark.asyncio + async def test_finished_streaming_request_without_runner_output_notifies_waiter(self) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = RequestScheduler() + engine.scheduler.initialize(SimpleNamespace()) + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._out_queue_streaming = {} + engine.main_loop = asyncio.get_running_loop() + + req_id = engine.scheduler.add_request(_make_request("stream-abort")) + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_queue_streaming[req_id] = queue + engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED) + + engine._handle_empty_streaming_requests({req_id}) + + output = await asyncio.wait_for(queue.get(), timeout=1) + assert output.aborted is True + assert output.finished is True + assert engine.scheduler.get_request_state(req_id) is None + def test_initializes_step_scheduler_when_step_execution_enabled( self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, ) -> None: - od_config = SimpleNamespace(model_class_name="mock_model") + od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False) od_config.step_execution = True fake_executor = mocker.Mock() fake_executor_cls = mocker.Mock(return_value=fake_executor) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index a21813d28..6d26d073d 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -249,6 +249,7 @@ def _make_runner(): runner.od_config = SimpleNamespace( cache_backend=None, parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.device = torch.device("cpu") runner.pipeline = _StepPipeline() @@ -267,6 +268,7 @@ def _make_distributed_runner(mode: str, device: torch.device): runner.od_config = SimpleNamespace( cache_backend=None, parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.device = device runner.pipeline = _DistributedStepPipeline(mode=mode, device=device) @@ -327,7 +329,7 @@ def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_id def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine: engine = object.__new__(DiffusionEngine) - engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline") + engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline", streaming_output=False) engine.pre_process_func = None engine.post_process_func = None engine.scheduler = scheduler @@ -563,6 +565,7 @@ class TestRunner: step_execution=True, model_class_name="RequestOnlyPipeline", parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.device = torch.device("cpu") runner.pipeline = None @@ -710,6 +713,7 @@ class TestExecutor: def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture): executor = object.__new__(MultiprocDiffusionExecutor) + executor.od_config = SimpleNamespace(streaming_output=False) executor._ensure_open = lambda: None expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None) executor.collective_rpc = mocker.Mock(return_value=expected) diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py new file mode 100644 index 000000000..50ed772c1 --- /dev/null +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration tests for diffusion pipeline streaming output mode. + +Covers: +- Mock pipeline -> StageDiffusionClient (ZMQ subprocess path) +- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni +- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni -> `/v1/realtime/video` +""" + +import asyncio +import contextlib +import queue +import time +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +import zmq +from fastapi import FastAPI, WebSocket +from pytest_mock import MockerFixture +from starlette.testclient import TestClient +from vllm.utils.network_utils import get_open_zmq_ipc_path + +import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module +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.stage_diffusion_client import StageDiffusionClient +from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, + OmniMsgpackEncoder, +) +from vllm_omni.engine.async_omni_engine import StageRuntimeInfo +from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni +from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _StepStreamingPipeline: + supports_step_execution = True + + def __init__(self, outputs: list[DiffusionOutput]) -> None: + self.outputs = outputs + self.requests: list[Any] = [] + + def step_outputs(self, request): + self.requests.append(request) + return list(self.outputs) + + +class _FailingStreamingPipeline: + supports_step_execution = True + + def __init__(self) -> None: + self.requests: list[Any] = [] + + def step_outputs(self, request): + self.requests.append(request) + raise RuntimeError("stream failed midway") + + +class _PipelineBackedEngine: + """DiffusionEngine stand-in that exposes step-streaming chunks from a mock pipeline.""" + + def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None: + self.pipeline = pipeline + self.executor = SimpleNamespace( + register_failure_callback=MagicMock(), + check_health=MagicMock(), + ) + + async def step_streaming(self, request): + try: + outputs = self.pipeline.step_outputs(request) + except Exception as exc: + outputs = [ + DiffusionOutput(custom_output={"chunk": 0}, finished=False), + DiffusionOutput(error=str(exc), finished=True), + ] + for output in outputs: + if output.error is not None: + yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)] + continue + custom_output = output.custom_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 abort(self, request_id: str) -> None: + del request_id + + +class TestPipelineStreamingOutputToStageDiffusionClient: + """Streaming pipeline output over ZMQ into ``StageDiffusionClient``.""" + + @pytest.mark.asyncio + async def test_streaming_output_reaches_stage_client_from_mock_pipeline(self) -> None: + """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags.""" + pipeline = _StepStreamingPipeline( + [ + DiffusionOutput(custom_output={"chunk": 0}, finished=False), + DiffusionOutput(custom_output={"chunk": 1}, finished=True), + ] + ) + + outputs = await self._run_streaming_client_proc(pipeline, request_id="req-stream") + + assert [output.request_id for output in outputs] == ["req-stream", "req-stream"] + assert [output.finished for output in outputs] == [False, True] + assert [output.custom_output["chunk"] for output in outputs] == [0, 1] + assert pipeline.requests[0].request_id == "req-stream" + + @pytest.mark.asyncio + async def test_streaming_midway_error_reaches_stage_client_after_prior_chunk(self) -> None: + """A pipeline error after the first chunk still delivers that chunk then a terminal error output.""" + pipeline = _FailingStreamingPipeline() + + outputs = await self._run_streaming_client_proc(pipeline, request_id="req-error") + + assert len(outputs) == 2 + assert outputs[0].request_id == "req-error" + assert outputs[0].finished is False + assert outputs[0].custom_output["chunk"] == 0 + assert outputs[1].request_id == "req-error" + assert outputs[1].finished is True + assert outputs[1].error is not None + assert "stream failed midway" in outputs[1].error + + @classmethod + async def _run_streaming_client_proc( + cls, + pipeline: _StepStreamingPipeline | _FailingStreamingPipeline, + *, + request_id: str, + ) -> list[OmniRequestOutput]: + request_address = get_open_zmq_ipc_path() + response_address = get_open_zmq_ipc_path() + proc = object.__new__(StageDiffusionProc) + proc._od_config = SimpleNamespace(streaming_output=True) + proc._engine = _PipelineBackedEngine(pipeline) + proc._closed = False + client = cls._make_client(request_address, response_address) + proc_task = asyncio.create_task(proc.run_loop(request_address, response_address)) + try: + await asyncio.sleep(0.05) + await client.add_request_async( + request_id, + "prompt", + OmniDiffusionSamplingParams(), + ) + deadline = asyncio.get_running_loop().time() + 2.0 + outputs: list[OmniRequestOutput] = [] + while (not outputs or not outputs[-1].finished) and asyncio.get_running_loop().time() < deadline: + output = client.get_diffusion_output_nowait() + if output is not None: + outputs.append(output) + else: + await asyncio.sleep(0.01) + return outputs + finally: + try: + client._request_socket.send(client._encoder.encode({"type": "shutdown"})) + except Exception: + pass + try: + await asyncio.wait_for(proc_task, timeout=2) + finally: + cls._close_client(client) + + @staticmethod + def _make_client(request_address: str, response_address: str) -> StageDiffusionClient: + client = object.__new__(StageDiffusionClient) + client.stage_id = 0 + client.replica_id = 0 + client.final_output = True + client.final_output_type = "image" + client.default_sampling_params = {} + client.requires_multimodal_data = False + client.custom_process_input_func = None + client.engine_input_source = [] + client._proc = None + client._proc_manager = None + client._owns_process = False + client._zmq_ctx = zmq.Context() + client._request_socket = client._zmq_ctx.socket(zmq.PUSH) + client._request_socket.bind(request_address) + client._response_socket = client._zmq_ctx.socket(zmq.PULL) + client._response_socket.bind(response_address) + client._encoder = OmniMsgpackEncoder() + client._decoder = OmniMsgpackDecoder() + client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() + client._rpc_results = {} + client._pending_rpcs = set() + client._tasks = {} + client._shutting_down = False + client._engine_dead = False + return client + + @staticmethod + def _close_client(client: StageDiffusionClient) -> None: + client._request_socket.close(linger=0) + client._response_socket.close(linger=0) + client._zmq_ctx.term() + + +class TestPipelineStreamingOutputToEntrypoint: + """Streaming pipeline output through inline stage, orchestrator, and entrypoints.""" + + @pytest.mark.asyncio + async def test_streaming_output_reaches_async_omni_from_mock_pipeline(self) -> None: + """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator.""" + pipeline = _StepStreamingPipeline( + [ + DiffusionOutput(custom_output={"chunk": 0}, finished=False), + DiffusionOutput(custom_output={"chunk": 1}, finished=True), + ] + ) + inline_client = self._make_inline_pipeline_client(pipeline) + fixture = _build_harness([inline_client]) + omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture)) + + try: + 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) + + await _wait_for(lambda: len(pipeline.requests) == 1) + assert pipeline.requests[0].request_id.startswith("req-omni-") + assert [output.custom_output["chunk"] for output in outputs] == [0, 1] + assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True] + finally: + await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client) + + @pytest.mark.asyncio + async def test_streaming_output_reaches_api_websocket_from_mock_pipeline(self, mocker: MockerFixture) -> None: + """Mock pipeline streaming chunks reach the video WebSocket handler through AsyncOmni.generate().""" + frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)] + pipeline = _StepStreamingPipeline( + [ + DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False), + DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True), + ] + ) + inline_client = self._make_inline_pipeline_client(pipeline) + fixture = _build_harness([inline_client]) + bridge_engine = self._OrchestratorBridgeEngine(fixture) + omni = self._make_async_omni(bridge_engine) + + encoded = [b"pipeline-fmp4-0", b"pipeline-fmp4-1"] + + class FakeStreamingVideoEncoder: + encode_calls: list[int] = [] + + def encode(self, video): + del video + idx = len(self.encode_calls) + self.encode_calls.append(idx) + return encoded[idx] + + def close(self): + return b"" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", + lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", + return_value="diffusion", + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + + handler = OmniStreamingVideoOutputHandler( + engine_client=omni, + model_name="test-model", + stage_configs=bridge_engine.stage_configs, + stall_timeout=5.0, + start_timeout=5.0, + ) + app = FastAPI() + + @app.websocket("/v1/realtime/video") + async def ws_endpoint(websocket: WebSocket): + await handler.handle_session(websocket) + + try: + with TestClient(app) as client: + 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_bytes() == b"pipeline-fmp4-0" + assert ws.receive_bytes() == b"pipeline-fmp4-1" + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["chunks"] == 2 + + await _wait_for(lambda: len(pipeline.requests) == 1) + assert pipeline.requests[0].request_id + finally: + await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client) + + 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 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 + + @classmethod + def _make_inline_pipeline_client(cls, pipeline: _StepStreamingPipeline) -> InlineStageDiffusionClient: + """Build an inline diffusion stage client whose engine runs the given mock pipeline.""" + 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, + ) + pipeline_engine = _PipelineBackedEngine(pipeline) + with patch.object(InlineStageDiffusionClient, "_enrich_config"): + with patch( + "vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine", + return_value=pipeline_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 + 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._is_sleeping = False + omni.prom_metrics = MagicMock() + omni.mod_metrics = MagicMock() + omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params + 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 + omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None) + 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) + + +def _make_vllm_config(): + @contextlib.contextmanager + def set_priority(*args, **kwargs): + del args, kwargs + yield + + return SimpleNamespace( + kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)), + compilation_config=SimpleNamespace(ir_enable_torch_wrap=True), + ) + + +class TestSupportedPipelines: + """Streaming-output protocol checks for supported pipelines.""" + + def test_helios_supports_step_execution_for_streaming_output(self) -> None: + from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline + from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution + + # Avoid loading model weights; protocol membership depends on the class contract. + pipeline = object.__new__(HeliosPipeline) + + assert pipeline.supports_step_execution is True + assert supports_step_execution(pipeline) is True + assert isinstance(pipeline, SupportsStepExecution) is True + + def test_load_model_rejects_streaming_output_without_step_execution(self, monkeypatch) -> None: + class _NoStepPipeline: + def forward(self): ... + + class _FakeDiffusersPipelineLoader: + def __init__(self, *args, **kwargs): ... + + def load_model(self, **kwargs): + return _NoStepPipeline() + + class _FakeDeviceMemoryProfiler: + consumed_memory = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = _make_vllm_config() + runner.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enforce_eager=True, + cache_backend=None, + cache_config=None, + streaming_output=True, + step_execution=False, + model_class_name="NoStepPipeline", + parallel_config=SimpleNamespace(use_hsdp=False), + ) + runner.device = torch.device("cpu") + runner.pipeline = None + runner.cache_backend = None + runner.offload_backend = None + runner.kv_transfer_manager = SimpleNamespace() + + monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader) + monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler) + monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) + + with pytest.raises(ValueError, match="NoStepPipeline"): + DiffusionModelRunner.load_model(runner) diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 9c10d46b4..b94f86db0 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -46,6 +46,7 @@ def client(mock_engine): ) with patch.object(InlineStageDiffusionClient, "_enrich_config"): od_config = MagicMock(spec=OmniDiffusionConfig) + od_config.streaming_output = False c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) yield c c.shutdown() @@ -72,6 +73,35 @@ async def test_inline_dispatch_request_success(client, mock_engine): mock_engine.step.assert_called_once() +@pytest.mark.asyncio +async def test_inline_dispatch_request_streaming_success(client, mock_engine): + """Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine""" + chunks = [ + OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=False), + OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=True), + ] + + async def _step_streaming(_request): + for chunk in chunks: + yield [chunk] + + client.od_config.streaming_output = True + mock_engine.step_streaming = _step_streaming + + await client.add_request_async("req-stream", "A test prompt", OmniDiffusionSamplingParams()) + + outputs = [] + for _ in range(20): + output = client.get_diffusion_output_nowait() + if output is not None: + outputs.append(output) + if output.finished: + break + await asyncio.sleep(0.01) + + assert outputs == chunks + + @pytest.mark.asyncio async def test_inline_dispatch_request_error(client, mock_engine): # Setup mock engine step to raise an exception diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 1b015c02a..5bfb8d738 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -21,6 +21,7 @@ from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE from vllm_omni.diffusion.sched import RequestScheduler from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc +from vllm_omni.diffusion.worker.utils import RunnerOutput from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -49,7 +50,7 @@ def _make_executor(num_gpus: int = 1): Returns ``(executor, request_queue, result_queue)``. """ - od_cfg = SimpleNamespace(num_gpus=num_gpus) + od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None) executor = MultiprocDiffusionExecutor(od_cfg) @@ -75,6 +76,7 @@ def _make_engine(num_gpus: int = 1): """Create a lightweight ``DiffusionEngine`` wired to mocked executor.""" executor, req_q, res_q = _make_executor(num_gpus) engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(streaming_output=False) sched = RequestScheduler() sched.initialize(SimpleNamespace()) engine.scheduler = sched @@ -581,6 +583,58 @@ class TestMultiprocExecutorRaisesEngineDeadError: ) +class TestMultiprocExecutorStepStreamingOutput: + """Streaming output uses step execution and one worker reply per step.""" + + def test_execute_step_allows_streaming_output_mode(self): + executor, req_q, res_q = _make_executor() + executor.od_config = SimpleNamespace(streaming_output=True) # pyright: ignore[reportAttributeAccessIssue] + runner_outputs = [ + RunnerOutput( + request_id="sched-stream", + step_index=1, + finished=False, + result=DiffusionOutput(output={"chunk": 0}, finished=False, chunk_index=0, total_chunks=2), + ), + RunnerOutput( + request_id="sched-stream", + step_index=2, + finished=True, + result=DiffusionOutput(output={"chunk": 1}, finished=True, chunk_index=1, total_chunks=2), + ), + ] + scheduler_output = SimpleNamespace( + scheduled_request_ids=["sched-stream"], + ) + + def _worker(): + for runner_output in runner_outputs: + req_q.get(timeout=10) + res_q.put(runner_output) + + thread = threading.Thread(target=_worker, daemon=True) + thread.start() + + first: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType] + executor, + scheduler_output, # pyright: ignore[reportArgumentType] + ) + second: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType] + executor, + scheduler_output, # pyright: ignore[reportArgumentType] + ) + + assert first is runner_outputs[0] + assert first.result is not None + assert first.result.output == {"chunk": 0} + assert first.finished is False + assert second is runner_outputs[1] + assert second.result is not None + assert second.result.output == {"chunk": 1} + assert second.finished is True + thread.join(timeout=2) + + class TestDiffusionEngineDeadErrorPassthrough: """``DiffusionEngine.add_req_and_wait_for_response`` re-raises ``EngineDeadError`` from executor and wraps other errors.""" diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index dcefafb8d..3038d4e90 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -11,6 +11,7 @@ import pytest from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -54,6 +55,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info(): proc = object.__new__(StageDiffusionProc) proc._engine = SimpleNamespace(step=step) + proc._od_config = SimpleNamespace(streaming_output=False) proc._executor = ThreadPoolExecutor(max_workers=1) try: @@ -102,6 +104,40 @@ class MockDiffusionEngine: return [MockOmniRequestOutput(request_id=request.request_id)] +@pytest.mark.asyncio +async def test_proc_streaming_request_yields_each_engine_chunk(): + """Ensure that the streaming output chunks from DiffusionEngine reaches StageDiffusionProc""" + captured = {} + chunks = [ + OmniRequestOutput.from_diffusion(request_id="", images=[], finished=False), + OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True), + ] + + class _StreamingEngine: + async def step_streaming(self, request): + captured["request"] = request + for chunk in chunks: + yield [chunk] + + stage_proc = object.__new__(StageDiffusionProc) + stage_proc._engine = _StreamingEngine() + + outputs = [ + output + async for output in stage_proc._process_streaming_request( + request_id="req-stream", + prompt="prompt", + sampling_params_dict=asdict(OmniDiffusionSamplingParams()), + kv_sender_info={0: {"host": "127.0.0.1"}}, + ) + ] + + assert outputs == chunks + assert [output.request_id for output in outputs] == ["req-stream", "req-stream"] + assert [output.finished for output in outputs] == [False, True] + assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}} + + @pytest.mark.asyncio async def test_proc_process_request_with_batching_async_output(): stage_proc = object.__new__(StageDiffusionProc) diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index 593a9960a..ebe635ea8 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/tests/e2e/accuracy/helpers.py @@ -305,7 +305,7 @@ def assert_video_similarity_metrics( offline_path: Path, ssim_threshold: float, psnr_threshold: float, -) -> None: +) -> tuple[float, float]: ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path) psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path) ssim_score = parse_ssim_score(ssim_output) @@ -339,6 +339,7 @@ def assert_video_similarity_metrics( f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. " f"online={online_path} offline={offline_path}" ) + return ssim_score, psnr_score def assert_similarity( diff --git a/tests/e2e/accuracy/test_video_streaming_output_similarity.py b/tests/e2e/accuracy/test_video_streaming_output_similarity.py new file mode 100644 index 000000000..799e6e125 --- /dev/null +++ b/tests/e2e/accuracy/test_video_streaming_output_similarity.py @@ -0,0 +1,143 @@ +""" +Online serving smoke test for streaming video output similarity. +""" + +import json +from pathlib import Path + +import pytest + +from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_video +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +MODEL = "BestWishYsh/Helios-Distilled" +PROMPT = "A serene lakeside sunrise with mist over the water." +WIDTH = 640 +HEIGHT = 384 +FPS = 16 +NUM_FRAMES = 65 +NUM_INFERENCE_STEPS = 5 +SEED = 42 +MIN_PSNR_DB = 28.0 +MIN_SSIM = 0.95 + +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + +HELIOS_DISTILLED_EXTRA_PARAMS = { + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": True, +} + + +def _server_args(*, streaming_output: bool) -> list[str]: + args = [ + "--stage-init-timeout", + "600", + "--init-timeout", + "900", + "--log-stats", + ] + if streaming_output: + args.append("--diffusion-streaming-output") + return args + + +def _request_config(model: str, *, serialize_extra_params: bool) -> dict: + extra_params = ( + json.dumps(HELIOS_DISTILLED_EXTRA_PARAMS) if serialize_extra_params else HELIOS_DISTILLED_EXTRA_PARAMS + ) + return { + "model": model, + "form_data": { + "prompt": PROMPT, + "width": WIDTH, + "height": HEIGHT, + "fps": FPS, + "num_frames": NUM_FRAMES, + "num_inference_steps": NUM_INFERENCE_STEPS, + "guidance_scale": 1.0, + "seed": SEED, + "extra_params": extra_params, + }, + } + + +def _client_for(server: OmniServer, *, run_level: str) -> OpenAIClientHandler: + return OpenAIClientHandler( + host=server.host, + port=server.port, + api_key="EMPTY", + run_level=run_level, + log_stats=server.log_stats, + ) + + +def _write_video_artifact(response: DiffusionResponse, output_path: Path) -> Path: + assert response.videos, "No video artifact returned" + output_path.write_bytes(response.videos[0]) + return output_path + + +@pytest.mark.parametrize( + "model", + [ + pytest.param( + MODEL, + id="helios_distilled", + marks=SINGLE_CARD_MARKS, + ) + ], +) +def test_helios_streaming_output_matches_non_streaming( + model: str, + model_prefix: str, + run_level: str, + tmp_path: Path, + record_property, +): + model = model_prefix + model + + with OmniServer(model, _server_args(streaming_output=True), use_omni=True) as server: + client = _client_for(server, run_level=run_level) + streaming_response = client.send_streaming_video_diffusion_request( + _request_config(server.model, serialize_extra_params=False) + )[0] + + streaming_path = _write_video_artifact(streaming_response, tmp_path / "helios_streaming.mp4") + record_property("helios_streaming_video", str(streaming_path)) + record_property("helios_streaming_latency_s", streaming_response.e2e_latency) + + with OmniServer(model, _server_args(streaming_output=False), use_omni=True) as server: + client = _client_for(server, run_level=run_level) + non_streaming_response = client.send_video_diffusion_request( + _request_config(server.model, serialize_extra_params=True) + )[0] + + non_streaming_path = _write_video_artifact(non_streaming_response, tmp_path / "helios_non_streaming.mp4") + record_property("helios_non_streaming_video", str(non_streaming_path)) + record_property("helios_non_streaming_latency_s", non_streaming_response.e2e_latency) + + streaming_metadata = probe_video(streaming_path) + non_streaming_metadata = probe_video(non_streaming_path) + assert streaming_metadata == non_streaming_metadata, ( + f"Video metadata mismatch for helios_streaming_output:\n" + f"streaming={streaming_metadata}\n" + f"non_streaming={non_streaming_metadata}\n" + f"streaming_path={streaming_path}\n" + f"non_streaming_path={non_streaming_path}" + ) + + ssim, psnr = assert_video_similarity_metrics( + label="helios_streaming_output", + online_path=streaming_path, + offline_path=non_streaming_path, + ssim_threshold=MIN_SSIM, + psnr_threshold=MIN_PSNR_DB, + ) + record_property("helios_streaming_vs_non_streaming_psnr_db", psnr) + record_property("helios_streaming_vs_non_streaming_ssim_percent", ssim * 100.0) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index f73504074..23f8cfaa2 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -523,6 +523,67 @@ async def test_run_single_stage_diffusion(orchestrator_factory) -> None: await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +async def test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks(orchestrator_factory) -> None: + """Intermediate diffusion chunks (finished=False) reach the frontend before the final chunk.""" + stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + orchestrator_fixture = orchestrator_factory([stage0]) + params = OmniDiffusionSamplingParams() + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-stream", + prompt={"prompt": "draw a cat"}, + original_prompt={"prompt": "draw a cat"}, + sampling_params_list=[params], + final_stage_id=0, + ) + + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + stage0.push_diffusion_output( + OmniRequestOutput.from_diffusion( + request_id="req-stream", + images=[], + final_output_type="image", + custom_output={"chunk": 0}, + finished=False, + ) + ) + stage0.push_diffusion_output( + OmniRequestOutput.from_diffusion( + request_id="req-stream", + images=[], + final_output_type="image", + custom_output={"chunk": 1}, + finished=True, + ) + ) + + output_msgs: list[OutputMessage] = [] + deadline = time.monotonic() + 2.0 + while not output_msgs or not output_msgs[-1].finished: + if time.monotonic() >= deadline: + raise AssertionError( + f"Timed out waiting for finished orchestrator output, got {len(output_msgs)} message(s)" + ) + try: + msg = orchestrator_fixture.output_sync_q.get_nowait() + except queue.Empty: + await asyncio.sleep(0.01) + continue + if isinstance(msg, OutputMessage): + output_msgs.append(msg) + + assert [msg.request_id for msg in output_msgs] == ["req-stream", "req-stream"] + assert [msg.finished for msg in output_msgs] == [False, True] + assert [msg.engine_outputs.finished for msg in output_msgs] == [False, True] + assert [msg.engine_outputs.custom_output["chunk"] for msg in output_msgs] == [0, 1] + await _wait_for(lambda: "req-stream" not in orchestrator_fixture.orchestrator.request_states) + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + @pytest.mark.asyncio async def test_run_llm_to_diffusion(orchestrator_factory) -> None: stage0 = FakeStageClient(stage_type="llm", final_output=False) diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py new file mode 100644 index 000000000..72d293853 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -0,0 +1,446 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for `/v1/realtime/video` WebSocket video output streaming.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncGenerator, Callable +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +from fastapi import FastAPI, WebSocket +from pytest_mock import MockerFixture +from starlette.testclient import TestClient + +from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +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 _build_test_app( + *, + mocker: MockerFixture, + streaming_chunks: list[tuple[bytes, bool]], + mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]], + final_chunk: bytes = b"", + stall_timeout: float = 5.0, +) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]: + engine_client = mocker.MagicMock() + engine_client.abort = mocker.AsyncMock() + engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()] + engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] + engine_client.generate = mock_generate + + class FakeStreamingVideoEncoder: + encode_calls: list[int] = [] + + def encode(self, video): + del video + idx = len(self.encode_calls) + self.encode_calls.append(idx) + return streaming_chunks[idx][0] + + def close(self): + return final_chunk + + def _make_encoder(*, output_format, fps, video_codec_options=None): + del output_format, fps, video_codec_options + return FakeStreamingVideoEncoder() + + encoder_factory = mocker.MagicMock(side_effect=_make_encoder) + engine_client.streaming_encoder_factory = encoder_factory + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", + encoder_factory, + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", + return_value="diffusion", + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + + handler = OmniStreamingVideoOutputHandler( + engine_client=engine_client, + model_name="test-model", + stage_configs=engine_client.stage_configs, + stall_timeout=stall_timeout, + start_timeout=5.0, + ) + app = FastAPI() + + @app.websocket("/v1/realtime/video") + async def ws_endpoint(websocket: WebSocket): + await handler.handle_session(websocket) + + return app, handler, engine_client + + +class TestStreamingVideoOutputWebSocket: + """WebSocket protocol tests for streaming generated video output.""" + + def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture): + """A full session delivers video.start, binary chunks, then session.done.""" + + async def mock_generate(*_args, **_kwargs): + for frames, finished in [ + (_fake_video_frames(2), False), + (_fake_video_frames(3), True), + ]: + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[frames], + final_output_type="image", + finished=finished, + ) + + 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": "A cat walking in the rain"}) + + start = ws.receive_json() + assert start["type"] == "video.start" + assert start["format"] == "m4s" + 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" + + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["chunks"] == 2 + assert done["stopped"] is False + + engine_client.abort.assert_not_called() + engine_client.streaming_encoder_factory.assert_called_once() + assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s" + + def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture): + 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.start", "prompt": "hello", "format": "webm"}) + err = ws.receive_json() + assert err["type"] == "error" + assert "format" in err["message"] + + engine_client.abort.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.""" + + async def mock_generate(*_args, **_kwargs): + 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, + final_chunk=b"mp4-trailer", + ) + + with TestClient(app) as client: + 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" + + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["chunks"] == 2 + assert done["stopped"] is False + + engine_client.abort.assert_not_called() + + def test_generation_error_output_returns_websocket_error(self, mocker: MockerFixture): + """An erroneous OmniRequestOutput from generate is sent to the client as a WebSocket error.""" + error_message = "diffusion engine exploded" + + 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, + ) + yield OmniRequestOutput.from_error("req-ws", error_message) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", False)], + 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": "hello"}) + assert ws.receive_json()["type"] == "video.start" + assert ws.receive_bytes() == b"mp4-chunk-0" + err = ws.receive_json() + assert err["type"] == "error" + assert error_message in err["message"] + + engine_client.abort.assert_not_called() + + +class _MockVideoOutputWebSocket: + """Feeds scripted client messages and records server sends.""" + + def __init__(self, messages: list[str]) -> None: + self._messages: asyncio.Queue[str] = asyncio.Queue() + for message in messages: + self._messages.put_nowait(message) + self.sent: list[dict | bytes] = [] + + def enqueue(self, message: str) -> None: + self._messages.put_nowait(message) + + async def accept(self) -> None: + return None + + async def receive_text(self) -> str: + return await self._messages.get() + + async def send_json(self, data: dict) -> None: + self.sent.append(data) + + async def send_bytes(self, data: bytes) -> None: + self.sent.append(data) + + @property + def sent_json(self) -> list[dict]: + return [m for m in self.sent if isinstance(m, dict)] + + +def _build_handler_for_async_tests( + mocker: MockerFixture, + *, + mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]], + streaming_chunks: list[tuple[bytes, bool]] | None = None, + stall_timeout: float = 5.0, +) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]: + engine_client = mocker.MagicMock() + engine_client.abort = mocker.AsyncMock() + engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()] + engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] + engine_client.generate = mock_generate + + chunks = streaming_chunks or [(b"mp4-chunk-0", True)] + + class FakeStreamingVideoEncoder: + def __init__(self) -> None: + self.encode_calls = 0 + + def encode(self, video): + del video + idx = self.encode_calls + self.encode_calls += 1 + return chunks[idx][0] + + def close(self): + return b"" + + encoder_factory = mocker.MagicMock(side_effect=lambda **_: FakeStreamingVideoEncoder()) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", + encoder_factory, + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", + return_value="diffusion", + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + + handler = OmniStreamingVideoOutputHandler( + engine_client=engine_client, + model_name="test-model", + stage_configs=engine_client.stage_configs, + stall_timeout=stall_timeout, + start_timeout=5.0, + ) + start_msg = json.dumps({"type": "session.start", "prompt": "async test"}) + ws = _MockVideoOutputWebSocket([start_msg]) + return handler, engine_client, ws + + +class TestStreamingVideoOutputStallTimeout: + """Stall-timeout and session.ping behavior.""" + + def test_long_generation_without_client_messages_succeeds(self, mocker: MockerFixture): + """Engine progress between yields keeps the session alive without client pings.""" + + 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, + stall_timeout=0.5, + ) + + with TestClient(app) as client: + 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" + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["stopped"] is False + + engine_client.abort.assert_not_called() + + def test_stall_timeout_aborts_when_engine_silent(self, mocker: MockerFixture): + """No engine output for longer than stall_timeout aborts the request.""" + + async def mock_generate(*_args, **_kwargs): + await asyncio.sleep(1.0) + 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, + stall_timeout=0.2, + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "stalled"}) + assert ws.receive_json()["type"] == "video.start" + err = ws.receive_json() + assert err["type"] == "error" + assert "Stall timeout" in err["message"] + + engine_client.abort.assert_called() + + @pytest.mark.asyncio + async def test_session_ping_returns_pong(self, mocker: MockerFixture): + """session.ping during generation receives session.pong.""" + + async def mock_generate(*_args, **_kwargs): + await asyncio.sleep(0.5) + 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, + stall_timeout=5.0, + ) + + async def _send_ping_after_delay() -> None: + await asyncio.sleep(0.05) + ws.enqueue(json.dumps({"type": "session.ping"})) + + ping_task = asyncio.create_task(_send_ping_after_delay()) + await handler.handle_session(ws) # type: ignore[arg-type] + await ping_task + + assert any(m.get("type") == "session.pong" for m in ws.sent_json) + + @pytest.mark.asyncio + async def test_ping_prevents_stall_during_silent_engine(self, mocker: MockerFixture): + """Client pings refresh the stall clock while the engine produces no output.""" + + 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, + stall_timeout=0.3, + ) + + async def _pinger() -> None: + while not release.is_set(): + ws.enqueue(json.dumps({"type": "session.ping"})) + await asyncio.sleep(0.1) + + pinger_task = asyncio.create_task(_pinger()) + session_task = asyncio.create_task(handler.handle_session(ws)) # type: ignore[arg-type] + await asyncio.sleep(0.8) + release.set() + await session_task + await pinger_task + + errors = [m for m in ws.sent_json if m.get("type") == "error"] + assert not any("Stall timeout" in m.get("message", "") for m in errors) + 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() diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 9e732403f..722819682 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -48,6 +48,107 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): assert mux_calls[0]["audio"] is None +def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch): + muxers = [] + + class FakeFragmentedMP4Muxer: + def __init__(self, *, width, height, fps, video_codec_options=None): + self.calls = [] + muxers.append( + { + "width": width, + "height": height, + "fps": fps, + "video_codec_options": video_codec_options, + "instance": self, + } + ) + + def mux_video_frames(self, frames): + self.calls.append(frames) + return f"fragment-{len(self.calls)}".encode() + + def close(self): + return b"trailer" + + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer", + FakeFragmentedMP4Muxer, + ) + + encoder = video_api_utils.FragmentedMP4VideoEncoder( + fps=12, + video_codec_options={"preset": "ultrafast"}, + ) + assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1" + assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2" + assert encoder.close() == b"trailer" + + assert len(muxers) == 1 + assert muxers[0]["width"] == 6 + assert muxers[0]["height"] == 4 + assert muxers[0]["fps"] == 12.0 + assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"} + assert len(muxers[0]["instance"].calls) == 2 + + +def test_create_streaming_video_encoder_selects_requested_format(): + assert isinstance( + video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12), + video_api_utils.FragmentedMP4VideoEncoder, + ) + + +def test_finalize_streaming_mp4_bytes_produces_progressive_mp4(): + """Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes.""" + import numpy as np + + from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes + + def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]: + import io + + import av + + with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container: + stream = container.streams.video[0] + frame_count = sum(1 for _ in container.decode(stream)) + assert stream.average_rate is not None + fps = float(stream.average_rate) + duration_sec = None + if stream.duration is not None: + assert stream.time_base is not None + duration_sec = float(stream.duration * stream.time_base) + return frame_count, fps, duration_sec, stream.width, stream.height + + width = 32 + height = 32 + fps = 16.0 + input_frame_count = 2 + + muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps) + frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8) + streamed = muxer.mux_video_frames(frames) + muxer.close() + + finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps) + assert finalized + assert finalized != streamed + + streamed_info = _read_mp4_video_info(streamed) + final_info = _read_mp4_video_info(finalized) + expected_duration = input_frame_count / fps + + assert streamed_info[0] == input_frame_count + assert final_info[0] == input_frame_count + assert streamed_info[1] == pytest.approx(fps) + assert final_info[1] == pytest.approx(fps) + assert streamed_info[3:] == (width, height) + assert final_info[3:] == (width, height) + + assert final_info[2] == pytest.approx(expected_duration, rel=0.05) + assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05) + + def test_rife_model_inference_runs_on_dummy_tensors(): model = rife_interpolator.Model().eval() img0 = torch.rand(1, 3, 32, 32) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 7e6fa2a73..94b17b2b8 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -1334,7 +1334,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture) extra_params = { "is_enable_stage2": True, "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [20, 20, 20], + "pyramid_num_inference_steps_list": [1, 1, 1], "use_cfg_zero_star": True, } response = test_client.post( @@ -1352,7 +1352,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture) captured = engine.captured_sampling_params_list[0] assert captured.extra_args["is_enable_stage2"] is True assert captured.extra_args["pyramid_num_stages"] == 3 - assert captured.extra_args["pyramid_num_inference_steps_list"] == [20, 20, 20] + assert captured.extra_args["pyramid_num_inference_steps_list"] == [1, 1, 1] assert captured.extra_args["use_cfg_zero_star"] is True diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index b8637c1e9..0d7700e27 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -8,6 +8,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model] @@ -198,6 +199,48 @@ def test_generate_accepts_request_after_repeated_cancellations(): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_yields_streaming_diffusion_chunks_before_final(): + """AsyncOmni.generate yields every intermediate diffusion chunk before the final one.""" + + async def streaming_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts): + del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts + yield OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + final_output_type="image", + custom_output={"chunk": 0}, + finished=False, + ) + yield OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + final_output_type="image", + custom_output={"chunk": 1}, + finished=True, + ) + + async def run_test(): + omni = get_async_omni_instance() + omni._process_orchestrator_results = streaming_process_results + + outputs = [] + async for output in AsyncOmni.generate( + omni, + prompt={"prompt": "a cat"}, + request_id="req-stream", + sampling_params_list=[SimpleNamespace()], + output_modalities=["image"], + ): + outputs.append(output) + + assert len(outputs) == 2 + assert [output.finished for output in outputs] == [False, True] + assert [output.custom_output["chunk"] for output in outputs] == [0, 1] + + asyncio.run(run_test()) + + @pytest.mark.cpu @pytest.mark.parametrize( "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE] diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index da62ceab2..de77722b3 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -115,6 +115,11 @@ def assert_video_diffusion_response( expected_height = _maybe_int(form_data.get("height")) expected_fps = _maybe_int(form_data.get("fps")) + # Skip num_frames assertion for Helios models because they round up frames + model = request_config.get("model", "") + if "Helios" in model: + expected_frames = None + for vid_bytes in response.videos: assert_video_valid( vid_bytes, diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index c12154f86..08be653fd 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -19,7 +19,7 @@ from collections.abc import Generator from dataclasses import asdict, dataclass from io import BytesIO from pathlib import Path -from typing import Any, NamedTuple +from typing import Any, NamedTuple, cast from urllib.parse import quote import psutil @@ -2168,6 +2168,102 @@ class OpenAIClientHandler: timeout=timeout, ) + def send_streaming_video_diffusion_request( + self, + request_config: dict[str, Any], + request_num: int = 1, + *, + timeout_seconds: float = 600.0, + ) -> list[DiffusionResponse]: + """ + Send a native ``/v1/realtime/video`` WebSocket request and return one + finalized MP4 artifact assembled from the streamed binary fragments. + """ + if request_num != 1: + raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented") + + response = asyncio.run( + self._send_streaming_video_diffusion_request_once( + request_config, + timeout_seconds=timeout_seconds, + ) + ) + assert_diffusion_response(response, request_config, run_level=self.run_level) + if response.e2e_latency is not None: + self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s") + else: + self._print_client_stat("[diffusion.stream] request#1 completed") + return [response] + + async def _send_streaming_video_diffusion_request_once( + self, + request_config: dict[str, Any], + *, + timeout_seconds: float, + ) -> DiffusionResponse: + form_data = request_config.get("form_data") + if not isinstance(form_data, dict): + raise ValueError("Video request_config must contain 'form_data'") + payload: dict[str, Any] = { + "type": "session.start", + **{key: value for key, value in form_data.items() if value is not None}, + } + model = request_config.get("model") + if model is not None: + payload["model"] = model + payload.setdefault("format", "m4s") + + fps = float(payload.get("fps") or 16) + stream_format = payload["format"] + url = self._build_ws_url("/v1/realtime/video") + + result = DiffusionResponse() + chunks: list[bytes] = [] + start_time = time.perf_counter() + deadline = start_time + timeout_seconds + + import websockets + + async with websockets.connect(url, max_size=None) as websocket: + await websocket.send(json.dumps(payload)) + + while True: + remaining = deadline - time.perf_counter() + if remaining <= 0: + raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s") + + message = await asyncio.wait_for(websocket.recv(), timeout=remaining) + if isinstance(message, bytes): + chunks.append(message) + continue + + msg = json.loads(message) + msg_type = msg.get("type") + if msg_type == "video.start": + stream_format = msg.get("format") or stream_format + continue + if msg_type == "session.done": + break + if msg_type == "error": + raise RuntimeError(str(msg.get("message", msg))) + + from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes + from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat + + streamed_bytes = b"".join(chunks) + if not streamed_bytes: + raise RuntimeError("Streaming video request completed without binary video chunks") + result.videos = [ + finalize_streaming_video_bytes( + streamed_bytes, + input_format=cast(StreamingVideoFormat, stream_format), + fps=fps, + ) + ] + result.e2e_latency = time.perf_counter() - start_time + result.success = True + return result + def _wait_until_video_completed( self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300 ) -> None: diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 6a25c6dab..421a42ed0 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -689,6 +689,9 @@ class OmniDiffusionConfig: # Step mode settings step_execution: bool = False + # Streaming mode settings + streaming_output: bool = False # Start (video) generation with initial prompt, but streaming output in chunks + # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 @@ -1164,6 +1167,11 @@ class DiffusionOutput: # logged timings info, directly from Req.timings # timings: Optional["RequestTimings"] = None + # Streaming info (the defaults should make sense for non-streaming mode) + finished: bool = True + chunk_index: int = 0 + total_chunks: int = 1 + # logged duration of stages stage_durations: dict[str, float] = field(default_factory=dict) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index a7869b140..1d8753682 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -9,7 +9,7 @@ import inspect import queue import threading import time -from collections.abc import Iterable +from collections.abc import AsyncGenerator, Iterable from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any @@ -126,9 +126,14 @@ class DiffusionEngine: self.action_post_process_func, "custom_output" ) + self.step_execution = bool(getattr(od_config, "step_execution", False)) + if self.od_config.streaming_output and not self.step_execution: + logger.warning("streaming_output=True requires step_execution=True; enabling step execution.") + self.od_config.step_execution = True + self.step_execution = True + executor_class = DiffusionExecutor.get_class(od_config) self.executor = executor_class(od_config) - self.step_execution = bool(getattr(od_config, "step_execution", False)) self.scheduler: SchedulerInterface = scheduler or ( StepScheduler() if self.step_execution else RequestScheduler() ) @@ -149,11 +154,15 @@ class DiffusionEngine: self._rpc_lock = threading.RLock() self._cv = threading.Condition(self._rpc_lock) self._out_queue: dict[str, asyncio.Future] = {} + self._out_queue_streaming: dict[str, asyncio.Queue[DiffusionOutput]] = {} self._closed = False self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() - self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request + if self.step_execution: + self.execute_fn = self.executor.execute_step + else: + self.execute_fn = self.executor.execute_request try: self._dummy_run() @@ -197,7 +206,37 @@ class DiffusionEngine: exec_start_time = time.perf_counter() output = await self.async_add_req_and_wait_for_response(request) exec_total_time = time.perf_counter() - exec_start_time + return self.postprocess_output(request, output, diffusion_engine_start_time, preprocess_time, exec_total_time) + async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[list[OmniRequestOutput], None]: + await self._check_and_start_background_loop() + + diffusion_engine_start_time = time.perf_counter() + + preprocess_time = 0.0 + if self.pre_process_func is not None: + preprocess_start_time = time.perf_counter() + request = self.pre_process_func(request) + preprocess_time = time.perf_counter() - preprocess_start_time + logger.debug("Pre-processing completed in %.4f seconds", preprocess_time) + + exec_start_time = time.perf_counter() + generator = self.async_add_req_and_stream_response(request) + async for output in generator: + exec_total_time = time.perf_counter() - exec_start_time + yield self.postprocess_output( + request, output, diffusion_engine_start_time, preprocess_time, exec_total_time + ) + + def postprocess_output( + self, + request: OmniDiffusionRequest, + output: DiffusionOutput, + diffusion_engine_start_time: float, + preprocess_time: float, + exec_total_time: float, + ) -> list[OmniRequestOutput]: + """Convert a DiffusionOutput to a list of OmniRequestOutput, attaching profiling metrics.""" if output.aborted: raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.") if output.error: @@ -212,7 +251,7 @@ class DiffusionEngine: if output.output is None: logger.warning("Output is None, returning empty OmniRequestOutput") - return format_empty_diffusion_outputs(request) + return format_empty_diffusion_outputs(request, finished=output.finished) # When CPU offload is enabled, move output to CPU before # post-processing to avoid device OOM — model weights may still @@ -312,11 +351,14 @@ class DiffusionEngine: sched_output = self.scheduler.schedule() if sched_output.is_empty: - self._handle_finished_requests(sched_output.finished_req_ids, None) + if self.od_config.streaming_output: + self._handle_empty_streaming_requests(sched_output.finished_req_ids) + else: + self._handle_finished_requests(sched_output.finished_req_ids, None) continue try: - runner_output = self.execute_fn(sched_output) + runner_output: BaseRunnerOutput = self.execute_fn(sched_output) # pyright: ignore[reportAssignmentType] except Exception as exc: logger.error( "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True @@ -336,7 +378,14 @@ class DiffusionEngine: self._process_aborts_queue() self._process_rpc_queue() finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output) - self._handle_finished_requests(finished_req_ids, runner_output) + if self.od_config.streaming_output: + self._handle_step_streaming_runner_output( + finished_req_ids, + sched_output.scheduled_request_ids, + runner_output, + ) + else: + self._handle_finished_requests(finished_req_ids, runner_output) # Engine is stopping: fail any RPCs still queued so callers don't hang. self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down.")) @@ -411,6 +460,57 @@ class DiffusionEngine: out = self._finalize_finished_request(rid, _output, missing_result_error) self._complete_future(fut, out) + def _handle_empty_streaming_requests( + self, + finished_ids: set[str], + missing_result_error: str = "Diffusion streaming request finished without execution output.", + ) -> None: + """Mirrors `_handle_finished_requests()` in non-streaming mode when used for empty scheduler output.""" + for rid in finished_ids: + out = self._finalize_finished_request(rid, None, missing_result_error=missing_result_error) + self._put_streaming_output_with_cv(rid, out) + + def _handle_step_streaming_runner_output( + self, + finished_req_ids: set[str], + scheduled_request_ids: list[str], + runner_output: BaseRunnerOutput, + ) -> None: + """ + Deliver partial step-execution outputs in streaming mode. + + Step execution returns one ``RunnerOutput`` per scheduled request per + engine tick. Most denoise steps have ``result=None``; chunk boundaries + return a ``DiffusionOutput`` that must be delivered even before the + scheduler marks the request finished. + """ + delivered_finished_req_ids: set[str] = set() + + # finished_ids may have some requests that are not scheduler in this round. + # First handle this-round requests. + for request_id in scheduled_request_ids: + req_output = runner_output.get_request_output(request_id) + if request_id in finished_req_ids: + # This entire request is finished (this is the last chunk) + out = self._finalize_finished_request( + request_id, + req_output, + missing_result_error="Diffusion streaming execution finished without a final output.", + ) + self._put_streaming_output_with_cv(request_id, out) + delivered_finished_req_ids.add(request_id) + elif req_output is not None and req_output.result is not None: + # This is a non-terminal chunk. So it is not in scheduler's finished_req_ids, but still need delivering. + self._put_streaming_output_with_cv(request_id, req_output.result) + + # Then handle other requests that are finished in this round. + for request_id in finished_req_ids - delivered_finished_req_ids: + out = self._finalize_finished_request( + request_id, + missing_result_error="Diffusion streaming request finished without execution output.", + ) + self._put_streaming_output_with_cv(request_id, out) + @staticmethod def make_engine( config: OmniDiffusionConfig, @@ -430,9 +530,14 @@ class DiffusionEngine: with self._cv: if self._closed: raise RuntimeError("DiffusionEngine is closed.") - fut = self.main_loop.create_future() - request_id = self.scheduler.add_request(request) - self._out_queue[request_id] = fut + if not self.od_config.streaming_output: + fut = self.main_loop.create_future() + request_id = self.scheduler.add_request(request) + self._out_queue[request_id] = fut + else: + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + request_id = self.scheduler.add_request(request) + self._out_queue_streaming[request_id] = queue self._cv.notify_all() return request_id @@ -448,12 +553,39 @@ class DiffusionEngine: logger.error(f"Wait for response failed: {e}") raise + async def get_streaming_result(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]: + """Mirrors `get_result()` in non-streaming mode.""" + with self._cv: + queue = self._out_queue_streaming.get(request_id) + if queue is None: + raise RuntimeError(f"Request {request_id} not found in output queue.") + try: + while True: + output: DiffusionOutput = await queue.get() + yield output + if output.finished: + break + except Exception as e: + logger.error(f"Wait for response failed: {e}") + raise + finally: + # In streaming mode, an output queue is maintained until the terminal chunk is met. + # So unlike the non-streaming mode where output Future is popped in `_handle_finished_requests` (immediately + # after the request is returned), the streaming mode needs to pop the output queue here (one layer above). + with self._cv: + if self._out_queue_streaming.get(request_id) is queue: + self._out_queue_streaming.pop(request_id, None) + async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: # No lock needed: add_request is already protected by self._cv, and # all executor calls are serialized inside the busy loop. request_id = self.add_request(request) return await self.get_result(request_id) + def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]: + request_id = self.add_request(request) + return self.get_streaming_result(request_id) + def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: with self._rpc_lock: if self._closed: @@ -475,7 +607,7 @@ class DiffusionEngine: # within _dummy_run, only one request will be scheduled request_id = sched_output.scheduled_request_ids[0] try: - runner_output = self.execute_fn(sched_output) + runner_output: BaseRunnerOutput = self.execute_fn(sched_output) # pyright: ignore[reportAssignmentType] except EngineDeadError: raise except Exception as exc: @@ -495,10 +627,10 @@ class DiffusionEngine: if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1: raise ValueError("Sync func should receive one result at one time") if target_request_id in finished_req_ids: - runner_output = runner_output.get_request_output(target_request_id) + req_output = runner_output.get_request_output(target_request_id) return self._finalize_finished_request( target_request_id, - runner_output=runner_output, + runner_output=req_output, missing_result_error="Diffusion execution finished without a final output.", ) @@ -693,8 +825,28 @@ class DiffusionEngine: else: _set_result() + def _put_streaming_queue_output( + self, + queue: asyncio.Queue[DiffusionOutput], + output: DiffusionOutput, + ) -> None: + """Append to streaming output queue in a safe event loop. Mirrors `_complete_future()` in non-streaming mode.""" + loop = self.main_loop + if loop is not None and loop.is_running(): + loop.call_soon_threadsafe(queue.put_nowait, output) + else: + queue.put_nowait(output) + + def _put_streaming_output_with_cv(self, request_id: str, output: DiffusionOutput) -> None: + with self._cv: + queue = self._out_queue_streaming.get(request_id) + if queue is None: + return + self._put_streaming_queue_output(queue, output) + def close(self) -> None: pending_futures: list[asyncio.Future] = [] + pending_streaming_queues: list[asyncio.Queue[DiffusionOutput]] = [] with self._cv: if self._closed and self._shutdown_complete: return @@ -703,12 +855,16 @@ class DiffusionEngine: if self.stop_event is not None: self.stop_event.set() pending_futures = list(self._out_queue.values()) + pending_streaming_queues = list(self._out_queue_streaming.values()) self._out_queue.clear() + self._out_queue_streaming.clear() self._cv.notify_all() closed_output = DiffusionOutput(error="DiffusionEngine is closed.") for fut in pending_futures: self._complete_future(fut, closed_output) + for streaming_queue in pending_streaming_queues: + self._put_streaming_queue_output(streaming_queue, closed_output) worker_thread = self.worker_thread if worker_thread is not None: diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 137f2b887..a54ce8b18 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -134,12 +134,18 @@ class InlineStageDiffusionClient(StageClientBase): kv_sender_info=kv_sender_info, ) - results = await self._engine.step(request) - result = results[0] - if not result.request_id: - result.request_id = request_id - - self._output_queue.put_nowait(result) + if self.od_config.streaming_output: + async for results in self._engine.step_streaming(request): + result = results[0] + if not result.request_id: + result.request_id = request_id + self._output_queue.put_nowait(result) + else: + results = await self._engine.step(request) + result = results[0] + if not result.request_id: + result.request_id = request_id + self._output_queue.put_nowait(result) except DiffusionRequestAbortedError as e: logger.info("request_id: %s aborted: %s", request_id, str(e)) except Exception as e: diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 9b625b6b9..9acd13b53 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -8,7 +8,7 @@ import json import logging import math import os -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING, Any, ClassVar import numpy as np @@ -33,9 +33,10 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.base_config import ( - QuantizationConfig, - ) + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import DiffusionRequestState logger = logging.getLogger(__name__) @@ -161,6 +162,7 @@ class HeliosPipeline( Implements chunked video generation with multi-term memory history context. """ + supports_step_execution: ClassVar[bool] = True _dit_modules: ClassVar[list[str]] = ["transformer"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -267,6 +269,640 @@ class HeliosPipeline( def current_timestep(self): return self._current_timestep + def _stage1_sigmas(self, num_steps: int) -> np.ndarray: + # DMD drops the last timestep in set_timesteps(). Only compensate for + # the one-step dummy-run edge case; otherwise preserve the historical + # stage-1 schedule used by request-level forward(). + sigma_count = num_steps + 2 if self.is_distilled and num_steps == 1 else num_steps + 1 + return np.linspace(0.999, 0.0, sigma_count)[:-1] + + def prepare_encode( + self, + state: DiffusionRequestState, + **kwargs: Any, + ) -> DiffusionRequestState: + """Initialize Helios request state for chunk-wise step execution.""" + del kwargs + req = OmniDiffusionRequest( + prompts=state.prompts or [], + sampling_params=state.sampling, + request_id=state.request_id, + ) + extra = getattr(state.sampling, "extra_args", {}) or {} + + history_sizes = sorted(extra.get("history_sizes", [16, 2, 1]), reverse=True) + num_latent_frames_per_chunk = int(extra.get("num_latent_frames_per_chunk", 9)) + keep_first_frame = bool(extra.get("keep_first_frame", True)) + frame_num = int(extra.get("frame_num", 132)) + height = (int(state.sampling.height or extra.get("height", 384)) // 16) * 16 + width = (int(state.sampling.width or extra.get("width", 640)) // 16) * 16 + num_frames = max(int(state.sampling.num_frames or frame_num), 1) + num_steps = int(state.sampling.num_inference_steps or extra.get("num_inference_steps", 50)) + output_type = extra.get("output_type", "np") + + image = extra.get("image") + video = extra.get("video") + if image is not None and video is not None: + raise ValueError("image and video cannot be provided simultaneously") + if len(req.prompts) > 1: + raise ValueError("Helios step execution supports a single prompt, not a batched request.") + + prompt = None + negative_prompt = None + if len(req.prompts) == 1: + prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") + negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") + if prompt is None: + raise ValueError("Prompt is required for Helios generation.") + + guidance_scale = float(extra.get("guidance_scale", 5.0)) + if state.sampling.guidance_scale_provided: + guidance_scale = state.sampling.guidance_scale + self._guidance_scale = guidance_scale + + device = self.device + dtype = self.transformer.dtype + generator = state.sampling.generator + if generator is None and state.sampling.seed is not None: + generator = torch.Generator(device=device).manual_seed(state.sampling.seed) + state.sampling.generator = generator + + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=state.sampling.num_outputs_per_prompt or 1, + max_sequence_length=state.sampling.max_sequence_length or 226, + device=device, + dtype=dtype, + ) + batch_size = prompt_embeds.shape[0] + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(self.vae.device, self.vae.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + self.vae.device, self.vae.dtype + ) + + add_noise_to_image_latents = bool(extra.get("add_noise_to_image_latents", True)) + image_noise_sigma_min = float(extra.get("image_noise_sigma_min", 0.111)) + image_noise_sigma_max = float(extra.get("image_noise_sigma_max", 0.135)) + add_noise_to_video_latents = bool(extra.get("add_noise_to_video_latents", True)) + video_noise_sigma_min = float(extra.get("video_noise_sigma_min", 0.111)) + video_noise_sigma_max = float(extra.get("video_noise_sigma_max", 0.135)) + + fake_image_latents = None + image_latents = None + if image is not None: + image_latents, fake_image_latents = self.prepare_image_latents( + image, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + ) + + if image_latents is not None and add_noise_to_image_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + fake_image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + fake_image_latents = ( + fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + + (1 - fake_image_noise_sigma) * fake_image_latents + ) + + video_latents = None + if video is not None: + image_latents, video_latents = self.prepare_video_latents( + video, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + ) + + if video_latents is not None and add_noise_to_video_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + noisy_latents_chunks = [] + num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk + for i in range(num_latent_chunks): + chunk_start = i * num_latent_frames_per_chunk + chunk_end = chunk_start + num_latent_frames_per_chunk + latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] + chunk_frames = latent_chunk.shape[2] + frame_sigmas = ( + torch.rand(chunk_frames, device=device, generator=generator) + * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) + noisy_chunk = ( + frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + + (1 - frame_sigmas) * latent_chunk + ) + noisy_latents_chunks.append(noisy_chunk) + video_latents = torch.cat(noisy_latents_chunks, dim=2) + + num_channels_latents = self.transformer.config.in_channels + window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) + num_history_latent_frames = sum(history_sizes) + total_generated_latent_frames = 0 + if not keep_first_frame: + history_sizes[-1] = history_sizes[-1] + 1 + + history_latents = torch.zeros( + batch_size, + num_channels_latents, + num_history_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + if fake_image_latents is not None: + history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) + total_generated_latent_frames += 1 + if video_latents is not None: + history_frames = history_latents.shape[2] + video_frames = video_latents.shape[2] + if video_frames < history_frames: + keep_frames = history_frames - video_frames + history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) + else: + history_latents = video_latents + total_generated_latent_frames += video_latents.shape[2] + + if keep_first_frame: + indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) + else: + indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) + + state.prompt_embeds = prompt_embeds + state.negative_prompt_embeds = negative_prompt_embeds + state.do_true_cfg = self.do_classifier_free_guidance + state.chunk_index = 0 + state.step_in_chunk = 0 + state.total_chunks = num_latent_chunk + state.extra.update( + { + "attention_kwargs": extra.get("attention_kwargs", {}) or {}, + "batch_size": batch_size, + "dtype": dtype, + "generator": generator, + "guidance_scale": guidance_scale, + "height": height, + "history_latents": history_latents, + "history_sizes": history_sizes, + "history_video": None, + "image_latents": image_latents, + "indices_hidden_states": indices_hidden_states.unsqueeze(0), + "indices_latents_history_long": indices_latents_history_long.unsqueeze(0), + "indices_latents_history_mid": indices_latents_history_mid.unsqueeze(0), + "indices_latents_history_short": indices_latents_history_short.unsqueeze(0), + "is_amplify_first_chunk": bool(extra.get("is_amplify_first_chunk", False)), + "is_enable_stage2": bool(extra.get("is_enable_stage2", False)), + "is_skip_first_chunk": bool(extra.get("is_skip_first_chunk", False)), + "keep_first_frame": keep_first_frame, + "latents_mean": latents_mean, + "latents_std": latents_std, + "num_channels_latents": num_channels_latents, + "num_history_latent_frames": num_history_latent_frames, + "num_latent_frames_per_chunk": num_latent_frames_per_chunk, + "num_steps": num_steps, + "output_type": output_type, + "pyramid_num_inference_steps_list": extra.get("pyramid_num_inference_steps_list", [10, 10, 10]), + "pyramid_num_stages": int(extra.get("pyramid_num_stages", 3)), + "total_generated_latent_frames": total_generated_latent_frames, + "use_cfg_zero_star": bool(extra.get("use_cfg_zero_star", False)), + "use_zero_init": bool(extra.get("use_zero_init", True)), + "vae_dtype": self.vae.dtype, + "width": width, + "window_num_frames": window_num_frames, + "zero_steps": int(extra.get("zero_steps", 1)), + } + ) + self._prepare_next_chunk(state) + return state + + def _prepare_next_chunk(self, state: DiffusionRequestState) -> None: + extra = state.extra + k = state.chunk_index + is_first_chunk = k == 0 + is_second_chunk = k == 1 + history_sizes = extra["history_sizes"] + history_latents = extra["history_latents"] + keep_first_frame = extra["keep_first_frame"] + image_latents = extra["image_latents"] + num_history_latent_frames = extra["num_history_latent_frames"] + + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + if image_latents is None and is_first_chunk: + latents_prefix = torch.zeros( + ( + extra["batch_size"], + extra["num_channels_latents"], + 1, + latents_history_1x.shape[-2], + latents_history_1x.shape[-1], + ), + device=latents_history_1x.device, + dtype=latents_history_1x.dtype, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + + state.latents = self.prepare_latents( + extra["batch_size"], + extra["num_channels_latents"], + extra["height"], + extra["width"], + extra["window_num_frames"], + dtype=torch.float32, + device=self.device, + generator=extra["generator"], + latents=None, + ) + extra["stage1_start_latents"] = state.latents + extra.update( + { + "is_first_chunk": is_first_chunk, + "is_second_chunk": is_second_chunk, + "latents_history_long": latents_history_long, + "latents_history_mid": latents_history_mid, + "latents_history_short": latents_history_short, + } + ) + + if not extra["is_enable_stage2"]: + patch_size = self.transformer.config.patch_size + image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + sigmas = self._stage1_sigmas(extra["num_steps"]) + mu = calculate_shift(image_seq_len) + self.scheduler.set_timesteps(extra["num_steps"], device=self.device, sigmas=sigmas, mu=mu) + state.timesteps = self.scheduler.timesteps + state.chunk_num_steps = len(state.timesteps) + else: + self._prepare_stage2_chunk(state) + + state.step_in_chunk = 0 + state.step_index = 0 + self._num_timesteps = state.chunk_num_steps + + def _prepare_stage2_chunk(self, state: DiffusionRequestState) -> None: + extra = state.extra + batch_size, num_channel, num_frames_lat, height, width = state.latents.shape + latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames_lat, num_channel, height, width + ) + for _ in range(extra["pyramid_num_stages"] - 1): + height //= 2 + width //= 2 + latents_flat = F.interpolate(latents_flat, size=(height, width), mode="bilinear") * 2 + state.latents = latents_flat.reshape(batch_size, num_frames_lat, num_channel, height, width).permute( + 0, 2, 1, 3, 4 + ) + extra["stage2_height"] = height + extra["stage2_width"] = width + extra["stage2_start_point_list"] = [state.latents] if self.is_distilled else None + extra["stage_index"] = 0 + extra["stage_step_index"] = 0 + amplify_first_chunk = extra["is_amplify_first_chunk"] and extra["is_first_chunk"] + state.chunk_num_steps = sum( + self._stage2_effective_num_steps(int(num_steps), amplify_first_chunk) + for num_steps in extra["pyramid_num_inference_steps_list"] + ) + self._set_stage2_timesteps(state) + + def _stage2_effective_num_steps(self, num_steps: int, is_amplify_first_chunk: bool) -> int: + if self.is_distilled and is_amplify_first_chunk: + return num_steps * 2 + return num_steps + + def _set_stage2_timesteps(self, state: DiffusionRequestState) -> None: + extra = state.extra + patch_size = self.transformer.config.patch_size + image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + mu = calculate_shift(image_seq_len) + stage_index = extra["stage_index"] + self.scheduler.set_timesteps( + extra["pyramid_num_inference_steps_list"][stage_index], + stage_index=stage_index, + device=self.device, + mu=mu, + is_amplify_first_chunk=extra["is_amplify_first_chunk"] and extra["is_first_chunk"], + ) + state.timesteps = self.scheduler.timesteps + state.step_index = extra["stage_step_index"] + + def denoise_step( + self, + input_batch: InputBatch, + states: Sequence[DiffusionRequestState], + **kwargs: Any, + ) -> torch.Tensor | None: + del kwargs + if len(states) != 1: + raise ValueError("Helios step execution supports a single request, not a batched request.") + state = states[0] + if state.extra["is_enable_stage2"]: + return self._denoise_stage2_step(state) + return self._denoise_stage1_step(state, input_batch.latents, input_batch.timesteps) + + def _denoise_stage1_step( + self, + state: DiffusionRequestState, + latents: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + extra = state.extra + batch_size = latents.shape[0] + t = timesteps + self._current_timestep = t + timestep = t.expand(batch_size) + transformer_kwargs = { + "hidden_states": latents.to(extra["dtype"]), + "timestep": timestep, + "indices_hidden_states": extra["indices_hidden_states"], + "indices_latents_history_short": extra["indices_latents_history_short"], + "indices_latents_history_mid": extra["indices_latents_history_mid"], + "indices_latents_history_long": extra["indices_latents_history_long"], + "latents_history_short": extra["latents_history_short"].to(extra["dtype"]), + "latents_history_mid": extra["latents_history_mid"].to(extra["dtype"]), + "latents_history_long": extra["latents_history_long"].to(extra["dtype"]), + "attention_kwargs": extra["attention_kwargs"], + "return_dict": False, + } + if extra["use_cfg_zero_star"] and state.do_true_cfg: + noise_pred = self.transformer( + encoder_hidden_states=state.prompt_embeds, + **transformer_kwargs, + )[0] + noise_uncond = self.transformer( + encoder_hidden_states=state.negative_prompt_embeds, + **transformer_kwargs, + )[0] + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))).to(noise_pred.dtype) + if (state.step_in_chunk <= extra["zero_steps"]) and extra["use_zero_init"]: + return noise_pred * 0.0 + return noise_uncond * alpha_cfg + extra["guidance_scale"] * (noise_pred - noise_uncond * alpha_cfg) + + positive_kwargs = { + "encoder_hidden_states": state.prompt_embeds, + **transformer_kwargs, + } + negative_kwargs = ( + { + "encoder_hidden_states": state.negative_prompt_embeds, + **transformer_kwargs, + } + if state.do_true_cfg + else None + ) + return self.predict_noise_maybe_with_cfg( + do_true_cfg=state.do_true_cfg, + true_cfg_scale=extra["guidance_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + def _denoise_stage2_step(self, state: DiffusionRequestState) -> torch.Tensor: + extra = state.extra + latents = state.latents + assert latents is not None + batch_size = latents.shape[0] + t = state.current_timestep + assert t is not None + self._current_timestep = t + timestep = t.expand(batch_size).to(torch.int64) + transformer_kwargs = { + "hidden_states": latents.to(extra["dtype"]), + "timestep": timestep, + "indices_hidden_states": extra["indices_hidden_states"], + "indices_latents_history_short": extra["indices_latents_history_short"], + "indices_latents_history_mid": extra["indices_latents_history_mid"], + "indices_latents_history_long": extra["indices_latents_history_long"], + "latents_history_short": extra["latents_history_short"].to(extra["dtype"]), + "latents_history_mid": extra["latents_history_mid"].to(extra["dtype"]), + "latents_history_long": extra["latents_history_long"].to(extra["dtype"]), + "attention_kwargs": extra["attention_kwargs"], + "return_dict": False, + } + noise_pred = self.transformer(encoder_hidden_states=state.prompt_embeds, **transformer_kwargs)[0] + if state.do_true_cfg: + noise_uncond = self.transformer(encoder_hidden_states=state.negative_prompt_embeds, **transformer_kwargs)[0] + if extra["use_cfg_zero_star"]: + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))).to(noise_pred.dtype) + if ( + extra["stage_index"] == 0 + and extra["stage_step_index"] <= extra["zero_steps"] + and extra["use_zero_init"] + ): + noise_pred = noise_pred * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + extra["guidance_scale"] * ( + noise_pred - noise_uncond * alpha_cfg + ) + else: + noise_pred = noise_uncond + extra["guidance_scale"] * (noise_pred - noise_uncond) + return noise_pred + + def step_scheduler( + self, + state: DiffusionRequestState, + noise_pred: torch.Tensor, + **kwargs: Any, + ) -> None: + del kwargs + if state.extra["is_enable_stage2"]: + self._step_scheduler_stage2(state, noise_pred) + else: + t = state.current_timestep + assert t is not None + if self.is_distilled: + state.latents = self.scheduler.step( + noise_pred, + t, + state.latents, + return_dict=False, + cur_sampling_step=state.step_in_chunk, + dmd_noisy_tensor=state.extra["stage1_start_latents"], + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=state.timesteps, + )[0] + else: + state.latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, state.latents, state.do_true_cfg) + state.step_in_chunk += 1 + state.step_index = state.step_in_chunk + + def _step_scheduler_stage2(self, state: DiffusionRequestState, noise_pred: torch.Tensor) -> None: + extra = state.extra + t = state.current_timestep + assert t is not None and state.latents is not None + start_points = extra["stage2_start_point_list"] + state.latents = self.scheduler.step( + noise_pred, + t, + state.latents, + return_dict=False, + cur_sampling_step=extra["stage_step_index"], + dmd_noisy_tensor=start_points[extra["stage_index"]] if start_points is not None else None, + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=state.timesteps, + )[0] + extra["stage_step_index"] += 1 + state.step_in_chunk += 1 + state.step_index = extra["stage_step_index"] + + if extra["stage_step_index"] < len(state.timesteps): + return + if extra["stage_index"] >= extra["pyramid_num_stages"] - 1: + return + + extra["stage_index"] += 1 + extra["stage_step_index"] = 0 + extra["stage2_height"] *= 2 + extra["stage2_width"] *= 2 + batch_size, num_channel, num_frames_cur = state.latents.shape[:3] + latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames_cur, + num_channel, + extra["stage2_height"] // 2, + extra["stage2_width"] // 2, + ) + latents_flat = F.interpolate(latents_flat, size=(extra["stage2_height"], extra["stage2_width"]), mode="nearest") + state.latents = latents_flat.reshape( + batch_size, + num_frames_cur, + num_channel, + extra["stage2_height"], + extra["stage2_width"], + ).permute(0, 2, 1, 3, 4) + + ori_sigma = 1 - self.scheduler.ori_start_sigmas[extra["stage_index"]] + gamma = self.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + patch_size = self.transformer.config.patch_size + noise = self.sample_block_noise( + batch_size, + num_channel, + state.latents.shape[2], + extra["stage2_height"], + extra["stage2_width"], + patch_size, + generator=extra["generator"], + ).to(device=self.device, dtype=extra["dtype"]) + state.latents = alpha * state.latents + beta * noise + if self.is_distilled and extra["stage2_start_point_list"] is not None: + extra["stage2_start_point_list"].append(state.latents) + self._set_stage2_timesteps(state) + + def post_decode( + self, + state: DiffusionRequestState, + **kwargs: Any, + ) -> DiffusionOutput: + del kwargs + extra = state.extra + is_first_chunk = extra["is_first_chunk"] + is_second_chunk = extra["is_second_chunk"] + if extra["keep_first_frame"] and ( + (is_first_chunk and extra["image_latents"] is None) or (extra["is_skip_first_chunk"] and is_second_chunk) + ): + extra["image_latents"] = state.latents[:, :, 0:1, :, :] + + extra["total_generated_latent_frames"] += state.latents.shape[2] + extra["history_latents"] = torch.cat([extra["history_latents"], state.latents], dim=2) + real_history_latents = extra["history_latents"][:, :, -extra["total_generated_latent_frames"] :] + current_latents = ( + real_history_latents[:, :, -extra["num_latent_frames_per_chunk"] :].to(extra["vae_dtype"]) + / extra["latents_std"] + + extra["latents_mean"] + ) + current_video = self.vae.decode(current_latents, return_dict=False)[0] + if extra["history_video"] is None: + extra["history_video"] = current_video + else: + extra["history_video"] = torch.cat([extra["history_video"], current_video], dim=2) + + output = current_latents if extra["output_type"] == "latent" else current_video + completed_chunk_index = state.chunk_index + state.chunk_index += 1 + finished = state.request_denoise_completed + if not finished: + self._prepare_next_chunk(state) + else: + self._current_timestep = None + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + return DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else {}, + chunk_index=completed_chunk_index, + total_chunks=state.total_chunks, + finished=finished, + ) + def forward( self, req: OmniDiffusionRequest, diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index c9d013ff3..657a45bf8 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -56,15 +56,19 @@ class SupportsStepExecution(Protocol): def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState: """Prepare request-level inputs and return initialized state.""" + ... def denoise_step(self, input_batch: InputBatch, **kwargs: Any) -> torch.Tensor | None: """Run one denoise forward on the runner-assembled batch.""" + ... def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: """Run one scheduler step.""" + ... def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput: - """Decode output after denoise loop.""" + """Decode output after denoise loop or at a partial chunk boundary.""" + ... @runtime_checkable diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 8d00b8f85..26cfa992e 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -70,7 +70,11 @@ def normalize_diffusion_postprocess_output( ) -def format_empty_diffusion_outputs(request: OmniDiffusionRequest) -> list[OmniRequestOutput]: +def format_empty_diffusion_outputs( + request: OmniDiffusionRequest, + *, + finished: bool = True, +) -> list[OmniRequestOutput]: return [ OmniRequestOutput.from_diffusion( request_id=request.request_id, @@ -78,6 +82,7 @@ def format_empty_diffusion_outputs(request: OmniDiffusionRequest) -> list[OmniRe prompt=prompt, metrics={}, latents=None, + finished=finished, ) for prompt in request.prompts ] @@ -125,6 +130,7 @@ def format_diffusion_outputs( is_text_output=is_text_output, is_audio_output=is_audio_output, audio_sample_rate=audio_sample_rate, + finished=diffusion_output.finished, ) return _format_multi_prompt_outputs( @@ -135,6 +141,7 @@ def format_diffusion_outputs( postprocess_output=postprocess_output, is_audio_output=is_audio_output, audio_sample_rate=audio_sample_rate, + finished=diffusion_output.finished, ) @@ -185,6 +192,7 @@ def _format_single_prompt_output( is_text_output: bool, is_audio_output: bool, audio_sample_rate: int | None, + finished: bool = True, ) -> list[OmniRequestOutput]: prompt = request.prompts[0] request_id = request.request_id @@ -202,6 +210,7 @@ def _format_single_prompt_output( final_output_type="text", stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, ), ] @@ -227,6 +236,7 @@ def _format_single_prompt_output( final_output_type="audio", stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, ), ] @@ -245,6 +255,7 @@ def _format_single_prompt_output( multimodal_output=mm_output, stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, ), ] @@ -258,6 +269,7 @@ def _format_multi_prompt_outputs( postprocess_output: DiffusionPostprocessOutput, is_audio_output: bool, audio_sample_rate: int | None, + finished: bool = True, ) -> list[OmniRequestOutput]: results = [] output_idx = 0 @@ -298,6 +310,7 @@ def _format_multi_prompt_outputs( final_output_type="audio", stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, ), ) continue @@ -337,6 +350,7 @@ def _format_multi_prompt_outputs( multimodal_output=mm_output, stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, ), ) diff --git a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py index f67e9f692..8d4bf41de 100644 --- a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py +++ b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py @@ -81,7 +81,9 @@ class DiffusionPipelineProfilerMixin: _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"] def setup_diffusion_pipeline_profiler( - self, profiler_targets: list[str] | None = None, enable_diffusion_pipeline_profiler: bool = False + self, + profiler_targets: list[str] | None = None, + enable_diffusion_pipeline_profiler: bool = False, ) -> None: self.enable_diffusion_pipeline_profiler = enable_diffusion_pipeline_profiler if not enable_diffusion_pipeline_profiler: diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 7530e61c9..aa058c8d0 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -10,6 +10,7 @@ import asyncio import contextlib import multiprocessing.connection import signal +from collections.abc import AsyncGenerator from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from typing import TYPE_CHECKING, Any @@ -172,6 +173,29 @@ class StageDiffusionProc: result.request_id = request_id return result + async def _process_streaming_request( + self, + request_id: str, + prompt: Any, + sampling_params_dict: dict, + kv_sender_info: dict[str, Any] | None = None, + ) -> AsyncGenerator[OmniRequestOutput, None]: + """Process a streaming diffusion request and yield the results from DiffusionEngine.step_streaming().""" + sampling_params = self._reconstruct_sampling_params(sampling_params_dict) + + request = OmniDiffusionRequest( + prompts=[prompt], + sampling_params=sampling_params, + request_id=request_id, + kv_sender_info=kv_sender_info, + ) + + async for results in self._engine.step_streaming(request): # pyright: ignore[reportOptionalMemberAccess] + result = results[0] + if not result.request_id: + result.request_id = request_id + yield result + async def _process_batch_request( self, request_id: str, @@ -186,6 +210,9 @@ class StageDiffusionProc: ``images`` list contains every generated image, matching the contract expected by the orchestrator and tests. """ + if self._od_config.streaming_output: + raise NotImplementedError("Streaming output is not supported for batched requests") + sampling_params = self._reconstruct_sampling_params(sampling_params_dict) request = OmniDiffusionRequest( @@ -388,13 +415,22 @@ class StageDiffusionProc: ) -> None: """Process a single diffusion request and send the response.""" try: - result = await self._process_request( - request_id, - prompt, - sampling_params_dict, - kv_sender_info=kv_sender_info, - ) - await response_socket.send(encoder.encode({"type": "result", "output": result})) + if not self._od_config.streaming_output: + result = await self._process_request( + request_id, + prompt, + sampling_params_dict, + kv_sender_info=kv_sender_info, + ) + await response_socket.send(encoder.encode({"type": "result", "output": result})) + else: + async for result in self._process_streaming_request( + request_id, + prompt, + sampling_params_dict, + kv_sender_info=kv_sender_info, + ): + await response_socket.send(encoder.encode({"type": "result", "output": result})) except DiffusionRequestAbortedError as e: logger.info( "request_id: %s aborted: %s", diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index a09cd4595..314a9d2bf 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -6,11 +6,124 @@ from __future__ import annotations import io from fractions import Fraction +from typing import Any, cast import av import numpy as np +class FragmentedMP4Muxer: + """Incrementally mux video frames into one fragmented MP4 byte stream.""" + + def __init__( + self, + *, + width: int, + height: int, + fps: float = 25.0, + video_codec: str = "h264", + crf: str = "18", + video_codec_options: dict[str, str] | None = None, + ) -> None: + self._buf = io.BytesIO() + self._closed = False + self._container = av.open( + self._buf, + mode="w", + format="mp4", + options={"movflags": "+frag_every_frame+empty_moov+default_base_moof"}, + ) + + self._stream: av.VideoStream = cast( + av.VideoStream, + self._container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)), + ) + self._stream.width = width + self._stream.height = height + self._stream.pix_fmt = "yuv420p" + + options: dict[str, object] = {"crf": str(crf)} + if video_codec_options: + options.update(video_codec_options) + self._stream.options = options + + try: + self._stream.codec_context.max_b_frames = 0 + except AttributeError: + pass + + def mux_video_frames(self, video_frames: np.ndarray) -> bytes: + """Mux a batch of ``uint8`` RGB frames and return newly written MP4 bytes.""" + if self._closed: + raise RuntimeError("Cannot mux frames after FragmentedMP4Muxer.close().") + if video_frames.ndim != 4 or video_frames.shape[-1] != 3: + raise ValueError("video_frames must have shape (T, H, W, 3).") + if video_frames.dtype != np.uint8: + raise ValueError("video_frames must be uint8.") + if video_frames.shape[1] != self._stream.height or video_frames.shape[2] != self._stream.width: + raise ValueError("All fragmented MP4 chunks in a session must use the same frame size.") + + for frame_data in video_frames: + frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") + for packet in self._stream.encode(frame): + self._container.mux(packet) + return self._read_new_bytes() + + def close(self) -> bytes: + """Flush delayed encoder packets, close the container, and return final bytes.""" + if self._closed: + return b"" + for packet in self._stream.encode(): + self._container.mux(packet) + self._container.close() + self._closed = True + return self._read_new_bytes() + + def _read_new_bytes(self) -> bytes: + """Return newly muxed bytes in the current video container, + then clear the buffer to prepare for the next chunk.""" + chunk = self._buf.getvalue() + self._buf.seek(0) + self._buf.truncate() + return chunk + + +def finalize_streaming_video_bytes( + video_bytes: bytes, + *, + input_format: str, + fps: float = 25.0, + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Convert streamed video bytes into a progressive MP4 for local playback.""" + if not video_bytes: + return video_bytes + + normalized_format = input_format.lower() + if normalized_format == "m4s": + demux_format = "mp4" + else: + raise ValueError(f"Unsupported streaming video format: {input_format}") + + try: + with cast(Any, av.open(io.BytesIO(video_bytes), format=demux_format)) as container: + stream = container.streams.video[0] + frame_arrays = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)] + except Exception: + return video_bytes + + if not frame_arrays: + return video_bytes + + frames_u8 = np.ascontiguousarray(np.stack(frame_arrays, axis=0), dtype=np.uint8) + return mux_video_audio_bytes( + frames_u8, + None, + fps=float(fps), + video_codec_options=video_codec_options, + ) + + def mux_video_audio_bytes( video_frames: np.ndarray, audio_waveform: np.ndarray | None = None, @@ -39,17 +152,19 @@ def mux_video_audio_bytes( buf = io.BytesIO() container = av.open(buf, mode="w", format="mp4") - v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)) + v_stream = cast(av.VideoStream, container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))) v_stream.width = video_frames.shape[2] v_stream.height = video_frames.shape[1] v_stream.pix_fmt = "yuv420p" - options = {"crf": str(crf)} + options: dict[str, object] = {"crf": str(crf)} if video_codec_options: options.update(video_codec_options) v_stream.options = options - a_stream = None + a_stream: av.AudioStream | None = None + samples: np.ndarray | None = None + layout: str | None = None if audio_waveform is not None: samples = audio_waveform.astype(np.float32) if samples.ndim == 1: @@ -58,7 +173,7 @@ def mux_video_audio_bytes( samples = np.ascontiguousarray(samples.T) num_channels = samples.shape[0] layout = "stereo" if num_channels >= 2 else "mono" - a_stream = container.add_stream(audio_codec, rate=audio_sample_rate) + a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=audio_sample_rate)) a_stream.layout = layout for frame_data in video_frames: @@ -69,6 +184,8 @@ def mux_video_audio_bytes( container.mux(packet) if a_stream is not None and audio_waveform is not None: + if samples is None or layout is None: + raise ValueError("Audio samples were not prepared for muxing.") audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout) audio_frame.sample_rate = audio_sample_rate for packet in a_stream.encode(audio_frame): diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 288973637..9f3598abe 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -152,12 +152,21 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) logger.info("Model runner: Model loaded successfully.") + if self.od_config.streaming_output and not getattr(self.od_config, "step_execution", False): + logger.warning("streaming_output=True requires step_execution=True; enabling step execution.") + self.od_config.step_execution = True + if getattr(self.od_config, "step_execution", False) and not self.supports_step_mode(): raise ValueError( "step_execution=True requires a pipeline implementing " "prepare_encode(), denoise_step(), step_scheduler(), and post_decode(); " f"{self.od_config.model_class_name} does not support that contract." ) + if self.od_config.streaming_output and not self.supports_step_mode(): + raise ValueError( + "streaming_output=True requires step execution support; " + f"{self.od_config.model_class_name} does not support that contract." + ) # Apply CPU offloading self.offload_backend = get_offload_backend(self.od_config, device=self.device) @@ -357,7 +366,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): and self.od_config.enable_cache_dit_summary ): cache_summary(self.pipeline, details=True) - return output # ------------------------------------------------------------------ @@ -457,7 +465,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): scatter_latents(states, input_batch) for state in states: - if interrupted or state.denoise_completed: + if interrupted or state.request_denoise_completed: self.state_cache.pop(state.request_id, None) def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any: @@ -492,7 +500,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): omni_diffusion_config=self.od_config, attn_metadata=attn_metadata, ): - noise_pred = self.pipeline.denoise_step(input_batch) + noise_pred = self.pipeline.denoise_step(input_batch, states=states) runner_output_list = [] pipeline_interrupted = getattr(self.pipeline, "interrupt", False) @@ -515,15 +523,24 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): req, noise_pred[offset : offset + row_num] if noise_pred is not None else None ) offset = offset + row_num - if req.denoise_completed: + if self.od_config.streaming_output: + should_decode = req.chunk_denoise_completed + else: + should_decode = req.denoise_completed + + if should_decode: result = self.pipeline.post_decode(req) else: result = None + # finished should be computed after post_decode() advanced chunk_index + finished = ( + req.request_denoise_completed if self.od_config.streaming_output else req.denoise_completed + ) runner_output_list.append( RunnerOutput( request_id=req.request_id, step_index=req.step_index, - finished=req.denoise_completed, + finished=finished, result=result, ) ) diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index a0bfcea9b..02cb0b14a 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -55,6 +55,12 @@ class DiffusionRequestState: timesteps: torch.Tensor | list[torch.Tensor] | None = None step_index: int = 0 + # ── Optional chunked streaming progress ── + chunk_index: int = 0 + step_in_chunk: int = 0 + total_chunks: int = 1 + chunk_num_steps: int | None = None + # ── Per-request scheduler instance (set once by prepare_encode) ── scheduler: Any | None = None @@ -103,6 +109,18 @@ class DiffusionRequestState: return False return self.step_index >= total_steps + @property + def chunk_denoise_completed(self) -> bool: + if self.chunk_num_steps is None: + return False + return self.step_in_chunk >= self.chunk_num_steps + + @property + def request_denoise_completed(self) -> bool: + if self.chunk_num_steps is None: + return self.denoise_completed + return self.chunk_index >= self.total_chunks + @property def new_request(self) -> bool: # TODO: this is only an approximation for current stepwise mode. diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index ac6140291..3cee3ccd8 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1027,6 +1027,7 @@ class AsyncOmniEngine: **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), + "streaming_output": kwargs.get("diffusion_streaming_output", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), "extras": { "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index e45681b5c..3902b11ba 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -662,6 +662,14 @@ class OmniServeCommand(CLISubcommand): type=int, help="Maximum generated image size in pixels (height * width).", ) + # Diffusion model (mainly video generation models) streaming output mode + omni_config_group.add_argument( + "--diffusion-streaming-output", + dest="diffusion_streaming_output", + action="store_true", + default=False, + help="Enable chunked streaming output for diffusion (mainly video generation) models that support it.", + ) # TTS-specific parameters omni_config_group.add_argument( diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index fc22493fc..c9f6a1d93 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -603,6 +603,7 @@ class OmniBase(PDDisaggregationMixin): metrics=response_metrics, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, + finished=finished, ) def shutdown(self, timeout: float | None = None) -> None: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 7de919fc8..405803d9c 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -132,6 +132,7 @@ from vllm_omni.entrypoints.openai.serving_video import ( ReferenceImage, ReferenceVideo, ) +from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler from vllm_omni.entrypoints.openai.serving_video_stream import create_streaming_video_handler from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, @@ -774,6 +775,11 @@ async def omni_init_app_state( model_name=model_name, stage_configs=diffusion_stage_configs, ) + state.openai_streaming_video_output = OmniStreamingVideoOutputHandler( + engine_client=engine_client, + model_name=model_name, + stage_configs=diffusion_stage_configs, + ) state.openai_serving_speech = OmniOpenAIServingSpeech.for_diffusion( diffusion_engine=engine_client, @@ -1594,6 +1600,23 @@ async def streaming_video_chat(websocket: WebSocket): await handler.handle_session(websocket) +@router.websocket("/v1/realtime/video") +async def streaming_video_output(websocket: WebSocket): + """WebSocket endpoint for streaming generated video output chunks.""" + handler = getattr(websocket.app.state, "openai_streaming_video_output", None) + if handler is None: + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "message": "Streaming video generation is not available", + } + ) + await websocket.close() + return + await handler.handle_session(websocket) + + @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py new file mode 100644 index 000000000..f713df07c --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -0,0 +1,519 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""WebSocket handler for streaming generated video chunks. + +Protocol: + Client -> Server: + {"type": "session.start", "model": "...", "prompt": "...", "format": "m4s", ...} + {"type": "session.prompt_update", "prompt": "..."} # currently unsupported + {"type": "session.stop"} + {"type": "session.ping"} # optional; refreshes stall clock + + Server -> Client: + {"type": "video.start", "request_id": "...", "config": {...}, "format": "m4s"} + + {"type": "session.done", ...} + {"type": "session.pong"} # reply to session.ping + {"type": "error", "message": "..."} +""" + +from __future__ import annotations + +import asyncio +import copy +import inspect +import json +import time +import uuid +from collections.abc import AsyncGenerator +from http import HTTPStatus +from typing import Any, cast + +from fastapi import HTTPException, WebSocket, WebSocketDisconnect +from pydantic import ValidationError +from vllm.engine.protocol import EngineClient +from vllm.logger import init_logger + +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 + +logger = init_logger(__name__) + +_DEFAULT_STALL_TIMEOUT = 60.0 +_DEFAULT_START_TIMEOUT = 10.0 +_CONTROL_POLL_INTERVAL = 1.0 +_MAX_START_MESSAGE_SIZE = 4 * 1024 * 1024 +_MAX_CONTROL_MESSAGE_SIZE = 128 * 1024 + + +class _SessionProgress: + """Tracks last server-side progress for stall-timeout enforcement.""" + + def __init__(self) -> None: + self._last_at = time.monotonic() + + def touch(self) -> None: + self._last_at = time.monotonic() + + def stalled_for(self, stall_timeout: float) -> bool: + return time.monotonic() - self._last_at > stall_timeout + + +class OmniStreamingVideoOutputHandler: + """Handles `/v1/realtime/video` sessions for generated video output.""" + + def __init__( + self, + engine_client: EngineClient, + model_name: str | None = None, + stage_configs: list[Any] | None = None, + stall_timeout: float = _DEFAULT_STALL_TIMEOUT, + start_timeout: float = _DEFAULT_START_TIMEOUT, + ) -> None: + self._engine_client = engine_client + self._model_name = model_name + self._stage_configs = stage_configs + self._stall_timeout = stall_timeout + self._start_timeout = start_timeout + + async def handle_session(self, websocket: WebSocket) -> None: + await websocket.accept() + + request_id: str | None = None + stop_event = asyncio.Event() + send_lock = asyncio.Lock() + control_task: asyncio.Task[None] | None = None + stopped = False + chunks_sent = 0 + + try: + start = await self._receive_start(websocket) + if start is None: + return + request, output_format = start + + request_id = f"video_stream-{uuid.uuid4().hex}" + progress = _SessionProgress() + async with send_lock: + await websocket.send_json( + { + "type": "video.start", + "request_id": request_id, + "format": output_format, + "config": request.model_dump(exclude_none=True), + } + ) + progress.touch() + + control_task = asyncio.create_task( + self._control_loop(websocket, request_id, stop_event, send_lock, progress) + ) + async for chunk in self._stream_video_bytes(request, request_id, output_format, progress): + if stop_event.is_set(): + stopped = True + break + async with send_lock: + await websocket.send_bytes(chunk) + chunks_sent += 1 + progress.touch() + + async with send_lock: + await websocket.send_json( + { + "type": "session.done", + "request_id": request_id, + "chunks": chunks_sent, + "stopped": stopped or stop_event.is_set(), + } + ) + except WebSocketDisconnect: + logger.info("Streaming video output: client disconnected") + if request_id is not None: + await self._abort_request(request_id) + except Exception as e: + if stop_event.is_set(): + try: + async with send_lock: + await websocket.send_json( + { + "type": "session.done", + "request_id": request_id, + "chunks": chunks_sent, + "stopped": True, + } + ) + except Exception: + logger.debug("Failed to send stopped session.done", exc_info=True) + return + logger.exception("Streaming video output session error: %s", e) + await self._send_error(websocket, str(e)) + finally: + if control_task is not None: + control_task.cancel() + await asyncio.gather(control_task, return_exceptions=True) + + async def _receive_start(self, websocket: WebSocket) -> tuple[VideoGenerationRequest, StreamingVideoFormat] | None: + try: + raw = await asyncio.wait_for(websocket.receive_text(), timeout=self._start_timeout) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.start") + return None + + if len(raw) > _MAX_START_MESSAGE_SIZE: + await self._send_error(websocket, "session.start message too large") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.start") + return None + + if not isinstance(msg, dict): + await self._send_error(websocket, "session.start must be a JSON object") + return None + + if msg.get("type") != "session.start": + await self._send_error(websocket, f"Expected session.start, got: {msg.get('type')}") + return None + + try: + output_format = self._coerce_output_format(msg.pop("format", "m4s")) + except ValueError as e: + await self._send_error(websocket, str(e)) + return None + + request_data = self._coerce_request_data(msg) + try: + request = VideoGenerationRequest(**request_data) + except ValidationError as e: + await self._send_error(websocket, f"Invalid session.start: {e}") + return None + + configured_model = self._model_name + if request.model is not None and configured_model is not None and request.model != configured_model: + await self._send_error( + websocket, + f"Model mismatch: request specifies '{request.model}' but server is running '{configured_model}'.", + ) + return None + + return request, output_format + + @staticmethod + def _coerce_request_data(msg: dict[str, Any]) -> dict[str, Any]: + request_data = {k: v for k, v in msg.items() if k not in {"type", "format", "sampling_params"}} + sampling_params = msg.get("sampling_params") + if isinstance(sampling_params, dict): + extra_params = dict(request_data.get("extra_params") or {}) + for key, value in sampling_params.items(): + if key in VideoGenerationRequest.model_fields and key not in request_data: + request_data[key] = value + else: + extra_params[key] = value + if extra_params: + request_data["extra_params"] = extra_params + return request_data + + @staticmethod + def _coerce_output_format(raw_format: Any) -> StreamingVideoFormat: + if raw_format == "m4s": + return cast(StreamingVideoFormat, raw_format) + raise ValueError("Invalid session.start format: expected 'm4s'.") + + async def _control_loop( + self, + websocket: WebSocket, + request_id: str, + stop_event: asyncio.Event, + send_lock: asyncio.Lock, + progress: _SessionProgress, + ) -> None: + while not stop_event.is_set(): + if progress.stalled_for(self._stall_timeout): + await self._send_error( + websocket, + "Stall timeout: no generation progress", + send_lock=send_lock, + ) + await self._abort_request(request_id) + stop_event.set() + return + + poll_timeout = min(_CONTROL_POLL_INTERVAL, self._stall_timeout) + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=poll_timeout, + ) + except asyncio.TimeoutError: + continue + + if len(raw) > _MAX_CONTROL_MESSAGE_SIZE: + await self._send_error(websocket, "control message too large", send_lock=send_lock) + continue + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON message", send_lock=send_lock) + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "WebSocket messages must be JSON objects", send_lock=send_lock) + continue + + msg_type = msg.get("type") + if msg_type == "session.stop": + await self._abort_request(request_id) + stop_event.set() + return + if msg_type == "session.ping": + progress.touch() + try: + async with send_lock: + await websocket.send_json({"type": "session.pong"}) + except Exception: + pass + continue + if msg_type == "session.prompt_update": + await self._send_error( + websocket, + "session.prompt_update is not supported by this backend yet", + send_lock=send_lock, + ) + 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 + except Exception: + logger.debug("Failed to abort streaming video request %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]: + """Yield encoded video bytes from diffusion streaming outputs.""" + prompt, gen_params, vp = 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"] + + output_fps = vp.fps or gen_params.resolved_frame_rate or 16 + encoder = create_streaming_video_encoder( + output_format=output_format, + fps=output_fps, + video_codec_options=video_codec_options, + ) + 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)) + videos = self._extract_video_outputs(result) + for video in videos: + chunk = encoder.encode(video) + if chunk: + yield chunk + + completed = True + final_chunk = encoder.close() + if final_chunk: + yield final_chunk + finally: + if not completed: + encoder.close() + + def _build_prompt_and_sampling_params( + self, + request: VideoGenerationRequest, + ) -> tuple[OmniTextPrompt, OmniDiffusionSamplingParams, VideoParams]: + """Build text-only diffusion inputs for the streaming endpoint.""" + prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt) + if request.negative_prompt is not None: + prompt["negative_prompt"] = request.negative_prompt + + gen_params = self._resolve_default_sampling_params() + vp = request.resolve_video_params() + if vp.width is not None and vp.height is not None: + gen_params.width = vp.width + gen_params.height = vp.height + if vp.num_frames is not None: + gen_params.num_frames = vp.num_frames + if vp.fps is not None: + gen_params.fps = vp.fps + gen_params.frame_rate = float(vp.fps) + + provided_fields = request.model_fields_set + if "num_inference_steps" in provided_fields and request.num_inference_steps is not None: + gen_params.num_inference_steps = request.num_inference_steps + if "guidance_scale" in provided_fields and request.guidance_scale is not None: + gen_params.guidance_scale = request.guidance_scale + if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None: + gen_params.guidance_scale_2 = request.guidance_scale_2 + if "true_cfg_scale" in provided_fields and request.true_cfg_scale is not None: + gen_params.true_cfg_scale = request.true_cfg_scale + if "seed" in provided_fields and request.seed is not None: + gen_params.seed = request.seed + if "boundary_ratio" in provided_fields and request.boundary_ratio is not None: + gen_params.boundary_ratio = request.boundary_ratio + if "flow_shift" in provided_fields and request.flow_shift is not None: + gen_params.extra_args["flow_shift"] = request.flow_shift + + if request.extra_params is not None: + if not isinstance(request.extra_params, dict): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="extra_params must be a JSON object/dict.", + ) + gen_params.extra_args.update(request.extra_params) + + self._apply_lora(request.lora, gen_params) + return prompt, gen_params, vp + + 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) + return OmniDiffusionSamplingParams() + + @staticmethod + def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None: + try: + lora_request, lora_scale = parse_lora_request(lora_body) + except ValueError as e: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=str(e), + ) from e + + if lora_request is None: + return + + gen_params.lora_request = lora_request + if lora_scale is not None: + gen_params.lora_scale = lora_scale + + async def _iter_generation_outputs( + self, + prompt: OmniTextPrompt, + gen_params: OmniDiffusionSamplingParams, + request_id: str, + progress: _SessionProgress, + ) -> AsyncGenerator[Any, None]: + stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None) + if not stage_configs: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="Stage configs not found. Start server with an omni diffusion model.", + ) + + for stage in stage_configs: + stage_type = get_stage_type(stage) + if stage_type != "diffusion": + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + 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), + diffusion_params=gen_params, + replace_diffusion_params=True, + ) + + async for output in engine_client.generate( + prompt=prompt, + request_id=request_id, + sampling_params_list=sampling_params_list, + ): + progress.touch() + yield output + + @staticmethod + def _normalize_video_outputs(videos: Any) -> list[Any]: + if videos is None: + return [] + if hasattr(videos, "ndim") 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] = [] + for item in videos: + if hasattr(item, "ndim") 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: + 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"): + 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: + videos = result.multimodal_output.get("video") + + normalized = self._normalize_video_outputs(videos) + if not normalized: + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + detail="No video outputs found in generation result.", + ) + return normalized + + @staticmethod + async def _send_error( + websocket: WebSocket, + message: str, + code: int | 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 send_lock is None: + await websocket.send_json(payload) + else: + async with send_lock: + await websocket.send_json(payload) + except Exception: + pass diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index fbb37c36b..aa6479189 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -423,16 +423,8 @@ def _coerce_audio_to_numpy(audio: Any) -> np.ndarray: return arr.astype(np.float32) -def _encode_video_bytes( - video: Any, - fps: int, - audio: Any | None = None, - audio_sample_rate: int | None = None, - video_codec_options: dict[str, str] | None = None, -) -> bytes: - """Encode a video payload into MP4 bytes, optionally muxing audio.""" - from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes - +def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: + """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3).""" frames = _coerce_video_to_frames(video) if not frames: raise ValueError("No frames found to encode.") @@ -448,13 +440,23 @@ def _encode_video_bytes( frames_np *= 255.0 frames_u8 = np.round(frames_np).astype(np.uint8) - # Ensure contiguous memory layout for faster PyAV muxing - frames_u8 = np.ascontiguousarray(frames_u8) + return np.ascontiguousarray(frames_u8) + + +def _encode_video_bytes( + video: Any, + fps: int, + audio: Any | None = None, + audio_sample_rate: int | None = None, + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Encode a video payload into MP4 bytes, optionally muxing audio.""" + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None return mux_video_audio_bytes( - frames_u8, + _coerce_video_to_uint8_frames(video), audio_np, fps=float(fps), audio_sample_rate=audio_sample_rate or 24000, @@ -462,6 +464,55 @@ def _encode_video_bytes( ) +class FragmentedMP4VideoEncoder: + """Normalize video chunks and append them to one fragmented MP4 stream.""" + + def __init__( + self, + *, + fps: int | float, + video_codec_options: dict[str, str] | None = None, + ) -> None: + self._fps = float(fps) + self._video_codec_options = video_codec_options + self._muxer: Any | None = None + + def encode(self, video: Any) -> bytes: + """Encode one generated video chunk and return newly emitted fMP4 bytes.""" + from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer + + frames_u8 = _coerce_video_to_uint8_frames(video) + if self._muxer is None: + self._muxer = FragmentedMP4Muxer( + width=frames_u8.shape[2], + height=frames_u8.shape[1], + fps=self._fps, + video_codec_options=self._video_codec_options, + ) + return self._muxer.mux_video_frames(frames_u8) + + def close(self) -> bytes: + """Close the underlying fMP4 muxer and return trailing bytes, if any.""" + if self._muxer is None: + return b"" + return self._muxer.close() + + +StreamingVideoFormat = Literal["m4s"] + + +def create_streaming_video_encoder( + *, + output_format: StreamingVideoFormat, + fps: int | float, + video_codec_options: dict[str, str] | None = None, +) -> FragmentedMP4VideoEncoder: + """Create an incremental encoder for the requested WebSocket video format.""" + if output_format == "m4s": + return FragmentedMP4VideoEncoder(fps=fps, video_codec_options=video_codec_options) + raise ValueError(f"Unsupported streaming video format: {output_format}") + + def encode_video_base64( video: Any, fps: int, diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 2dfa1699f..b3a597365 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -175,6 +175,7 @@ class OmniRequestOutput: final_output_type: str = "image", stage_durations: dict[str, float] | None = None, peak_memory_mb: float = 0.0, + finished: bool = True, ) -> "OmniRequestOutput": """Create output from diffusion model. @@ -211,7 +212,7 @@ class OmniRequestOutput: _custom_output=custom_output or {}, stage_durations=stage_durations or {}, peak_memory_mb=peak_memory_mb, - finished=True, + finished=finished, ) @property