Spaces:
Paused
Paused
File size: 2,151 Bytes
2d9b22b |
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 |
import os
import torch
import torch.nn as nn
import numpy as np
import librosa
from .stft import STFT, TorchSTFT
USING_TORCH_STFT = os.getenv("USING_TORCH_STFT") is not None
class MelSpectrogram(nn.Module):
def __init__(
self,
n_mel_channels: int,
sampling_rate: int,
win_length: int,
hop_length: int,
n_fft: int = None,
mel_fmin: int = 0,
mel_fmax: int = None,
clamp: float = 1e-5,
):
super().__init__()
n_fft = win_length if n_fft is None else n_fft
mel_basis = librosa.filters.mel(
sr=sampling_rate,
n_fft=n_fft,
n_mels=n_mel_channels,
fmin=mel_fmin,
fmax=mel_fmax,
htk=True,
)
mel_basis = torch.from_numpy(mel_basis).float()
self.register_buffer("mel_basis", mel_basis, persistent=False)
self.n_fft = n_fft
self.hop_length = hop_length
self.win_length = win_length
self.sampling_rate = sampling_rate
self.n_mel_channels = n_mel_channels
self.clamp = clamp
self.keyshift = 0
self.speed = 1
self.factor = 2 ** (self.keyshift / 12)
self.n_fft_new = int(np.round(self.n_fft * self.factor))
self.win_length_new = int(np.round(self.win_length * self.factor))
self.hop_length_new = int(np.round(self.hop_length * self.speed))
if USING_TORCH_STFT:
self.stft = TorchSTFT(
filter_length=self.n_fft_new,
hop_length=self.hop_length_new,
win_length=self.win_length_new,
window="hann",
)
else:
self.stft = STFT(
filter_length=self.n_fft_new,
hop_length=self.hop_length_new,
win_length=self.win_length_new,
window="hann",
)
def forward(self, audio: torch.Tensor):
magnitude = self.stft(audio)
mel_output = torch.matmul(self.mel_basis, magnitude)
log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
return log_mel_spec
|