Nick CaoandClaude
2bc183f6f0
Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Claude <noreply@anthropic.com>
62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Unit tests for vllm_omni.utils.audio.mel_filter_bank."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.utils.audio import mel_filter_bank
|
|
|
|
# Parameter combinations used across the codebase.
|
|
_PARAM_SETS = [
|
|
# Qwen3-TTS talker / speaker encoder (sr=24000)
|
|
dict(sr=24000, n_fft=1024, n_mels=128, fmin=0, fmax=12000),
|
|
# CosyVoice3 whisper encoder, Qwen3-TTS 25Hz tokenizer (sr=16000, 80 mels)
|
|
dict(sr=16000, n_fft=400, n_mels=80),
|
|
# CosyVoice3 whisper encoder (sr=16000, 128 mels)
|
|
dict(sr=16000, n_fft=400, n_mels=128),
|
|
]
|
|
|
|
_parametrize_params = pytest.mark.parametrize(
|
|
"params", _PARAM_SETS, ids=lambda p: f"{p['sr']}_{p['n_fft']}_{p['n_mels']}"
|
|
)
|
|
|
|
|
|
class TestMelFilterBank:
|
|
@_parametrize_params
|
|
def test_output_shape(self, params):
|
|
fb = mel_filter_bank(**params)
|
|
n_freqs = params["n_fft"] // 2 + 1
|
|
assert fb.shape == (params["n_mels"], n_freqs)
|
|
|
|
@_parametrize_params
|
|
def test_non_negative(self, params):
|
|
fb = mel_filter_bank(**params)
|
|
assert (fb >= 0).all()
|
|
|
|
def test_dtype_is_float(self):
|
|
fb = mel_filter_bank(sr=16000, n_fft=400, n_mels=80)
|
|
assert fb.dtype == torch.float32
|
|
|
|
def test_fmax_defaults_to_nyquist(self):
|
|
"""When fmax is omitted it should equal sr / 2."""
|
|
fb_default = mel_filter_bank(sr=16000, n_fft=400, n_mels=80)
|
|
fb_explicit = mel_filter_bank(sr=16000, n_fft=400, n_mels=80, fmax=8000.0)
|
|
torch.testing.assert_close(fb_default, fb_explicit)
|
|
|
|
def test_each_mel_band_has_nonzero_energy(self):
|
|
"""Every mel band should have at least one nonzero frequency bin."""
|
|
fb = mel_filter_bank(sr=24000, n_fft=1024, n_mels=128, fmin=0, fmax=12000)
|
|
for i in range(fb.shape[0]):
|
|
assert fb[i].sum() > 0, f"mel band {i} is all zeros"
|
|
|
|
def test_higher_fmax_extends_coverage(self):
|
|
"""A higher fmax should produce nonzero weights at higher frequency bins."""
|
|
fb_low = mel_filter_bank(sr=24000, n_fft=1024, n_mels=128, fmin=0, fmax=6000)
|
|
fb_high = mel_filter_bank(sr=24000, n_fft=1024, n_mels=128, fmin=0, fmax=12000)
|
|
# The highest nonzero column should be larger for fb_high.
|
|
last_nonzero_low = (fb_low.sum(dim=0) > 0).nonzero()[-1].item()
|
|
last_nonzero_high = (fb_high.sum(dim=0) > 0).nonzero()[-1].item()
|
|
assert last_nonzero_high > last_nonzero_low
|