| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from io import BytesIO |
| 4 | from pathlib import Path |
| 5 | from unittest.mock import patch |
| 6 | |
| 7 | import numpy as np |
| 8 | import pybase64 as base64 |
| 9 | import pytest |
| 10 | import soundfile as sf |
| 11 | |
| 12 | from vllm.multimodal.media import AudioMediaIO |
| 13 | from vllm.multimodal.media import audio as audio_module |
| 14 | from vllm.multimodal.media.audio import ( |
| 15 | load_audio, |
| 16 | load_audio_soundfile, |
| 17 | load_audio_torchcodec, |
| 18 | ) |
| 19 | |
| 20 | from ...conftest import AudioTestAssets |
| 21 | |
| 22 | pytestmark = pytest.mark.cpu_test |
| 23 | |
| 24 | ASSETS_DIR = Path(__file__).parent.parent / "assets" |
| 25 | assert ASSETS_DIR.exists() |
| 26 | |
| 27 | |
| 28 | @pytest.fixture |
| 29 | def dummy_audio(): |
| 30 | return np.array([0.0, 0.1, 0.2, 0.3, 0.4], dtype=float) |
| 31 | |
| 32 | |
| 33 | @pytest.fixture |
| 34 | def dummy_audio_bytes(audio_assets: AudioTestAssets): |
| 35 | with open(audio_assets[0].get_local_path(), "rb") as f: |
| 36 | return f.read() |
| 37 | |
| 38 | |
| 39 | def test_audio_media_io_load_bytes(dummy_audio_bytes): |
| 40 | audio_io = AudioMediaIO() |
| 41 | out = audio_io.load_bytes(dummy_audio_bytes) |
| 42 | assert isinstance(out[0], np.ndarray) |
| 43 | assert out[1] == 16000 |
| 44 | |
| 45 | |
| 46 | def test_audio_media_io_load_base64(dummy_audio_bytes): |
| 47 | audio_io = AudioMediaIO() |
| 48 | encoded = base64.b64encode(dummy_audio_bytes).decode("utf-8") |
| 49 | out = audio_io.load_base64("audio/wav", encoded) |
| 50 | assert isinstance(out[0], np.ndarray) |
| 51 | assert out[1] == 16000 |
| 52 | |
| 53 | |
| 54 | def test_audio_media_io_load_base64_rejects_malformed(dummy_audio_bytes): |
| 55 | """Malformed base64 must surface as a ValueError so the server answers 400. |
| 56 | Without strict decoding the bad characters are dropped, the garbage reaches |
| 57 | libsndfile, and the client gets a 500 instead.""" |
| 58 | encoded = base64.b64encode(dummy_audio_bytes).decode("utf-8") |
| 59 | malformed = encoded[:8] + "!!!@@@###" + encoded[8:] |
| 60 | |
| 61 | audio_io = AudioMediaIO() |
| 62 | with pytest.raises(ValueError): |
| 63 | audio_io.load_base64("audio/wav", malformed) |
| 64 | |
| 65 | |
| 66 | def test_audio_media_io_load_file(audio_assets: AudioTestAssets): |
| 67 | audio_io = AudioMediaIO() |
| 68 | path = audio_assets[0].get_local_path() |
| 69 | out = audio_io.load_file(path) |
| 70 | assert isinstance(out[0], np.ndarray) |
| 71 | assert out[1] == 16000 |
| 72 | |
| 73 | |
| 74 | def test_audio_media_io_encode_base64(dummy_audio): |
| 75 | audio_io = AudioMediaIO() |
| 76 | media = (dummy_audio, 16000) |
| 77 | with patch("soundfile.write") as mock_write: |
| 78 | |
| 79 | def write_to_buffer(buffer, *_args, **_kwargs): |
| 80 | buffer.write(b"dummy_wav_data") |
| 81 | |
| 82 | mock_write.side_effect = write_to_buffer |
| 83 | |
| 84 | out = audio_io.encode_base64(media) |
| 85 | decoded = base64.b64decode(out) |
| 86 | assert decoded == b"dummy_wav_data" |
| 87 | mock_write.assert_called_once() |
| 88 | |
| 89 | |
| 90 | def test_load_audio_max_duration_respected(dummy_audio_bytes): |
| 91 | """Valid audio within the duration limit should load successfully.""" |
| 92 | y, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, max_duration_s=3600) |
| 93 | assert isinstance(y, np.ndarray) |
| 94 | assert len(y) > 0 |
| 95 | |
| 96 | |
| 97 | def test_load_audio_max_duration_rejected(dummy_audio_bytes): |
| 98 | """Audio exceeding the duration limit must be rejected during decode.""" |
| 99 | with pytest.raises(ValueError, match="exceeds maximum allowed duration"): |
| 100 | load_audio(BytesIO(dummy_audio_bytes), sr=None, max_duration_s=0.0001) |
| 101 | |
| 102 | |
| 103 | def test_audio_media_io_from_video(video_assets): |
| 104 | audio_io = AudioMediaIO() |
| 105 | video_path = video_assets[0].video_path |
| 106 | with open(video_path, "rb") as f: |
| 107 | audio, sr = audio_io.load_bytes(f.read()) |
| 108 | audio_ref, sr_ref = load_audio(video_path, sr=None) |
| 109 | assert sr == sr_ref |
| 110 | np.testing.assert_allclose(audio_ref, audio, atol=1e-4) |
| 111 | |
| 112 | |
| 113 | def _make_flac_bytes(frames: int, channels: int, samplerate: int) -> bytes: |
| 114 | """Create a minimal FLAC file in memory for testing.""" |
| 115 | data = np.zeros((frames, channels), dtype=np.int16) |
| 116 | buf = BytesIO() |
| 117 | sf.write(buf, data, samplerate, format="FLAC") |
| 118 | return buf.getvalue() |
| 119 | |
| 120 | |
| 121 | def test_small_file_passes_memory_guard(): |
| 122 | """A small valid file should pass both duration and memory guards.""" |
| 123 | payload = _make_flac_bytes(frames=16000, channels=1, samplerate=16000) |
| 124 | y, sr = load_audio_soundfile( |
| 125 | BytesIO(payload), |
| 126 | sr=None, |
| 127 | max_duration_s=600, |
| 128 | max_decode_bytes=256 * 1024 * 1024, |
| 129 | ) |
| 130 | assert isinstance(y, np.ndarray) |
| 131 | assert len(y) == 16000 |
| 132 | |
| 133 | |
| 134 | def test_memory_guard_rejects_large_allocation(): |
| 135 | """A file whose frames*channels*4 exceeds the byte limit must be |
| 136 | rejected before allocating the buffer.""" |
| 137 | # 100_000 frames * 8 channels * 4 bytes = 3.2 MB |
| 138 | payload = _make_flac_bytes(frames=100_000, channels=8, samplerate=48000) |
| 139 | # Set limit to 1 MiB — should reject |
| 140 | with pytest.raises(ValueError, match="VLLM_MAX_AUDIO_DECODE_BYTES"): |
| 141 | load_audio_soundfile( |
| 142 | BytesIO(payload), |
| 143 | sr=None, |
| 144 | max_duration_s=600, |
| 145 | max_decode_bytes=1 * 1024 * 1024, |
| 146 | ) |
| 147 | |
| 148 | |
| 149 | def test_forged_samplerate_rejected_by_memory_guard(): |
| 150 | """The PoC scenario: high sample rate fools the duration guard but |
| 151 | the memory guard catches the large frame*channel allocation.""" |
| 152 | # Forged high sample rate: 655350 Hz, 8 channels, 1M frames |
| 153 | # Duration guard sees: 1_000_000 / 655_350 = 1.5s → passes |
| 154 | # Memory: 1_000_000 * 8 * 4 = 32 MB |
| 155 | payload = _make_flac_bytes(frames=1_000_000, channels=8, samplerate=655350) |
| 156 | # Set memory limit to 16 MiB — below the 32 MB allocation |
| 157 | with pytest.raises(ValueError, match="VLLM_MAX_AUDIO_DECODE_BYTES"): |
| 158 | load_audio_soundfile( |
| 159 | BytesIO(payload), |
| 160 | sr=None, |
| 161 | max_duration_s=600, |
| 162 | max_decode_bytes=16 * 1024 * 1024, |
| 163 | ) |
| 164 | |
| 165 | |
| 166 | def test_load_audio_threads_max_decode_bytes(): |
| 167 | """Verify load_audio passes max_decode_bytes through to backend.""" |
| 168 | # 50_000 frames * 4 channels * 4 bytes = 800 KB |
| 169 | payload = _make_flac_bytes(frames=50_000, channels=4, samplerate=44100) |
| 170 | # Limit of 512 KB should reject |
| 171 | with pytest.raises(ValueError, match="VLLM_MAX_AUDIO_DECODE_BYTES"): |
| 172 | load_audio( |
| 173 | BytesIO(payload), |
| 174 | sr=None, |
| 175 | max_duration_s=600, |
| 176 | max_decode_bytes=512 * 1024, |
| 177 | ) |
| 178 | |
| 179 | |
| 180 | @pytest.mark.parametrize("backend", ["soundfile", "pyav", "torchcodec"]) |
| 181 | def test_load_audio_backend_matches_default(backend, dummy_audio_bytes): |
| 182 | """Every explicit backend must decode to the same samples as `auto`.""" |
| 183 | if backend == "torchcodec": |
| 184 | pytest.importorskip("torchcodec") |
| 185 | ref_audio, ref_sr = load_audio(BytesIO(dummy_audio_bytes), sr=None) |
| 186 | audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, backend=backend) |
| 187 | assert sr == ref_sr |
| 188 | if backend == "torchcodec": |
| 189 | # Vorbis encoder padding is trimmed, so torchcodec must match the |
| 190 | # soundfile reference exactly, including the waveform length. |
| 191 | assert audio.shape == ref_audio.shape |
| 192 | np.testing.assert_allclose(ref_audio, audio, atol=1e-4) |
| 193 | return |
| 194 | # Decoders disagree only on codec encoder-delay/padding, so PyAV may |
| 195 | # emit a few extra trailing samples (e.g. ~192 for Ogg Vorbis). Compare |
| 196 | # the overlapping region, which must agree to float32 precision. |
| 197 | n = min(ref_audio.shape[-1], audio.shape[-1]) |
| 198 | assert n > 0 |
| 199 | np.testing.assert_allclose(ref_audio[:n], audio[:n], atol=1e-4) |
| 200 | |
| 201 | |
| 202 | def test_load_audio_default_preserves_vorbis_length(dummy_audio_bytes): |
| 203 | """Default decoding must not append Vorbis padding to the speech waveform.""" |
| 204 | expected, expected_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) |
| 205 | audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None) |
| 206 | assert sr == expected_sr |
| 207 | assert audio.shape == expected.shape |
| 208 | |
| 209 | |
| 210 | def test_load_audio_torchcodec_preserves_vorbis_length(dummy_audio_bytes): |
| 211 | """Torchcodec must trim Vorbis trailing padding even on FFmpeg < 5.0, |
| 212 | matching the soundfile reference exactly.""" |
| 213 | pytest.importorskip("torchcodec") |
| 214 | expected, expected_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) |
| 215 | audio, sr = load_audio_torchcodec(BytesIO(dummy_audio_bytes), sr=None) |
| 216 | assert sr == expected_sr |
| 217 | assert audio.shape == expected.shape |
| 218 | np.testing.assert_allclose(audio, expected, atol=1e-4) |
| 219 | |
| 220 | |
| 221 | def test_load_audio_unknown_backend_rejected(dummy_audio_bytes): |
| 222 | """An unknown backend must fail loudly instead of silently degrading.""" |
| 223 | with pytest.raises(ValueError, match="Unknown audio backend"): |
| 224 | load_audio(BytesIO(dummy_audio_bytes), sr=None, backend="not_a_backend") |
| 225 | |
| 226 | |
| 227 | def test_load_audio_auto_falls_back_without_torchcodec(dummy_audio_bytes): |
| 228 | """`auto` must fall back to the soundfile → PyAV chain when torchcodec |
| 229 | is not importable.""" |
| 230 | ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) |
| 231 | with patch.object(audio_module, "load_audio_torchcodec", side_effect=ImportError): |
| 232 | audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, backend="auto") |
| 233 | assert sr == ref_sr |
| 234 | np.testing.assert_array_equal(ref_audio, audio) |
| 235 | |
| 236 | |
| 237 | def test_load_audio_auto_falls_back_without_ffmpeg(dummy_audio_bytes): |
| 238 | """Torchcodec installed but system ffmpeg missing (`AudioDecoder is None`) |
| 239 | must surface as ImportError so `auto` falls back to soundfile → PyAV.""" |
| 240 | ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) |
| 241 | with patch.object(audio_module, "AudioDecoder", None): |
| 242 | with pytest.raises(ImportError, match="torchcodec audio backend"): |
| 243 | load_audio_torchcodec(BytesIO(dummy_audio_bytes), sr=None) |
| 244 | audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, backend="auto") |
| 245 | assert sr == ref_sr |
| 246 | np.testing.assert_array_equal(ref_audio, audio) |
| 247 | |
| 248 | |
| 249 | def test_load_audio_auto_falls_back_when_libtorchcodec_unloadable( |
| 250 | dummy_audio_bytes, |
| 251 | ): |
| 252 | """Torchcodec loads its ffmpeg-backed core lazily at decoder construction; |
| 253 | a "Could not load libtorchcodec" RuntimeError there (no system ffmpeg) |
| 254 | must also surface as ImportError so `auto` falls back.""" |
| 255 | |
| 256 | class _UnloadableDecoder: |
| 257 | def __init__(self, *args, **kwargs): |
| 258 | raise RuntimeError("Could not load libtorchcodec. Likely causes: ...") |
| 259 | |
| 260 | ref_audio, ref_sr = load_audio_soundfile(BytesIO(dummy_audio_bytes), sr=None) |
| 261 | with patch.object(audio_module, "AudioDecoder", _UnloadableDecoder): |
| 262 | with pytest.raises(ImportError, match="torchcodec audio backend"): |
| 263 | load_audio_torchcodec(BytesIO(dummy_audio_bytes), sr=None) |
| 264 | audio, sr = load_audio(BytesIO(dummy_audio_bytes), sr=None, backend="auto") |
| 265 | assert sr == ref_sr |
| 266 | np.testing.assert_array_equal(ref_audio, audio) |
| 267 | |
| 268 | |
| 269 | def test_audio_media_io_audio_backend_kwarg(dummy_audio_bytes): |
| 270 | """`audio_backend` selects the backend; unknown values fail at init.""" |
| 271 | audio, sr = AudioMediaIO(audio_backend="pyav").load_bytes(dummy_audio_bytes) |
| 272 | assert isinstance(audio, np.ndarray) |
| 273 | assert sr == 16000 |
| 274 | with pytest.raises(ValueError, match="Unknown audio_backend"): |
| 275 | AudioMediaIO(audio_backend="not_a_backend") |
| 276 | |
| 277 | |
| 278 | def test_torchcodec_max_duration_rejected(dummy_audio_bytes): |
| 279 | """The decompression-bomb duration guard must hold for torchcodec too.""" |
| 280 | pytest.importorskip("torchcodec") |
| 281 | with pytest.raises(ValueError, match="exceeds maximum allowed duration"): |
| 282 | load_audio_torchcodec( |
| 283 | BytesIO(dummy_audio_bytes), sr=None, max_duration_s=0.0001 |
| 284 | ) |
| 285 | |
| 286 | |
| 287 | def test_torchcodec_max_decode_bytes_rejected(dummy_audio_bytes): |
| 288 | """The decompression-bomb memory guard must hold for torchcodec too.""" |
| 289 | pytest.importorskip("torchcodec") |
| 290 | audio, _ = load_audio_torchcodec(BytesIO(dummy_audio_bytes), sr=None) |
| 291 | with pytest.raises(ValueError, match="VLLM_MAX_AUDIO_DECODE_BYTES"): |
| 292 | load_audio_torchcodec( |
| 293 | BytesIO(dummy_audio_bytes), |
| 294 | sr=None, |
| 295 | max_decode_bytes=audio.nbytes - 1, |
| 296 | ) |