[feat] Streaming diffusion video generation output (#3737)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: Zeyu Huang | 黃澤宇 <11222265+fhfuih@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

.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:

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]}'
+```

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

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

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

@@ -0,0 +1,11 @@
+<!-- This file and the corresponding JS 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. -->
+<div id="vllm-streaming-video-view" style="display:flex; flex-direction:column; gap:10px;">
+  <video id="vllm-streaming-video-player" autoplay muted controls playsinline
+    style="width:100%; min-height:320px; background:#111; border-radius:8px;"></video>
+  <div style="display:flex; align-items:center; gap:10px;">
+    <span id="vllm-streaming-video-status" style="font-weight:600;">Ready</span>
+  </div>
+  <pre id="vllm-streaming-video-log" style="height:220px; overflow:auto; margin:0; padding:10px; border-radius:8px;
+              background:#0f172a; color:#dbeafe; font-size:12px; white-space:pre-wrap;">Ready.</pre>
+</div>

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");
+})();

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",

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

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

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)

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)

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)

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

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."""

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)

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(

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)

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)

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

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)

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
 
 

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]

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,

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:

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)
 

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:

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:

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,

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

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

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:

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",

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

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

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.

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

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(

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:

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."""

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"}
+        <binary frame: fragmented MP4 bytes>
+        {"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

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,

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