From aff0af97546f5350493304c9d7a4f112782d3e2d Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Mon, 15 Jun 2026 22:50:50 +0800 Subject: [PATCH] [Refactor] Extract diffusion output formatting boundary (#4407) Signed-off-by: hsliuustc0106 --- .../test_diffusion_engine_dummy_run.py | 18 +- .../test_diffusion_engine_metrics.py | 54 +-- .../test_diffusion_output_formatter.py | 307 +++++++++++++++ tests/diffusion/test_diffusion_scheduler.py | 2 +- .../test_diffusers_backend_similarity.py | 2 +- vllm_omni/diffusion/diffusion_engine.py | 346 ++++------------- vllm_omni/diffusion/io_support.py | 77 ++++ vllm_omni/diffusion/output_formatter.py | 357 ++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 2 +- vllm_omni/entrypoints/openai/serving_chat.py | 2 +- 10 files changed, 852 insertions(+), 315 deletions(-) create mode 100644 tests/diffusion/test_diffusion_output_formatter.py create mode 100644 vllm_omni/diffusion/io_support.py create mode 100644 vllm_omni/diffusion/output_formatter.py diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 71cb47512..50dde3083 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index dd8db9fdf..2e8eae242 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/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" diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py new file mode 100644 index 000000000..49c07d98f --- /dev/null +++ b/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] == [{}, {}] diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index f5a4f7f77..503ea0842 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/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 diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index ad4dd5726..529152290 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/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( diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b3405e3bb..3b4843af0 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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(): diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py new file mode 100644 index 000000000..9c7b2b027 --- /dev/null +++ b/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())) diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py new file mode 100644 index 000000000..8d00b8f85 --- /dev/null +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4a298feeb..8df60b9e5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 2a7b7e269..24998436e 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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