audio.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. import os
  2. from functools import lru_cache
  3. from subprocess import CalledProcessError, run
  4. from typing import Optional, Union
  5. import numpy as np
  6. import torch
  7. import torch.nn.functional as F
  8. from .utils import exact_div
  9. # hard-coded audio hyperparameters
  10. SAMPLE_RATE = 16000
  11. N_FFT = 400
  12. N_MELS = 80
  13. HOP_LENGTH = 160
  14. CHUNK_LENGTH = 30
  15. N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk
  16. N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH) # 3000 frames in a mel spectrogram input
  17. N_SAMPLES_PER_TOKEN = HOP_LENGTH * 2 # the initial convolutions has stride 2
  18. FRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH) # 10ms per audio frame
  19. TOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN) # 20ms per audio token
  20. def load_audio(file: str, sr: int = SAMPLE_RATE):
  21. """
  22. Open an audio file and read as mono waveform, resampling as necessary
  23. Parameters
  24. ----------
  25. file: str
  26. The audio file to open
  27. sr: int
  28. The sample rate to resample the audio if necessary
  29. Returns
  30. -------
  31. A NumPy array containing the audio waveform, in float32 dtype.
  32. """
  33. # This launches a subprocess to decode audio while down-mixing
  34. # and resampling as necessary. Requires the ffmpeg CLI in PATH.
  35. # fmt: off
  36. cmd = [
  37. "ffmpeg",
  38. "-nostdin",
  39. "-threads", "0",
  40. "-i", file,
  41. "-f", "s16le",
  42. "-ac", "1",
  43. "-acodec", "pcm_s16le",
  44. "-ar", str(sr),
  45. "-"
  46. ]
  47. # fmt: on
  48. try:
  49. out = run(cmd, capture_output=True, check=True).stdout
  50. except CalledProcessError as e:
  51. raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
  52. return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
  53. def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
  54. """
  55. Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
  56. """
  57. if torch.is_tensor(array):
  58. if array.shape[axis] > length:
  59. array = array.index_select(
  60. dim=axis, index=torch.arange(length, device=array.device)
  61. )
  62. if array.shape[axis] < length:
  63. pad_widths = [(0, 0)] * array.ndim
  64. pad_widths[axis] = (0, length - array.shape[axis])
  65. array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
  66. else:
  67. if array.shape[axis] > length:
  68. array = array.take(indices=range(length), axis=axis)
  69. if array.shape[axis] < length:
  70. pad_widths = [(0, 0)] * array.ndim
  71. pad_widths[axis] = (0, length - array.shape[axis])
  72. array = np.pad(array, pad_widths)
  73. return array
  74. @lru_cache(maxsize=None)
  75. def mel_filters(device, n_mels: int = N_MELS) -> torch.Tensor:
  76. """
  77. load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
  78. Allows decoupling librosa dependency; saved using:
  79. np.savez_compressed(
  80. "mel_filters.npz",
  81. mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
  82. )
  83. """
  84. assert n_mels == 80, f"Unsupported n_mels: {n_mels}"
  85. with np.load(
  86. os.path.join(os.path.dirname(__file__), "assets", "mel_filters.npz")
  87. ) as f:
  88. return torch.from_numpy(f[f"mel_{n_mels}"]).to(device)
  89. def log_mel_spectrogram(
  90. audio: Union[str, np.ndarray, torch.Tensor],
  91. n_mels: int = N_MELS,
  92. padding: int = 0,
  93. device: Optional[Union[str, torch.device]] = None,
  94. ):
  95. """
  96. Compute the log-Mel spectrogram of
  97. Parameters
  98. ----------
  99. audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
  100. The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
  101. n_mels: int
  102. The number of Mel-frequency filters, only 80 is supported
  103. padding: int
  104. Number of zero samples to pad to the right
  105. device: Optional[Union[str, torch.device]]
  106. If given, the audio tensor is moved to this device before STFT
  107. Returns
  108. -------
  109. torch.Tensor, shape = (80, n_frames)
  110. A Tensor that contains the Mel spectrogram
  111. """
  112. if not torch.is_tensor(audio):
  113. if isinstance(audio, str):
  114. audio = load_audio(audio)
  115. audio = torch.from_numpy(audio)
  116. if device is not None:
  117. audio = audio.to(device)
  118. if padding > 0:
  119. audio = F.pad(audio, (0, padding))
  120. window = torch.hann_window(N_FFT).to(audio.device)
  121. stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
  122. magnitudes = stft[..., :-1].abs() ** 2
  123. filters = mel_filters(audio.device, n_mels)
  124. mel_spec = filters @ magnitudes
  125. log_spec = torch.clamp(mel_spec, min=1e-10).log10()
  126. log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
  127. log_spec = (log_spec + 4.0) / 4.0
  128. return log_spec