[Refactor] Extract diffusion output formatting boundary (#4407)

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from vllm_omni.diffusion import diffusion_engine
+from vllm_omni.diffusion import io_support
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -13,12 +13,12 @@ def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.Mo
         dummy_run_num_frames = 2
 
     monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
+        io_support.DiffusionModelRegistry,
         "_try_load_model_cls",
         lambda model_class_name: JointAudioVideoModel,
     )
 
-    assert diffusion_engine.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2
+    assert io_support.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2
 
 
 def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -26,12 +26,12 @@ def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.Mon
         support_audio_output = True
 
     monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
+        io_support.DiffusionModelRegistry,
         "_try_load_model_cls",
         lambda model_class_name: AudioOutputModel,
     )
 
-    assert diffusion_engine.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2
+    assert io_support.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2
 
 
 def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -39,19 +39,19 @@ def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.Monke
         pass
 
     monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
+        io_support.DiffusionModelRegistry,
         "_try_load_model_cls",
         lambda model_class_name: VideoOnlyModel,
     )
 
-    assert diffusion_engine.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1
+    assert io_support.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1
 
 
 def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
+        io_support.DiffusionModelRegistry,
         "_try_load_model_cls",
         lambda model_class_name: None,
     )
 
-    assert diffusion_engine.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
+    assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2

tests/diffusion/test_diffusion_engine_metrics.py

@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Source-level regression tests for DiffusionEngine.
+"""Source-level regression tests for diffusion output/engine helpers.
 
 These tests verify naming conventions and patterns by inspecting source code
 at the function level using AST. They are intentionally coupled to the source
-layout and should be updated whenever the metrics construction code is
-refactored.
+layout and should be updated whenever the inspected helper code is refactored.
 """
 
 from __future__ import annotations
@@ -23,10 +22,20 @@ _ENGINE_PATH = os.path.normpath(
         "diffusion_engine.py",
     )
 )
+_FORMATTER_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "diffusion",
+        "output_formatter.py",
+    )
+)
 
 
-def _read_engine_source() -> str:
-    with open(_ENGINE_PATH) as f:
+def _read_source(path: str) -> str:
+    with open(path) as f:
         return f.read()
 
 
@@ -60,37 +69,38 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) ->
 
 
 class TestMetricKeys:
-    """Verify metric naming conventions in DiffusionEngine.step()."""
+    """Verify metric naming conventions in diffusion output formatting."""
 
     def test_no_duplicate_preprocess_key(self) -> None:
-        """step() should not contain 'preprocessing_time_ms' (duplicate of 'preprocess_time_ms')."""
-        source = _read_engine_source()
-        step_source = _get_function_source(source, "DiffusionEngine", "step")
-        assert "preprocessing_time_ms" not in step_source, (
-            "Found duplicate key 'preprocessing_time_ms' in step() — should only use 'preprocess_time_ms'"
+        """format_diffusion_outputs() should not duplicate 'preprocess_time_ms'."""
+        source = _read_source(_FORMATTER_PATH)
+        formatter_source = _get_function_source(source, None, "format_diffusion_outputs")
+        assert "preprocessing_time_ms" not in formatter_source, (
+            "Found duplicate key 'preprocessing_time_ms' in "
+            "format_diffusion_outputs() — should only use 'preprocess_time_ms'"
         )
 
     def test_metric_key_naming_consistency(self) -> None:
-        """In step(): exec_time should use exec_total_time, total_time should use step_total_ms."""
-        source = _read_engine_source()
-        step_source = _get_function_source(source, "DiffusionEngine", "step")
-        lines = step_source.split("\n")
+        """Metric keys should map to the explicit timing fields."""
+        source = _read_source(_FORMATTER_PATH)
+        formatter_source = _get_function_source(source, None, "format_diffusion_outputs")
+        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 "exec_total_time" in line, (
-                    "diffusion_engine_exec_time_ms should measure executor time only (exec_total_time)"
+                assert "timings.exec_time_s" in line, (
+                    "diffusion_engine_exec_time_ms should measure executor time only (timings.exec_time_s)"
                 )
             if '"diffusion_engine_total_time_ms"' in line:
                 found_total = True
-                assert "step_total_ms" in line, (
-                    "diffusion_engine_total_time_ms should measure full step time (step_total_ms)"
+                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 step()"
-        assert found_total, "diffusion_engine_total_time_ms key not found in step()"
+        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()"
 
 
 class TestDummyRunAllocation:
@@ -98,7 +108,7 @@ class TestDummyRunAllocation:
 
     def test_no_oversized_allocation(self) -> None:
         """_dummy_run should not allocate more audio than needed."""
-        source = _read_engine_source()
+        source = _read_source(_ENGINE_PATH)
         dummy_source = _get_function_source(source, "DiffusionEngine", "_dummy_run")
         assert "audio_sr * audio_duration_sec" not in dummy_source, (
             "_dummy_run should generate exact-sized audio, not allocate and slice"

tests/diffusion/test_diffusion_output_formatter.py

@@ -0,0 +1,307 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+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,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _request(
+    prompts: list[str] | None = None,
+    *,
+    num_outputs_per_prompt: int = 1,
+) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=prompts or ["prompt"],
+        request_id="req-1",
+        sampling_params=OmniDiffusionSamplingParams(
+            num_inference_steps=1,
+            num_outputs_per_prompt=num_outputs_per_prompt,
+            resolution=512,
+        ),
+    )
+
+
+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_custom_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    custom_output = {"from_output": "kept"}
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "video": ["frame-0"],
+            "audio": "audio-0",
+            "actions": "action-0",
+            "custom_output": {"from_postprocess": "merged"},
+            "audio_sample_rate": 48000,
+            "fps": 24.0,
+        },
+        custom_output,
+    )
+
+    results = format_diffusion_outputs(
+        request=_request(["prompt-0"]),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(
+            output=None,
+            custom_output=custom_output,
+            stage_durations={"execute": 1.25},
+            peak_memory_mb=321.0,
+        ),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert len(results) == 1
+    result = results[0]
+    assert result.images == ["frame-0"]
+    assert result.prompt == "prompt-0"
+    assert result.final_output_type == "image"
+    assert result.custom_output == {
+        "from_output": "kept",
+        "from_postprocess": "merged",
+    }
+    assert custom_output == {"from_output": "kept"}
+    assert result.multimodal_output == {
+        "audio": "audio-0",
+        "audio_sample_rate": 48000,
+        "fps": 24.0,
+        "actions": "action-0",
+    }
+    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,
+        "postprocess_time_ms": 30.0,
+    }
+
+
+def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        "caption",
+        {"text_output": "caption"},
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request(["describe this"]),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output="caption"),
+        output_data="caption",
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == []
+    assert result.prompt == "describe this"
+    assert result.final_output_type == "text"
+    assert result.custom_output == {"text_output": "caption"}
+
+
+def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    class AudioModel:
+        audio_sample_rate = 44100
+
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True)
+    monkeypatch.setattr(
+        output_formatter.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda _: AudioModel,
+    )
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform"], {})
+
+    [result] = format_diffusion_outputs(
+        request=_request(["speak"]),
+        od_config=_config("audio_model"),
+        diffusion_output=DiffusionOutput(output=["waveform"]),
+        output_data=["waveform"],
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == []
+    assert result.prompt == "speak"
+    assert result.final_output_type == "audio"
+    assert result.multimodal_output == {
+        "audio": "waveform",
+        "audio_sample_rate": 44100,
+    }
+
+
+def test_formatter_preserves_audio_model_video_audio_and_actions(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "video": ["frame-0"],
+            "audio": "audio-0",
+            "actions": "action-0",
+            "audio_sample_rate": 16000,
+            "fps": 30.0,
+        },
+        {},
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request(["watch and listen"]),
+        od_config=_config("audio_video_model"),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == ["frame-0"]
+    assert result.prompt == "watch and listen"
+    assert result.final_output_type == "image"
+    assert result.multimodal_output == {
+        "audio": "audio-0",
+        "audio_sample_rate": 16000,
+        "fps": 30.0,
+        "actions": "action-0",
+    }
+
+
+def test_formatter_preserves_audio_only_postprocess_dict(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "audio": "waveform",
+            "audio_sample_rate": 24000,
+        },
+        {},
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request(["speak"]),
+        od_config=_config("audio_model"),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == []
+    assert result.prompt == "speak"
+    assert result.final_output_type == "audio"
+    assert result.multimodal_output == {
+        "audio": "waveform",
+        "audio_sample_rate": 24000,
+    }
+
+
+def test_formatter_preserves_single_prompt_multiple_audio_outputs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True)
+    monkeypatch.setattr(
+        output_formatter.DiffusionModelRegistry,
+        "_try_load_model_cls",
+        lambda _: None,
+    )
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"], {})
+
+    [result] = format_diffusion_outputs(
+        request=_request(["speak"], num_outputs_per_prompt=2),
+        od_config=_config("audio_model"),
+        diffusion_output=DiffusionOutput(output=["waveform-0", "waveform-1"]),
+        output_data=["waveform-0", "waveform-1"],
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == []
+    assert result.prompt == "speak"
+    assert result.final_output_type == "audio"
+    assert result.multimodal_output == {"audio": ["waveform-0", "waveform-1"]}
+
+
+def test_formatter_preserves_multi_prompt_audio_and_action_slicing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "video": ["frame-0", "frame-1"],
+            "audio": ["audio-0", "audio-1"],
+            "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+            "custom_output": {"shared": True},
+            "fps": 12.5,
+        },
+        {},
+    )
+
+    results = format_diffusion_outputs(
+        request=_request(["prompt-0", "prompt-1"]),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert len(results) == 2
+    assert results[0].images == ["frame-0"]
+    assert results[1].images == ["frame-1"]
+    assert results[0].custom_output == {"shared": True}
+    assert results[1].custom_output == {"shared": True}
+    assert results[0].multimodal_output["audio"] == "audio-0"
+    assert results[1].multimodal_output["audio"] == "audio-1"
+    assert results[0].multimodal_output["fps"] == 12.5
+    assert results[1].multimodal_output["fps"] == 12.5
+    torch.testing.assert_close(
+        results[0].multimodal_output["actions"],
+        torch.tensor([1.0, 2.0]),
+    )
+    torch.testing.assert_close(
+        results[1].multimodal_output["actions"],
+        torch.tensor([3.0, 4.0]),
+    )
+
+
+def test_format_empty_diffusion_outputs_preserves_empty_response_shape() -> None:
+    results = format_empty_diffusion_outputs(_request(["prompt-0", "prompt-1"]))
+
+    assert len(results) == 2
+    assert [result.prompt for result in results] == ["prompt-0", "prompt-1"]
+    assert [result.images for result in results] == [[], []]
+    assert [result.metrics for result in results] == [{}, {}]

tests/diffusion/test_diffusion_scheduler.py

@@ -550,7 +550,7 @@ class TestDiffusionEngine:
             request_id="req-batch",
         )
 
-        mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False)
+        mocker.patch("vllm_omni.diffusion.output_formatter.supports_audio_output", return_value=False)
         outputs = await engine.step(request)
 
         assert len(outputs) == 2

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -338,7 +338,7 @@ def test_diffusers_backend_i2v_matches_diffusers(
 
 
 def _diffusers_dummy_run(pipe: DiffusionPipeline) -> None:
-    from vllm_omni.diffusion.diffusion_engine import supports_multimodal_input
+    from vllm_omni.diffusion.io_support import supports_multimodal_input
 
     supports_image_input, supports_audio_input = supports_multimodal_input(
         SimpleNamespace(

vllm_omni/diffusion/diffusion_engine.py

@@ -10,8 +10,8 @@ import queue
 import threading
 import time
 from collections.abc import Iterable
-from dataclasses import dataclass, field
-from typing import Any
+from dataclasses import dataclass, field, replace
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import PIL.Image
@@ -25,8 +25,21 @@ from vllm_omni.diffusion.data import (
     OmniDiffusionConfig,
 )
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.io_support import (
+    get_dummy_run_num_frames,
+    get_extra_body_params,
+    get_extra_output_params,
+    image_color_format,
+    supports_audio_output,
+    supports_multimodal_input,
+)
+from vllm_omni.diffusion.output_formatter import (
+    DiffusionStepTimings,
+    format_diffusion_outputs,
+    format_empty_diffusion_outputs,
+    normalize_diffusion_postprocess_output,
+)
 from vllm_omni.diffusion.registry import (
-    DiffusionModelRegistry,
     get_diffusion_action_post_process_func,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
@@ -37,52 +50,23 @@ from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
 from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
 from vllm_omni.errors import client_error_from_metadata, is_client_error_status
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
-from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
-
-@dataclass
-class _RpcTask:
-    """A pending collective_rpc invocation queued for the busy loop."""
-
-    method: str
-    args: tuple
-    kwargs: dict | None
-    deadline: float | None
-    unique_reply_rank: int | None
-    future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future)
-
-
-def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]:
-    if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None:
-        signature = inspect.signature(pipe_cls.__call__)
-        support_image_input = "image" in signature.parameters
-        support_audio_input = (
-            "audio" in signature.parameters or "audio_latents" in signature.parameters
-        )  # ref. LTX-2 format
-        return support_image_input, support_audio_input
-
-    supports_image_input = False
-    supports_audio_input = False
-
-    model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
-    if model_cls is not None:
-        supports_image_input = bool(getattr(model_cls, "support_image_input", False))
-        supports_audio_input = bool(getattr(model_cls, "support_audio_input", False))
-    return supports_image_input, supports_audio_input
-
-
-def image_color_format(model_class_name: str) -> str:
-    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
-    return getattr(model_cls, "color_format", "RGB")
-
-
-def supports_audio_output(model_class_name: str) -> bool:
-    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
-    if model_cls is None:
-        return False
-    return bool(getattr(model_cls, "support_audio_output", False))
+__all__ = [
+    "DiffusionEngine",
+    "_RpcTask",
+    "_move_tensor_tree_to_cpu",
+    "get_dummy_run_num_frames",
+    "get_extra_body_params",
+    "get_extra_output_params",
+    "image_color_format",
+    "supports_audio_output",
+    "supports_multimodal_input",
+]
 
 
 def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool:
@@ -106,41 +90,16 @@ def _move_tensor_tree_to_cpu(value: object) -> object:
     return value
 
 
-def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int:
-    """Get num_frames for the dummy warmup run. Returns 0 to skip warmup."""
+@dataclass
+class _RpcTask:
+    """A pending collective_rpc invocation queued for the busy loop."""
 
-    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
-    if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
-        return int(getattr(model_cls, "dummy_run_num_frames"))
-    return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
-
-
-def get_extra_body_params(model_class_name: str) -> frozenset[str]:
-    """Return the set of extra_body keys accepted by a pipeline.
-
-    Each pipeline can declare ``EXTRA_BODY_PARAMS: ClassVar[frozenset[str]]``
-    to advertise which request-level parameters should be forwarded from
-    ``extra_body`` to ``OmniDiffusionSamplingParams.extra_args``.
-    Returns an empty frozenset when the pipeline does not declare any.
-    """
-    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
-    if model_cls is None:
-        return frozenset()
-    return frozenset(getattr(model_cls, "EXTRA_BODY_PARAMS", frozenset()))
-
-
-def get_extra_output_params(model_class_name: str) -> frozenset[str]:
-    """Return the set of custom_output keys to expose in API response metrics.
-
-    Each pipeline can declare ``EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]]``
-    to advertise which ``DiffusionOutput.custom_output`` keys should be
-    copied into the response ``metrics`` dict.
-    Returns an empty frozenset when the pipeline does not declare any.
-    """
-    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
-    if model_cls is None:
-        return frozenset()
-    return frozenset(getattr(model_cls, "EXTRA_OUTPUT_PARAMS", frozenset()))
+    method: str
+    args: tuple
+    kwargs: dict | None
+    deadline: float | None
+    unique_reply_rank: int | None
+    future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future)
 
 
 class DiffusionEngine:
@@ -257,16 +216,7 @@ class DiffusionEngine:
 
         if output.output is None:
             logger.warning("Output is None, returning empty OmniRequestOutput")
-            return [
-                OmniRequestOutput.from_diffusion(
-                    request_id=request.request_id,
-                    images=[],
-                    prompt=prompt,
-                    metrics={},
-                    latents=None,
-                )
-                for i, prompt in enumerate(request.prompts)
-            ]
+            return format_empty_diffusion_outputs(request)
 
         # When CPU offload is enabled, move output to CPU before
         # post-processing to avoid device OOM — model weights may still
@@ -291,21 +241,17 @@ class DiffusionEngine:
                 outputs = self.post_process_func(output_data)
         else:
             outputs = output_data
-        audio_payload = None
-        model_audio_sample_rate = None
-        model_fps = None
-        if isinstance(outputs, dict):
-            audio_payload = outputs.get("audio")
-            action_payload = outputs.get("actions")
-            custom_output.update(outputs.get("custom_output") or {})
-            model_audio_sample_rate = outputs.get("audio_sample_rate")
-            model_fps = outputs.get("fps")
-            outputs = outputs.get("video", outputs)
+
+        postprocess_output = normalize_diffusion_postprocess_output(outputs, custom_output)
+        custom_output = postprocess_output.custom_output
+        action_payload = postprocess_output.action_payload
         if action_payload is None:
             action_payload = custom_output.get("actions")
+            if action_payload is not None:
+                postprocess_output = replace(postprocess_output, action_payload=action_payload)
         action_post_process_func = getattr(self, "action_post_process_func", None)
         if action_payload is None and action_post_process_func is not None:
-            raw_action_payload = custom_output.get("action", action_payload)
+            raw_action_payload = custom_output.get("action")
             if raw_action_payload is not None:
                 action_kwargs: dict[str, Any] = {}
                 if getattr(self, "_action_post_process_accepts_custom_output", False):
@@ -313,7 +259,12 @@ class DiffusionEngine:
                 if getattr(self, "_action_post_process_accepts_sampling_params", False):
                     action_kwargs["sampling_params"] = request.sampling_params
                 action_payload = action_post_process_func(raw_action_payload, **action_kwargs)
-                custom_output["actions"] = action_payload
+                custom_output = {**custom_output, "actions": action_payload}
+                postprocess_output = replace(
+                    postprocess_output,
+                    custom_output=custom_output,
+                    action_payload=action_payload,
+                )
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
 
@@ -327,184 +278,19 @@ class DiffusionEngine:
             step_total_ms,
         )
 
-        # Convert to OmniRequestOutput format
-        # Ensure outputs is a list
-        if not isinstance(outputs, list):
-            outputs = [outputs] if outputs is not None else []
-
-        metrics = {
-            "preprocess_time_ms": preprocess_time * 1000,
-            "diffusion_engine_exec_time_ms": exec_total_time * 1000,
-            "diffusion_engine_total_time_ms": step_total_ms,
-            "image_num": int(request.sampling_params.num_outputs_per_prompt),
-            "resolution": int(request.sampling_params.resolution),
-            "postprocess_time_ms": postprocess_time * 1000,
-        }
-
-        # Detect text output: when the pipeline returns a string (e.g.,
-        # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it
-        # as a text-type response instead of an image.
-        is_text_output = isinstance(output_data, str) and custom_output.get("text_output") is not None
-
-        # Handle single request or multiple requests
-        is_audio_output = supports_audio_output(self.od_config.model_class_name)
-        if is_audio_output and model_audio_sample_rate is None:
-            model_cls = DiffusionModelRegistry._try_load_model_cls(self.od_config.model_class_name)
-            model_audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
-
-        def _audio_mm(payload: Any) -> dict[str, Any]:
-            mm: dict[str, Any] = {"audio": payload}
-            if model_audio_sample_rate is not None:
-                mm["audio_sample_rate"] = model_audio_sample_rate
-            return mm
-
-        if len(request.prompts) == 1:
-            # Single request: return single OmniRequestOutput
-            prompt = request.prompts[0]
-            request_id = request.request_id
-
-            if is_text_output:
-                return [
-                    OmniRequestOutput.from_diffusion(
-                        request_id=request_id,
-                        images=[],
-                        prompt=prompt,
-                        metrics=metrics,
-                        custom_output=custom_output,
-                        final_output_type="text",
-                        stage_durations=output.stage_durations,
-                        peak_memory_mb=output.peak_memory_mb,
-                    ),
-                ]
-            if is_audio_output:
-                request_audio_payload = outputs[0] if len(outputs) == 1 else outputs
-                return [
-                    OmniRequestOutput.from_diffusion(
-                        request_id=request_id,
-                        images=[],
-                        prompt=prompt,
-                        metrics=metrics,
-                        latents=output.trajectory_latents,
-                        trajectory_latents=output.trajectory_latents,
-                        trajectory_timesteps=output.trajectory_timesteps,
-                        trajectory_log_probs=output.trajectory_log_probs,
-                        trajectory_decoded=output.trajectory_decoded,
-                        multimodal_output=_audio_mm(request_audio_payload),
-                        final_output_type="audio",
-                        stage_durations=output.stage_durations,
-                        peak_memory_mb=output.peak_memory_mb,
-                    ),
-                ]
-            else:
-                mm_output = {}
-                if audio_payload is not None:
-                    mm_output["audio"] = audio_payload
-                if model_audio_sample_rate is not None:
-                    mm_output["audio_sample_rate"] = model_audio_sample_rate
-                if model_fps is not None:
-                    mm_output["fps"] = model_fps
-                if action_payload is not None:
-                    mm_output["actions"] = action_payload
-                return [
-                    OmniRequestOutput.from_diffusion(
-                        request_id=request_id,
-                        images=outputs,
-                        prompt=prompt,
-                        metrics=metrics,
-                        latents=output.trajectory_latents,
-                        trajectory_latents=output.trajectory_latents,
-                        trajectory_timesteps=output.trajectory_timesteps,
-                        trajectory_log_probs=output.trajectory_log_probs,
-                        trajectory_decoded=output.trajectory_decoded,
-                        custom_output=custom_output,
-                        multimodal_output=mm_output,
-                        stage_durations=output.stage_durations,
-                        peak_memory_mb=output.peak_memory_mb,
-                    ),
-                ]
-        else:
-            # Multiple requests: return list of OmniRequestOutput
-            # Split images based on num_outputs_per_prompt for each request
-            results = []
-            output_idx = 0
-            request_id = request.request_id
-
-            for i, prompt in enumerate(request.prompts):
-                # Get images for this request
-                num_outputs = request.sampling_params.num_outputs_per_prompt
-                start_idx = output_idx
-                end_idx = start_idx + num_outputs
-                request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else []
-                output_idx = end_idx
-
-                if is_audio_output:
-                    request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs
-                    results.append(
-                        OmniRequestOutput.from_diffusion(
-                            request_id=request_id,
-                            images=[],
-                            prompt=prompt,
-                            metrics=metrics,
-                            latents=output.trajectory_latents,
-                            trajectory_latents=output.trajectory_latents,
-                            trajectory_timesteps=output.trajectory_timesteps,
-                            trajectory_log_probs=output.trajectory_log_probs,
-                            trajectory_decoded=output.trajectory_decoded,
-                            multimodal_output=_audio_mm(request_audio_payload),
-                            final_output_type="audio",
-                            stage_durations=output.stage_durations,
-                            peak_memory_mb=output.peak_memory_mb,
-                        ),
-                    )
-                else:
-                    mm_output = {}
-                    if audio_payload is not None:
-                        sliced_audio = audio_payload
-                        if isinstance(audio_payload, (list, tuple)):
-                            sliced_audio = audio_payload[start_idx:end_idx]
-                            if len(sliced_audio) == 1:
-                                sliced_audio = sliced_audio[0]
-                        elif hasattr(audio_payload, "shape") and getattr(audio_payload, "shape", None) is not None:
-                            if len(audio_payload.shape) > 0 and audio_payload.shape[0] >= end_idx:
-                                sliced_audio = audio_payload[start_idx:end_idx]
-                                if num_outputs == 1:
-                                    sliced_audio = sliced_audio[0]
-                        mm_output["audio"] = sliced_audio
-                    if model_audio_sample_rate is not None:
-                        mm_output["audio_sample_rate"] = model_audio_sample_rate
-                    if model_fps is not None:
-                        mm_output["fps"] = model_fps
-                    if action_payload is not None:
-                        sliced_actions = action_payload
-                        if isinstance(action_payload, (list, tuple)):
-                            sliced_actions = action_payload[start_idx:end_idx]
-                            if len(sliced_actions) == 1:
-                                sliced_actions = sliced_actions[0]
-                        elif hasattr(action_payload, "shape") and getattr(action_payload, "shape", None) is not None:
-                            if len(action_payload.shape) > 0 and action_payload.shape[0] >= end_idx:
-                                sliced_actions = action_payload[start_idx:end_idx]
-                                if num_outputs == 1:
-                                    sliced_actions = sliced_actions[0]
-                        mm_output["actions"] = sliced_actions
-                    results.append(
-                        OmniRequestOutput.from_diffusion(
-                            request_id=request_id,
-                            images=request_outputs,
-                            prompt=prompt,
-                            metrics=metrics,
-                            latents=output.trajectory_latents,
-                            trajectory_latents=output.trajectory_latents,
-                            trajectory_timesteps=output.trajectory_timesteps,
-                            trajectory_log_probs=output.trajectory_log_probs,
-                            trajectory_decoded=output.trajectory_decoded,
-                            custom_output=custom_output,
-                            multimodal_output=mm_output,
-                            stage_durations=output.stage_durations,
-                            peak_memory_mb=output.peak_memory_mb,
-                        ),
-                    )
-
-            return results
+        return format_diffusion_outputs(
+            request=request,
+            od_config=self.od_config,
+            diffusion_output=output,
+            output_data=output_data,
+            postprocess_output=postprocess_output,
+            timings=DiffusionStepTimings(
+                preprocess_time_s=preprocess_time,
+                exec_time_s=exec_total_time,
+                postprocess_time_s=postprocess_time,
+                total_time_ms=step_total_ms,
+            ),
+        )
 
     def _busy_loop(self):
         while not self.stop_event.is_set():

vllm_omni/diffusion/io_support.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import inspect
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+
+def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]:
+    if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None:
+        signature = inspect.signature(pipe_cls.__call__)
+        support_image_input = "image" in signature.parameters
+        support_audio_input = (
+            "audio" in signature.parameters or "audio_latents" in signature.parameters
+        )  # ref. LTX-2 format
+        return support_image_input, support_audio_input
+
+    supports_image_input = False
+    supports_audio_input = False
+
+    model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
+    if model_cls is not None:
+        supports_image_input = bool(getattr(model_cls, "support_image_input", False))
+        supports_audio_input = bool(getattr(model_cls, "support_audio_input", False))
+    return supports_image_input, supports_audio_input
+
+
+def image_color_format(model_class_name: str) -> str:
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    return getattr(model_cls, "color_format", "RGB")
+
+
+def supports_audio_output(model_class_name: str) -> bool:
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is None:
+        return False
+    return bool(getattr(model_cls, "support_audio_output", False))
+
+
+def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int:
+    """Get num_frames for the dummy warmup run. Returns 0 to skip warmup."""
+
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
+        return int(getattr(model_cls, "dummy_run_num_frames"))
+    return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
+
+
+def get_extra_body_params(model_class_name: str) -> frozenset[str]:
+    """Return the set of extra_body keys accepted by a pipeline.
+
+    Each pipeline can declare ``EXTRA_BODY_PARAMS: ClassVar[frozenset[str]]``
+    to advertise which request-level parameters should be forwarded from
+    ``extra_body`` to ``OmniDiffusionSamplingParams.extra_args``.
+    Returns an empty frozenset when the pipeline does not declare any.
+    """
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is None:
+        return frozenset()
+    return frozenset(getattr(model_cls, "EXTRA_BODY_PARAMS", frozenset()))
+
+
+def get_extra_output_params(model_class_name: str) -> frozenset[str]:
+    """Return the set of custom_output keys to expose in API response metrics.
+
+    Each pipeline can declare ``EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]]``
+    to advertise which ``DiffusionOutput.custom_output`` keys should be
+    copied into the response ``metrics`` dict.
+    Returns an empty frozenset when the pipeline does not declare any.
+    """
+    model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
+    if model_cls is None:
+        return frozenset()
+    return frozenset(getattr(model_cls, "EXTRA_OUTPUT_PARAMS", frozenset()))

vllm_omni/diffusion/output_formatter.py

@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.io_support import supports_audio_output
+from vllm_omni.diffusion.registry import DiffusionModelRegistry
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.outputs import OmniRequestOutput
+
+
+@dataclass(frozen=True)
+class DiffusionStepTimings:
+    preprocess_time_s: float
+    exec_time_s: float
+    postprocess_time_s: float
+    total_time_ms: float
+
+
+@dataclass(frozen=True)
+class DiffusionPostprocessOutput:
+    outputs: Any
+    custom_output: dict[str, Any]
+    audio_payload: Any | None = None
+    action_payload: Any | None = None
+    audio_sample_rate: int | None = None
+    fps: float | None = None
+    has_video_payload: bool = False
+
+
+def normalize_diffusion_postprocess_output(
+    outputs: Any,
+    custom_output: dict[str, Any],
+) -> DiffusionPostprocessOutput:
+    """Normalize the legacy postprocess dict shape used by diffusion models.
+
+    The returned envelope owns a shallow merged copy of ``custom_output`` so
+    postprocess values keep legacy override precedence without mutating the
+    caller-owned dict.
+    """
+
+    audio_payload = None
+    action_payload = None
+    audio_sample_rate = None
+    fps = None
+    has_video_payload = False
+    merged_custom_output = dict(custom_output)
+
+    if isinstance(outputs, dict):
+        has_video_payload = "video" in outputs
+        audio_payload = outputs.get("audio")
+        action_payload = outputs.get("actions")
+        merged_custom_output.update(outputs.get("custom_output") or {})
+        audio_sample_rate = outputs.get("audio_sample_rate")
+        fps = outputs.get("fps")
+        outputs = outputs.get("video", outputs)
+
+    return DiffusionPostprocessOutput(
+        outputs=outputs,
+        custom_output=merged_custom_output,
+        audio_payload=audio_payload,
+        action_payload=action_payload,
+        audio_sample_rate=audio_sample_rate,
+        fps=fps,
+        has_video_payload=has_video_payload,
+    )
+
+
+def format_empty_diffusion_outputs(request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+    return [
+        OmniRequestOutput.from_diffusion(
+            request_id=request.request_id,
+            images=[],
+            prompt=prompt,
+            metrics={},
+            latents=None,
+        )
+        for prompt in request.prompts
+    ]
+
+
+def format_diffusion_outputs(
+    *,
+    request: OmniDiffusionRequest,
+    od_config: OmniDiffusionConfig,
+    diffusion_output: DiffusionOutput,
+    output_data: Any,
+    postprocess_output: DiffusionPostprocessOutput,
+    timings: DiffusionStepTimings,
+) -> list[OmniRequestOutput]:
+    """Convert a finished diffusion model output into API-facing outputs."""
+
+    outputs = _ensure_list(postprocess_output.outputs)
+    metrics = {
+        "preprocess_time_ms": timings.preprocess_time_s * 1000,
+        "diffusion_engine_exec_time_ms": timings.exec_time_s * 1000,
+        "diffusion_engine_total_time_ms": timings.total_time_ms,
+        "image_num": int(request.sampling_params.num_outputs_per_prompt),
+        "resolution": int(request.sampling_params.resolution),
+        "postprocess_time_ms": timings.postprocess_time_s * 1000,
+    }
+
+    # Detect text output: when the pipeline returns a string (e.g.,
+    # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it
+    # as a text-type response instead of an image.
+    is_text_output = isinstance(output_data, str) and postprocess_output.custom_output.get("text_output") is not None
+
+    is_audio_output = supports_audio_output(od_config.model_class_name)
+    audio_sample_rate = postprocess_output.audio_sample_rate
+    if is_audio_output and audio_sample_rate is None:
+        model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
+        audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
+
+    if len(request.prompts) == 1:
+        return _format_single_prompt_output(
+            request=request,
+            diffusion_output=diffusion_output,
+            outputs=outputs,
+            metrics=metrics,
+            postprocess_output=postprocess_output,
+            is_text_output=is_text_output,
+            is_audio_output=is_audio_output,
+            audio_sample_rate=audio_sample_rate,
+        )
+
+    return _format_multi_prompt_outputs(
+        request=request,
+        diffusion_output=diffusion_output,
+        outputs=outputs,
+        metrics=metrics,
+        postprocess_output=postprocess_output,
+        is_audio_output=is_audio_output,
+        audio_sample_rate=audio_sample_rate,
+    )
+
+
+def _ensure_list(outputs: Any) -> list[Any]:
+    if isinstance(outputs, list):
+        return outputs
+    return [outputs] if outputs is not None else []
+
+
+def _format_audio_multimodal_output(payload: Any, audio_sample_rate: int | None) -> dict[str, Any]:
+    mm_output: dict[str, Any] = {"audio": payload}
+    if audio_sample_rate is not None:
+        mm_output["audio_sample_rate"] = audio_sample_rate
+    return mm_output
+
+
+def _has_non_audio_postprocess_payload(postprocess_output: DiffusionPostprocessOutput) -> bool:
+    return (
+        postprocess_output.has_video_payload
+        or postprocess_output.action_payload is not None
+        or postprocess_output.fps is not None
+    )
+
+
+def _build_multimodal_output(
+    postprocess_output: DiffusionPostprocessOutput,
+    audio_sample_rate: int | None,
+) -> dict[str, Any]:
+    mm_output: dict[str, Any] = {}
+    if postprocess_output.audio_payload is not None:
+        mm_output["audio"] = postprocess_output.audio_payload
+    if audio_sample_rate is not None:
+        mm_output["audio_sample_rate"] = audio_sample_rate
+    if postprocess_output.fps is not None:
+        mm_output["fps"] = postprocess_output.fps
+    if postprocess_output.action_payload is not None:
+        mm_output["actions"] = postprocess_output.action_payload
+    return mm_output
+
+
+def _format_single_prompt_output(
+    *,
+    request: OmniDiffusionRequest,
+    diffusion_output: DiffusionOutput,
+    outputs: list[Any],
+    metrics: dict[str, Any],
+    postprocess_output: DiffusionPostprocessOutput,
+    is_text_output: bool,
+    is_audio_output: bool,
+    audio_sample_rate: int | None,
+) -> list[OmniRequestOutput]:
+    prompt = request.prompts[0]
+    request_id = request.request_id
+    mm_output = _build_multimodal_output(postprocess_output, audio_sample_rate)
+
+    if is_text_output:
+        return [
+            OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=[],
+                prompt=prompt,
+                metrics=metrics,
+                custom_output=postprocess_output.custom_output,
+                multimodal_output=mm_output,
+                final_output_type="text",
+                stage_durations=diffusion_output.stage_durations,
+                peak_memory_mb=diffusion_output.peak_memory_mb,
+            ),
+        ]
+
+    if is_audio_output and not _has_non_audio_postprocess_payload(postprocess_output):
+        request_audio_payload = postprocess_output.audio_payload
+        if request_audio_payload is None:
+            request_audio_payload = outputs[0] if len(outputs) == 1 else outputs
+        return [
+            OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=[],
+                prompt=prompt,
+                metrics=metrics,
+                latents=diffusion_output.trajectory_latents,
+                trajectory_latents=diffusion_output.trajectory_latents,
+                trajectory_timesteps=diffusion_output.trajectory_timesteps,
+                trajectory_log_probs=diffusion_output.trajectory_log_probs,
+                trajectory_decoded=diffusion_output.trajectory_decoded,
+                multimodal_output=_format_audio_multimodal_output(
+                    request_audio_payload,
+                    audio_sample_rate,
+                ),
+                final_output_type="audio",
+                stage_durations=diffusion_output.stage_durations,
+                peak_memory_mb=diffusion_output.peak_memory_mb,
+            ),
+        ]
+
+    return [
+        OmniRequestOutput.from_diffusion(
+            request_id=request_id,
+            images=outputs,
+            prompt=prompt,
+            metrics=metrics,
+            latents=diffusion_output.trajectory_latents,
+            trajectory_latents=diffusion_output.trajectory_latents,
+            trajectory_timesteps=diffusion_output.trajectory_timesteps,
+            trajectory_log_probs=diffusion_output.trajectory_log_probs,
+            trajectory_decoded=diffusion_output.trajectory_decoded,
+            custom_output=postprocess_output.custom_output,
+            multimodal_output=mm_output,
+            stage_durations=diffusion_output.stage_durations,
+            peak_memory_mb=diffusion_output.peak_memory_mb,
+        ),
+    ]
+
+
+def _format_multi_prompt_outputs(
+    *,
+    request: OmniDiffusionRequest,
+    diffusion_output: DiffusionOutput,
+    outputs: list[Any],
+    metrics: dict[str, Any],
+    postprocess_output: DiffusionPostprocessOutput,
+    is_audio_output: bool,
+    audio_sample_rate: int | None,
+) -> list[OmniRequestOutput]:
+    results = []
+    output_idx = 0
+    request_id = request.request_id
+
+    for prompt in request.prompts:
+        num_outputs = request.sampling_params.num_outputs_per_prompt
+        start_idx = output_idx
+        end_idx = start_idx + num_outputs
+        request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else []
+        output_idx = end_idx
+
+        if is_audio_output and not _has_non_audio_postprocess_payload(postprocess_output):
+            if postprocess_output.audio_payload is not None:
+                request_audio_payload = _slice_batch_payload(
+                    postprocess_output.audio_payload,
+                    start_idx,
+                    end_idx,
+                    num_outputs,
+                )
+            else:
+                request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs
+            results.append(
+                OmniRequestOutput.from_diffusion(
+                    request_id=request_id,
+                    images=[],
+                    prompt=prompt,
+                    metrics=metrics,
+                    latents=diffusion_output.trajectory_latents,
+                    trajectory_latents=diffusion_output.trajectory_latents,
+                    trajectory_timesteps=diffusion_output.trajectory_timesteps,
+                    trajectory_log_probs=diffusion_output.trajectory_log_probs,
+                    trajectory_decoded=diffusion_output.trajectory_decoded,
+                    multimodal_output=_format_audio_multimodal_output(
+                        request_audio_payload,
+                        audio_sample_rate,
+                    ),
+                    final_output_type="audio",
+                    stage_durations=diffusion_output.stage_durations,
+                    peak_memory_mb=diffusion_output.peak_memory_mb,
+                ),
+            )
+            continue
+
+        mm_output: dict[str, Any] = {}
+        if postprocess_output.audio_payload is not None:
+            mm_output["audio"] = _slice_batch_payload(
+                postprocess_output.audio_payload,
+                start_idx,
+                end_idx,
+                num_outputs,
+            )
+        if audio_sample_rate is not None:
+            mm_output["audio_sample_rate"] = audio_sample_rate
+        if postprocess_output.fps is not None:
+            mm_output["fps"] = postprocess_output.fps
+        if postprocess_output.action_payload is not None:
+            mm_output["actions"] = _slice_batch_payload(
+                postprocess_output.action_payload,
+                start_idx,
+                end_idx,
+                num_outputs,
+            )
+
+        results.append(
+            OmniRequestOutput.from_diffusion(
+                request_id=request_id,
+                images=request_outputs,
+                prompt=prompt,
+                metrics=metrics,
+                latents=diffusion_output.trajectory_latents,
+                trajectory_latents=diffusion_output.trajectory_latents,
+                trajectory_timesteps=diffusion_output.trajectory_timesteps,
+                trajectory_log_probs=diffusion_output.trajectory_log_probs,
+                trajectory_decoded=diffusion_output.trajectory_decoded,
+                custom_output=postprocess_output.custom_output,
+                multimodal_output=mm_output,
+                stage_durations=diffusion_output.stage_durations,
+                peak_memory_mb=diffusion_output.peak_memory_mb,
+            ),
+        )
+
+    return results
+
+
+def _slice_batch_payload(payload: Any, start_idx: int, end_idx: int, num_outputs: int) -> Any:
+    sliced = payload
+    if isinstance(payload, (list, tuple)):
+        sliced = payload[start_idx:end_idx]
+        if len(sliced) == 1:
+            sliced = sliced[0]
+    elif hasattr(payload, "shape") and getattr(payload, "shape", None) is not None:
+        if len(payload.shape) > 0 and payload.shape[0] >= end_idx:
+            sliced = payload[start_idx:end_idx]
+            if num_outputs == 1:
+                sliced = sliced[0]
+    return sliced

vllm_omni/engine/async_omni_engine.py

@@ -35,8 +35,8 @@ from vllm.v1.engine.input_processor import InputProcessor
 
 from vllm_omni.config.stage_config import strip_parent_engine_args
 from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
-from vllm_omni.diffusion.diffusion_engine import supports_audio_output
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
+from vllm_omni.diffusion.io_support import supports_audio_output
 from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
 from vllm_omni.diffusion.stage_diffusion_proc import (
     complete_diffusion_handshake,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -23,7 +23,7 @@ from vllm.entrypoints.chat_utils import (
     make_tool_call_id,
 )
 
-from vllm_omni.diffusion.diffusion_engine import get_extra_body_params, get_extra_output_params
+from vllm_omni.diffusion.io_support import get_extra_body_params, get_extra_output_params
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
 from vllm_omni.entrypoints.utils import coerce_param_message_types