| import torch |
| import torch.nn.functional as F |
| from torch.autograd import Variable |
| 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 audio.audio_processing import dynamic_range_compression |
| from audio.audio_processing import dynamic_range_decompression |
| from audio.audio_processing import 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) |
| |
| 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() |
|
|
| |
| 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): |
| num_batches = input_data.size(0) |
| num_samples = input_data.size(1) |
|
|
| self.num_samples = num_samples |
|
|
| |
| 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.cuda(), |
| Variable(self.forward_basis, requires_grad=False).cuda(), |
| 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): |
| recombine_magnitude_phase = torch.cat( |
| [magnitude*torch.cos(phase), magnitude*torch.sin(phase)], dim=1) |
|
|
| inverse_transform = F.conv_transpose1d( |
| recombine_magnitude_phase, |
| 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) |
| |
| 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.cuda() if magnitude.is_cuda else window_sum |
| inverse_transform[:, :, |
| approx_nonzero_indices] /= window_sum[approx_nonzero_indices] |
|
|
| |
| 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=0.0, |
| mel_fmax=8000.0): |
| 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): |
| output = dynamic_range_compression(magnitudes) |
| return output |
|
|
| def spectral_de_normalize(self, magnitudes): |
| output = dynamic_range_decompression(magnitudes) |
| return output |
|
|
| def mel_spectrogram(self, y): |
| """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) |
| assert(torch.max(y.data) <= 1) |
|
|
| 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) |
| energy = torch.norm(magnitudes, dim=1) |
| |
| return mel_output, energy |
|
|
