chore: import upstream snapshot with attribution

modules/aux_decoder/__init__.py

@@ -0,0 +1,70 @@
+import torch.nn
+from torch import nn
+
+from .convnext import ConvNeXtDecoder
+from utils import filter_kwargs
+
+AUX_DECODERS = {
+    'convnext': ConvNeXtDecoder
+}
+AUX_LOSSES = {
+    'convnext': nn.L1Loss
+}
+
+
+def build_aux_decoder(
+        in_dims: int, out_dims: int,
+        aux_decoder_arch: str, aux_decoder_args: dict
+) -> torch.nn.Module:
+    decoder_cls = AUX_DECODERS[aux_decoder_arch]
+    kwargs = filter_kwargs(aux_decoder_args, decoder_cls)
+    return AUX_DECODERS[aux_decoder_arch](in_dims, out_dims, **kwargs)
+
+
+def build_aux_loss(aux_decoder_arch):
+    return AUX_LOSSES[aux_decoder_arch]()
+
+
+class AuxDecoderAdaptor(nn.Module):
+    def __init__(self, in_dims: int, out_dims: int, num_feats: int,
+                 spec_min: list, spec_max: list,
+                 aux_decoder_arch: str, aux_decoder_args: dict):
+        super().__init__()
+        self.decoder = build_aux_decoder(
+            in_dims=in_dims, out_dims=out_dims * num_feats,
+            aux_decoder_arch=aux_decoder_arch,
+            aux_decoder_args=aux_decoder_args
+        )
+        self.out_dims = out_dims
+        self.n_feats = num_feats
+        if spec_min is not None and spec_max is not None:
+            # spec: [B, T, M] or [B, F, T, M]
+            # spec_min and spec_max: [1, 1, M] or [1, 1, F, M] => transpose(-3, -2) => [1, 1, M] or [1, F, 1, M]
+            spec_min = torch.FloatTensor(spec_min)[None, None, :].transpose(-3, -2)
+            spec_max = torch.FloatTensor(spec_max)[None, None, :].transpose(-3, -2)
+            self.register_buffer('spec_min', spec_min, persistent=False)
+            self.register_buffer('spec_max', spec_max, persistent=False)
+
+    def norm_spec(self, x):
+        k = (self.spec_max - self.spec_min) / 2.
+        b = (self.spec_max + self.spec_min) / 2.
+        return (x - b) / k
+
+    def denorm_spec(self, x):
+        k = (self.spec_max - self.spec_min) / 2.
+        b = (self.spec_max + self.spec_min) / 2.
+        return x * k + b
+
+    def forward(self, condition, infer=False):
+        x = self.decoder(condition, infer=infer)  # [B, T, F x C]
+
+        if self.n_feats > 1:
+            # This is the temporary solution since PyTorch 1.13
+            # does not support exporting aten::unflatten to ONNX
+            # x = x.unflatten(dim=2, sizes=(self.n_feats, self.in_dims))
+            x = x.reshape(-1, x.shape[1], self.n_feats, self.out_dims)  # [B, T, F, C]
+            x = x.transpose(1, 2)  # [B, F, T, C]
+        if infer:
+            x = self.denorm_spec(x)
+
+        return x  # [B, T, C] or [B, F, T, C]

modules/aux_decoder/convnext.py

@@ -0,0 +1,89 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from modules.commons.common_layers import AdamWConv1d
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+            self,
+            dim: int,
+            intermediate_dim: int,
+            layer_scale_init_value: Optional[float] = None, drop_out: float = 0.0
+
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+        # self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path = nn.Identity()
+        self.dropout = nn.Dropout(drop_out) if drop_out > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor, ) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+        x = self.dropout(x)
+
+        x = residual + self.drop_path(x)
+        return x
+
+
+class ConvNeXtDecoder(nn.Module):
+    def __init__(
+            self, in_dims, out_dims, /, *,
+            num_channels=512, num_layers=6, kernel_size=7, dropout_rate=0.1
+    ):
+        super().__init__()
+        self.inconv = nn.Conv1d(
+            in_dims, num_channels, kernel_size,
+            stride=1, padding=(kernel_size - 1) // 2
+        )
+        self.conv = nn.ModuleList(
+            ConvNeXtBlock(
+                dim=num_channels, intermediate_dim=num_channels * 4,
+                layer_scale_init_value=1e-6, drop_out=dropout_rate
+            ) for _ in range(num_layers)
+        )
+        self.outconv = AdamWConv1d(
+            num_channels, out_dims, kernel_size,
+            stride=1, padding=(kernel_size - 1) // 2
+        )
+
+    # noinspection PyUnusedLocal
+    def forward(self, x, infer=False):
+        x = x.transpose(1, 2)
+        x = self.inconv(x)
+        for conv in self.conv:
+            x = conv(x)
+        x = self.outconv(x)
+        x = x.transpose(1, 2)
+        return x