Spaces:
Running on Zero
Running on Zero
File size: 2,684 Bytes
cafad09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 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)
|