|
import torch |
|
import numpy as np |
|
import torch.nn.functional as F |
|
from torch.autograd import Variable |
|
from scipy.signal import get_window |
|
from librosa.util import pad_center, tiny |
|
from maha_tts.utils.audio import window_sumsquare |
|
|
|
|
|
class STFT(torch.nn.Module): |
|
"""adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft""" |
|
def __init__(self, filter_length=800, hop_length=200, win_length=800, |
|
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, size = 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, |
|
Variable(self.forward_basis, requires_grad=False), |
|
stride=self.hop_length, |
|
padding=0) |
|
|
|
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 |
