Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| class SubbandSTFT: | |
| def __init__(self, config): | |
| self.n_fft = config.n_fft | |
| self.hop_length = config.hop_length | |
| self.window = torch.hann_window(window_length=self.n_fft, periodic=True) | |
| self.dim_f = config.dim_f | |
| def __call__(self, x): | |
| window = self.window.to(x.device) | |
| batch_dims = x.shape[:-2] | |
| channels, length = x.shape[-2:] | |
| x = torch.stft( | |
| x.reshape(-1, length), | |
| n_fft=self.n_fft, | |
| hop_length=self.hop_length, | |
| window=window, | |
| center=True, | |
| return_complex=True, | |
| ) | |
| x = torch.view_as_real(x).permute(0, 3, 1, 2) | |
| x = x.reshape(*batch_dims, channels * 2, -1, x.shape[-1]) | |
| return x[..., : self.dim_f, :] | |
| def inverse(self, x): | |
| window = self.window.to(x.device) | |
| batch_dims = x.shape[:-3] | |
| channels, freq_bins, time_bins = x.shape[-3:] | |
| full_freq_bins = self.n_fft // 2 + 1 | |
| f_pad = torch.zeros([*batch_dims, channels, full_freq_bins - freq_bins, time_bins]).to(x.device) | |
| x = torch.cat([x, f_pad], -2) | |
| x = x.reshape(-1, 2, full_freq_bins, time_bins).permute(0, 2, 3, 1) | |
| x = x[..., 0] + x[..., 1] * 1.0j | |
| x = torch.istft( | |
| x, | |
| n_fft=self.n_fft, | |
| hop_length=self.hop_length, | |
| window=window, | |
| center=True, | |
| ) | |
| return x.reshape([*batch_dims, 2, -1]) | |
| def get_activation(act_type): | |
| if act_type == "gelu": | |
| return nn.GELU() | |
| if act_type == "relu": | |
| return nn.ReLU() | |
| if act_type[:3] == "elu": | |
| alpha = float(act_type.replace("elu", "")) | |
| return nn.ELU(alpha) | |
| raise Exception | |
| def cac_to_cws(x, num_subbands): | |
| batch, channels, freq_bins, time_bins = x.shape | |
| return x.reshape(batch, channels * num_subbands, freq_bins // num_subbands, time_bins) | |
| def cws_to_cac(x, num_subbands): | |
| batch, channels, freq_bins, time_bins = x.shape | |
| return x.reshape(batch, channels // num_subbands, freq_bins * num_subbands, time_bins) | |
| def forward_subband_mask_model(module, x, core_fn): | |
| x = module.stft(x) | |
| mix = x = cac_to_cws(x, module.num_subbands) | |
| first_conv_out = x = module.first_conv(x) | |
| x = core_fn(x.transpose(-1, -2)).transpose(-1, -2) | |
| x = x * first_conv_out | |
| x = module.final_conv(torch.cat([mix, x], 1)) | |
| x = cws_to_cac(x, module.num_subbands) | |
| if module.num_target_instruments > 1: | |
| batch, channels, freq_bins, time_bins = x.shape | |
| x = x.reshape(batch, module.num_target_instruments, -1, freq_bins, time_bins) | |
| return module.stft.inverse(x) | |