lt-tensor 0.0.1a14__py3-none-any.whl → 0.0.1a16__py3-none-any.whl

lt_tensor/model_zoo/features.py

@@ -323,8 +323,11 @@ class AudioEncoder(Model):
 
     def __init__(
         self,
-        channels: int = 80,
+        channels: int,
         alpha: float = 4.0,
+        feat_channels: int = 64,
+        out_features: Optional[int] = None,
+        out_channels: int = 1,
         interp_mode: Literal[
             "nearest",
             "linear",
@@ -338,16 +341,60 @@ class AudioEncoder(Model):
 
         self.net = nn.Sequential(
             nn.Conv1d(
-                channels, channels, kernel_size=3, stride=2, padding=5, groups=channels
+                channels, feat_channels, kernel_size=3, stride=1, padding=5, groups=1
             ),
             nn.LeakyReLU(0.1),
-            nn.Conv1d(channels, channels, kernel_size=7, stride=1, padding=1, groups=1),
+            nn.Conv1d(
+                feat_channels,
+                feat_channels,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                groups=feat_channels,
+            ),
+            nn.LeakyReLU(0.1),
+            nn.Conv1d(
+                feat_channels,
+                feat_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=feat_channels // 8,
+            ),
+            nn.LeakyReLU(0.1),
+            nn.Conv1d(
+                feat_channels,
+                feat_channels,
+                kernel_size=7,
+                stride=1,
+                padding=1,
+                groups=1,
+            ),
         )
-        self.fc = nn.Linear(channels, channels)
+        self.fc = nn.Linear(feat_channels, channels)
+        self.feat_channels = feat_channels
         self.activation = activation
         self.channels = channels
         self.mode = interp_mode
         self.alpha = alpha
+        self.post_conv = nn.Conv1d(
+            channels,
+            out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            dilation=1,
+            groups=1,
+            bias=True,
+        )
+        if out_features is not None:
+            self.format_out = lambda tensor: F.interpolate(
+                tensor,
+                size=out_features,
+                mode=interp_mode,
+            )
+        else:
+            self.format_out = nn.Identity()
 
     def forward(self, mels: Tensor, cr_audio: Tensor):
         sin = torch.asin(cr_audio)
@@ -367,14 +414,20 @@ class AudioEncoder(Model):
             .contiguous()
         )
         x = self.activation(x)
-        return self.fc(x).transpose(-1, -2)
+
+        xt = self.fc(x).transpose(-1, -2)
+        out = self.post_conv(xt)
+        return self.format_out(out)
 
 
 class AudioEncoderAttn(Model):
     def __init__(
         self,
-        channels: int = 80,
+        channels: int,
+        feat_channels: int = 64,
         alpha: float = 4.0,
+        out_channels: Optional[int] = None,
+        out_features: int = 1,
         interp_mode: Literal[
             "nearest",
             "linear",
@@ -388,16 +441,54 @@ class AudioEncoderAttn(Model):
 
         self.net = nn.Sequential(
             nn.Conv1d(
-                channels, channels, kernel_size=3, stride=2, padding=5, groups=channels
+                channels, feat_channels, kernel_size=3, stride=1, padding=1, groups=1
+            ),
+            nn.LeakyReLU(0.1),
+            nn.Conv1d(
+                feat_channels,
+                feat_channels,
+                kernel_size=3,
+                stride=2,
+                padding=5,
+                groups=feat_channels,
             ),
             nn.LeakyReLU(0.1),
-            nn.Conv1d(channels, channels, kernel_size=7, stride=1, padding=1, groups=1),
+            nn.Conv1d(
+                feat_channels,
+                feat_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=feat_channels // 8,
+            ),
+            nn.LeakyReLU(0.1),
+            nn.Conv1d(
+                feat_channels, channels, kernel_size=7, stride=1, padding=1, groups=1
+            ),
         )
         self.fusion = CrossAttentionFusion(channels, channels, 2, d_model=channels)
         self.channels = channels
         self.mode = interp_mode
         self.alpha = alpha
         self.activation = activation
+        self.post_conv = nn.Conv1d(
+            channels,
+            out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            dilation=1,
+            groups=1,
+            bias=True,
+        )
+        if out_features is not None:
+            self.format_out = lambda tensor: F.interpolate(
+                tensor,
+                size=out_features,
+                mode=interp_mode,
+            )
+        else:
+            self.format_out = nn.Identity()
 
     def forward(self, mels: Tensor, cr_audio: Tensor):
         sin = torch.asin(cr_audio)
@@ -408,9 +499,9 @@ class AudioEncoderAttn(Model):
         )
         x = self.activation(self.net(mod))
         x = F.interpolate(x, size=mels.shape[-1], mode=self.mode)
-
-        # Ensure contiguous before transpose
         x_t = x.transpose(-2, -1).contiguous()
         mels_t = mels.transpose(-2, -1).contiguous()
 
-        return self.fusion(x_t, mels_t).transpose(-2, -1)
+        xt = self.fusion(x_t, mels_t).transpose(-2, -1)
+        out = self.post_conv(xt)
+        return self.format_out(out)

lt_tensor/model_zoo/residual.py

@@ -5,7 +5,9 @@ __all__ = [
     "ResBlock2D",
     "ResBlock1DShuffled",
     "AdaResBlock1D",
-    "ResBlocks",
+    "ResBlocks1D",
+    "ResBlock1D2",
+    "ShuffleBlock2D",
 ]
 import math
 from lt_utils.common import *
@@ -45,38 +47,8 @@ class ConvNets(Model):
             m.weight.data.normal_(mean, std)
 
 
-class ResBlocks(ConvNets):
-    def __init__(
-        self,
-        channels: int,
-        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
-        resblock_dilation_sizes: List[Union[int, List[int]]] = [
-            [1, 3, 5],
-            [1, 3, 5],
-            [1, 3, 5],
-        ],
-        activation: nn.Module = nn.LeakyReLU(0.1),
-    ):
-        super().__init__()
-        self.num_kernels = len(resblock_kernel_sizes)
-        self.rb = nn.ModuleList()
-        self.activation = activation
-
-        for k, j in zip(resblock_kernel_sizes, resblock_dilation_sizes):
-            self.rb.append(ResBlock1D(channels, k, j, activation))
-
-        self.rb.apply(self.init_weights)
-
-    def forward(self, x: torch.Tensor):
-        xs = None
-        for i, block in enumerate(self.rb):
-            if i == 0:
-                xs = block(x)
-            else:
-                xs += block(x)
-        x = xs / self.num_kernels
-        return x
-
+def get_padding(ks, d):
+    return int((ks * d - d) / 2)
 
 
 class ResBlock1D(ConvNets):
@@ -92,14 +64,13 @@ class ResBlock1D(ConvNets):
         self.conv_nets = nn.ModuleList(
             [
                 self._get_conv_layer(i, channels, kernel_size, 1, dilation, activation)
-                for i in range(3)
+                for i in range(len(dilation))
             ]
         )
         self.conv_nets.apply(self.init_weights)
         self.last_index = len(self.conv_nets) - 1
 
     def _get_conv_layer(self, id, ch, k, stride, d, actv):
-        get_padding = lambda ks, d: int((ks * d - d) / 2)
         return nn.Sequential(
             actv,  # 1
             weight_norm(
@@ -126,16 +97,11 @@ class ResBlock1DShuffled(ConvNets):
         kernel_size=3,
         dilation=(1, 3, 5),
         activation: nn.Module = nn.LeakyReLU(0.1),
-        add_channel_shuffle: bool = False,  # requires pytorch 2.7.0 +
         channel_shuffle_groups=1,
     ):
         super().__init__()
 
-        self.channel_shuffle = (
-            nn.ChannelShuffle(channel_shuffle_groups)
-            if add_channel_shuffle
-            else nn.Identity()
-        )
+        self.channel_shuffle = nn.ChannelShuffle(channel_shuffle_groups)
 
         self.conv_nets = nn.ModuleList(
             [
@@ -171,29 +137,67 @@ class ResBlock1DShuffled(ConvNets):
 class ResBlock2D(Model):
     def __init__(
         self,
-        in_channels,
-        out_channels,
-        downsample=False,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        hidden_dim: int = 32,
+        downscale: bool = False,
+        activation: nn.Module = nn.LeakyReLU(0.2),
     ):
         super().__init__()
-        stride = 2 if downsample else 1
+        stride = 2 if downscale else 1
+        if out_channels is None:
+            out_channels = in_channels
 
         self.block = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, 3, stride, 1),
-            nn.LeakyReLU(0.2),
-            nn.Conv2d(out_channels, out_channels, 3, 1, 1),
+            nn.Conv2d(in_channels, hidden_dim, 3, stride, 1),
+            activation,
+            nn.Conv2d(hidden_dim, hidden_dim, 7, 1, 3),
+            activation,
+            nn.Conv2d(hidden_dim, out_channels, 3, 1, 1),
         )
 
         self.skip = nn.Identity()
-        if downsample or in_channels != out_channels:
+        if downscale or in_channels != out_channels:
             self.skip = spectral_norm_select(
                 nn.Conv2d(in_channels, out_channels, 1, stride)
             )
-        # on less to be handled every cicle
+        # on less to be handled every cycle
         self.sqrt_2 = math.sqrt(2)
 
     def forward(self, x: Tensor):
-        return (self.block(x) + self.skip(x)) / self.sqrt_2
+        return x + ((self.block(x) + self.skip(x)) / self.sqrt_2)
+
+
+class ShuffleBlock2D(ConvNets):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        hidden_dim: int = 32,
+        downscale: bool = False,
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+        if out_channels is None:
+            out_channels = channels
+        self.shuffle = nn.ChannelShuffle(groups=2)
+        self.ch_split = lambda tensor: torch.split(tensor, 1, dim=1)
+        self.activation = activation
+        self.resblock_2d = ResBlock2D(
+            channels, out_channels, hidden_dim, downscale, activation
+        )
+
+    def shuffle_channels(self, tensor: torch.Tensor):
+        with torch.no_grad():
+            x = F.channel_shuffle(tensor.transpose(1, -1), tensor.shape[1]).transpose(
+                -1, 1
+            )
+        return self.ch_split(x)
+
+    def forward(self, x: torch.Tensor):
+        ch1, ch2 = self.shuffle_channels(x)
+        ch2 = self.resblock_2d(ch2)
+        return torch.cat((ch1, ch2), dim=1)
 
 
 class AdaResBlock1D(ConvNets):
@@ -207,46 +211,111 @@ class AdaResBlock1D(ConvNets):
     ):
         super().__init__()
 
+        self.alpha1 = nn.ModuleList()
+        self.alpha2 = nn.ModuleList()
         self.conv_nets = nn.ModuleList(
             [
                 self._get_conv_layer(
-                    i,
+                    d,
                     res_block_channels,
                     ada_channel_in,
                     kernel_size,
-                    1,
-                    dilation,
                 )
-                for i in range(3)
+                for d in dilation
             ]
         )
         self.conv_nets.apply(self.init_weights)
         self.last_index = len(self.conv_nets) - 1
         self.activation = activation
 
-    def _get_conv_layer(self, id, ch, ada_ch, k, stride, d):
-        get_padding = lambda ks, d: int((ks * d - d) / 2)
+    def _get_conv_layer(self, d, ch, ada_ch, k):
+        self.alpha1.append(nn.Parameter(torch.ones(1, ada_ch, 1)))
+        self.alpha2.append(nn.Parameter(torch.ones(1, ada_ch, 1)))
         return nn.ModuleDict(
             dict(
                 norm1=AdaFusion1D(ada_ch, ch),
                 norm2=AdaFusion1D(ada_ch, ch),
-                alpha1=nn.Parameter(torch.ones(1, ada_ch, 1)),
-                alpha2=nn.Parameter(torch.ones(1, ada_ch, 1)),
                 conv1=weight_norm(
                     nn.Conv1d(
-                        ch, ch, k, stride, dilation=d[id], padding=get_padding(k, d[id])
+                        ch, ch, k, 1, dilation=d, padding=get_padding(k, d)
                     )
                 ),  # 2
                 conv2=weight_norm(
-                    nn.Conv1d(ch, ch, k, stride, dilation=1, padding=get_padding(k, 1))
+                    nn.Conv1d(ch, ch, k, 1, dilation=1, padding=get_padding(k, 1))
                 ),  # 4
             )
         )
 
     def forward(self, x: torch.Tensor, y: torch.Tensor):
-        for cnn in self.conv_nets:
-            xt = self.activation(cnn["norm1"](x, y, cnn["alpha1"]))
+        for i, cnn in enumerate(self.conv_nets):
+            xt = self.activation(cnn["norm1"](x, y, self.alpha1[i]))
             xt = cnn["conv1"](xt)
-            xt = self.activation(cnn["norm2"](xt, y, cnn["alpha2"]))
+            xt = self.activation(cnn["norm2"](xt, y, self.alpha2[i]))
             x = cnn["conv2"](xt) + x
         return x
+
+
+class ResBlock1D2(ConvNets):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in range(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
+
+
+class ResBlocks1D(ConvNets):
+    def __init__(
+        self,
+        channels: int,
+        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
+        resblock_dilation_sizes: List[Union[int, List[int]]] = [
+            [1, 3, 5],
+            [1, 3, 5],
+            [1, 3, 5],
+        ],
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        block: Union[ResBlock1D, ResBlock1D2] = ResBlock1D,
+    ):
+        super().__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.rb = nn.ModuleList()
+        self.activation = activation
+
+        for k, j in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+            self.rb.append(block(channels, k, j, activation))
+
+        self.rb.apply(self.init_weights)
+
+    def forward(self, x: torch.Tensor):
+        xs = None
+        for i, block in enumerate(self.rb):
+            if i == 0:
+                xs = block(x)
+            else:
+                xs += block(x)
+        return xs / self.num_kernels

{lt_tensor-0.0.1a14.dist-info → lt_tensor-0.0.1a16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a14
+Version: 0.0.1a16
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336

{lt_tensor-0.0.1a14.dist-info → lt_tensor-0.0.1a16.dist-info}/RECORD

@@ -4,29 +4,29 @@ lt_tensor/losses.py,sha256=zvkCOnE5XpF3v6ymivRIdqPTsMM5zc94ZMom7YDi3zM,4946
 lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
 lt_tensor/math_ops.py,sha256=TkD4WQG42KsQ9Fg7FXOjf8f-ixtW0apf2XjaooecVx4,2257
 lt_tensor/misc_utils.py,sha256=S57M5XuGsIuaOKnEGZJsY3B2dTmggpdhsqQr51CQsYo,28754
-lt_tensor/model_base.py,sha256=lxzRXfPlR_t_6LfgRw2dct55evrtmwTiDqZGAe3jLro,20026
+lt_tensor/model_base.py,sha256=qqqIVpYz6nv01MnZuuAj1dxq4_NN-zSivP1GaegA9TI,21597
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
 lt_tensor/noise_tools.py,sha256=wFeAsHhLhSlEc5XU5LbFKaXoHeVxrWjiMeljjGdIKyM,11363
 lt_tensor/torch_commons.py,sha256=fntsEU8lhBQo0ebonI1iXBkMbWMN3HpBsG13EWlP5s8,718
 lt_tensor/transform.py,sha256=dZm8T_ov0blHMQu6nGiehsdG1VSB7bZBUVmTkT-PBdc,13257
 lt_tensor/datasets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lt_tensor/datasets/audio.py,sha256=j73oRyXt-AK4tWWYWjH-3c5RYouQBgDSCTuWHmyG8kQ,7450
-lt_tensor/model_zoo/__init__.py,sha256=RzG7fltZLyiIU_Za4pgfBPli5uPITiJkq4sTCd4uA_0,319
-lt_tensor/model_zoo/basic.py,sha256=_26H_jJk5Ld3DZiNpIhGosGfMxoFDZrI8bpDAYUOYno,10660
-lt_tensor/model_zoo/discriminator.py,sha256=dS5UmJZV5MxIFiaBlIXfgGLDdUT3y0Vuv9lDGHsjJE8,5849
-lt_tensor/model_zoo/features.py,sha256=CTFMidzza31pqQjwPfp_g0BNVfuQ8Dlo5JnxpYpKgag,13144
+lt_tensor/datasets/audio.py,sha256=5Wvz1BJ7xXkLYpVLLw9RY3X3RgMdPPeGiN0-MmJDQy0,8045
+lt_tensor/model_zoo/__init__.py,sha256=ltVTvmOlbOCfDc5Trvg0-Ta_Ujgkw0UVF9V5rqHx-RI,378
+lt_tensor/model_zoo/basic.py,sha256=pI8HyiHK-cmWcEEaVY_EduUJOjZW6HOtXvJd8Rbhq30,15452
+lt_tensor/model_zoo/features.py,sha256=DO8dlE0kmPKTNC1Xkv9wKegOOYkQa_rkxM4hhcNwJWA,15655
 lt_tensor/model_zoo/fusion.py,sha256=usC1bcjQRNivDc8xzkIS5T1glm78OLcs2V_tPqfp-eI,5422
 lt_tensor/model_zoo/pos_encoder.py,sha256=3d1EYLinCU9UAy-WuEWeYMGhMqaGknCiQ5qEmhw_UYM,4487
-lt_tensor/model_zoo/residual.py,sha256=3tc2fJaz6SxtKYAsxndahhwIxlN6oLk5tcdIXtUKaQc,7357
+lt_tensor/model_zoo/residual.py,sha256=i5V4ju7DB3WesKBVm6KH_LyPoKGDUOyo2Usfs-PyP58,9394
 lt_tensor/model_zoo/transformer.py,sha256=HUFoFFh7EQJErxdd9XIxhssdjvNVx2tNGDJOTUfwG2A,4301
-lt_tensor/model_zoo/istft/__init__.py,sha256=SV96w9WUWfHMee8Vjgn2MP0igKft7_mLTju9rFVYGHY,102
-lt_tensor/model_zoo/istft/generator.py,sha256=wWHUfLFIItN-tB3pWkc1r9aTWpHYBFg7UfvLN4_cD78,3179
-lt_tensor/model_zoo/istft/sg.py,sha256=EaEi3otw_uY5QfqDBNIWBWTJSg3KnwzzR4FBr0u09C0,4838
-lt_tensor/model_zoo/istft/trainer.py,sha256=KZXsAptOJeLYlr6t-DPX1qxgN526-2EBKoQQlcsHp8Y,21054
+lt_tensor/model_zoo/audio_models/__init__.py,sha256=CLoLqvbA_ltqm3KOg5AH3A0co0HtsLfFPUBsxxLSCgI,39
+lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=aFSmr8PYpmOfbe15lhNoj-ZzP5ChrZcikovKLZKg7nw,140
+lt_tensor/model_zoo/audio_models/diffwave/model.py,sha256=kHo76bxLJtTBn1m0gq5KKrUsjm9ASsCCwf8MvWaB1R8,6901
+lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=BOBZSK2HFOdMcFyjrzwZi_TeAtBGIcpb8pQxiGlwLEE,12302
+lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=o7Ie1qI22u_g9t1252PX4vl4uF6JHynAJryuz2lAZE0,12920
 lt_tensor/processors/__init__.py,sha256=4b9MxAJolXiJfSm20ZEspQTDm1tgLazwlPWA_jB1yLM,63
 lt_tensor/processors/audio.py,sha256=SMqNSl4Den-x1awTCQ8-TcR-0jPiv5lDaUpU93SRRaw,14749
-lt_tensor-0.0.1a14.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
-lt_tensor-0.0.1a14.dist-info/METADATA,sha256=mxwJTAo51GfGEEW87lT-Tp1AHtoRvuKCmcPxAyqJxLQ,1033
-lt_tensor-0.0.1a14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a14.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a14.dist-info/RECORD,,
+lt_tensor-0.0.1a16.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
+lt_tensor-0.0.1a16.dist-info/METADATA,sha256=uxk1cMeQkLniYUIgEjHD2eJ8_JGwAKS2minrCmAJfMo,1033
+lt_tensor-0.0.1a16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a16.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a16.dist-info/RECORD,,