[BugFix] Remove pydub dependency for Python 3.13 compatibility (#4035)

Signed-off-by: FED4 <feiyyfrank@outlook.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

requirements/common.txt

@@ -22,4 +22,3 @@ prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
-pydub

tests/helpers/media.py

@@ -588,15 +588,32 @@ def cosine_similarity_text(text1, text2, n: int = 3):
     return cosine * length_harmony
 
 
-def _merge_base64_audio_to_segment(base64_list: list[str]):
-    from pydub import AudioSegment
+class _AudioBuffer:
+    """Minimal replacement for pydub.AudioSegment used by test helpers."""
 
-    merged = None
+    def __init__(self, data: np.ndarray, sample_rate: int):
+        self.data = data
+        self.sample_rate = sample_rate
+
+    def export(self, buf: io.BytesIO, format: str = "wav"):
+        sf.write(buf, self.data, self.sample_rate, format=format.upper())
+        buf.seek(0)
+
+
+def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
+    from vllm.multimodal.media.audio import AudioMediaIO
+
+    io_ = AudioMediaIO()
+    chunks: list[np.ndarray] = []
+    sample_rate: int | None = None
     for b64 in base64_list:
         raw = base64.b64decode(b64.split(",", 1)[-1])
-        seg = AudioSegment.from_file(io.BytesIO(raw))
-        merged = seg if merged is None else merged + seg
-    return merged
+        waveform, sr = io_.load_bytes(raw)
+        if sample_rate is None:
+            sample_rate = int(sr)
+        chunks.append(waveform)
+    merged = np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
+    return _AudioBuffer(merged, sample_rate or 16000)
 
 
 @contextmanager

tests/test_pydub_removal.py

@@ -0,0 +1,182 @@
+"""Functional tests for pydub removal.
+
+Verifies that the new av/numpy/torchaudio replacements produce
+identical results to the old pydub-based code for all three paths:
+  1. WAV PCM → duration & frame count (pure arithmetic)
+  2. Non-WAV (mp3/ogg) → duration & frame count (av.open)
+  3. TTS WER PCM export: resample to 24000Hz mono s16
+"""
+
+import io
+import wave
+
+import av
+import numpy as np
+import soundfile as sf
+
+
+def _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=1.0):
+    """Generate a WAV file with a sine wave, return (wav_bytes, pcm_bytes, params)."""
+    n_samples = int(sample_rate * duration_sec)
+    t = np.linspace(0, duration_sec, n_samples, endpoint=False)
+    samples = (np.sin(2 * np.pi * 440 * t) * 32767 * 0.5).astype(np.int16)
+    if channels == 2:
+        samples = np.column_stack([samples, samples]).flatten()
+
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(channels)
+        wf.setsampwidth(sample_width)
+        wf.setframerate(sample_rate)
+        wf.writeframes(samples.tobytes())
+    wav_bytes = buf.getvalue()
+
+    # Extract raw PCM (skip WAV header)
+    with wave.open(io.BytesIO(wav_bytes), "rb") as wr:
+        pcm_bytes = wr.readframes(wr.getnframes())
+        params = (wr.getnchannels(), wr.getsampwidth(), wr.getframerate())
+
+    return wav_bytes, pcm_bytes, params
+
+
+def _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0):
+    """Generate an OGG Vorbis file with a sine wave, return bytes."""
+    n_samples = int(sample_rate * duration_sec)
+    t = np.linspace(0, duration_sec, n_samples, endpoint=False)
+    samples = (np.sin(2 * np.pi * 440 * t) * 0.5).astype(np.float32)
+    buf = io.BytesIO()
+    sf.write(buf, samples, sample_rate, format="OGG", subtype="VORBIS")
+    return buf.getvalue()
+
+
+def test_wav_path_duration_and_frames():
+    """Path 1: WAV PCM buffer → duration & frames via arithmetic."""
+    for sr in [16000, 24000, 44100]:
+        for ch in [1, 2]:
+            for dur in [0.5, 1.0, 2.5]:
+                _, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes(
+                    sample_rate=sr, channels=ch, duration_sec=dur
+                )
+
+                # New code logic (from patch.py)
+                frame_width = sample_width * channels
+                audio_frames = len(pcm_bytes) // frame_width
+                audio_duration_sec = audio_frames / frame_rate
+
+                assert abs(audio_duration_sec - dur) < 0.001, f"sr={sr} ch={ch} dur={dur}: got {audio_duration_sec}"
+                expected_frames = int(sr * dur)
+                assert audio_frames == expected_frames, (
+                    f"sr={sr} ch={ch} dur={dur}: frames {audio_frames} != {expected_frames}"
+                )
+
+    print("[PASS] test_wav_path_duration_and_frames")
+
+
+def test_non_wav_path_duration_av():
+    """Path 2: non-WAV audio → duration & frames via av."""
+    for dur in [0.5, 1.0, 2.0]:
+        ogg_bytes = _make_ogg_bytes(sample_rate=16000, duration_sec=dur)
+
+        container = av.open(io.BytesIO(ogg_bytes))
+        stream = container.streams.audio[0]
+        audio_frames = sum(f.samples for f in container.decode(stream))
+        audio_duration_sec = audio_frames / stream.rate if stream.rate else 0.0
+        container.close()
+
+        # OGG encoding may add padding, allow small tolerance
+        assert abs(audio_duration_sec - dur) < 0.05, f"dur={dur}: got {audio_duration_sec}"
+
+    print("[PASS] test_non_wav_path_duration_av")
+
+
+def test_tts_pcm_export_wav_path():
+    """Path 3a: WAV PCM → resample to 24000Hz mono s16."""
+    import torch
+    import torchaudio
+
+    _, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes(
+        sample_rate=16000, channels=2, sample_width=2, duration_sec=1.0
+    )
+
+    # New code logic
+    pcm_array = np.frombuffer(pcm_bytes, dtype=np.int16)
+    if channels > 1:
+        pcm_array = pcm_array.reshape(-1, channels).mean(axis=1).astype(np.int16)
+    if frame_rate != 24000:
+        t = torch.from_numpy(pcm_array.astype(np.float32)).unsqueeze(0)
+        t = torchaudio.functional.resample(t, frame_rate, 24000)
+        pcm_array = t.squeeze(0).to(torch.int16).numpy()
+
+    result_bytes = pcm_array.tobytes()
+    # Should be mono, 24000Hz, 2 bytes per sample
+    expected_samples = int(24000 * 1.0)
+    actual_samples = len(result_bytes) // 2
+    assert abs(actual_samples - expected_samples) < 10, f"expected ~{expected_samples} samples, got {actual_samples}"
+    assert len(pcm_array.shape) == 1, "should be mono (1D array)"
+
+    print("[PASS] test_tts_pcm_export_wav_path")
+
+
+def test_tts_pcm_export_non_wav_path():
+    """Path 3b: non-WAV → resample to 24000Hz mono s16 via av.AudioResampler."""
+    ogg_bytes = _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0)
+
+    container = av.open(io.BytesIO(ogg_bytes))
+    resampler = av.AudioResampler(format="s16", layout="mono", rate=24000)
+    frames = []
+    for frame in container.decode(audio=0):
+        for rf in resampler.resample(frame):
+            frames.append(rf.to_ndarray().flatten())
+    container.close()
+
+    result = np.concatenate(frames)
+    result_bytes = result.tobytes()
+    expected_samples = int(24000 * 1.0)
+    actual_samples = len(result_bytes) // 2
+    assert abs(actual_samples - expected_samples) < 500, f"expected ~{expected_samples} samples, got {actual_samples}"
+
+    print("[PASS] test_tts_pcm_export_non_wav_path")
+
+
+def test_media_merge_and_export():
+    """media.py: _merge_base64_audio_to_segment decode + concat + export WAV."""
+    import base64 as b64mod
+    import sys
+
+    sys.path.insert(0, ".")
+    from tests.helpers.media import _merge_base64_audio_to_segment
+
+    # Create two WAV chunks as base64
+    chunks_b64 = []
+    for dur in [0.5, 0.3]:
+        wav_bytes, _, _ = _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=dur)
+        chunks_b64.append(b64mod.b64encode(wav_bytes).decode())
+
+    merged = _merge_base64_audio_to_segment(chunks_b64)
+
+    # Test export to WAV
+    buf = io.BytesIO()
+    merged.export(buf, format="wav")
+    wav_out = buf.getvalue()
+    assert len(wav_out) > 44, f"WAV output too small: {len(wav_out)} bytes"
+
+    # Verify merged duration (~0.8s)
+    out_data, out_sr = sf.read(io.BytesIO(wav_out))
+    actual_dur = len(out_data) / out_sr
+    assert abs(actual_dur - 0.8) < 0.05, f"expected ~0.8s merged duration, got {actual_dur}"
+
+    print(f"[PASS] test_media_merge_and_export (duration={actual_dur:.3f}s, wav_size={len(wav_out)} bytes)")
+
+
+if __name__ == "__main__":
+    test_wav_path_duration_and_frames()
+    test_non_wav_path_duration_av()
+    import importlib.util
+
+    if importlib.util.find_spec("torch") and importlib.util.find_spec("torchaudio"):
+        test_tts_pcm_export_wav_path()
+    else:
+        print("[SKIP] test_tts_pcm_export_wav_path (torch/torchaudio not available)")
+    test_tts_pcm_export_non_wav_path()
+    test_media_merge_and_export()
+    print("\nAll tests passed.")

vllm_omni/benchmarks/patch/patch.py

@@ -16,8 +16,8 @@ from datetime import datetime
 from typing import Any, Literal
 
 import aiohttp
+import numpy as np
 import pybase64 as base64
-from pydub import AudioSegment
 from tqdm.asyncio import tqdm
 from vllm.benchmarks import datasets
 from vllm.benchmarks.datasets import SampleRequest
@@ -653,9 +653,8 @@ async def async_request_openai_chat_omni_completions(
         # Reset per-attempt state so that retries do not mix partial
         # outputs or metrics from previous attempts.
         generated_text = ""
-        generated_audio = None
         # For wav responses, accumulate decoded PCM bytes per chunk
-        # to avoid repeated AudioSegment decode/concat.
+        # to avoid repeated decode/concat.
         wav_pcm_buffer = bytearray()
         wav_audio_params: tuple[int, int, int] | None = None
         wav_inconsistent_chunk_count = 0
@@ -795,30 +794,33 @@ async def async_request_openai_chat_omni_completions(
 
                     output.latency = timestamp - st
                     output.generated_text = generated_text
+                    audio_duration_sec = 0.0
+                    audio_frames = 0
                     if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
                         channels, sample_width, frame_rate = wav_audio_params
-                        generated_audio = AudioSegment(
-                            data=bytes(wav_pcm_buffer),
-                            sample_width=sample_width,
-                            frame_rate=frame_rate,
-                            channels=channels,
-                        )
+                        frame_width = sample_width * channels
+                        if frame_width > 0:
+                            audio_frames = len(wav_pcm_buffer) // frame_width
+                            audio_duration_sec = audio_frames / frame_rate
+                        else:
+                            logger.warning("Audio frame width is zero")
                     elif audio_bytes_buffer:
                         try:
-                            generated_audio = AudioSegment.from_file(
+                            from vllm.multimodal.audio import get_audio_duration
+                            from vllm.multimodal.media.audio import load_audio
+
+                            waveform, sr = load_audio(
                                 io.BytesIO(bytes(audio_bytes_buffer)),
-                                format=response_format,
+                                sr=None,
+                                mono=False,
                             )
+                            audio_duration_sec = get_audio_duration(y=waveform, sr=sr)
+                            audio_frames = int(audio_duration_sec * sr)
                         except Exception as ex:
                             logger.warning("Failed to decode accumulated audio bytes: %s", ex)
-                    if generated_audio is not None:
-                        output.audio_duration = len(generated_audio) / 1000.0
-                        frame_width = generated_audio.frame_width
-                        if frame_width > 0:
-                            output.audio_frames = len(generated_audio.raw_data) // frame_width
-                        else:
-                            output.audio_frames = 0
-                            logger.warning("Audio frame width is zero")
+                    if audio_duration_sec > 0 or audio_frames > 0:
+                        output.audio_duration = audio_duration_sec
+                        output.audio_frames = audio_frames
                         audio_duration = output.audio_duration
                         if audio_duration > 0:
                             output.audio_rtf = audio_generate_time / output.audio_duration
@@ -827,8 +829,29 @@ async def async_request_openai_chat_omni_completions(
                             logger.warning("Audio duration is zero")
                         if _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False):
                             try:
-                                seg = generated_audio.set_frame_rate(24000).set_channels(1).set_sample_width(2)
-                                output.tts_output_pcm_bytes = bytes(seg.raw_data)
+                                if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
+                                    from vllm.multimodal.audio import AudioResampler
+
+                                    pcm_channels, pcm_sw, pcm_rate = wav_audio_params
+                                    pcm = np.frombuffer(
+                                        bytes(wav_pcm_buffer), dtype=np.int16 if pcm_sw == 2 else np.float32
+                                    )
+                                    if pcm_channels > 1:
+                                        pcm = pcm.reshape(-1, pcm_channels).mean(axis=1).astype(pcm.dtype)
+                                    pcm_f32 = pcm.astype(np.float32) / 32767.0 if pcm.dtype == np.int16 else pcm
+                                    if pcm_rate != 24000:
+                                        resampler = AudioResampler(target_sr=24000)
+                                        pcm_f32 = resampler.resample(pcm_f32, orig_sr=pcm_rate)
+                                    output.tts_output_pcm_bytes = (pcm_f32 * 32767).astype(np.int16).tobytes()
+                                elif audio_bytes_buffer:
+                                    from vllm.multimodal.media.audio import load_audio
+
+                                    waveform, _ = load_audio(
+                                        io.BytesIO(bytes(audio_bytes_buffer)),
+                                        sr=24000,
+                                        mono=True,
+                                    )
+                                    output.tts_output_pcm_bytes = (waveform * 32767).astype(np.int16).tobytes()
                             except Exception as ex:
                                 logger.warning("seed_tts WER PCM export failed: %s", ex)
                     output.success = True