lt-tensor 0.0.1a34__py3-none-any.whl → 0.0.1a35__py3-none-any.whl

lt_tensor/model_zoo/audio_models/istft/__init__.py

@@ -2,10 +2,9 @@ __all__ = ["iSTFTNetGenerator", "iSTFTNetConfig"]
 from lt_utils.common import *
 from lt_tensor.torch_commons import *
 from lt_tensor.model_zoo.convs import ConvNets
-from torch.nn import functional as F
 from lt_tensor.config_templates import ModelConfig
-from lt_tensor.misc_utils import get_config, get_weights
-from lt_utils.file_ops import load_json, is_file, is_dir, is_path_valid
+from lt_utils.file_ops import is_file
+from lt_tensor.model_zoo.audio_models.resblocks import ResBlock1, ResBlock2
 
 
 class iSTFTNetConfig(ModelConfig):
@@ -22,6 +21,7 @@ class iSTFTNetConfig(ModelConfig):
     ]
 
     activation: nn.Module = nn.LeakyReLU(0.1)
+    resblock_activation: nn.Module = nn.LeakyReLU(0.1)
     resblock: int = 0
     gen_istft_n_fft: int = 16
     sampling_rate: Number = 24000
@@ -39,7 +39,8 @@ class iSTFTNetConfig(ModelConfig):
             [1, 3, 5],
         ],
         activation: nn.Module = nn.LeakyReLU(0.1),
-        resblock: Union[int, str] = "1",
+        resblock_activation: nn.Module = nn.LeakyReLU(0.1),
+        resblock: int = 0,
         gen_istft_n_fft: int = 16,
         sampling_rate: Number = 24000,
         *args,
@@ -53,6 +54,7 @@ class iSTFTNetConfig(ModelConfig):
             "resblock_kernel_sizes": resblock_kernel_sizes,
             "resblock_dilation_sizes": resblock_dilation_sizes,
             "activation": activation,
+            "resblock_activation": resblock_activation,
             "resblock": resblock,
             "gen_istft_n_fft": gen_istft_n_fft,
             "sampling_rate": sampling_rate,
@@ -62,134 +64,6 @@ class iSTFTNetConfig(ModelConfig):
     def post_process(self):
         if isinstance(self.resblock, str):
             self.resblock = 0 if self.resblock == "1" else 1
-            
-            
-def get_padding(ks, d):
-    return int((ks * d - d) / 2)
-
-
-class ResBlock1(ConvNets):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
-        super().__init__()
-        self.convs1 = nn.ModuleList(
-            [
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[0],
-                        padding=get_padding(kernel_size, dilation[0]),
-                    )
-                ),
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[1],
-                        padding=get_padding(kernel_size, dilation[1]),
-                    )
-                ),
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[2],
-                        padding=get_padding(kernel_size, dilation[2]),
-                    )
-                ),
-            ]
-        )
-        self.convs1.apply(self.init_weights)
-
-        self.convs2 = nn.ModuleList(
-            [
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=get_padding(kernel_size, 1),
-                    )
-                ),
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=get_padding(kernel_size, 1),
-                    )
-                ),
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=get_padding(kernel_size, 1),
-                    )
-                ),
-            ]
-        )
-        self.activation = nn.LeakyReLU(0.1)
-        self.convs2.apply(self.init_weights)
-
-    def forward(self, x):
-        for c1, c2 in zip(self.convs1, self.convs2):
-            xt = self.activation(x)
-            xt = c1(xt)
-            xt = self.activation(xt)
-            xt = c2(xt)
-            x = xt + x
-        return x
-
-
-class ResBlock2(ConvNets):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
-        super().__init__()
-        self.convs = nn.ModuleList(
-            [
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[0],
-                        padding=get_padding(kernel_size, dilation[0]),
-                    )
-                ),
-                weight_norm(
-                    nn.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[1],
-                        padding=get_padding(kernel_size, dilation[1]),
-                    )
-                ),
-            ]
-        )
-        self.activation = nn.LeakyReLU(0.1)
-        self.convs.apply(self.init_weights)
-
-    def forward(self, x):
-        for c in self.convs:
-            xt = self.activation(x)
-            xt = c(xt)
-            x = xt + x
-        return x
 
 
 class iSTFTNetGenerator(ConvNets):
@@ -241,7 +115,7 @@ class iSTFTNetGenerator(ConvNets):
             for j, (k, d) in enumerate(
                 zip(cfg.resblock_kernel_sizes, cfg.resblock_dilation_sizes)
             ):
-                self.resblocks.append(resblock(ch, k, d))
+                self.resblocks.append(resblock(ch, k, d, cfg.resblock_activation))
 
         self.post_n_fft = cfg.gen_istft_n_fft
         self.conv_post = weight_norm(
@@ -252,7 +126,7 @@ class iSTFTNetGenerator(ConvNets):
         self.activation = cfg.activation
         self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         x = self.conv_pre(x)
         for i in range(self.num_upsamples):
             x = self.activation(x)
@@ -283,7 +157,7 @@ class iSTFTNetGenerator(ConvNets):
         **pickle_load_args,
     ):
         try:
-            incompatible_keys = super().load_weights(
+            return super().load_weights(
                 path,
                 raise_if_not_exists,
                 strict,
@@ -292,18 +166,6 @@ class iSTFTNetGenerator(ConvNets):
                 mmap,
                 **pickle_load_args,
             )
-            if incompatible_keys:
-                self.remove_norms()
-                incompatible_keys = super().load_weights(
-                    path,
-                    raise_if_not_exists,
-                    strict,
-                    assign,
-                    weights_only,
-                    mmap,
-                    **pickle_load_args,
-                )
-            return incompatible_keys
         except RuntimeError:
             self.remove_norms()
             return super().load_weights(
@@ -322,6 +184,7 @@ class iSTFTNetGenerator(ConvNets):
         model_file: PathLike,
         model_config: Union[iSTFTNetConfig, Dict[str, Any]],
         *,
+        remove_norms: bool = False,
         strict: bool = False,
         map_location: str = "cpu",
         weights_only: bool = False,
@@ -339,11 +202,11 @@ class iSTFTNetGenerator(ConvNets):
             h = iSTFTNetConfig(**model_config)
 
         model = cls(h)
+        if remove_norms:
+            model.remove_norms()
         try:
-            incompatible_keys = model.load_state_dict(model_state_dict, strict=strict)
-            if incompatible_keys:
-                model.remove_norms()
-                model.load_state_dict(model_state_dict, strict=strict)
+            model.load_state_dict(model_state_dict, strict=strict)
+            return model
         except RuntimeError:
             print(
                 f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!"
@@ -351,201 +214,3 @@ class iSTFTNetGenerator(ConvNets):
             model.remove_norms()
             model.load_state_dict(model_state_dict, strict=strict)
         return model
-
-
-class DiscriminatorP(ConvNets):
-    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
-        super().__init__()
-        self.period = period
-        self.activation = nn.LeakyReLU(0.1)
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-        self.convs = nn.ModuleList(
-            [
-                norm_f(
-                    nn.Conv2d(
-                        1,
-                        32,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(5, 1), 0),
-                    )
-                ),
-                norm_f(
-                    nn.Conv2d(
-                        32,
-                        128,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(5, 1), 0),
-                    )
-                ),
-                norm_f(
-                    nn.Conv2d(
-                        128,
-                        512,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(5, 1), 0),
-                    )
-                ),
-                norm_f(
-                    nn.Conv2d(
-                        512,
-                        1024,
-                        (kernel_size, 1),
-                        (stride, 1),
-                        padding=(get_padding(5, 1), 0),
-                    )
-                ),
-                norm_f(nn.Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
-            ]
-        )
-        self.conv_post = norm_f(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
-
-    def forward(self, x):
-        fmap = []
-
-        # 1d to 2d
-        b, c, t = x.shape
-        if t % self.period != 0:  # pad first
-            n_pad = self.period - (t % self.period)
-            x = F.pad(x, (0, n_pad), "reflect")
-            t = t + n_pad
-        x = x.view(b, c, t // self.period, self.period)
-
-        for l in self.convs:
-            x = l(x)
-            x = self.activation(x)
-            fmap.append(x)
-        x = self.conv_post(x)
-        fmap.append(x)
-        x = torch.flatten(x, 1, -1)
-
-        return x, fmap
-
-
-class MultiPeriodDiscriminator(ConvNets):
-    def __init__(self):
-        super().__init__()
-        self.discriminators = nn.ModuleList(
-            [
-                DiscriminatorP(2),
-                DiscriminatorP(3),
-                DiscriminatorP(5),
-                DiscriminatorP(7),
-                DiscriminatorP(11),
-            ]
-        )
-
-    def forward(self, y, y_hat):
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-        for i, d in enumerate(self.discriminators):
-            y_d_r, fmap_r = d(y)
-            y_d_g, fmap_g = d(y_hat)
-            y_d_rs.append(y_d_r)
-            fmap_rs.append(fmap_r)
-            y_d_gs.append(y_d_g)
-            fmap_gs.append(fmap_g)
-
-        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
-
-
-class DiscriminatorS(ConvNets):
-    def __init__(self, use_spectral_norm=False):
-        super().__init__()
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-        self.convs = nn.ModuleList(
-            [
-                norm_f(nn.Conv1d(1, 128, 15, 1, padding=7)),
-                norm_f(nn.Conv1d(128, 128, 41, 2, groups=4, padding=20)),
-                norm_f(nn.Conv1d(128, 256, 41, 2, groups=16, padding=20)),
-                norm_f(nn.Conv1d(256, 512, 41, 4, groups=16, padding=20)),
-                norm_f(nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
-                norm_f(nn.Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
-                norm_f(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
-            ]
-        )
-        self.activation = nn.LeakyReLU(0.1)
-        self.conv_post = norm_f(nn.Conv1d(1024, 1, 3, 1, padding=1))
-
-    def forward(self, x):
-        fmap = []
-        for l in self.convs:
-            x = self.activation(l(x))
-
-            fmap.append(x)
-        x = self.conv_post(x)
-        fmap.append(x)
-        x = torch.flatten(x, 1, -1)
-
-        return x, fmap
-
-
-class MultiScaleDiscriminator(ConvNets):
-    def __init__(self):
-        super().__init__()
-        self.discriminators = nn.ModuleList(
-            [
-                DiscriminatorS(use_spectral_norm=True),
-                DiscriminatorS(),
-                DiscriminatorS(),
-            ]
-        )
-        self.meanpools = nn.ModuleList(
-            [nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)]
-        )
-
-    def forward(self, y, y_hat):
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-        for i, d in enumerate(self.discriminators):
-            if i != 0:
-                y = self.meanpools[i - 1](y)
-                y_hat = self.meanpools[i - 1](y_hat)
-            y_d_r, fmap_r = d(y)
-            y_d_g, fmap_g = d(y_hat)
-            y_d_rs.append(y_d_r)
-            fmap_rs.append(fmap_r)
-            y_d_gs.append(y_d_g)
-            fmap_gs.append(fmap_g)
-
-        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
-
-
-def feature_loss(fmap_r, fmap_g):
-    loss = 0
-    for dr, dg in zip(fmap_r, fmap_g):
-        for rl, gl in zip(dr, dg):
-            loss += torch.mean(torch.abs(rl - gl))
-
-    return loss * 2
-
-
-def discriminator_loss(disc_real_outputs, disc_generated_outputs):
-    loss = 0
-    r_losses = []
-    g_losses = []
-    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
-        r_loss = torch.mean((1 - dr) ** 2)
-        g_loss = torch.mean(dg**2)
-        loss += r_loss + g_loss
-        r_losses.append(r_loss.item())
-        g_losses.append(g_loss.item())
-
-    return loss, r_losses, g_losses
-
-
-def generator_loss(disc_outputs):
-    loss = 0
-    gen_losses = []
-    for dg in disc_outputs:
-        l = torch.mean((1 - dg) ** 2)
-        gen_losses.append(l)
-        loss += l
-
-    return loss, gen_losses

lt_tensor/model_zoo/audio_models/resblocks.py

@@ -0,0 +1,248 @@
+from lt_utils.common import *
+from torch.nn import functional as F
+from lt_tensor.torch_commons import *
+from lt_tensor.model_zoo.convs import ConvNets
+from lt_tensor.model_zoo.activations import snake, alias_free
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class ResBlock1(ConvNets):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs1.apply(self.init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in dilation
+            ]
+        )
+        self.convs2.apply(self.init_weights)
+        self.activation = activation
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = c1(self.activation(x))
+            x = c2(self.activation(xt)) + x
+        return x
+
+
+class ResBlock2(ConvNets):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs.apply(self.init_weights)
+        self.activation = activation
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = c(self.activation(x))
+            x = xt + x
+        return x
+
+
+def get_snake(name: Literal["snake", "snakebeta"] = "snake"):
+    assert name.lower() in [
+        "snake",
+        "snakebeta",
+    ], f"'{name}' is not a valid snake activation! use 'snake' or 'snakebeta'"
+    if name.lower() == "snake":
+        return snake.Snake
+    return snake.SnakeBeta
+
+
+class AMPBlock1(ConvNets):
+    """Modified from 'https://github.com/NVIDIA/BigVGAN/blob/main/bigvgan.py' under MIT license, found in 'bigvgan/LICENSE'
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
+
+    Args:
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        snake_logscale: (bool): to use logscale with snake activation. Default to True.
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Defaults to 'snakebeta'.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        snake_logscale: bool = True,
+        activation: Literal["snake", "snakebeta"] = "snakebeta",
+    ):
+        super().__init__()
+        actv = get_snake(activation)
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs1.apply(self.init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in range(len(dilation))
+            ]
+        )
+        self.convs2.apply(self.init_weights)
+
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # Total number of conv layers
+
+        # Activation functions
+        self.activations = nn.ModuleList(
+            [
+                alias_free.Activation1d(
+                    activation=actv(channels, alpha_logscale=snake_logscale)
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+
+class AMPBlock2(ConvNets):
+    """Modified from 'https://github.com/NVIDIA/BigVGAN/blob/main/bigvgan.py' under MIT license, found in 'bigvgan/LICENSE'
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1
+
+    Args:
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        snake_logscale: (bool): to use logscale with snake activation. Default to True.
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Defaults to 'snakebeta'.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        snake_logscale: bool = True,
+        activation: Literal["snake", "snakebeta"] = "snakebeta",
+    ):
+        super().__init__()
+        actv = get_snake(activation)
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs.apply(self.init_weights)
+
+        self.num_layers = len(self.convs)  # Total number of conv layers
+
+        # Activation functions
+        self.activations = nn.ModuleList(
+            [
+                alias_free.Activation1d(
+                    activation=actv(channels, alpha_logscale=snake_logscale)
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+    def forward(self, x):
+        for c, a in zip(self.convs, self.activations):
+            xt = a(x)
+            xt = c(xt)
+            x = xt + x
+        return x