import torch from torch import nn import torch.nn.functional as F from uvr5_pack.lib_v5 import layers from uvr5_pack.lib_v5 import spec_utils class BaseASPPNet(nn.Module): def __init__(self, nin, ch, dilations=(4, 8, 16)): super(BaseASPPNet, self).__init__() self.enc1 = layers.Encoder(nin, ch, 3, 2, 1) self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1) self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1) self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1) self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations) self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1) self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1) self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1) self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1) def __call__(self, x): h, e1 = self.enc1(x) h, e2 = self.enc2(h) h, e3 = self.enc3(h) h, e4 = self.enc4(h) h = self.aspp(h) h = self.dec4(h, e4) h = self.dec3(h, e3) h = self.dec2(h, e2) h = self.dec1(h, e1) return h class CascadedASPPNet(nn.Module): def __init__(self, n_fft): super(CascadedASPPNet, self).__init__() self.stg1_low_band_net = BaseASPPNet(2, 16) self.stg1_high_band_net = BaseASPPNet(2, 16) self.stg2_bridge = layers.Conv2DBNActiv(18, 8, 1, 1, 0) self.stg2_full_band_net = BaseASPPNet(8, 16) self.stg3_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0) self.stg3_full_band_net = BaseASPPNet(16, 32) self.out = nn.Conv2d(32, 2, 1, bias=False) self.aux1_out = nn.Conv2d(16, 2, 1, bias=False) self.aux2_out = nn.Conv2d(16, 2, 1, bias=False) self.max_bin = n_fft // 2 self.output_bin = n_fft // 2 + 1 self.offset = 128 def forward(self, x, aggressiveness=None): mix = x.detach() x = x.clone() x = x[:, :, :self.max_bin] bandw = x.size()[2] // 2 aux1 = torch.cat([ self.stg1_low_band_net(x[:, :, :bandw]), self.stg1_high_band_net(x[:, :, bandw:]) ], dim=2) h = torch.cat([x, aux1], dim=1) aux2 = self.stg2_full_band_net(self.stg2_bridge(h)) h = torch.cat([x, aux1, aux2], dim=1) h = self.stg3_full_band_net(self.stg3_bridge(h)) mask = torch.sigmoid(self.out(h)) mask = F.pad( input=mask, pad=(0, 0, 0, self.output_bin - mask.size()[2]), mode='replicate') if self.training: aux1 = torch.sigmoid(self.aux1_out(aux1)) aux1 = F.pad( input=aux1, pad=(0, 0, 0, self.output_bin - aux1.size()[2]), mode='replicate') aux2 = torch.sigmoid(self.aux2_out(aux2)) aux2 = F.pad( input=aux2, pad=(0, 0, 0, self.output_bin - aux2.size()[2]), mode='replicate') return mask * mix, aux1 * mix, aux2 * mix else: if aggressiveness: mask[:, :, :aggressiveness['split_bin']] = torch.pow(mask[:, :, :aggressiveness['split_bin']], 1 + aggressiveness['value'] / 3) mask[:, :, aggressiveness['split_bin']:] = torch.pow(mask[:, :, aggressiveness['split_bin']:], 1 + aggressiveness['value']) return mask * mix def predict(self, x_mag, aggressiveness=None): h = self.forward(x_mag, aggressiveness) if self.offset > 0: h = h[:, :, :, self.offset:-self.offset] assert h.size()[3] > 0 return h