Spaces:
Running
Running
File size: 3,299 Bytes
e319a7b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Standalone pure-torch equivalent of torchaudio.transforms.MelSpectrogram.
Matches torchaudio with mel_scale='htk', norm=None, win_length == n_fft.
"""
import torch
from torch import nn
def _hz_to_mel_htk(freq: torch.Tensor) -> torch.Tensor:
return 2595.0 * torch.log10(1.0 + freq / 700.0)
def _mel_to_hz_htk(mel: torch.Tensor) -> torch.Tensor:
return 700.0 * (10 ** (mel / 2595.0) - 1.0)
def melscale_fbanks(
n_freqs: int,
f_min: float,
f_max: float,
n_mels: int,
sample_rate: int,
) -> torch.Tensor:
"""Triangular mel filterbank matching torchaudio.functional.melscale_fbanks
with mel_scale='htk' and norm=None. Returns a tensor of shape [n_freqs, n_mels]."""
all_freqs = torch.linspace(0, sample_rate // 2, n_freqs)
m_min = _hz_to_mel_htk(torch.tensor(float(f_min)))
m_max = _hz_to_mel_htk(torch.tensor(float(f_max)))
m_pts = torch.linspace(m_min.item(), m_max.item(), n_mels + 2)
f_pts = _mel_to_hz_htk(m_pts)
f_diff = f_pts[1:] - f_pts[:-1]
slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)
down_slopes = -slopes[:, :-2] / f_diff[:-1]
up_slopes = slopes[:, 2:] / f_diff[1:]
return torch.maximum(torch.zeros(()), torch.minimum(down_slopes, up_slopes))
class _Spectrogram(nn.Module):
"""Holds the STFT window so the safetensors key
`...mel_spec_transform.spectrogram.window` round-trips."""
def __init__(self, n_fft: int):
super().__init__()
self.register_buffer("window", torch.hann_window(n_fft))
class _MelScale(nn.Module):
"""Holds the mel filterbank so the safetensors key
`...mel_spec_transform.mel_scale.fb` round-trips."""
def __init__(self, fb: torch.Tensor):
super().__init__()
self.register_buffer("fb", fb)
class _MelSpectrogram(nn.Module):
"""Pure-torch equivalent of torchaudio.transforms.MelSpectrogram
(htk mel scale, no Slaney norm, win_length == n_fft)."""
def __init__(
self,
sample_rate: int,
n_fft: int,
hop_length: int,
n_mels: int,
power: float = 2.0,
center: bool = True,
pad_mode: str = "reflect",
):
super().__init__()
self.n_fft = n_fft
self.hop_length = hop_length
self.power = power
self.center = center
self.pad_mode = pad_mode
self.spectrogram = _Spectrogram(n_fft)
fb = melscale_fbanks(
n_freqs=n_fft // 2 + 1,
f_min=0.0,
f_max=sample_rate / 2.0,
n_mels=n_mels,
sample_rate=sample_rate,
)
self.mel_scale = _MelScale(fb)
def forward(self, x: torch.Tensor) -> torch.Tensor:
leading = x.shape[:-1]
x = x.reshape(-1, x.shape[-1])
spec = torch.stft(
x,
n_fft=self.n_fft,
hop_length=self.hop_length,
win_length=self.n_fft,
window=self.spectrogram.window,
center=self.center,
pad_mode=self.pad_mode,
return_complex=True,
normalized=False,
onesided=True,
)
spec = spec.abs() ** self.power
mel = torch.matmul(spec.transpose(-1, -2), self.mel_scale.fb).transpose(-1, -2)
return mel.reshape(*leading, *mel.shape[-2:])
|