Spaces:
Runtime error
Runtime error
File size: 3,002 Bytes
61e6f25 | 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 | import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
class MelodySpectrogram(torch.nn.Module):
def __init__(
self,
n_mel_channels=80,
sampling_rate=44100,
win_length=2048,
hop_length=512,
n_fft=None,
mel_fmin=0,
mel_fmax=None,
clamp=1e-5,
):
from librosa.filters import mel
super().__init__()
n_fft = win_length if n_fft is None else n_fft
self.hann_window = {}
mel_basis = 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)
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
def _mel_forward(self, audio, keyshift=0, speed=1, center=True):
factor = 2 ** (keyshift / 12)
n_fft_new = int(np.round(self.n_fft * factor))
win_length_new = int(np.round(self.win_length * factor))
hop_length_new = int(np.round(self.hop_length * speed))
keyshift_key = str(keyshift) + "_" + str(audio.device)
if keyshift_key not in self.hann_window:
self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
audio.device
)
fft = torch.stft(
audio,
n_fft=n_fft_new,
hop_length=hop_length_new,
win_length=win_length_new,
window=self.hann_window[keyshift_key],
center=center,
return_complex=True,
)
magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
if keyshift != 0:
size = self.n_fft // 2 + 1
resize = magnitude.size(1)
if resize < size:
magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
mel_output = torch.matmul(self.mel_basis, magnitude)
log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
return log_mel_spec
@torch.no_grad()
def forward(self, audio, sr, sil_len_to_end=None, keyshift=0, speed=1):
# audio, sr = torchaudio.load(audio_path)
if sil_len_to_end is not None:
silence = torch.zeros(audio.shape[0], int(sr * sil_len_to_end))
audio = torch.cat([audio, silence], dim=1)
if sr != self.sampling_rate:
audio = torchaudio.transforms.Resample(sr, self.sampling_rate)(audio)
if audio.shape[0] > 1:
audio = torch.mean(audio, dim=0, keepdim=True)
audio = audio.to(self.mel_basis.device)
return self._mel_forward(audio, keyshift=keyshift, speed=speed)
|