Spaces:
Running
on
A10G
Running
on
A10G
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| from scipy.signal import get_window | |
| from librosa.util import pad_center, tiny | |
| from librosa.filters import mel as librosa_mel_fn | |
| from audioldm.audio.audio_processing import ( | |
| dynamic_range_compression, | |
| dynamic_range_decompression, | |
| window_sumsquare, | |
| ) | |
| class STFT(torch.nn.Module): | |
| """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft""" | |
| def __init__(self, filter_length, hop_length, win_length, window="hann"): | |
| super(STFT, self).__init__() | |
| self.filter_length = filter_length | |
| self.hop_length = hop_length | |
| self.win_length = win_length | |
| self.window = window | |
| self.forward_transform = None | |
| scale = self.filter_length / self.hop_length | |
| fourier_basis = np.fft.fft(np.eye(self.filter_length)) | |
| cutoff = int((self.filter_length / 2 + 1)) | |
| fourier_basis = np.vstack( | |
| [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])] | |
| ) | |
| forward_basis = torch.FloatTensor(fourier_basis[:, None, :]) | |
| inverse_basis = torch.FloatTensor( | |
| np.linalg.pinv(scale * fourier_basis).T[:, None, :] | |
| ) | |
| if window is not None: | |
| assert filter_length >= win_length | |
| # get window and zero center pad it to filter_length | |
| fft_window = get_window(window, win_length, fftbins=True) | |
| fft_window = pad_center(fft_window, filter_length) | |
| fft_window = torch.from_numpy(fft_window).float() | |
| # window the bases | |
| forward_basis *= fft_window | |
| inverse_basis *= fft_window | |
| self.register_buffer("forward_basis", forward_basis.float()) | |
| self.register_buffer("inverse_basis", inverse_basis.float()) | |
| def transform(self, input_data): | |
| device = self.forward_basis.device | |
| input_data = input_data.to(device) | |
| num_batches = input_data.size(0) | |
| num_samples = input_data.size(1) | |
| self.num_samples = num_samples | |
| # similar to librosa, reflect-pad the input | |
| input_data = input_data.view(num_batches, 1, num_samples) | |
| input_data = F.pad( | |
| input_data.unsqueeze(1), | |
| (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0), | |
| mode="reflect", | |
| ) | |
| input_data = input_data.squeeze(1) | |
| forward_transform = F.conv1d( | |
| input_data, | |
| torch.autograd.Variable(self.forward_basis, requires_grad=False), | |
| stride=self.hop_length, | |
| padding=0, | |
| )#.cpu() | |
| cutoff = int((self.filter_length / 2) + 1) | |
| real_part = forward_transform[:, :cutoff, :] | |
| imag_part = forward_transform[:, cutoff:, :] | |
| magnitude = torch.sqrt(real_part**2 + imag_part**2) | |
| phase = torch.autograd.Variable(torch.atan2(imag_part.data, real_part.data)) | |
| return magnitude, phase | |
| def inverse(self, magnitude, phase): | |
| device = self.forward_basis.device | |
| magnitude, phase = magnitude.to(device), phase.to(device) | |
| recombine_magnitude_phase = torch.cat( | |
| [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1 | |
| ) | |
| inverse_transform = F.conv_transpose1d( | |
| recombine_magnitude_phase, | |
| torch.autograd.Variable(self.inverse_basis, requires_grad=False), | |
| stride=self.hop_length, | |
| padding=0, | |
| ) | |
| if self.window is not None: | |
| window_sum = window_sumsquare( | |
| self.window, | |
| magnitude.size(-1), | |
| hop_length=self.hop_length, | |
| win_length=self.win_length, | |
| n_fft=self.filter_length, | |
| dtype=np.float32, | |
| ) | |
| # remove modulation effects | |
| approx_nonzero_indices = torch.from_numpy( | |
| np.where(window_sum > tiny(window_sum))[0] | |
| ) | |
| window_sum = torch.autograd.Variable( | |
| torch.from_numpy(window_sum), requires_grad=False | |
| ) | |
| window_sum = window_sum | |
| inverse_transform[:, :, approx_nonzero_indices] /= window_sum[ | |
| approx_nonzero_indices | |
| ] | |
| # scale by hop ratio | |
| inverse_transform *= float(self.filter_length) / self.hop_length | |
| inverse_transform = inverse_transform[:, :, int(self.filter_length / 2) :] | |
| inverse_transform = inverse_transform[:, :, : -int(self.filter_length / 2) :] | |
| return inverse_transform | |
| def forward(self, input_data): | |
| self.magnitude, self.phase = self.transform(input_data) | |
| reconstruction = self.inverse(self.magnitude, self.phase) | |
| return reconstruction | |
| class TacotronSTFT(torch.nn.Module): | |
| def __init__( | |
| self, | |
| filter_length, | |
| hop_length, | |
| win_length, | |
| n_mel_channels, | |
| sampling_rate, | |
| mel_fmin, | |
| mel_fmax, | |
| ): | |
| super(TacotronSTFT, self).__init__() | |
| self.n_mel_channels = n_mel_channels | |
| self.sampling_rate = sampling_rate | |
| self.stft_fn = STFT(filter_length, hop_length, win_length) | |
| mel_basis = librosa_mel_fn( | |
| sampling_rate, filter_length, n_mel_channels, mel_fmin, mel_fmax | |
| ) | |
| mel_basis = torch.from_numpy(mel_basis).float() | |
| self.register_buffer("mel_basis", mel_basis) | |
| def spectral_normalize(self, magnitudes, normalize_fun): | |
| output = dynamic_range_compression(magnitudes, normalize_fun) | |
| return output | |
| def spectral_de_normalize(self, magnitudes): | |
| output = dynamic_range_decompression(magnitudes) | |
| return output | |
| def mel_spectrogram(self, y, normalize_fun=torch.log): | |
| """Computes mel-spectrograms from a batch of waves | |
| PARAMS | |
| ------ | |
| y: Variable(torch.FloatTensor) with shape (B, T) in range [-1, 1] | |
| RETURNS | |
| ------- | |
| mel_output: torch.FloatTensor of shape (B, n_mel_channels, T) | |
| """ | |
| assert torch.min(y.data) >= -1, torch.min(y.data) | |
| assert torch.max(y.data) <= 1, torch.max(y.data) | |
| magnitudes, phases = self.stft_fn.transform(y) | |
| magnitudes = magnitudes.data | |
| mel_output = torch.matmul(self.mel_basis, magnitudes) | |
| mel_output = self.spectral_normalize(mel_output, normalize_fun) | |
| energy = torch.norm(magnitudes, dim=1) | |
| log_magnitudes = self.spectral_normalize(magnitudes, normalize_fun) | |
| return mel_output, log_magnitudes, energy | |