[CI][Bugfix] Fix LTX audio-video warmup output typing (#3964)

Signed-off-by: david6666666 <530634352@qq.com>

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

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,

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"]

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,

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"]