anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
Hongsheng Liu authored and GitHub committed 2026-06-15 22:50:50 +08:00
1 parent b1a4b08a7e
commit aff0af9754
10 files changed
+852 -315
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_engine_dummy_run.py
+9 -9
View File
@@ -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