From c1ebb607472e9c3ab9df416e7016e92c56c32818 Mon Sep 17 00:00:00 2001 From: Dong W <89223086+sniper35@users.noreply.github.com> Date: Wed, 31 Dec 2025 20:12:03 -0800 Subject: [PATCH] [Bugfix]Fix Qwen 2.5 omni hardcoded max_mel_frames to resolve shape mismatch (#543) Signed-off-by: Dong Wang Co-authored-by: Hongsheng Liu --- .buildkite/scripts/simple_test.sh | 1 + .../models/qwen2_5_omni/test_audio_length.py | 53 ++++++++++++++ .../models/qwen2_5_omni/audio_length.py | 69 +++++++++++++++++++ .../qwen2_5_omni/qwen2_5_omni_token2wav.py | 24 +++++-- 4 files changed, 143 insertions(+), 4 deletions(-) create mode 100644 tests/model_executor/models/qwen2_5_omni/test_audio_length.py create mode 100644 vllm_omni/model_executor/models/qwen2_5_omni/audio_length.py diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh index bb2e1eca0..33248d99c 100755 --- a/.buildkite/scripts/simple_test.sh +++ b/.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 diff --git a/tests/model_executor/models/qwen2_5_omni/test_audio_length.py b/tests/model_executor/models/qwen2_5_omni/test_audio_length.py new file mode 100644 index 000000000..dd5f09817 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/audio_length.py b/vllm_omni/model_executor/models/qwen2_5_omni/audio_length.py new file mode 100644 index 000000000..4162f88e5 --- /dev/null +++ b/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) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py index 6ea444916..5ad90eff1 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py +++ b/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)