lt-tensor 0.0.1a26__py3-none-any.whl → 0.0.1a28__py3-none-any.whl

lt_tensor/model_base.py

@@ -80,6 +80,62 @@ class _Devices_Base(nn.Module):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
 
+    def freeze_all(self, exclude: list[str] = []):
+        for name, module in self.named_modules():
+            if name in exclude or not hasattr(module, "requires_grad"):
+                continue
+            try:
+                self.freeze_module(module)
+            except:
+                pass
+
+    def unfreeze_all(self, exclude: list[str] = []):
+        for name, module in self.named_modules():
+            if name in exclude or not hasattr(module, "requires_grad"):
+                continue
+            try:
+                self.unfreeze_module(module)
+            except:
+                pass
+
+    def freeze_module(
+        self, module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor]
+    ):
+        self._change_gradient_state(module_or_name, False)
+
+    def unfreeze_module(
+        self, module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor]
+    ):
+        self._change_gradient_state(module_or_name, True)
+
+    def _change_gradient_state(
+        self,
+        module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor],
+        new_state: bool,  # True = unfreeze
+    ):
+        assert isinstance(
+            module_or_name, (str, nn.Module, nn.Parameter, Model, Tensor)
+        ), f"Item '{module_or_name}' is not a valid module, parameter, tensor or a string."
+        if isinstance(module_or_name, (nn.Module, nn.Parameter, Model, Tensor)):
+            target = module_or_name
+        else:
+            target = getattr(self, module_or_name)
+
+        if isinstance(target, Tensor):
+            target.requires_grad = new_state
+        elif isinstance(target, nn.Parameter):
+            target.requires_grad = new_state
+        elif isinstance(target, Model):
+            target.freeze_all()
+        elif isinstance(target, nn.Module):
+            for param in target.parameters():
+                if hasattr(param, "requires_grad"):
+                    param.requires_grad = new_state
+        else:
+            raise ValueError(
+                f"Item '{module_or_name}' is not a valid module, parameter or tensor."
+            )
+
     def _apply_device(self):
         """Add here components that are needed to have device applied to them,
         that usually the '.to()' function fails to apply
@@ -182,20 +238,12 @@ class Model(_Devices_Base, ABC):
     """
 
     _autocast: bool = False
-    _is_unfrozen: bool = False
-    # list with modules that can be frozen or unfrozen
-    registered_freezable_modules: List[str] = []
-    is_frozen: bool = False
-    _can_be_frozen: bool = (
-        False  # to control if the module can or cannot be freezed by other modules from 'Model' class
-    )
+
     # this is to be used on the case of they module requires low-rank adapters
     _low_rank_lambda: Optional[Callable[[], nn.Module]] = (
         None  # Example: lambda: nn.Linear(32, 32, True)
     )
     low_rank_adapter: Union[nn.Identity, nn.Module, nn.Sequential] = nn.Identity()
-    # never freeze:
-    _never_freeze_modules: List[str] = ["low_rank_adapter"]
 
     # dont save list:
     _dont_save_items: List[str] = []
@@ -208,75 +256,6 @@ class Model(_Devices_Base, ABC):
     def autocast(self, value: bool):
         self._autocast = value
 
-    def freeze_all(self, exclude: Optional[List[str]] = None, force: bool = False):
-        no_exclusions = not exclude
-        no_exclusions = not exclude
-        results = []
-        for name, module in self.named_modules():
-            if (
-                name in self._never_freeze_modules
-                or not force
-                and name not in self.registered_freezable_modules
-            ):
-                results.append(
-                    (
-                        name,
-                        "Unregistered module, to freeze/unfreeze it add its name into 'registered_freezable_modules'.",
-                    )
-                )
-                continue
-            if no_exclusions:
-                self.change_frozen_state(True, module)
-            elif not any(exclusion in name for exclusion in exclude):
-                results.append((name, self.change_frozen_state(True, module)))
-            else:
-                results.append((name, "excluded"))
-        return results
-
-    def unfreeze_all(self, exclude: Optional[list[str]] = None, force: bool = False):
-        """Unfreezes all model parameters except specified layers."""
-        no_exclusions = not exclude
-        results = []
-        for name, module in self.named_modules():
-            if (
-                name in self._never_freeze_modules
-                or not force
-                and name not in self.registered_freezable_modules
-            ):
-
-                results.append(
-                    (
-                        name,
-                        "Unregistered module, to freeze/unfreeze it add it into 'registered_freezable_modules'.",
-                    )
-                )
-                continue
-            if no_exclusions:
-                self.change_frozen_state(False, module)
-            elif not any(exclusion in name for exclusion in exclude):
-                results.append((name, self.change_frozen_state(False, module)))
-            else:
-                results.append((name, "excluded"))
-        return results
-
-    def change_frozen_state(self, freeze: bool, module: nn.Module):
-        assert isinstance(module, nn.Module)
-        if module.__class__.__name__ in self._never_freeze_modules:
-            return "Not Allowed"
-        try:
-            if isinstance(module, Model):
-                if module._can_be_frozen:
-                    if freeze:
-                        return module.freeze_all()
-                    return module.unfreeze_all()
-                else:
-                    return "Not Allowed"
-            else:
-                module.requires_grad_(not freeze)
-                return not freeze
-        except Exception as e:
-            return e
-
     def trainable_parameters(self, module_name: Optional[str] = None):
         """Gets the number of trainable parameters from either the entire model or from a specific module."""
         if module_name is not None:

lt_tensor/model_zoo/__init__.py

@@ -9,6 +9,7 @@ __all__ = [
     "audio_models",
     "hifigan",
     "istft",
+    "losses",
 ]
 from .audio_models import hifigan, istft
 from . import (
@@ -19,4 +20,5 @@ from . import (
     pos_encoder,
     residual,
     transformer,
+    losses,
 )

lt_tensor/model_zoo/losses/__init__.py

@@ -0,0 +1,3 @@
+from . import discriminators
+
+__all__ = ["discriminators"]

lt_tensor/model_zoo/losses/discriminators.py

@@ -0,0 +1,610 @@
+from lt_tensor.model_zoo.audio_models.hifigan import ConvNets
+from lt_utils.common import *
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from lt_tensor.model_zoo.convs import ConvNets
+from torch.nn import functional as F
+from torchaudio import transforms as T
+
+MULTI_DISC_OUT_TYPE: TypeAlias = Tuple[
+    List[Tensor],
+    List[Tensor],
+    List[List[Tensor]],
+    List[List[Tensor]],
+]
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class MultiDiscriminatorWrapper(ConvNets):
+    """Base for all multi-steps type of discriminators"""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.leaky_relu = nn.LeakyReLU(kwargs.get("negative_slope", 0.1))
+
+    def forward(self, y: Tensor, y_hat: Tensor) -> MULTI_DISC_OUT_TYPE:
+        pass
+
+    # for type hinting
+    def __call__(self, *args, **kwds) -> MULTI_DISC_OUT_TYPE:
+        return super().__call__(*args, **kwds)
+
+    def gen_step(self, y: Tensor, y_hat: Tensor) -> tuple[Tensor, Tensor, List[float]]:
+        """For generator loss step [feature loss, generator loss, list of generator losses (float)]"""
+        _, y_hat_gen, feat_map_real, feat_map_gen = self.train_step(y, y_hat)
+        loss_feat = self.feature_loss(feat_map_real, feat_map_gen)
+        loss_generator, losses_gen_s = self.generator_loss(y_hat_gen)
+        return loss_feat, loss_generator, losses_gen_s
+
+    def disc_step(
+        self, y: Tensor, y_hat: Tensor
+    ) -> tuple[Tensor, tuple[List[float], List[float]]]:
+        """For discriminator loss step [discriminator loss, (disc losses real, disc losses generated)]"""
+        y_hat_real, y_hat_gen, _, _ = self.train_step(y, y_hat)
+
+        loss_disc, losses_disc_real, losses_disc_generated = self.discriminator_loss(
+            y_hat_real, y_hat_gen
+        )
+        return loss_disc, (losses_disc_real, losses_disc_generated)
+
+    @staticmethod
+    def discriminator_loss(
+        disc_real_outputs, disc_generated_outputs
+    ) -> Tuple[Tensor, List[float], List[float]]:
+        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
+
+    @staticmethod
+    def feature_loss(fmap_r, fmap_g) -> Tensor:
+        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
+
+    @staticmethod
+    def generator_loss(disc_outputs) -> Tuple[Tensor, List[float]]:
+        loss = 0
+        gen_losses = []
+        for dg in disc_outputs:
+            l = torch.mean((1 - dg) ** 2)
+            gen_losses.append(l.item())
+            loss += l
+
+        return loss, gen_losses
+
+
+class DiscriminatorP(ConvNets):
+    def __init__(
+        self,
+        period: List[int],
+        discriminator_channel_mult: Number = 1,
+        kernel_size: int = 5,
+        stride: int = 3,
+        use_spectral_norm: bool = False,
+    ):
+        super().__init__()
+        self.period = period
+        norm_f = weight_norm if not use_spectral_norm else spectral_norm
+        dsc = lambda x: int(x * discriminator_channel_mult)
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    nn.Conv2d(
+                        1,
+                        dsc(32),
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        dsc(32),
+                        dsc(128),
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        dsc(128),
+                        dsc(512),
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        dsc(512),
+                        dsc(1024),
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        dsc(1024),
+                        dsc(1024),
+                        (kernel_size, 1),
+                        1,
+                        padding=(2, 0),
+                    )
+                ),
+            ]
+        )
+        self.conv_post = norm_f(nn.Conv2d(dsc(1024), 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
+        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 = F.leaky_relu(x, 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return x.flatten(1, -1), fmap
+
+
+class MultiPeriodDiscriminator(MultiDiscriminatorWrapper):
+    def __init__(
+        self,
+        discriminator_channel_mult: Number = 1,
+        mpd_reshapes: list[int] = [2, 3, 5, 7, 11],
+        use_spectral_norm: bool = False,
+    ):
+        super().__init__()
+        self.mpd_reshapes = mpd_reshapes
+        print(f"mpd_reshapes: {self.mpd_reshapes}")
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorP(
+                    rs,
+                    use_spectral_norm=use_spectral_norm,
+                    discriminator_channel_mult=discriminator_channel_mult,
+                )
+                for rs in self.mpd_reshapes
+            ]
+        )
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor) -> MULTI_DISC_OUT_TYPE:
+        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 EnvelopeExtractor(nn.Module):
+    """Extracts the amplitude envelope of the audio signal."""
+
+    def __init__(self, kernel_size=101):
+        super().__init__()
+        # Lowpass filter for smoothing envelope (moving average)
+        self.kernel_size = kernel_size
+        self.register_buffer("kernel", torch.ones(1, 1, kernel_size) / kernel_size)
+
+    def forward(self, x):
+        # x: (B, 1, T) -> abs(x)
+        envelope = torch.abs(x)
+        # Apply low-pass smoothing (via conv1d)
+        envelope = F.pad(
+            envelope, (self.kernel_size // 2, self.kernel_size // 2), mode="reflect"
+        )
+        envelope = F.conv1d(envelope, self.kernel)
+        return envelope
+
+
+class DiscriminatorEnvelope(ConvNets):
+    def __init__(self, use_spectral_norm=False):
+        super().__init__()
+        norm_f = weight_norm if not use_spectral_norm else spectral_norm
+        self.extractor = EnvelopeExtractor(kernel_size=101)
+        self.convs = nn.ModuleList(
+            [
+                norm_f(nn.Conv1d(1, 64, 15, stride=1, padding=7)),
+                norm_f(nn.Conv1d(64, 128, 41, stride=2, groups=4, padding=20)),
+                norm_f(nn.Conv1d(128, 256, 41, stride=2, groups=16, padding=20)),
+                norm_f(nn.Conv1d(256, 512, 41, stride=4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(512, 512, 41, stride=4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(512, 512, 5, stride=1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(nn.Conv1d(512, 1, 3, stride=1, padding=1))
+        self.activation = nn.LeakyReLU(0.1)
+
+    def forward(self, x):
+        # Input: raw audio (B, 1, T)
+        x = self.extractor(x)
+        fmap = []
+        for layer in self.convs:
+            x = self.activation(layer(x))
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return x.flatten(1), fmap
+
+
+class MultiEnvelopeDiscriminator(MultiDiscriminatorWrapper):
+    def __init__(self, use_spectral_norm: bool = False):
+        super().__init__()
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorEnvelope(use_spectral_norm),  # raw envelope
+                DiscriminatorEnvelope(use_spectral_norm),  # downsampled once
+                DiscriminatorEnvelope(use_spectral_norm),  # downsampled twice
+            ]
+        )
+        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)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorB(ConvNets):
+    """
+    Multi-band multi-scale STFT discriminator, with the architecture based on https://github.com/descriptinc/descript-audio-codec.
+    and the modified code adapted from https://github.com/gemelo-ai/vocos.
+    """
+
+    def __init__(
+        self,
+        window_length: int,
+        channels: int = 32,
+        hop_factor: float = 0.25,
+        bands: Tuple[Tuple[float, float], ...] = (
+            (0.0, 0.1),
+            (0.1, 0.25),
+            (0.25, 0.5),
+            (0.5, 0.75),
+            (0.75, 1.0),
+        ),
+    ):
+        super().__init__()
+        self.window_length = window_length
+        self.hop_factor = hop_factor
+        self.spec_fn = T.Spectrogram(
+            n_fft=window_length,
+            hop_length=int(window_length * hop_factor),
+            win_length=window_length,
+            power=None,
+        )
+        n_fft = window_length // 2 + 1
+        bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
+        self.bands = bands
+        convs = lambda: nn.ModuleList(
+            [
+                weight_norm(nn.Conv2d(2, channels, (3, 9), (1, 1), padding=(1, 4))),
+                weight_norm(
+                    nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))
+                ),
+                weight_norm(
+                    nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))
+                ),
+                weight_norm(
+                    nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))
+                ),
+                weight_norm(
+                    nn.Conv2d(channels, channels, (3, 3), (1, 1), padding=(1, 1))
+                ),
+            ]
+        )
+        self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
+
+        self.conv_post = weight_norm(
+            nn.Conv2d(channels, 1, (3, 3), (1, 1), padding=(1, 1))
+        )
+
+    def spectrogram(self, x: Tensor) -> List[Tensor]:
+        # Remove DC offset
+        x = x - x.mean(dim=-1, keepdims=True)
+        # Peak normalize the volume of input audio
+        x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
+        x = self.spec_fn(x)
+        x = torch.view_as_real(x)
+        x = x.permute(0, 3, 2, 1)  # [B, F, T, C] -> [B, C, T, F]
+        # Split into bands
+        x_bands = [x[..., b[0] : b[1]] for b in self.bands]
+        return x_bands
+
+    def forward(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
+        x_bands = self.spectrogram(x.squeeze(1))
+        fmap = []
+        x = []
+
+        for band, stack in zip(x_bands, self.band_convs):
+            for i, layer in enumerate(stack):
+                band = layer(band)
+                band = torch.nn.functional.leaky_relu(band, 0.1)
+                if i > 0:
+                    fmap.append(band)
+            x.append(band)
+
+        x = torch.cat(x, dim=-1)
+        x = self.conv_post(x)
+        fmap.append(x)
+
+        return x, fmap
+
+
+class MultiBandDiscriminator(MultiDiscriminatorWrapper):
+    """
+    Multi-band multi-scale STFT discriminator, with the architecture based on https://github.com/descriptinc/descript-audio-codec.
+    and the modified code adapted from https://github.com/gemelo-ai/vocos.
+    """
+
+    def __init__(
+        self,
+        mbd_fft_sizes: list[int] = [2048, 1024, 512],
+    ):
+        super().__init__()
+        self.fft_sizes = mbd_fft_sizes
+        self.discriminators = nn.ModuleList(
+            [DiscriminatorB(window_length=w) for w in self.fft_sizes]
+        )
+
+    def forward(self, y: Tensor, y_hat: Tensor) -> MULTI_DISC_OUT_TYPE:
+
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+
+        for d in self.discriminators:
+
+            y_d_r, fmap_r = d(x=y)
+            y_d_g, fmap_g = d(x=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 DiscriminatorR(ConvNets):
+    def __init__(
+        self,
+        resolution: List[int],
+        use_spectral_norm: bool = False,
+        discriminator_channel_mult: int = 1,
+    ):
+        super().__init__()
+
+        self.resolution = resolution
+        assert (
+            len(self.resolution) == 3
+        ), f"MRD layer requires list with len=3, got {self.resolution}"
+        self.lrelu_slope = 0.1
+
+        self.register_buffer("window", torch.hann_window(self.resolution[-1]))
+
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    nn.Conv2d(
+                        1, int(32 * discriminator_channel_mult), (3, 9), padding=(1, 4)
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_mult),
+                        (3, 9),
+                        stride=(1, 2),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_mult),
+                        (3, 9),
+                        stride=(1, 2),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_mult),
+                        (3, 9),
+                        stride=(1, 2),
+                        padding=(1, 4),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_mult),
+                        (3, 3),
+                        padding=(1, 1),
+                    )
+                ),
+            ]
+        )
+        self.conv_post = norm_f(
+            nn.Conv2d(int(32 * discriminator_channel_mult), 1, (3, 3), padding=(1, 1))
+        )
+
+    def forward(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
+        fmap = []
+        x = self.spectrogram(x)
+        x = x.unsqueeze(1)
+        for l in self.convs:
+            x = l(x)
+            x = F.leaky_relu(x, self.lrelu_slope)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+    def spectrogram(self, x: Tensor) -> Tensor:
+        n_fft, hop_length, win_length = self.resolution
+        x = F.pad(
+            x,
+            (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
+            mode="reflect",
+        )
+        x = x.squeeze(1)
+        x = torch.stft(
+            x,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            center=False,
+            return_complex=True,
+            window=self.window,
+        )
+        x = torch.view_as_real(x)  # [B, F, TT, 2]
+        mag = torch.norm(x, p=2, dim=-1)  # [B, F, TT]
+
+        return mag
+
+
+class MultiResolutionDiscriminator(MultiDiscriminatorWrapper):
+    def __init__(
+        self,
+        use_spectral_norm: bool = False,
+        discriminator_channel_mult: int = 1,
+        resolutions: List[List[int]] = [
+            [1024, 120, 600],
+            [2048, 240, 1200],
+            [512, 50, 240],
+        ],
+    ):
+        super().__init__()
+        self.resolutions = resolutions
+        assert (
+            len(self.resolutions) == 3
+        ), f"MRD requires list of list with len=3, each element having a list with len=3. Got {self.resolutions}, type: {type(self.resolutions)}"
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorR(
+                    resolution, use_spectral_norm, discriminator_channel_mult
+                )
+                for resolution in self.resolutions
+            ]
+        )
+
+    def forward(self, y: Tensor, y_hat: Tensor) -> MULTI_DISC_OUT_TYPE:
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for disc in self.discriminators:
+            y_d_r, fmap_r = disc(x=y)
+            y_d_g, fmap_g = disc(x=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 MultiDiscriminatorStep(Model):
+    def __init__(self, list_discriminator: List[MultiDiscriminatorWrapper]):
+        super().__init__()
+        self.disc: Sequence[MultiDiscriminatorWrapper] = nn.ModuleList(
+            list_discriminator
+        )
+        self.total = len(self.disc)
+
+    def forward(
+        self,
+        y: Tensor,
+        y_hat: Tensor,
+        step_type: Literal["discriminator", "generator"],
+    ) -> Union[
+        Tuple[Tensor, Tensor, List[float]], Tuple[Tensor, List[float], List[float]]
+    ]:
+        """
+        It returns the content based on the choice of "step_type", being it a
+        'discriminator' or 'generator'
+
+        For generator it returns:
+        Tuple[Tensor, Tensor, List[float]]
+        "gen_loss, feat_loss, all_g_losses"
+
+        For 'discriminator' it returns:
+        Tuple[Tensor, List[float], List[float]]
+        "disc_loss, disc_real_losses, disc_gen_losses"
+        """
+        if step_type == "generator":
+            all_g_losses: List[float] = []
+            feat_loss: Tensor = 0
+            gen_loss: Tensor = 0
+        else:
+            disc_loss: Tensor = 0
+            disc_real_losses: List[float] = []
+            disc_gen_losses: List[float] = []
+
+        for disc in self.disc:
+            if step_type == "generator":
+                #  feature loss, generator loss, list of generator losses (float)]
+                f_loss, g_loss, g_losses = disc.gen_step(y, y_hat)
+                gen_loss += g_loss
+                feat_loss += f_loss
+                all_g_losses.extend(g_losses)
+            else:
+                # [discriminator loss, (disc losses real, disc losses generated)]
+                d_loss, (d_real_losses, d_gen_losses) = disc.disc_step(y, y_hat)
+                disc_loss += d_loss
+                disc_real_losses.extend(d_real_losses)
+                disc_gen_losses.extend(d_gen_losses)
+
+        if step_type == "generator":
+            return gen_loss, feat_loss, all_g_losses
+        return disc_loss, disc_real_losses, disc_gen_losses

lt_tensor/processors/audio.py

@@ -105,7 +105,6 @@ class AudioProcessor(Model):
             onesided=self.cfg.onesided,
             normalized=self.cfg.normalized,
         )
-        self.griffin_lm_iters = 32
         self.mel_rscale = torchaudio.transforms.InverseMelScale(
             n_stft=self.cfg.n_stft,
             n_mels=self.cfg.n_mels,
@@ -114,21 +113,19 @@ class AudioProcessor(Model):
             f_max=self.cfg.f_max,
             mel_scale=self.cfg.mel_scale,
         )
-        self.giffin_lim = torchaudio.transforms.GriffinLim(
-            n_fft=self.cfg.n_fft,
-            win_length=self.cfg.win_length,
-            hop_length=self.cfg.hop_length,
-        )
+
         self.register_buffer(
             "window",
             (torch.hann_window(self.cfg.win_length) if window is None else window),
         )
 
     def _apply_device(self):
-        print(f"Audio Processor Device: {self.device.type}")
-        self.giffin_lim.to(device=self.device)
         self._mel_spec.to(device=self.device)
         self.mel_rscale.to(device=self.device)
+        try:
+            self.window.to(device=self.device)
+        except:
+            pass
 
     def from_numpy(
         self,
@@ -173,7 +170,9 @@ class AudioProcessor(Model):
         )
 
         if audio is None and mel is not None:
-            return self.from_numpy(librosa.feature.rms(S=mel, **rms_kwargs)[0])
+            return self.from_numpy(
+                librosa.feature.rms(S=mel, **rms_kwargs)[0]
+            ).squeeze()
         default_dtype = audio.dtype
         default_device = audio.device
         if audio.ndim > 1:
@@ -192,8 +191,12 @@ class AudioProcessor(Model):
         audio = self.to_numpy_safe(audio)
         if B == 1:
             if mel is None:
-                return self.from_numpy(librosa.feature.rms(y=audio, **rms_kwargs)[0])
-            return self.from_numpy(librosa.feature.rms(y=audio, S=mel, **rms_kwargs)[0])
+                return self.from_numpy(
+                    librosa.feature.rms(y=audio, **rms_kwargs)[0]
+                ).squeeze()
+            return self.from_numpy(
+                librosa.feature.rms(y=audio, S=mel, **rms_kwargs)[0]
+            ).squeeze()
         else:
             rms_ = []
             for i in range(B):
@@ -201,7 +204,7 @@ class AudioProcessor(Model):
                     0
                 ]
                 rms_.append(_r)
-            return self.from_numpy_batch(rms_, default_device, default_dtype)
+            return self.from_numpy_batch(rms_, default_device, default_dtype).squeeze()
 
     def compute_pitch(
         self,
@@ -250,7 +253,7 @@ class AudioProcessor(Model):
             for i in range(B):
                 f0_.append(librosa.yin(self.to_numpy_safe(audio[i, :]), **yn_kwargs))
             f0 = self.from_numpy_batch(f0_, default_device, default_dtype)
-        return f0
+        return f0.squeeze()
 
     def compute_pitch_torch(
         self,
@@ -273,7 +276,7 @@ class AudioProcessor(Model):
             win_length=win_length,
             freq_low=fmin,
             freq_high=fmax,
-        )
+        ).squeeze()
 
     def interpolate(
         self,
@@ -312,7 +315,7 @@ class AudioProcessor(Model):
             antialias=antialias,
         )
 
-    def inverse_transform(
+    def istft(
         self,
         spec: Tensor,
         phase: Tensor,
@@ -320,6 +323,10 @@ class AudioProcessor(Model):
         hop_length: Optional[int] = None,
         win_length: Optional[int] = None,
         length: Optional[int] = None,
+        center: Optional[bool] = None,
+        normalized: Optional[bool] = None,
+        onesided: Optional[bool] = None,
+        return_complex: bool = False,
         *,
         _recall: bool = False,
     ):
@@ -331,25 +338,25 @@ class AudioProcessor(Model):
         try:
             return torch.istft(
                 spec * torch.exp(phase * 1j),
-                n_fft=n_fft or self.cfg.n_fft,
-                hop_length=hop_length or self.cfg.hop_length,
-                win_length=win_length or self.cfg.win_length,
+                n_fft=default(n_fft, self.cfg.n_fft),
+                hop_length=default(hop_length, self.cfg.hop_length),
+                win_length=default(win_length, self.cfg.win_length),
                 window=window,
-                center=self.cfg.center,
-                normalized=self.cfg.normalized,
-                onesided=self.cfg.onesided,
+                center=default(center, self.cfg.center),
+                normalized=default(normalized, self.cfg.normalized),
+                onesided=default(onesided, self.cfg.onesided),
                 length=length,
-                return_complex=False,
+                return_complex=return_complex,
             )
         except RuntimeError as e:
             if not _recall and spec.device != self.window.device:
                 self.window = self.window.to(spec.device)
-                return self.inverse_transform(
+                return self.istft(
                     spec, phase, n_fft, hop_length, win_length, length, _recall=True
                 )
             raise e
 
-    def normalize_audio(
+    def istft_norm(
         self,
         wave: Tensor,
         length: Optional[int] = None,
@@ -389,7 +396,7 @@ class AudioProcessor(Model):
         except RuntimeError as e:
             if not _recall and wave.device != self.window.device:
                 self.window = self.window.to(wave.device)
-                return self.normalize_audio(wave, length, _recall=True)
+                return self.istft_norm(wave, length, _recall=True)
             raise e
 
     def compute_mel(
@@ -407,11 +414,7 @@ class AudioProcessor(Model):
                 mel_tensor = (
                     torch.log(eps + mel_tensor.unsqueeze(0)) - self.cfg.mean
                 ) / self.cfg.std
-            if mel_tensor.ndim == 4:
-                return mel_tensor.squeeze()
-            elif mel_tensor.ndim == 2:
-                return mel_tensor.unsqueeze(0)
-            return mel_tensor
+            return mel_tensor.squeeze()
 
         except RuntimeError as e:
             if not _recall:
@@ -419,14 +422,6 @@ class AudioProcessor(Model):
                 return self.compute_mel(wave, raw_mel_only, eps, _recall=True)
             raise e
 
-    def inverse_mel_spectogram(self, mel: Tensor, n_iter: Optional[int] = None):
-        if isinstance(n_iter, int) and n_iter != self.griffin_lm_iters:
-            self.giffin_lim.n_iter = n_iter
-            self.griffin_lm_iters = n_iter
-        return self.giffin_lim.forward(
-            self.mel_rscale(mel),
-        )
-
     def load_audio(
         self,
         path: PathLike,
@@ -510,14 +505,9 @@ class AudioProcessor(Model):
             maximum,
         )
 
-    def stft_loss(
-        self,
-        signal: Tensor,
-        ground: Tensor,
-    ):
-        with torch.no_grad():
-            ground = F.interpolate(ground, signal.shape[-1]).to(signal.device)
-        return F.l1_loss(signal, ground)
+    def stft_loss(self, signal: Tensor, ground: Tensor, magnitude: float = 1.0):
+        ground = F.interpolate(ground, signal.shape[-1]).to(signal.device)
+        return F.l1_loss(signal.squeeze(), ground.squeeze()) * magnitude
 
     def forward(
         self,

{lt_tensor-0.0.1a26.dist-info → lt_tensor-0.0.1a28.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a26
+Version: 0.0.1a28
 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.1a26.dist-info → lt_tensor-0.0.1a28.dist-info}/RECORD

@@ -4,12 +4,12 @@ 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=N2r3UmxC4RM2BZBQhpjDZ_BKLrzsyIlKzopTzJbnjFU,28962
-lt_tensor/model_base.py,sha256=GvmQdt97ZSfOObBpBIq7UUTwpIE1g-aBm23za36YA0M,18431
+lt_tensor/model_base.py,sha256=DTg44N6eTXLmpIAj_ac29-M5dI_iY_sC0yA_K3E13GI,17446
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
 lt_tensor/noise_tools.py,sha256=wFeAsHhLhSlEc5XU5LbFKaXoHeVxrWjiMeljjGdIKyM,11363
 lt_tensor/torch_commons.py,sha256=8l0bxmrAzwvyqjivCIVISXlbvKarlg4DdE0BOGSnMuQ,812
 lt_tensor/transform.py,sha256=dZm8T_ov0blHMQu6nGiehsdG1VSB7bZBUVmTkT-PBdc,13257
-lt_tensor/model_zoo/__init__.py,sha256=ltVTvmOlbOCfDc5Trvg0-Ta_Ujgkw0UVF9V5rqHx-RI,378
+lt_tensor/model_zoo/__init__.py,sha256=yPUVchgVhU2nAJ2ocA4HFfG7IMEiBu8qOi8I1KWTTkU,404
 lt_tensor/model_zoo/basic.py,sha256=pI8HyiHK-cmWcEEaVY_EduUJOjZW6HOtXvJd8Rbhq30,15452
 lt_tensor/model_zoo/convs.py,sha256=YQRxek75Qpsha8nfc7wLhmJS9XxPeCa4WxuftLg6IcE,3927
 lt_tensor/model_zoo/features.py,sha256=DO8dlE0kmPKTNC1Xkv9wKegOOYkQa_rkxM4hhcNwJWA,15655
@@ -26,10 +26,12 @@ lt_tensor/model_zoo/audio_models/__init__.py,sha256=MoG9YjxLyvscq_6njK1ljGBletK9
 lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=PDuDYN1omD1RoAXcmxH3tEgfAuM3ZHAWzimD6ElMqEQ,9073
 lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=7GJqKLw7-juXpfp5IFzjASLut0uouDhjZ1CQknf3H68,16533
 lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=ltIuD9t1gmS3bTmCqZIwJHKrhC6DYya3OaXlskWX9kw,17606
+lt_tensor/model_zoo/losses/__init__.py,sha256=B9RAUxBiOZwooztnij1oLeRwZ7_MjnN3mPoum7saD6s,59
+lt_tensor/model_zoo/losses/discriminators.py,sha256=yYh7HzRTUtr0RVTG7cWpcYsJZsRCz6yzg6Loq8FtyOk,20405
 lt_tensor/processors/__init__.py,sha256=Pvxhh0KR65zLCgUd53_k5Z0y5JWWcO0ZBXFK9rv0o5w,109
-lt_tensor/processors/audio.py,sha256=WkumFNx8OXGQkTEU5Rkede9NLMrsGaTGY37Ti784Wv8,17028
-lt_tensor-0.0.1a26.dist-info/licenses/LICENSE,sha256=tQHc38scHOba4kDBNG4U0U6PpObaloiZG-FvKSgv2b0,11336
-lt_tensor-0.0.1a26.dist-info/METADATA,sha256=2STSK6jgD_qECwz9WygTXNDwfapEAR2mpHiS14bi9tQ,1062
-lt_tensor-0.0.1a26.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a26.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a26.dist-info/RECORD,,
+lt_tensor/processors/audio.py,sha256=rsnnNi8MtxPq9vAYoiRQ7lGjorfJIpRvrKEe3zA8YJk,16668
+lt_tensor-0.0.1a28.dist-info/licenses/LICENSE,sha256=tQHc38scHOba4kDBNG4U0U6PpObaloiZG-FvKSgv2b0,11336
+lt_tensor-0.0.1a28.dist-info/METADATA,sha256=2LLguzaCAM2bcAdy_D66j4PS9Oh5PU3ZnA9qy7xcx0w,1062
+lt_tensor-0.0.1a28.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a28.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a28.dist-info/RECORD,,