Spaces:
Running
Running
| """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:]) | |