optimize(uvr5): remove redundant files

infer/lib/uvr5_pack/lib_v5/layers.py

@@ -26,40 +26,17 @@ class Conv2DBNActiv(nn.Module):
         return self.conv(x)
 
 
-class SeperableConv2DBNActiv(nn.Module):
-    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
-        super(SeperableConv2DBNActiv, self).__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(
-                nin,
-                nin,
-                kernel_size=ksize,
-                stride=stride,
-                padding=pad,
-                dilation=dilation,
-                groups=nin,
-                bias=False,
-            ),
-            nn.Conv2d(nin, nout, kernel_size=1, bias=False),
-            nn.BatchNorm2d(nout),
-            activ(),
-        )
-
-    def __call__(self, x):
-        return self.conv(x)
-
-
 class Encoder(nn.Module):
     def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
         super(Encoder, self).__init__()
-        self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
-        self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ)
+        self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
 
     def __call__(self, x):
-        skip = self.conv1(x)
-        h = self.conv2(skip)
+        h = self.conv1(x)
+        h = self.conv2(h)
 
-        return h, skip
+        return h
 
 
 class Decoder(nn.Module):
@@ -67,15 +44,19 @@ class Decoder(nn.Module):
         self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False
     ):
         super(Decoder, self).__init__()
-        self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        # self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
         self.dropout = nn.Dropout2d(0.1) if dropout else None
 
     def __call__(self, x, skip=None):
         x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
+
         if skip is not None:
             skip = spec_utils.crop_center(skip, x)
             x = torch.cat([x, skip], dim=1)
-        h = self.conv(x)
+
+        h = self.conv1(x)
+        # h = self.conv2(h)
 
         if self.dropout is not None:
             h = self.dropout(h)
@@ -84,25 +65,24 @@ class Decoder(nn.Module):
 
 
 class ASPPModule(nn.Module):
-    def __init__(self, nin, nout, dilations=(4, 8, 16), activ=nn.ReLU):
+    def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False):
         super(ASPPModule, self).__init__()
         self.conv1 = nn.Sequential(
             nn.AdaptiveAvgPool2d((1, None)),
-            Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
+            Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ),
         )
-        self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
-        self.conv3 = SeperableConv2DBNActiv(
-            nin, nin, 3, 1, dilations[0], dilations[0], activ=activ
+        self.conv2 = Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
+        self.conv3 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[0], dilations[0], activ=activ
         )
-        self.conv4 = SeperableConv2DBNActiv(
-            nin, nin, 3, 1, dilations[1], dilations[1], activ=activ
+        self.conv4 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[1], dilations[1], activ=activ
         )
-        self.conv5 = SeperableConv2DBNActiv(
-            nin, nin, 3, 1, dilations[2], dilations[2], activ=activ
-        )
-        self.bottleneck = nn.Sequential(
-            Conv2DBNActiv(nin * 5, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1)
+        self.conv5 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[2], dilations[2], activ=activ
         )
+        self.bottleneck = Conv2DBNActiv(nout * 5, nout, 1, 1, 0, activ=activ)
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
 
     def forward(self, x):
         _, _, h, w = x.size()
@@ -114,5 +94,32 @@ class ASPPModule(nn.Module):
         feat4 = self.conv4(x)
         feat5 = self.conv5(x)
         out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
-        bottle = self.bottleneck(out)
-        return bottle
+        out = self.bottleneck(out)
+
+        if self.dropout is not None:
+            out = self.dropout(out)
+
+        return out
+
+
+class LSTMModule(nn.Module):
+    def __init__(self, nin_conv, nin_lstm, nout_lstm):
+        super(LSTMModule, self).__init__()
+        self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
+        self.lstm = nn.LSTM(
+            input_size=nin_lstm, hidden_size=nout_lstm // 2, bidirectional=True
+        )
+        self.dense = nn.Sequential(
+            nn.Linear(nout_lstm, nin_lstm), nn.BatchNorm1d(nin_lstm), nn.ReLU()
+        )
+
+    def forward(self, x):
+        N, _, nbins, nframes = x.size()
+        h = self.conv(x)[:, 0]  # N, nbins, nframes
+        h = h.permute(2, 0, 1)  # nframes, N, nbins
+        h, _ = self.lstm(h)
+        h = self.dense(h.reshape(-1, h.size()[-1]))  # nframes * N, nbins
+        h = h.reshape(nframes, N, 1, nbins)
+        h = h.permute(1, 2, 3, 0)
+
+        return h