[Bugfix]Fix Qwen 2.5 omni hardcoded max_mel_frames to resolve shape mismatch (#543)

Signed-off-by: Dong Wang <dongw2019@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/scripts/simple_test.sh

@@ -51,3 +51,4 @@ VENV_PYTHON="${VENV_DIR}/bin/python"
 "${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
 "${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
 "${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
+"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py

tests/model_executor/models/qwen2_5_omni/test_audio_length.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+
+def test_resolve_max_mel_frames_default():
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames
+
+    assert resolve_max_mel_frames(None, default=30000) == 30000
+    assert resolve_max_mel_frames(None, default=6000) == 6000
+
+
+def test_resolve_max_mel_frames_explicit():
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames
+
+    # Explicit argument always wins over default
+    assert resolve_max_mel_frames(123, default=30000) == 123
+    assert resolve_max_mel_frames(6000, default=30000) == 6000
+    assert resolve_max_mel_frames(0, default=30000) == 0
+
+
+@pytest.mark.parametrize("repeats", [2, 4])
+@pytest.mark.parametrize("code_len", [0, 1, 32768])
+@pytest.mark.parametrize("max_mel_frames", [None, -1, 0, 1, 6000, 30000])
+def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames):
+    """Guard that any max_mel_frames yields a mel length aligned to repeats, and
+    consistent with the truncated code length (prevents concat mismatch).
+    """
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length
+
+    target_code_len, target_mel_len = cap_and_align_mel_length(
+        code_len=code_len,
+        repeats=repeats,
+        max_mel_frames=max_mel_frames,
+    )
+
+    assert isinstance(target_code_len, int)
+    assert isinstance(target_mel_len, int)
+
+    if code_len == 0:
+        assert target_code_len == 0
+        assert target_mel_len == 0
+        return
+
+    assert target_code_len >= 1
+    assert target_mel_len >= repeats
+    assert target_mel_len % repeats == 0
+    assert target_mel_len == target_code_len * repeats
+    assert target_code_len <= code_len
+
+    if max_mel_frames is not None and int(max_mel_frames) > 0 and int(max_mel_frames) >= repeats:
+        assert target_mel_len <= int(max_mel_frames)

vllm_omni/model_executor/models/qwen2_5_omni/audio_length.py

@@ -0,0 +1,69 @@
+"""Audio length helpers (pure-Python).
+
+These utilities are used to keep code2wav (codec -> mel -> wav) lengths aligned
+when applying a mel-frame cap.
+"""
+
+from __future__ import annotations
+
+
+def resolve_max_mel_frames(max_mel_frames: int | None, *, default: int = 30000) -> int:
+    """Resolve max mel frames from an explicit value or default.
+
+    Args:
+        max_mel_frames: Explicit value to use. If None, uses `default`.
+        default: Default value to use when `max_mel_frames` is None.
+
+    Returns:
+        The resolved max mel frames value.
+    """
+    if max_mel_frames is not None:
+        return int(max_mel_frames)
+    return int(default)
+
+
+def cap_and_align_mel_length(
+    *,
+    code_len: int,
+    repeats: int,
+    max_mel_frames: int | None,
+    default_max_mel_frames: int = 30000,
+) -> tuple[int, int]:
+    """Compute a (target_code_len, target_mel_len) pair.
+
+    - `mel_len` is always a multiple of `repeats` (codec expansion factor).
+    - If `max_mel_frames` is None, uses `default_max_mel_frames`.
+    - If `max_mel_frames` <= 0, no cap is applied (mel_len == code_len * repeats).
+    - If `max_mel_frames` is smaller than `repeats` and code_len > 0, we still
+      return at least one codec token worth of mel frames (mel_len == repeats)
+      so downstream repeat-interleave stays valid.
+    """
+    code_len = int(code_len)
+    repeats = int(repeats)
+    if repeats <= 0:
+        raise ValueError(f"repeats must be > 0, got {repeats}")
+    if code_len <= 0:
+        return 0, 0
+
+    if max_mel_frames is None:
+        max_mel_frames = int(default_max_mel_frames)
+    else:
+        max_mel_frames = int(max_mel_frames)
+
+    maximum_duration = int(code_len * repeats)
+    if max_mel_frames > 0:
+        target_duration = min(maximum_duration, max_mel_frames)
+    else:
+        target_duration = maximum_duration
+
+    # Align down to repeats; then ensure we keep at least one codec token.
+    target_duration = (target_duration // repeats) * repeats
+    if target_duration <= 0:
+        target_duration = min(maximum_duration, repeats)
+
+    target_code_len = target_duration // repeats
+    if target_code_len <= 0:
+        target_code_len = 1
+        target_duration = repeats
+
+    return int(target_code_len), int(target_duration)

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py

@@ -32,6 +32,7 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
+from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames
 from vllm_omni.utils.platform_utils import is_npu
 
 
@@ -1261,12 +1262,24 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         num_steps=10,
         guidance_scale=0.5,
         sway_coefficient=-1.0,
+        max_mel_frames: int | None = None,
     ):
-        noise_initialization = torch.randn([1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype)
-        maximum_duration = quantized_code.shape[1] * self.repeats
-        initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device)
+        max_mel_frames = resolve_max_mel_frames(max_mel_frames, default=30000)
+        target_code_len, target_duration = cap_and_align_mel_length(
+            code_len=int(quantized_code.shape[1]),
+            repeats=int(self.repeats),
+            max_mel_frames=max_mel_frames,
+        )
+        if int(quantized_code.shape[1]) != target_code_len:
+            quantized_code = quantized_code[:, :target_code_len]
+
+        initial_state = torch.randn(
+            [1, target_duration, self.mel_dim],
+            dtype=reference_mel_spectrogram.dtype,
+            device=quantized_code.device,
+        )
         batch_size = reference_mel_spectrogram.shape[0]
-        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
+        conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, target_duration, 1)
 
         if batch_size != 1:
             raise ValueError("Only batch size = 1 is currently supported")
@@ -1281,6 +1294,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
                     time_step=time_step,
                     drop_audio_conditioning=False,
                     drop_code=False,
+                    apply_cfg=False,
                 )
                 return prediction
 
@@ -1461,6 +1475,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
         num_steps=10,
         guidance_scale=0.5,
         sway_coefficient=-1.0,
+        max_mel_frames: int | None = None,
         **kwargs,
     ):
         """Generates a waveform from input code and conditioning parameters."""
@@ -1472,6 +1487,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             sway_coefficient=sway_coefficient,
+            max_mel_frames=max_mel_frames,
         ).to(self.code2wav_bigvgan_model.dtype)
 
         waveform = self.code2wav_bigvgan_model(mel_spectrogram).to(self.dtype)