[Refact] Refactor diffusion output streaming (#5166)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/test_diffusion_engine.py

@@ -15,7 +15,11 @@ from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _move_tensor_tree_to_cpu
+from vllm_omni.diffusion.diffusion_engine import (
+    DiffusionEngine,
+    DiffusionExecutionMode,
+    _move_tensor_tree_to_cpu,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
@@ -221,7 +225,7 @@ class TestRequestBatchCapability:
             diffusion_engine_module.supports_request_batch(od_config)
         registry_load.assert_not_called()
 
-    def test_engine_disables_batch_dispatch_for_single_request_pipeline(
+    def test_engine_uses_request_batch_mode_for_single_request_pipeline(
         self,
         monkeypatch: pytest.MonkeyPatch,
         mocker: MockerFixture,
@@ -230,6 +234,7 @@ class TestRequestBatchCapability:
             model_class_name="SinglePipeline",
             custom_pipeline_args=None,
             streaming_output=False,
+            max_num_seqs=1,
         )
         fake_executor = SimpleNamespace(
             execute_request=mocker.Mock(return_value="per-request"),
@@ -255,15 +260,50 @@ class TestRequestBatchCapability:
             "_try_load_model_cls",
             lambda model_class_name: _SingleRequestPipeline,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
         engine = DiffusionEngine(od_config)
-        output = engine.execute_fn(_make_request_mode_sched_output("req-a", "req-b"))
+        output = engine.execute_fn(_make_request_mode_sched_output("req-a"))
 
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
         assert engine.supports_request_batch is False
-        assert output == "per-request"
-        fake_executor.execute_request.assert_called_once()
-        fake_executor.execute_batch.assert_not_called()
+        assert output == "batch"
+        fake_executor.execute_batch.assert_called_once()
+        fake_executor.execute_request.assert_not_called()
+
+    def test_engine_rejects_multi_seq_request_batch_for_single_request_pipeline(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            streaming_output=False,
+            max_num_seqs=2,
+        )
+        fake_executor_cls = mocker.Mock()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+
+        with pytest.raises(ValueError, match="max_num_seqs=1"):
+            DiffusionEngine(od_config)
+        fake_executor_cls.assert_not_called()
 
     @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
     def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
@@ -276,6 +316,7 @@ class TestRequestBatchCapability:
             model_class_name="BatchPipeline",
             custom_pipeline_args=None,
             streaming_output=False,
+            max_num_seqs=2,
         )
         fake_executor = SimpleNamespace(
             execute_request=mocker.Mock(return_value="per-request"),
@@ -301,16 +342,55 @@ class TestRequestBatchCapability:
             "_try_load_model_cls",
             lambda model_class_name: _BatchCapablePipeline,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
         engine = DiffusionEngine(od_config)
         output = engine.execute_fn(_make_request_mode_sched_output(*request_ids))
 
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
         assert engine.supports_request_batch is True
         assert output == "batch"
         fake_executor.execute_batch.assert_called_once()
         fake_executor.execute_request.assert_not_called()
 
+    def test_make_engine_runs_startup_warmup(self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            engine_backend="default",
+            streaming_output=False,
+        )
+        fake_executor = SimpleNamespace(
+            execute_request=mocker.Mock(),
+            execute_batch=mocker.Mock(),
+            execute_step=mocker.Mock(),
+        )
+        fake_executor_cls = mocker.Mock(return_value=fake_executor)
+        warmup = mocker.Mock()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+        monkeypatch.setattr(DiffusionEngine, "run_startup_warmup", warmup)
+
+        engine = DiffusionEngine.make_engine(od_config)
+
+        assert isinstance(engine, DiffusionEngine)
+        warmup.assert_called_once_with()
+
 
 class TestRequestBatchAdmission:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -470,11 +550,20 @@ def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
     assert moved["items"][1][0] == "keep"
 
 
+async def _consume_final_output(generator):
+    final_output = None
+    async for output in generator:
+        final_output = output
+    if final_output is None:
+        raise RuntimeError("Diffusion execution finished without output.")
+    return final_output
+
+
 @pytest.mark.asyncio
-async def test_async_add_req_and_wait_for_response():
+async def test_async_add_req_and_stream_response():
     engine = object.__new__(DiffusionEngine)
     engine.scheduler = MockScheduler()
-    engine._out_queue = {}
+    engine._out_streams = {}
     engine.abort_queue: queue.Queue[str] = queue.Queue()
     engine._rpc_queue = queue.Queue()
     engine._rpc_lock = threading.RLock()
@@ -485,6 +574,7 @@ async def test_async_add_req_and_wait_for_response():
     engine._loop_started = False
     engine.main_loop = None
     engine.supports_request_batch = False
+    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
 
     engine._finalize_finished_request = lambda rid, out, err: out.result
 
@@ -512,7 +602,7 @@ async def test_async_add_req_and_wait_for_response():
     async def run_task(rid):
         req = SimpleNamespace(request_id=rid)
         start = time.time()
-        res = await engine.async_add_req_and_wait_for_response(req)
+        res = await _consume_final_output(engine.async_add_req_and_stream_response(req))
         return rid, res, time.time() - start
 
     task_ids = [f"req_{i}" for i in range(5)]

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -5,15 +5,17 @@ import asyncio
 import queue
 import threading
 from types import SimpleNamespace
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 
+from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -33,8 +35,7 @@ def _make_engine() -> DiffusionEngine:
     engine.executor = SimpleNamespace(shutdown=Mock())
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
-    engine._out_queue = {}
-    engine._out_queue_streaming = {}
+    engine._out_streams = {}
     engine._closed = False
     engine._shutdown_complete = False
     engine.abort_queue = queue.Queue()
@@ -44,28 +45,13 @@ def _make_engine() -> DiffusionEngine:
     return engine
 
 
-def test_close_completes_pending_async_waiters() -> None:
-    engine = _make_engine()
-    event_loop = asyncio.new_event_loop()
-    try:
-        future = event_loop.create_future()
-        engine._out_queue["pending-req"] = future
-
-        engine.close()
-
-        assert future.done()
-        assert future.result().error == "DiffusionEngine is closed."
-    finally:
-        event_loop.close()
-
-
-def test_close_completes_pending_streaming_waiters() -> None:
+def test_close_completes_pending_output_streams() -> 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._out_streams["pending-stream"] = queue
 
         engine.close()
 
@@ -76,21 +62,101 @@ def test_close_completes_pending_streaming_waiters() -> None:
         event_loop.close()
 
 
-def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
+def test_emit_finished_outputs_finalizes_already_drained_waiter() -> None:
     class RacingOutQueue(dict):
-        def __contains__(self, key):
-            return True
-
-        def pop(self, key, default=None):
+        def get(self, key, default=None):
             return default
 
     engine = _make_engine()
-    engine._out_queue = RacingOutQueue()
-    engine._finalize_finished_request = Mock(side_effect=AssertionError("should not finalize drained waiters"))
+    request_id = engine.scheduler.add_request(_make_request("pending-req"))
+    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+    engine._out_streams = RacingOutQueue()
 
-    engine._handle_finished_requests({"pending-req"})
+    engine._emit_finished_outputs({request_id})
 
-    engine._finalize_finished_request.assert_not_called()
+    assert engine.scheduler.get_request_state(request_id) is None
+
+
+def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
+    engine = _make_engine()
+    engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
+    request_id = engine.scheduler.add_request(_make_request("step-drained"))
+    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+    engine._emit_outputs({request_id}, [request_id], SimpleNamespace(get_request_output=lambda _request_id: None))
+
+    assert engine.scheduler.get_request_state(request_id) is None
+
+
+def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
+    od_config = SimpleNamespace(
+        custom_pipeline_args=None,
+        model_class_name="CustomSchedulerPipeline",
+        streaming_output=False,
+    )
+    custom_scheduler = RequestScheduler()
+    fake_executor = SimpleNamespace(
+        execute_request=Mock(),
+        execute_batch=Mock(),
+        execute_step=Mock(),
+    )
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+        lambda *args, **kwargs: Mock(return_value=fake_executor),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.supports_request_batch",
+        lambda *args, **kwargs: False,
+    )
+
+    engine = DiffusionEngine(od_config, scheduler=custom_scheduler)
+
+    assert engine.scheduler is custom_scheduler
+
+
+@pytest.mark.asyncio
+async def test_step_compatibility_wrapper_returns_final_batch() -> None:
+    engine = _make_engine()
+    first = [OmniRequestOutput.from_diffusion(request_id="req", images=[], finished=False)]
+    final = [OmniRequestOutput.from_diffusion(request_id="req", images=[], finished=True)]
+
+    async def _step_streaming(_request):
+        yield first
+        yield final
+
+    engine.step_streaming = _step_streaming  # type: ignore[method-assign]
+    with patch.object(diffusion_engine_module.logger, "warning_once") as warning_once:
+        output = await engine.step(_make_request("req"))
+
+    assert output is final
+    warning_once.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_async_wait_compatibility_wrapper_returns_final_output() -> None:
+    engine = _make_engine()
+    first = DiffusionOutput(output="chunk", finished=False)
+    final = DiffusionOutput(output="final", finished=True)
+
+    async def _stream_response(_request):
+        yield first
+        yield final
+
+    engine.async_add_req_and_stream_response = _stream_response  # type: ignore[method-assign]
+    with patch.object(diffusion_engine_module.logger, "warning_once") as warning_once:
+        output = await engine.async_add_req_and_wait_for_response(_make_request("req"))
+
+    assert output is final
+    warning_once.assert_called_once()
 
 
 def test_abort_request_id_aborts_scheduler_request() -> None:

tests/diffusion/test_diffusion_engine_metrics.py

@@ -80,27 +80,23 @@ class TestMetricKeys:
             "format_diffusion_outputs() — should only use 'preprocess_time_ms'"
         )
 
-    def test_metric_key_naming_consistency(self) -> None:
-        """Metric keys should map to the explicit timing fields."""
+    def test_timing_metric_key_naming_consistency(self) -> None:
+        """Timing metrics should be attached by the engine orchestration path."""
         source = _read_source(_FORMATTER_PATH)
         formatter_source = _get_function_source(source, None, "format_diffusion_outputs")
+        engine_source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(engine_source, "DiffusionEngine", "step_streaming")
         lines = formatter_source.split("\n")
 
-        found_exec = False
-        found_total = False
         for line in lines:
             if '"diffusion_engine_exec_time_ms"' in line:
-                found_exec = True
-                assert "timings.exec_time_s" in line, (
-                    "diffusion_engine_exec_time_ms should measure executor time only (timings.exec_time_s)"
-                )
+                raise AssertionError("diffusion_engine_exec_time_ms should be attached in step_streaming()")
             if '"diffusion_engine_total_time_ms"' in line:
-                found_total = True
-                assert "timings.total_time_ms" in line, (
-                    "diffusion_engine_total_time_ms should measure full step time (timings.total_time_ms)"
-                )
-        assert found_exec, "diffusion_engine_exec_time_ms key not found in format_diffusion_outputs()"
-        assert found_total, "diffusion_engine_total_time_ms key not found in format_diffusion_outputs()"
+                raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
+
+        assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
+        assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source
+        assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source
 
 
 class TestDummyRunAllocation:

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -30,7 +30,7 @@ from typing import Any
 import pytest
 
 from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _RpcTask
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode, _RpcTask
 from vllm_omni.diffusion.sched import RequestScheduler
 from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey
 from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -91,7 +91,7 @@ class _ConcurrencyTrackingExecutor:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
                 tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
-                return DiffusionOutput(error=f"result_for_{tag}")
+                return DiffusionOutput(error=f"result_for_{tag}", finished=True)
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
         finally:
@@ -149,12 +149,12 @@ def _make_engine_with_loop(
     engine.scheduler = sched
     engine.step_execution = False
     engine.supports_request_batch = False
+    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
     engine.execute_fn = engine.executor.execute_request
 
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
-    engine._out_queue = {}
-    engine._out_queue_streaming = {}
+    engine._out_streams = {}
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
@@ -187,6 +187,15 @@ def _stop_engine(engine: DiffusionEngine) -> None:
     assert not engine.worker_thread.is_alive(), "Busy loop thread did not stop"
 
 
+async def _consume_final_output(generator):
+    final_output = None
+    async for output in generator:
+        final_output = output
+    if final_output is None:
+        raise RuntimeError("Diffusion execution finished without output.")
+    return final_output
+
+
 # ─────────────────────── single-thread invariant ───────────────────────────
 
 
@@ -199,7 +208,7 @@ async def test_executor_only_called_from_busy_loop_thread():
     busy_tid = engine.worker_thread.ident
     try:
         # Per-request path
-        await engine.async_add_req_and_wait_for_response(_make_request("req1"))
+        await _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
         # Raw RPC path (sync from a worker thread)
         result = await asyncio.to_thread(engine.collective_rpc, "ping", args=("a",), unique_reply_rank=0)
         assert result.error == "rpc_result_for_a"
@@ -228,7 +237,7 @@ async def test_executor_calls_never_overlap_under_load():
             return await engine.async_collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)
 
         async def _request(i: int):
-            return await engine.async_add_req_and_wait_for_response(_make_request(f"r{i}"))
+            return await _consume_final_output(engine.async_add_req_and_stream_response(_make_request(f"r{i}")))
 
         tasks = [_rpc(i) for i in range(20)] + [_request(i) for i in range(5)]
         results = await asyncio.gather(*tasks)
@@ -510,7 +519,9 @@ async def test_rpc_and_request_results_do_not_swap():
     engine = _make_engine_with_loop(loop, rpc_delay=0.02)
     try:
         rpc_task = asyncio.create_task(engine.async_collective_rpc("ping", args=("rpc1",), unique_reply_rank=0))
-        req_task = asyncio.create_task(engine.async_add_req_and_wait_for_response(_make_request("req1")))
+        req_task = asyncio.create_task(
+            _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
+        )
         rpc_res, req_res = await asyncio.gather(rpc_task, req_task)
     finally:
         _stop_engine(engine)

tests/diffusion/test_diffusion_model_runner.py

@@ -98,6 +98,29 @@ class _ChunkStepPipeline:
         return output
 
 
+class _FinalOnlyStepPipeline:
+    device = torch.device("cpu")
+    supports_step_execution = True
+
+    def prepare_encode(self, state):
+        state.prompt_embeds = torch.zeros(1, 1, 1)
+        state.latents = torch.zeros(1, 1)
+        state.timesteps = torch.tensor([1.0])
+        state.step_index = 0
+        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
+
+    def post_decode(self, state):
+        return DiffusionOutput(output=state.latents.clone())
+
+
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -228,6 +251,35 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     assert fourth.get_request_output("req").finished is True
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_stepwise_streaming_decodes_final_only_pipeline(monkeypatch):
+    """Step streaming still returns a final output when no chunk boundaries exist."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _FinalOnlyStepPipeline()
+    runner.od_config.streaming_output = True
+    runner.od_config.step_execution = True
+    req = _make_request()
+    req.request_id = "req"
+    req.sampling_params.num_inference_steps = 1
+
+    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=[]),
+    )
+
+    output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output).get_request_output("req")
+
+    assert output.finished is True
+    assert output.result is not None
+    assert torch.equal(output.result.output, torch.ones(1, 1))
+
+
 @pytest.mark.core_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):

tests/diffusion/test_diffusion_output_formatter.py

@@ -11,7 +11,6 @@ import torch
 from vllm_omni.diffusion import output_formatter
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.output_formatter import (
-    DiffusionStepTimings,
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
     normalize_diffusion_postprocess_output,
@@ -42,15 +41,6 @@ def _config(model_class_name: str = "mock_model") -> SimpleNamespace:
     return SimpleNamespace(model_class_name=model_class_name)
 
 
-def _timings() -> DiffusionStepTimings:
-    return DiffusionStepTimings(
-        preprocess_time_s=0.01,
-        exec_time_s=0.02,
-        postprocess_time_s=0.03,
-        total_time_ms=60.0,
-    )
-
-
 def test_formatter_preserves_single_video_audio_actions_and_metadata(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -80,7 +70,6 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
         ),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert len(results) == 1
@@ -102,14 +91,10 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
     assert result.stage_durations == {"execute": 1.25}
     assert result.peak_memory_mb == 321.0
     assert result.metrics == {
-        "preprocess_time_ms": 10.0,
-        "diffusion_engine_exec_time_ms": 20.0,
-        "diffusion_engine_total_time_ms": 60.0,
         "image_num": 1,
         "resolution": 512,
         "width": None,
         "height": None,
-        "postprocess_time_ms": 30.0,
     }
 
 
@@ -162,7 +147,6 @@ def test_formatter_normalizes_payload_metadata_envelope(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.multimodal_output == {
@@ -198,7 +182,6 @@ def test_formatter_maps_trajectory_payload_to_request_output(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == ["image-0"]
@@ -231,7 +214,6 @@ def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPa
         diffusion_output=DiffusionOutput(output=None),
         output_data={"payload": {"text": "caption"}},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -269,7 +251,6 @@ def test_formatter_preserves_image_with_reasoning_metadata(
         diffusion_output=DiffusionOutput(output=None),
         output_data=None,
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert postprocess_output.primary_key == "image"
@@ -298,7 +279,6 @@ def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
         diffusion_output=DiffusionOutput(output=["waveform"]),
         output_data=["waveform"],
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -330,7 +310,6 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == ["frame-0"]
@@ -365,7 +344,6 @@ def test_formatter_preserves_audio_only_postprocess_dict(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -397,7 +375,6 @@ def test_formatter_preserves_single_prompt_multiple_audio_outputs(
         diffusion_output=DiffusionOutput(output=["waveform-0", "waveform-1"]),
         output_data=["waveform-0", "waveform-1"],
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -425,7 +402,6 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert len(results) == 1

tests/diffusion/test_diffusion_scheduler.py

@@ -11,7 +11,7 @@ import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -532,13 +532,11 @@ class TestDiffusionEngine:
         assert output is runner_output.result
         engine.execute_fn.assert_called_once()
 
-    def test_initializes_injected_scheduler(
+    def test_initializes_default_request_scheduler(
         self,
         monkeypatch: pytest.MonkeyPatch,
         mocker: MockerFixture,
     ) -> None:
-        request = _make_request("init")
-        scheduler = _StubScheduler(request, DiffusionOutput(output=None))
         od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
         fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
 
@@ -554,11 +552,11 @@ class TestDiffusionEngine:
             "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
             lambda *args, **kwargs: fake_executor_cls,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
-        DiffusionEngine(od_config, scheduler=scheduler)
+        engine = DiffusionEngine(od_config)
 
-        assert scheduler.initialized_with is od_config
+        assert isinstance(engine.scheduler, RequestScheduler)
+        assert engine.scheduler.max_num_running_reqs == 1
         fake_executor_cls.assert_called_once_with(od_config)
 
     def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
@@ -573,16 +571,19 @@ class TestDiffusionEngine:
         assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
 
     @pytest.mark.asyncio
-    async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
+    async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._check_and_start_background_loop = mocker.AsyncMock()
         engine.pre_process_func = None
-        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
-            return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
-        )
+
+        async def _stream(_request):
+            yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
+
+        engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None))
 
         with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            await engine.step(_make_request("req-abort"))
+            async for _ in engine.step_streaming(_make_request("req-abort")):
+                pass
 
     def test_abort_queue_marks_request_finished_aborted(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
@@ -619,12 +620,13 @@ class TestDiffusionEngine:
         engine.scheduler.initialize(SimpleNamespace())
         engine._rpc_lock = threading.RLock()
         engine._cv = threading.Condition(engine._rpc_lock)
-        engine._out_queue_streaming = {}
+        engine._out_streams = {}
+        engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
         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
+        engine._out_streams[req_id] = queue
         sched_output = engine.scheduler.schedule()
 
         chunk = RunnerOutput(
@@ -634,7 +636,7 @@ class TestDiffusionEngine:
             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)
+        engine._emit_outputs(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"
@@ -648,7 +650,7 @@ class TestDiffusionEngine:
             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)
+        engine._emit_outputs(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"
@@ -662,15 +664,15 @@ class TestDiffusionEngine:
         engine.scheduler.initialize(SimpleNamespace())
         engine._rpc_lock = threading.RLock()
         engine._cv = threading.Condition(engine._rpc_lock)
-        engine._out_queue_streaming = {}
+        engine._out_streams = {}
         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._out_streams[req_id] = queue
         engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
 
-        engine._handle_empty_streaming_requests({req_id})
+        engine._emit_finished_outputs({req_id})
 
         output = await asyncio.wait_for(queue.get(), timeout=1)
         assert output.aborted is True
@@ -699,9 +701,9 @@ class TestDiffusionEngine:
             "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
             lambda *args, **kwargs: fake_executor_cls,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
         engine = DiffusionEngine(od_config)
 
+        assert engine.execution_mode == DiffusionExecutionMode.STEP_BATCH
         assert isinstance(engine.scheduler, StepScheduler)
         assert engine.execute_fn is fake_executor.execute_step
         fake_executor_cls.assert_called_once_with(od_config)

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -54,9 +54,12 @@ def client(mock_engine):
 
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_success(client, mock_engine):
-    # Setup mock engine step to return a successful result
     mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
-    mock_engine.step.return_value = [mock_result]
+
+    async def _step_streaming(_request):
+        yield [mock_result]
+
+    mock_engine.step_streaming = _step_streaming
 
     sampling_params = OmniDiffusionSamplingParams()
     await client.add_request_async("req-1", "A test prompt", sampling_params)
@@ -70,7 +73,6 @@ async def test_inline_dispatch_request_success(client, mock_engine):
 
     assert output is not None
     assert output.request_id == "req-1"
-    mock_engine.step.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -104,8 +106,11 @@ async def test_inline_dispatch_request_streaming_success(client, mock_engine):
 
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_error(client, mock_engine):
-    # Setup mock engine step to raise an exception
-    mock_engine.step.side_effect = RuntimeError("Engine failure")
+    async def _step_streaming(_request):
+        raise RuntimeError("Engine failure")
+        yield  # pragma: no cover
+
+    mock_engine.step_streaming = _step_streaming
 
     sampling_params = OmniDiffusionSamplingParams()
     await client.add_request_async("req-err", "A test prompt", sampling_params)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -91,7 +91,7 @@ def _make_engine(num_gpus: int = 1):
     engine._loop_started = False
     engine._rpc_queue = queue.Queue()
     engine.abort_queue = queue.Queue()
-    engine.execute_fn = executor.execute_request
+    engine.execute_fn = executor.execute_batch
     return engine, executor, req_q, res_q
 
 

tests/diffusion/test_stage_diffusion_proc.py

@@ -27,7 +27,7 @@ DELAY_BASE = 0.01
 
 
 class MockDiffusionEngine:
-    async def step(self, request):
+    async def step_streaming(self, request):
         def simulate_step_delay(height, width, num_inference_steps) -> float:
             return (height / BASE_HEIGHT) * (width / BASE_WIDTH) * (num_inference_steps / BASE_INFER_STEPS)
 
@@ -37,7 +37,7 @@ class MockDiffusionEngine:
         )
         delay = DELAY_BASE + delay_scale * DELAY_BASE
         await asyncio.sleep(delay)
-        return [MockOmniRequestOutput(request_id=request.request_id)]
+        yield [MockOmniRequestOutput(request_id=request.request_id)]
 
 
 @pytest.mark.asyncio