lt-tensor 0.0.1a10__py3-none-any.whl → 0.0.1a12__py3-none-any.whl

lt_tensor/model_zoo/disc.py

@@ -1,4 +1,4 @@
-from ..torch_commons import *
+from lt_tensor.torch_commons import *
 import torch.nn.functional as F
 from lt_tensor.model_base import Model
 from lt_utils.common import *
@@ -76,20 +76,6 @@ class PeriodDiscriminator(Model):
         return x.flatten(1, -1), f_map
 
 
-class MultiPeriodDiscriminator(Model):
-    def __init__(self, periods=[2, 3, 5, 7, 11]):
-        super().__init__()
-
-        self.discriminators = nn.ModuleList([PeriodDiscriminator(p) for p in periods])
-
-    def forward(self, x: torch.Tensor):
-        """
-        x: (B, T)
-        Returns: list of tuples of outputs from each period discriminator and the f_map.
-        """
-        return [d(x) for d in self.discriminators]
-
-
 class ScaleDiscriminator(nn.Module):
     def __init__(self, use_spectral_norm=False):
         super().__init__()
@@ -123,11 +109,11 @@ class ScaleDiscriminator(nn.Module):
 
 
 class MultiScaleDiscriminator(Model):
-    def __init__(self):
+    def __init__(self, layers: int = 3):
         super().__init__()
         self.pooling = nn.AvgPool1d(4, 2, padding=2)
         self.discriminators = nn.ModuleList(
-            [ScaleDiscriminator(i == 0) for i in range(3)]
+            [ScaleDiscriminator(i == 0) for i in range(layers)]
         )
 
     def forward(self, x: torch.Tensor):
@@ -136,57 +122,75 @@ class MultiScaleDiscriminator(Model):
         Returns: list of outputs from each scale discriminator
         """
         outputs = []
+        features = []
         for i, d in enumerate(self.discriminators):
             if i != 0:
                 x = self.pooling(x)
-            outputs.append(d(x))
-        return outputs
+            out, f_map = d(x)
+            outputs.append(out)
+            features.append(f_map)
+        return outputs, features
 
 
-class GeneralLossDescriminator(Model):
-    """TODO: build an unified loss for both mpd and msd here."""
-
-    def __init__(self):
+class MultiPeriodDiscriminator(Model):
+    def __init__(self, periods: List[int] = [2, 3, 5, 7, 11]):
         super().__init__()
-        self.mpd = MultiPeriodDiscriminator()
-        self.msd = MultiScaleDiscriminator()
-        self.print_trainable_parameters()
-
-    def _get_group_(self):
-        pass
+        self.discriminators = nn.ModuleList([PeriodDiscriminator(p) for p in periods])
 
-    def forward(self, x: Tensor, y_hat: Tensor):
-        return
+    def forward(self, x: torch.Tensor):
+        """
+        x: (B, T)
+        Returns: list of tuples of outputs from each period discriminator and the f_map.
+        """
+        # torch.log(torch.clip(x, min=clip_val))
+        out_map = []
+        feat_map = []
+        for d in self.discriminators:
+            out, feat = d(x)
+            out_map.append(out)
+            feat_map.append(feat)
+        return out_map, feat_map
 
 
-def discriminator_loss(d_outputs_real, d_outputs_fake):
+def discriminator_loss(real_out_map, fake_out_map):
     loss = 0.0
-    for real_out, fake_out in zip(d_outputs_real, d_outputs_fake):
-        real_score = real_out[0]
-        fake_score = fake_out[0]
-        loss += torch.mean(F.relu(1.0 - real_score)) + torch.mean(
-            F.relu(1.0 + fake_score)
-        )
-    return loss
+    rl, fl = [], []
+    for real_out, fake_out in zip(real_out_map, fake_out_map):
+        real_loss = torch.mean((1.0 - real_out) ** 2)
+        fake_loss = torch.mean(fake_out**2)
+        loss += real_loss + fake_loss
+        rl.append(real_loss.item())
+        fl.append(fake_loss.item())
+    return loss, sum(rl), sum(fl)
 
 
-def generator_adv_loss(d_outputs_fake):
+def generator_adv_loss(fake_disc_outputs: List[Tensor]):
     loss = 0.0
-    for fake_out in d_outputs_fake:
+    for fake_out in fake_disc_outputs:
         fake_score = fake_out[0]
         loss += -torch.mean(fake_score)
     return loss
 
 
-def feature_matching_loss(
-    d_outputs_real,
-    d_outputs_fake,
-    loss_fn: Callable[[Tensor, Tensor], Tensor] = F.mse_loss,
+def feature_loss(
+    fmap_r,
+    fmap_g,
+    weight=2.0,
+    loss_fn: Callable[[Tensor, Tensor], Tensor] = F.l1_loss,
 ):
     loss = 0.0
-    for real_out, fake_out in zip(d_outputs_real, d_outputs_fake):
-        real_feats = real_out[1]
-        fake_feats = fake_out[1]
-        for real_f, fake_f in zip(real_feats, fake_feats):
-            loss += loss_fn(fake_f, real_f)
-    return loss
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += loss_fn(rl - gl)
+    return loss * weight
+
+
+def generator_loss(disc_generated_outputs):
+    loss = 0.0
+    gen_losses = []
+    for dg in disc_generated_outputs:
+        l = torch.mean((1.0 - dg) ** 2)
+        gen_losses.append(l.item())
+        loss += l
+
+    return loss, gen_losses

lt_tensor/model_zoo/fsn.py

@@ -6,8 +6,8 @@ __all__ = [
     "GatedFusion",
 ]
 
-from ..torch_commons import *
-from ..model_base import Model
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
 
 
 class ConcatFusion(Model):

lt_tensor/model_zoo/gns.py

@@ -7,10 +7,10 @@ __all__ = [
     "NoisePredictor1D",
 ]
 
-from ..torch_commons import *
-from ..model_base import Model
-from .rsd import ResBlock1D, ResBlocks
-from ..misc_utils import log_tensor
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from lt_tensor.model_zoo.rsd import ResBlock1D
+from lt_tensor.misc_utils import log_tensor
 
 import torch.nn.functional as F
 

lt_tensor/model_zoo/istft/__init__.py

@@ -0,0 +1,5 @@
+from .generator import iSTFTGenerator
+from . import trainer
+
+
+__all__ = ["iSTFTGenerator", "trainer"]

lt_tensor/model_zoo/istft/generator.py

@@ -0,0 +1,150 @@
+__all__ = ["iSTFTGenerator", "ResBlocks"]
+import gc
+import math
+import itertools
+from lt_utils.common import *
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from lt_tensor.misc_utils import log_tensor
+from lt_tensor.model_zoo.rsd import ResBlock1D, ConvNets, get_weight_norm
+from lt_utils.misc_utils import log_traceback
+from lt_tensor.processors import AudioProcessor
+from lt_utils.type_utils import is_dir, is_pathlike
+from lt_tensor.misc_utils import set_seed, clear_cache
+from lt_tensor.model_zoo.disc import MultiPeriodDiscriminator, MultiScaleDiscriminator
+import torch.nn.functional as F
+from lt_tensor.config_templates import updateDict, ModelConfig
+
+
+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 self.activation(x)
+
+
+class iSTFTGenerator(ConvNets):
+    def __init__(
+        self,
+        in_channels: int = 80,
+        upsample_rates: List[Union[int, List[int]]] = [8, 8],
+        upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16],
+        upsample_initial_channel: int = 512,
+        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],
+        ],
+        n_fft: int = 16,
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        hop_length: int = 256,
+    ):
+        super().__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.hop_length = hop_length
+        self.conv_pre = weight_norm(
+            nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
+        )
+        self.blocks = nn.ModuleList()
+        self.activation = activation
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.blocks.append(
+                self._make_blocks(
+                    (i, k, u),
+                    upsample_initial_channel,
+                    resblock_kernel_sizes,
+                    resblock_dilation_sizes,
+                )
+            )
+
+        ch = upsample_initial_channel // (2 ** (i + 1))
+        self.post_n_fft = n_fft // 2 + 1
+        self.conv_post = weight_norm(nn.Conv1d(ch, n_fft + 2, 7, 1, padding=3))
+        self.conv_post.apply(self.init_weights)
+        self.reflection_pad = nn.ReflectionPad1d((1, 0))
+
+        self.phase = nn.Sequential(
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+        )
+        self.spec = nn.Sequential(
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+        )
+
+    def _make_blocks(
+        self,
+        state: Tuple[int, int, int],
+        upsample_initial_channel: int,
+        resblock_kernel_sizes: List[Union[int, List[int]]],
+        resblock_dilation_sizes: List[int | List[int]],
+    ):
+        i, k, u = state
+        channels = upsample_initial_channel // (2 ** (i + 1))
+        return nn.ModuleDict(
+            dict(
+                up=nn.Sequential(
+                    self.activation,
+                    weight_norm(
+                        nn.ConvTranspose1d(
+                            upsample_initial_channel // (2**i),
+                            channels,
+                            k,
+                            u,
+                            padding=(k - u) // 2,
+                        )
+                    ).apply(self.init_weights),
+                ),
+                residual=ResBlocks(
+                    channels,
+                    resblock_kernel_sizes,
+                    resblock_dilation_sizes,
+                    self.activation,
+                ),
+            )
+        )
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+        for block in self.blocks:
+            x = block["up"](x)
+            x = block["residual"](x)
+
+        x = self.reflection_pad(x)
+        x = self.conv_post(x)
+        spec = torch.exp(self.spec(x[:, : self.post_n_fft, :]))
+        phase = torch.sin(self.phase(x[:, self.post_n_fft :, :]))
+
+        return spec, phase