[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>

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)