| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from dataclasses import dataclass |
| 5 | from enum import Enum |
| 6 | from functools import lru_cache |
| 7 | from typing import Literal |
| 8 | |
| 9 | import numpy as np |
| 10 | import numpy.typing as npt |
| 11 | import torch |
| 12 | |
| 13 | from vllm.utils.import_utils import PlaceholderModule |
| 14 | from vllm.utils.torch_utils import set_default_torch_num_threads |
| 15 | |
| 16 | try: |
| 17 | import av as av |
| 18 | except ImportError: |
| 19 | av = PlaceholderModule("av") # type: ignore[assignment] |
| 20 | |
| 21 | try: |
| 22 | import scipy.signal as scipy_signal |
| 23 | except ImportError: |
| 24 | scipy_signal = PlaceholderModule("scipy").placeholder_attr("signal") # type: ignore[assignment] |
| 25 | |
| 26 | try: |
| 27 | import soxr as soxr |
| 28 | except ImportError: |
| 29 | soxr = PlaceholderModule("soxr") # type: ignore[assignment] |
| 30 | |
| 31 | try: |
| 32 | import torchaudio |
| 33 | except ImportError: |
| 34 | torchaudio = PlaceholderModule("torchaudio") # type: ignore[assignment] |
| 35 | |
| 36 | |
| 37 | # ============================================================ |
| 38 | # Aligned with `librosa.get_duration` function |
| 39 | def get_audio_duration(*, y: npt.NDArray[np.floating], sr: float = 22050) -> float: |
| 40 | """Get the duration of an audio array in seconds. |
| 41 | |
| 42 | Args: |
| 43 | y: Audio time series. Can be 1D (samples,) or 2D (channels, samples). |
| 44 | sr: Sample rate of the audio in Hz. |
| 45 | |
| 46 | Returns: |
| 47 | Duration of the audio in seconds. |
| 48 | |
| 49 | """ |
| 50 | n_samples = y.shape[-1] |
| 51 | return float(n_samples) / sr |
| 52 | |
| 53 | |
| 54 | class ChannelReduction(str, Enum): |
| 55 | """Method to reduce multi-channel audio to target channels.""" |
| 56 | |
| 57 | MEAN = "mean" # Average across channels (default, preserves energy balance) |
| 58 | FIRST = "first" # Take first channel only |
| 59 | MAX = "max" # Take max value across channels |
| 60 | SUM = "sum" # Sum across channels |
| 61 | |
| 62 | |
| 63 | @dataclass |
| 64 | class AudioSpec: |
| 65 | """Specification for target audio format. |
| 66 | |
| 67 | This dataclass defines the expected audio format for a model's feature |
| 68 | extractor. It is used to normalize audio data before processing. |
| 69 | |
| 70 | Attributes: |
| 71 | target_channels: Number of output channels. None means passthrough |
| 72 | (no normalization). 1 = mono, 2 = stereo, etc. |
| 73 | channel_reduction: Method to reduce channels when input has more |
| 74 | channels than target. Only used when reducing channels. |
| 75 | |
| 76 | """ |
| 77 | |
| 78 | target_channels: int | None = 1 |
| 79 | channel_reduction: ChannelReduction = ChannelReduction.MEAN |
| 80 | |
| 81 | @property |
| 82 | def needs_normalization(self) -> bool: |
| 83 | """Whether audio normalization is needed.""" |
| 84 | return self.target_channels is not None |
| 85 | |
| 86 | def __repr__(self) -> str: |
| 87 | if self.target_channels is None: |
| 88 | return "AudioSpec(passthrough)" |
| 89 | return ( |
| 90 | f"AudioSpec(channels={self.target_channels}, " |
| 91 | f"reduction={self.channel_reduction.value})" |
| 92 | ) |
| 93 | |
| 94 | |
| 95 | # Pre-defined specs for common use cases |
| 96 | MONO_AUDIO_SPEC = AudioSpec(target_channels=1, channel_reduction=ChannelReduction.MEAN) |
| 97 | PASSTHROUGH_AUDIO_SPEC = AudioSpec(target_channels=None) |
| 98 | |
| 99 | |
| 100 | def normalize_audio( |
| 101 | audio: npt.NDArray[np.floating] | torch.Tensor, |
| 102 | spec: AudioSpec, |
| 103 | ) -> npt.NDArray[np.floating] | torch.Tensor: |
| 104 | """Normalize audio to the specified format. |
| 105 | |
| 106 | This function handles channel reduction for multi-channel audio, |
| 107 | supporting both numpy arrays and torch tensors. |
| 108 | |
| 109 | Args: |
| 110 | audio: Input audio data. Can be: |
| 111 | - 1D array/tensor: (time,) - already mono |
| 112 | - 2D array/tensor: (channels, time) - standard format from torchaudio |
| 113 | - 2D array/tensor: (time, channels) - format from soundfile |
| 114 | (will be auto-detected and transposed if time > channels) |
| 115 | spec: AudioSpec defining the target format. |
| 116 | |
| 117 | Returns: |
| 118 | Normalized audio in the same type as input (numpy or torch). |
| 119 | For mono output (target_channels=1), returns 1D array/tensor. |
| 120 | |
| 121 | Raises: |
| 122 | ValueError: If audio has unsupported dimensions or channel expansion |
| 123 | is requested (e.g., mono to stereo). |
| 124 | |
| 125 | """ |
| 126 | if not spec.needs_normalization: |
| 127 | return audio |
| 128 | |
| 129 | # Handle 1D audio (already mono) |
| 130 | if audio.ndim == 1: |
| 131 | if spec.target_channels == 1: |
| 132 | return audio |
| 133 | raise ValueError(f"Cannot expand mono audio to {spec.target_channels} channels") |
| 134 | |
| 135 | # Handle 2D audio |
| 136 | if audio.ndim != 2: |
| 137 | raise ValueError(f"Unsupported audio shape: {audio.shape}. Expected 1D or 2D.") |
| 138 | |
| 139 | # Auto-detect format: if shape[0] > shape[1], assume (time, channels) |
| 140 | # This handles soundfile format where time dimension is typically much larger |
| 141 | if audio.shape[0] > audio.shape[1]: |
| 142 | # Transpose from (time, channels) to (channels, time) |
| 143 | audio = audio.T if isinstance(audio, np.ndarray) else audio.T |
| 144 | |
| 145 | num_channels = audio.shape[0] |
| 146 | |
| 147 | # No reduction needed if already at target |
| 148 | if num_channels == spec.target_channels: |
| 149 | return audio |
| 150 | |
| 151 | # Cannot expand channels |
| 152 | if num_channels < spec.target_channels: |
| 153 | raise ValueError( |
| 154 | f"Cannot expand {num_channels} channels to {spec.target_channels}" |
| 155 | ) |
| 156 | |
| 157 | # Reduce channels |
| 158 | is_numpy = isinstance(audio, np.ndarray) |
| 159 | |
| 160 | if spec.target_channels == 1: |
| 161 | # Reduce to mono |
| 162 | if spec.channel_reduction == ChannelReduction.MEAN: |
| 163 | result = np.mean(audio, axis=0) if is_numpy else audio.mean(dim=0) |
| 164 | elif spec.channel_reduction == ChannelReduction.FIRST: |
| 165 | result = audio[0] |
| 166 | elif spec.channel_reduction == ChannelReduction.MAX: |
| 167 | result = np.max(audio, axis=0) if is_numpy else audio.max(dim=0).values |
| 168 | elif spec.channel_reduction == ChannelReduction.SUM: |
| 169 | result = np.sum(audio, axis=0) if is_numpy else audio.sum(dim=0) |
| 170 | else: |
| 171 | raise ValueError(f"Unknown reduction method: {spec.channel_reduction}") |
| 172 | return result |
| 173 | else: |
| 174 | # Reduce to N channels (take first N and apply reduction if needed) |
| 175 | # For now, just take first N channels |
| 176 | return audio[: spec.target_channels] |
| 177 | |
| 178 | |
| 179 | # ============================================================ |
| 180 | # Audio Resampling |
| 181 | # ============================================================ |
| 182 | |
| 183 | |
| 184 | def resample_audio_pyav( |
| 185 | audio: npt.NDArray[np.floating], |
| 186 | *, |
| 187 | orig_sr: float, |
| 188 | target_sr: float, |
| 189 | ) -> npt.NDArray[np.floating]: |
| 190 | """Resample audio using PyAV (libswresample via FFmpeg). |
| 191 | |
| 192 | Args: |
| 193 | audio: Input audio. Can be: |
| 194 | - 1D array ``(samples,)``: mono audio |
| 195 | - 2D array ``(channels, samples)``: stereo audio |
| 196 | orig_sr: Original sample rate in Hz. |
| 197 | target_sr: Target sample rate in Hz. |
| 198 | |
| 199 | Returns: |
| 200 | Resampled audio with the same shape as the input (1D → 1D, 2D → 2D). |
| 201 | |
| 202 | """ |
| 203 | orig_sr_int = int(round(orig_sr)) |
| 204 | target_sr_int = int(round(target_sr)) |
| 205 | |
| 206 | if orig_sr_int == target_sr_int: |
| 207 | return audio |
| 208 | |
| 209 | if audio.ndim == 2: |
| 210 | # Resample each channel independently and re-stack. |
| 211 | return np.stack( |
| 212 | [ |
| 213 | resample_audio_pyav(ch, orig_sr=orig_sr, target_sr=target_sr) |
| 214 | for ch in audio |
| 215 | ], |
| 216 | axis=0, |
| 217 | ) |
| 218 | |
| 219 | expected_len = int(math.ceil(audio.shape[-1] * target_sr_int / orig_sr_int)) |
| 220 | |
| 221 | # from_ndarray expects shape (channels, samples) for planar formats. |
| 222 | # libswresample requires a minimum number of input samples to produce |
| 223 | # output frames; pad short inputs with zeros so we always get output, |
| 224 | # then trim to the expected output length. |
| 225 | _MIN_SAMPLES = 1024 |
| 226 | audio_f32 = np.asarray(audio, dtype=np.float32) |
| 227 | if len(audio_f32) < _MIN_SAMPLES: |
| 228 | audio_f32 = np.pad(audio_f32, (0, _MIN_SAMPLES - len(audio_f32))) |
| 229 | audio_f32 = audio_f32.reshape(1, -1) |
| 230 | |
| 231 | resampler = av.AudioResampler(format="fltp", layout="mono", rate=target_sr_int) |
| 232 | |
| 233 | frame = av.AudioFrame.from_ndarray(audio_f32, format="fltp", layout="mono") |
| 234 | frame.sample_rate = orig_sr_int |
| 235 | |
| 236 | out_frames = resampler.resample(frame) |
| 237 | out_frames.extend(resampler.resample(None)) # flush buffered samples |
| 238 | |
| 239 | result = np.concatenate([f.to_ndarray() for f in out_frames], axis=1).squeeze(0) |
| 240 | return result[:expected_len] |
| 241 | |
| 242 | |
| 243 | def resample_audio_scipy( |
| 244 | audio: npt.NDArray[np.floating], |
| 245 | *, |
| 246 | orig_sr: float, |
| 247 | target_sr: float, |
| 248 | ) -> npt.NDArray[np.floating]: |
| 249 | orig_sr_int = int(round(orig_sr)) |
| 250 | target_sr_int = int(round(target_sr)) |
| 251 | |
| 252 | if orig_sr_int == target_sr_int: |
| 253 | return audio |
| 254 | |
| 255 | gcd = math.gcd(orig_sr_int, target_sr_int) |
| 256 | return scipy_signal.resample_poly( |
| 257 | audio, |
| 258 | target_sr_int // gcd, |
| 259 | orig_sr_int // gcd, |
| 260 | axis=-1, |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | def resample_audio_soxr( |
| 265 | audio: npt.NDArray[np.floating], |
| 266 | *, |
| 267 | orig_sr: float, |
| 268 | target_sr: float, |
| 269 | ) -> npt.NDArray[np.floating]: |
| 270 | orig_sr_int = int(round(orig_sr)) |
| 271 | target_sr_int = int(round(target_sr)) |
| 272 | |
| 273 | if orig_sr_int == target_sr_int: |
| 274 | return audio |
| 275 | |
| 276 | if audio.ndim == 2: |
| 277 | return np.stack( |
| 278 | [ |
| 279 | resample_audio_soxr(ch, orig_sr=orig_sr, target_sr=target_sr) |
| 280 | for ch in audio |
| 281 | ], |
| 282 | axis=0, |
| 283 | ) |
| 284 | |
| 285 | return soxr.resample(audio, orig_sr_int, target_sr_int) |
| 286 | |
| 287 | |
| 288 | @lru_cache(maxsize=32) |
| 289 | def _get_torchaudio_resampler( |
| 290 | orig_sr: int, target_sr: int |
| 291 | ) -> "torchaudio.transforms.Resample": |
| 292 | # `torchaudio.transforms.Resample` precomputes its kernel for a fixed |
| 293 | # (orig_sr, target_sr) pair; cache instances so repeated requests at a |
| 294 | # common input rate skip the kernel rebuild. |
| 295 | return torchaudio.transforms.Resample(orig_sr, target_sr) |
| 296 | |
| 297 | |
| 298 | def resample_audio_torchaudio( |
| 299 | audio: npt.NDArray[np.floating], |
| 300 | *, |
| 301 | orig_sr: float, |
| 302 | target_sr: float, |
| 303 | ) -> npt.NDArray[np.floating]: |
| 304 | """Resample audio using torchaudio's bandlimited sinc interpolation. |
| 305 | |
| 306 | Unlike the PyAV resampler, this handles any input length without padding |
| 307 | and applies the kernel over the trailing axis, so 2D ``(channels, |
| 308 | samples)`` input needs no per-channel loop. |
| 309 | |
| 310 | Args: |
| 311 | audio: Input audio. Can be: |
| 312 | - 1D array ``(samples,)``: mono audio |
| 313 | - 2D array ``(channels, samples)``: stereo audio |
| 314 | orig_sr: Original sample rate in Hz. |
| 315 | target_sr: Target sample rate in Hz. |
| 316 | |
| 317 | Returns: |
| 318 | Resampled audio with the same shape as the input (1D → 1D, 2D → 2D). |
| 319 | |
| 320 | """ |
| 321 | orig_sr_int = int(round(orig_sr)) |
| 322 | target_sr_int = int(round(target_sr)) |
| 323 | |
| 324 | if orig_sr_int == target_sr_int: |
| 325 | return audio |
| 326 | |
| 327 | # The kernel is float32; cast the input to match (same coercion as the |
| 328 | # PyAV path). |
| 329 | tensor = torch.as_tensor(audio, dtype=torch.float32) |
| 330 | # Resampling runs in the API/server parent process. Keep it from |
| 331 | # touching OpenMP or oneDNN thread state: both poison subsequently |
| 332 | # forked engine-core processes, which then segfault on their first |
| 333 | # parallel CPU op. |
| 334 | with set_default_torch_num_threads(1), torch.backends.mkldnn.flags(enabled=False): |
| 335 | resampler = _get_torchaudio_resampler(orig_sr_int, target_sr_int) |
| 336 | return resampler(tensor).numpy() |
| 337 | |
| 338 | |
| 339 | class AudioResampler: |
| 340 | """Resample audio data to a target sample rate.""" |
| 341 | |
| 342 | _METHODS = ("pyav", "scipy", "soxr", "torchaudio") |
| 343 | |
| 344 | def __init__( |
| 345 | self, |
| 346 | target_sr: float | None = None, |
| 347 | method: Literal["pyav", "scipy", "soxr", "torchaudio"] = "torchaudio", |
| 348 | ): |
| 349 | # Eager validation so a bad method fails at construction rather than |
| 350 | # on the first audio request. |
| 351 | if method not in self._METHODS: |
| 352 | raise ValueError( |
| 353 | f"Invalid resampling method: {method!r}. " |
| 354 | f"Supported methods: {list(self._METHODS)}." |
| 355 | ) |
| 356 | self.target_sr = target_sr |
| 357 | self.method = method |
| 358 | |
| 359 | def resample( |
| 360 | self, |
| 361 | audio: npt.NDArray[np.floating], |
| 362 | *, |
| 363 | orig_sr: float, |
| 364 | ) -> npt.NDArray[np.floating]: |
| 365 | if self.target_sr is None: |
| 366 | raise RuntimeError( |
| 367 | "Audio resampling is not supported when `target_sr` is not provided" |
| 368 | ) |
| 369 | if math.isclose( |
| 370 | float(orig_sr), |
| 371 | float(self.target_sr), |
| 372 | rel_tol=0.0, |
| 373 | abs_tol=1e-6, |
| 374 | ): |
| 375 | return audio |
| 376 | if self.method == "pyav": |
| 377 | return resample_audio_pyav(audio, orig_sr=orig_sr, target_sr=self.target_sr) |
| 378 | elif self.method == "scipy": |
| 379 | return resample_audio_scipy( |
| 380 | audio, orig_sr=orig_sr, target_sr=self.target_sr |
| 381 | ) |
| 382 | elif self.method == "soxr": |
| 383 | return resample_audio_soxr(audio, orig_sr=orig_sr, target_sr=self.target_sr) |
| 384 | elif self.method == "torchaudio": |
| 385 | return resample_audio_torchaudio( |
| 386 | audio, orig_sr=orig_sr, target_sr=self.target_sr |
| 387 | ) |
| 388 | else: |
| 389 | raise ValueError( |
| 390 | f"Invalid resampling method: {self.method}. " |
| 391 | f"Supported methods are {list(self._METHODS)}." |
| 392 | ) |
| 393 | |
| 394 | |
| 395 | # ============================================================ |
| 396 | # Audio Chunking / Splitting |
| 397 | # ============================================================ |
| 398 | |
| 399 | |
| 400 | def split_audio( |
| 401 | audio_data: np.ndarray, |
| 402 | sample_rate: int, |
| 403 | max_clip_duration_s: float, |
| 404 | overlap_duration_s: float, |
| 405 | min_energy_window_size: int, |
| 406 | ) -> list[np.ndarray]: |
| 407 | """Split audio into chunks with intelligent split points. |
| 408 | |
| 409 | Splits long audio into smaller chunks at low-energy regions to minimize |
| 410 | cutting through speech. Uses overlapping windows to find quiet moments |
| 411 | for splitting. |
| 412 | |
| 413 | Args: |
| 414 | audio_data: 1D mono audio array to split. ASR models consume mono, so |
| 415 | callers must downmix before chunking. |
| 416 | sample_rate: Sample rate of the audio in Hz. |
| 417 | max_clip_duration_s: Maximum duration of each chunk in seconds. |
| 418 | overlap_duration_s: Overlap duration in seconds between consecutive chunks. |
| 419 | Used to search for optimal split points. |
| 420 | min_energy_window_size: Window size in samples for finding low-energy regions. |
| 421 | |
| 422 | Returns: |
| 423 | List of 1D audio chunks. |
| 424 | |
| 425 | Raises: |
| 426 | AssertionError: If ``audio_data`` is not 1D. |
| 427 | |
| 428 | Example: |
| 429 | >>> audio = np.random.randn(1040000) # 65 seconds at 16kHz |
| 430 | >>> chunks = split_audio( |
| 431 | ... audio_data=audio, |
| 432 | ... sample_rate=16000, |
| 433 | ... max_clip_duration_s=30.0, |
| 434 | ... overlap_duration_s=1.0, |
| 435 | ... min_energy_window_size=1600, |
| 436 | ... ) |
| 437 | >>> len(chunks) |
| 438 | 3 |
| 439 | |
| 440 | """ |
| 441 | if audio_data.ndim > 1: |
| 442 | raise ValueError( |
| 443 | f"split_audio expects mono audio, got shape {audio_data.shape}" |
| 444 | ) |
| 445 | |
| 446 | chunk_size = int(sample_rate * max_clip_duration_s) |
| 447 | overlap_size = int(sample_rate * overlap_duration_s) |
| 448 | chunks = [] |
| 449 | i = 0 |
| 450 | |
| 451 | while i < audio_data.shape[-1]: |
| 452 | if i + chunk_size >= audio_data.shape[-1]: |
| 453 | # Handle last chunk - take everything remaining |
| 454 | chunks.append(audio_data[..., i:]) |
| 455 | break |
| 456 | |
| 457 | # Find the best split point in the overlap region |
| 458 | search_start = i + chunk_size - overlap_size |
| 459 | search_end = min(i + chunk_size, audio_data.shape[-1]) |
| 460 | split_point = find_split_point( |
| 461 | audio_data, search_start, search_end, min_energy_window_size |
| 462 | ) |
| 463 | |
| 464 | # Guarantee forward progress: if split_point didn't advance, |
| 465 | # fall back to the hard chunk boundary. |
| 466 | if split_point <= i: |
| 467 | split_point = min(i + chunk_size, audio_data.shape[-1]) |
| 468 | |
| 469 | # Extract chunk up to the split point |
| 470 | chunks.append(audio_data[..., i:split_point]) |
| 471 | i = split_point |
| 472 | |
| 473 | return chunks |
| 474 | |
| 475 | |
| 476 | def find_split_point( |
| 477 | wav: np.ndarray, |
| 478 | start_idx: int, |
| 479 | end_idx: int, |
| 480 | min_energy_window: int, |
| 481 | ) -> int: |
| 482 | """Find the best point to split audio by looking for silence or low amplitude. |
| 483 | |
| 484 | Searches for the quietest region within a specified range by calculating |
| 485 | RMS energy in sliding windows. |
| 486 | |
| 487 | Args: |
| 488 | wav: 1D mono audio array. |
| 489 | start_idx: Start index of search region (inclusive). |
| 490 | end_idx: End index of search region (exclusive). |
| 491 | min_energy_window: Window size in samples for energy calculation. |
| 492 | |
| 493 | Returns: |
| 494 | Index of the quietest point within the search region. This is the |
| 495 | recommended split point to minimize audio artifacts. |
| 496 | |
| 497 | Example: |
| 498 | >>> audio = np.random.randn(32000) |
| 499 | >>> # Insert quiet region |
| 500 | >>> audio[16000:17600] = 0.01 |
| 501 | >>> split_idx = find_split_point( |
| 502 | ... wav=audio, |
| 503 | ... start_idx=0, |
| 504 | ... end_idx=32000, |
| 505 | ... min_energy_window=1600, |
| 506 | ... ) |
| 507 | >>> 16000 <= split_idx <= 17600 |
| 508 | True |
| 509 | |
| 510 | """ |
| 511 | segment = wav[start_idx:end_idx] |
| 512 | |
| 513 | # Calculate RMS energy in small windows |
| 514 | min_energy = math.inf |
| 515 | quietest_idx = start_idx |
| 516 | |
| 517 | for i in range(0, len(segment) - min_energy_window, min_energy_window): |
| 518 | window = segment[i : i + min_energy_window] |
| 519 | energy = (window**2).mean() ** 0.5 |
| 520 | if not math.isnan(energy) and energy < min_energy: |
| 521 | quietest_idx = i + start_idx |
| 522 | min_energy = energy |
| 523 | |
| 524 | return quietest_idx |