From 67bb9b4a124d1237bb33f371d5b0b92ca70db38a Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 29 May 2026 15:46:02 +0800 Subject: [PATCH] [CI][Bugfix] Fix LTX audio-video warmup output typing (#3964) Signed-off-by: david6666666 <530634352@qq.com> --- .../test_diffusion_engine_dummy_run.py | 57 +++++++++++++++++++ vllm_omni/diffusion/diffusion_engine.py | 12 ++-- .../diffusion/models/ltx2/pipeline_ltx2.py | 4 +- .../diffusion/models/ltx2/pipeline_ltx2_3.py | 2 +- .../models/ltx2/pipeline_ltx2_image2video.py | 1 + 5 files changed, 70 insertions(+), 6 deletions(-) create mode 100644 tests/diffusion/test_diffusion_engine_dummy_run.py diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py new file mode 100644 index 000000000..71cb47512 --- /dev/null +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm_omni.diffusion import diffusion_engine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None: + class JointAudioVideoModel: + dummy_run_num_frames = 2 + + monkeypatch.setattr( + diffusion_engine.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 + + +def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.MonkeyPatch) -> None: + class AudioOutputModel: + support_audio_output = True + + monkeypatch.setattr( + diffusion_engine.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 + + +def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.MonkeyPatch) -> None: + class VideoOnlyModel: + pass + + monkeypatch.setattr( + diffusion_engine.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 + + +def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + diffusion_engine.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: None, + ) + + assert diffusion_engine.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2 diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 8cb0d5296..b4bfecb77 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -83,6 +83,13 @@ def supports_audio_output(model_class_name: str) -> bool: return bool(getattr(model_cls, "support_audio_output", False)) +def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int: + 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. @@ -690,10 +697,7 @@ class DiffusionEngine: dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32) prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio - # Audio pipelines round audio token count from num_frames; the default - # of 1 yields seq_len=1 K/V which cuDNN SDPA refuses under torch.compile. - # 2 is the minimum that produces audio_num_frames > 1. - num_frames = 2 if supports_audio_input or supports_audio_output(self.od_config.model_class_name) else 1 + num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) req = OmniDiffusionRequest( prompts=[prompt], request_id=DUMMY_DIFFUSION_REQUEST_ID, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 29059a6d6..82f59e7ff 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -158,7 +158,7 @@ class _VideoAudioScheduler: class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): # Audio is diffused jointly with video; warmup must size audio tokens. - support_audio_output = True + dummy_run_num_frames = 2 def __init__( self, @@ -1121,6 +1121,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): """LTX2TwoStagesPipeline is for two stages image to video generation""" + dummy_run_num_frames = 2 + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index 3c18e70dc..e6985e186 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -123,7 +123,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): """ # Audio is diffused jointly with video; warmup must size audio tokens. - support_audio_output = True + dummy_run_num_frames = 2 def __init__( self, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index af4135d65..e93f4e744 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -735,6 +735,7 @@ class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsComponentDiscovery): """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation""" support_image_input = True + dummy_run_num_frames = 2 _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"]