[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