Spaces:
Running
Running
| #!/usr/bin/python3 | |
| # -*- coding: utf-8 -*- | |
| """ | |
| https://github.com/modelscope/modelscope/blob/master/modelscope/models/audio/ans/conv_stft.py | |
| """ | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from scipy.signal import get_window | |
| def init_kernels(nfft: int, win_size: int, hop_size: int, win_type: str = None, inverse=False): | |
| if win_type == "None" or win_type is None: | |
| window = np.ones(win_size) | |
| else: | |
| window = get_window(win_type, win_size, fftbins=True)**0.5 | |
| fourier_basis = np.fft.rfft(np.eye(nfft))[:win_size] | |
| real_kernel = np.real(fourier_basis) | |
| image_kernel = np.imag(fourier_basis) | |
| kernel = np.concatenate([real_kernel, image_kernel], 1).T | |
| if inverse: | |
| kernel = np.linalg.pinv(kernel).T | |
| kernel = kernel * window | |
| kernel = kernel[:, None, :] | |
| result = ( | |
| torch.from_numpy(kernel.astype(np.float32)), | |
| torch.from_numpy(window[None, :, None].astype(np.float32)) | |
| ) | |
| return result | |
| class ConvSTFT(nn.Module): | |
| def __init__(self, | |
| nfft: int, | |
| win_size: int, | |
| hop_size: int, | |
| win_type: str = "hamming", | |
| feature_type: str = "real", | |
| requires_grad: bool = False): | |
| super(ConvSTFT, self).__init__() | |
| if nfft is None: | |
| self.nfft = int(2**np.ceil(np.log2(win_size))) | |
| else: | |
| self.nfft = nfft | |
| kernel, _ = init_kernels(self.nfft, win_size, hop_size, win_type) | |
| self.weight = nn.Parameter(kernel, requires_grad=requires_grad) | |
| self.win_size = win_size | |
| self.hop_size = hop_size | |
| self.stride = hop_size | |
| self.dim = self.nfft | |
| self.feature_type = feature_type | |
| def forward(self, inputs: torch.Tensor): | |
| if inputs.dim() == 2: | |
| inputs = torch.unsqueeze(inputs, 1) | |
| outputs = F.conv1d(inputs, self.weight, stride=self.stride) | |
| if self.feature_type == "complex": | |
| return outputs | |
| else: | |
| dim = self.dim // 2 + 1 | |
| real = outputs[:, :dim, :] | |
| imag = outputs[:, dim:, :] | |
| mags = torch.sqrt(real**2 + imag**2) | |
| phase = torch.atan2(imag, real) | |
| return mags, phase | |
| class ConviSTFT(nn.Module): | |
| def __init__(self, | |
| win_size: int, | |
| hop_size: int, | |
| nfft: int = None, | |
| win_type: str = "hamming", | |
| feature_type: str = "real", | |
| requires_grad: bool = False): | |
| super(ConviSTFT, self).__init__() | |
| if nfft is None: | |
| self.nfft = int(2**np.ceil(np.log2(win_size))) | |
| else: | |
| self.nfft = nfft | |
| kernel, window = init_kernels(self.nfft, win_size, hop_size, win_type, inverse=True) | |
| self.weight = nn.Parameter(kernel, requires_grad=requires_grad) | |
| self.win_size = win_size | |
| self.hop_size = hop_size | |
| self.win_type = win_type | |
| self.stride = hop_size | |
| self.dim = self.nfft | |
| self.feature_type = feature_type | |
| self.register_buffer("window", window) | |
| self.register_buffer("enframe", torch.eye(win_size)[:, None, :]) | |
| def forward(self, | |
| inputs: torch.Tensor, | |
| phase: torch.Tensor = None): | |
| """ | |
| :param inputs: torch.Tensor, shape: [b, n+2, t] (complex spec) or [b, n//2+1, t] (mags) | |
| :param phase: torch.Tensor, shape: [b, n//2+1, t] | |
| :return: | |
| """ | |
| if phase is not None: | |
| real = inputs * torch.cos(phase) | |
| imag = inputs * torch.sin(phase) | |
| inputs = torch.cat([real, imag], 1) | |
| outputs = F.conv_transpose1d(inputs, self.weight, stride=self.stride) | |
| # this is from torch-stft: https://github.com/pseeth/torch-stft | |
| t = self.window.repeat(1, 1, inputs.size(-1))**2 | |
| coff = F.conv_transpose1d(t, self.enframe, stride=self.stride) | |
| outputs = outputs / (coff + 1e-8) | |
| return outputs | |
| def main(): | |
| stft = ConvSTFT(win_size=512, hop_size=200, feature_type="complex") | |
| istft = ConviSTFT(win_size=512, hop_size=200, feature_type="complex") | |
| mixture = torch.rand(size=(1, 8000*40), dtype=torch.float32) | |
| spec = stft.forward(mixture) | |
| # shape: [batch_size, freq_bins, time_steps] | |
| print(spec.shape) | |
| waveform = istft.forward(spec) | |
| # shape: [batch_size, channels, num_samples] | |
| print(waveform.shape) | |
| return | |
| if __name__ == "__main__": | |
| main() | |