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

lt_tensor/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.1a33"
+__version__ = "0.0.1a34"
 
 from . import (
     lr_schedulers,

lt_tensor/losses.py

@@ -6,19 +6,24 @@ __all__ = [
     "hybrid_loss",
     "diff_loss",
     "cosine_loss",
-    "gan_loss",
     "ft_n_loss",
+    "MultiMelScaleLoss",
 ]
 import math
 import random
 from lt_tensor.torch_commons import *
 from lt_utils.common import *
 import torch.nn.functional as F
+from lt_tensor.model_base import Model
+from lt_tensor.processors import AudioProcessor, AudioProcessorConfig
+from lt_tensor.math_ops import normalize_minmax, normalize_unit_norm, normalize_zscore
+
 
 def ft_n_loss(output: Tensor, target: Tensor, weight: Optional[Tensor] = None):
     if weight is not None:
-        return torch.mean((torch.abs(output - target) + weight) **0.5)
-    return torch.mean(torch.abs(output - target)**0.5)
+        return torch.mean((torch.abs(output - target) + weight) ** 0.5)
+    return torch.mean(torch.abs(output - target) ** 0.5)
+
 
 def adaptive_l1_loss(
     inp: Tensor,
@@ -58,50 +63,6 @@ def cosine_loss(inp, tgt):
     return 1 - cos.mean()  # Lower is better
 
 
-class GanLosses:
-    @staticmethod
-    def get_loss(
-        pred: Tensor,
-        target_is_real: bool,
-        loss_type: Literal["bce", "mse", "hinge", "wasserstein"] = "bce",
-    ) -> Tensor:
-        if loss_type == "bce":  # Standard GAN
-            target = torch.ones_like(pred) if target_is_real else torch.zeros_like(pred)
-            return F.binary_cross_entropy_with_logits(pred, target)
-
-        elif loss_type == "mse":  # LSGAN
-            target = torch.ones_like(pred) if target_is_real else torch.zeros_like(pred)
-            return F.mse_loss(torch.sigmoid(pred), target)
-
-        elif loss_type == "hinge":
-            if target_is_real:
-                return torch.mean(F.relu(1.0 - pred))
-            else:
-                return torch.mean(F.relu(1.0 + pred))
-
-        elif loss_type == "wasserstein":
-            return -pred.mean() if target_is_real else pred.mean()
-
-        else:
-            raise ValueError(f"Unknown loss_type: {loss_type}")
-
-    @staticmethod
-    def generator_loss(fake_pred: Tensor, loss_type: str = "bce") -> Tensor:
-        return GanLosses.get_loss(fake_pred, target_is_real=True, loss_type=loss_type)
-
-    @staticmethod
-    def discriminator_loss(
-        real_pred: Tensor, fake_pred: Tensor, loss_type: str = "bce"
-    ) -> Tensor:
-        real_loss = GanLosses.get_loss(
-            real_pred, target_is_real=True, loss_type=loss_type
-        )
-        fake_loss = GanLosses.get_loss(
-            fake_pred.detach(), target_is_real=False, loss_type=loss_type
-        )
-        return (real_loss + fake_loss) * 0.5
-
-
 def masked_cross_entropy(
     logits: torch.Tensor,  # [B, T, V]
     targets: torch.Tensor,  # [B, T]
@@ -157,3 +118,160 @@ def gan_d_loss(real_preds, fake_preds, use_lsgan=True):
                 torch.log(1 - fake + 1e-7)
             )
     return loss
+
+
+class MultiMelScaleLoss(Model):
+    def __init__(
+        self,
+        sample_rate: int,
+        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
+        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
+        n_ffts: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
+        hops: List[int] = [8, 16, 32, 64, 128, 256, 512],
+        f_min: float = [0, 0, 0, 0, 0, 0, 0],
+        f_max: Optional[float] = [None, None, None, None, None, None, None],
+        loss_fn: Callable = nn.L1Loss(),
+        center: bool = True,
+        power: float = 1.0,
+        normalized: bool = False,
+        pad_mode: str = "reflect",
+        onesided: Optional[bool] = None,
+        std: int = 4,
+        mean: int = -4,
+        use_istft_norm: bool = True,
+        use_pitch_loss: bool = True,
+        use_rms_loss: bool = True,
+        norm_pitch_fn: Callable[[Tensor], Tensor] = normalize_unit_norm,
+        norm_rms_fn: Callable[[Tensor], Tensor] = normalize_unit_norm,
+        lambda_mel: float = 1.0,
+        lambda_rms: float = 1.0,
+        lambda_pitch: float = 1.0,
+        weight: float = 1.0,
+    ):
+        super().__init__()
+        assert (
+            len(n_mels)
+            == len(window_lengths)
+            == len(n_ffts)
+            == len(hops)
+            == len(f_min)
+            == len(f_max)
+        )
+        self.loss_fn = loss_fn
+        self.lambda_mel = lambda_mel
+        self.weight = weight
+        self.use_istft_norm = use_istft_norm
+        self.use_pitch_loss = use_pitch_loss
+        self.use_rms_loss = use_rms_loss
+        self.lambda_pitch = lambda_pitch
+        self.lambda_rms = lambda_rms
+
+        self.norm_pitch_fn = norm_pitch_fn
+        self.norm_rms = norm_rms_fn
+
+        self._setup_mels(
+            sample_rate,
+            n_mels,
+            window_lengths,
+            n_ffts,
+            hops,
+            f_min,
+            f_max,
+            center,
+            power,
+            normalized,
+            pad_mode,
+            onesided,
+            std,
+            mean,
+        )
+
+    def _setup_mels(
+        self,
+        sample_rate: int,
+        n_mels: List[int],
+        window_lengths: List[int],
+        n_ffts: List[int],
+        hops: List[int],
+        f_min: List[float],
+        f_max: List[Optional[float]],
+        center: bool,
+        power: float,
+        normalized: bool,
+        pad_mode: str,
+        onesided: Optional[bool],
+        std: int,
+        mean: int,
+    ):
+        assert (
+            len(n_mels)
+            == len(window_lengths)
+            == len(n_ffts)
+            == len(hops)
+            == len(f_min)
+            == len(f_max)
+        )
+        _mel_kwargs = dict(
+            sample_rate=sample_rate,
+            center=center,
+            onesided=onesided,
+            normalized=normalized,
+            power=power,
+            pad_mode=pad_mode,
+            std=std,
+            mean=mean,
+        )
+        self.mel_spectrograms: List[AudioProcessor] = nn.ModuleList(
+            [
+                AudioProcessor(
+                    AudioProcessorConfig(
+                        **_mel_kwargs,
+                        n_mels=mel,
+                        n_fft=n_fft,
+                        win_length=win,
+                        hop_length=hop,
+                        f_min=fmin,
+                        f_max=fmax,
+                    )
+                )
+                for mel, win, n_fft, hop, fmin, fmax in zip(
+                    n_mels, window_lengths, n_ffts, hops, f_min, f_max
+                )
+            ]
+        )
+
+    def forward(
+        self, input_wave: torch.Tensor, target_wave: torch.Tensor
+    ) -> torch.Tensor:
+        assert self.use_istft_norm or input_wave.shape[-1] == target_wave.shape[-1]
+        target_wave = target_wave.to(input_wave.device)
+        losses = 0.0
+        for M in self.mel_spectrograms:
+            # Apply normalization if requested
+            if self.use_istft_norm:
+                input_proc = M.istft_norm(input_wave, length=target_wave.shape[-1])
+                target_proc = M.istft_norm(target_wave, length=target_wave.shape[-1])
+            else:
+                input_proc, target_proc = input_wave, target_wave
+
+            x_mels = M(input_proc)
+            y_mels = M(target_proc)
+
+            loss = self.loss_fn(x_mels.squeeze(), y_mels.squeeze())
+            losses += loss * self.lambda_mel
+
+            # pitch/f0 loss
+            if self.use_pitch_loss:
+                x_pitch = self.norm_pitch_fn(M.compute_pitch(input_proc))
+                y_pitch = self.norm_pitch_fn(M.compute_pitch(target_proc))
+                f0_loss = self.loss_fn(x_pitch, y_pitch)
+                losses += f0_loss * self.lambda_pitch
+
+            # energy/rms loss
+            if self.use_rms_loss:
+                x_rms = self.norm_rms(M.compute_rms(input_proc, x_mels))
+                y_rms = self.norm_rms(M.compute_rms(target_proc, y_mels))
+                rms_loss = self.loss_fn(x_rms, y_rms)
+                losses += rms_loss * self.lambda_rms
+
+        return losses * self.weight

lt_tensor/model_zoo/losses/discriminators.py

@@ -82,23 +82,6 @@ class MultiDiscriminatorWrapper(Model):
         return disc_loss, disc_real_losses, disc_gen_losses
 
 
-def normalize_unit_norm(x: torch.Tensor, eps: float = 1e-5):
-    norm = torch.norm(x, dim=-1, keepdim=True)
-    return x / (norm + eps)
-
-
-def normalize_minmax(x: torch.Tensor, eps: float = 1e-5):
-    min_val = x.amin(dim=-1, keepdim=True)
-    max_val = x.amax(dim=-1, keepdim=True)
-    return (x - min_val) / (max_val - min_val + eps)
-
-
-def normalize_zscore(x: torch.Tensor, eps: float = 1e-5):
-    mean = x.mean(dim=-1, keepdim=True)
-    std = x.std(dim=-1, keepdim=True)
-    return (x - mean) / (std + eps)
-
-
 def get_padding(kernel_size, dilation=1):
     return int((kernel_size * dilation - dilation) / 2)
 
@@ -637,179 +620,3 @@ class MultiResolutionDiscriminator(_MultiDiscriminatorT):
             y_d_gs.append(y_d_g)
             fmap_gs.append(fmap_g)
         return y_d_rs, y_d_gs, fmap_rs, fmap_gs
-
-
-class MultiMelScaleLoss(Model):
-    # TODO: Make the normalization an argument to be chosen by the dev
-    def __init__(
-        self,
-        sample_rate: int,
-        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
-        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
-        n_ffts: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
-        hops: List[int] = [8, 16, 32, 64, 128, 256, 512],
-        weight: float = 1.0,
-        lambda_mel: float = 1.0,
-        f_min: float = [0, 0, 0, 0, 0, 0, 0],
-        f_max: Optional[float] = [None, None, None, None, None, None, None],
-        loss_fn: Callable = nn.L1Loss(),
-        center: bool = True,
-        power: float = 1.0,
-        normalized: bool = False,
-        pad_mode: str = "reflect",
-        onesided: Optional[bool] = None,
-        std: int = 4,
-        mean: int = -4,
-        auto_interpolate: bool = True,
-        use_istft_norm: bool = True,
-        use_pitch_loss: bool = False,
-        use_rms_loss: bool = False,
-        lambda_pitch: float = 0.5,
-        lambda_rms: float = 0.5,
-    ):
-        super().__init__()
-        assert (
-            len(n_mels)
-            == len(window_lengths)
-            == len(n_ffts)
-            == len(hops)
-            == len(f_min)
-            == len(f_max)
-        )
-        self.loss_fn = loss_fn
-        self.lambda_mel = lambda_mel
-        self.weight = weight
-        self.use_istft_norm = use_istft_norm
-        self.auto_interpolate = auto_interpolate if not self.use_istft_norm else False
-        self.use_pitch_loss = use_pitch_loss
-        self.use_rms_loss = use_rms_loss
-        self.lambda_pitch = lambda_pitch
-        self.lambda_rms = lambda_rms
-
-        self._setup_mels(
-            sample_rate,
-            n_mels,
-            window_lengths,
-            n_ffts,
-            hops,
-            f_min,
-            f_max,
-            center,
-            power,
-            normalized,
-            pad_mode,
-            onesided,
-            std,
-            mean,
-        )
-
-    def _setup_mels(
-        self,
-        sample_rate: int,
-        n_mels: List[int],
-        window_lengths: List[int],
-        n_ffts: List[int],
-        hops: List[int],
-        f_min: List[float],
-        f_max: List[Optional[float]],
-        center: bool,
-        power: float,
-        normalized: bool,
-        pad_mode: str = "reflect",
-        onesided: Optional[bool] = None,
-        std: int = 4,
-        mean: int = -4,
-    ):
-        assert (
-            len(n_mels)
-            == len(window_lengths)
-            == len(n_ffts)
-            == len(hops)
-            == len(f_min)
-            == len(f_max)
-        )
-        _mel_kwargs = dict(
-            sample_rate=sample_rate,
-            center=center,
-            onesided=onesided,
-            normalized=normalized,
-            power=power,
-            pad_mode=pad_mode,
-            std=std,
-            mean=mean,
-        )
-        self.mel_spectrograms: List[AudioProcessor] = nn.ModuleList(
-            [
-                AudioProcessor(
-                    AudioProcessorConfig(
-                        **_mel_kwargs,
-                        n_mels=mel,
-                        n_fft=n_fft,
-                        win_length=win,
-                        hop_length=hop,
-                        f_min=fmin,
-                        f_max=fmax,
-                    )
-                )
-                for mel, win, n_fft, hop, fmin, fmax in zip(
-                    n_mels, window_lengths, n_ffts, hops, f_min, f_max
-                )
-            ]
-        )
-
-    def _process_tensor(
-        self,
-        input_wave: torch.Tensor,
-        target_wave: torch.Tensor,
-    ):
-        if input_wave.shape[-1] != target_wave.shape[-1]:
-            if input_wave.ndim < 3:
-                # To be compatible with interpolatin
-                if input_wave.ndim == 2:
-                    input_wave = input_wave.unsqueeze(1)
-                else:
-                    input_wave = input_wave.unsqueeze(0).unsqueeze(0)
-            input_wave = F.interpolate(input_wave, target_wave.shape[-1], mode="linear")
-        return input_wave
-
-    def forward(
-        self, input_wave: torch.Tensor, target_wave: torch.Tensor
-    ) -> torch.Tensor:
-        assert (
-            self.use_istft_norm
-            or self.auto_interpolate
-            or input_wave.shape[-1] == target_wave.shape[-1]
-        )
-        if self.auto_interpolate:
-            input_wave = self._process_tensor(input_wave, target_wave)
-
-        losses = 0.0
-        for M in self.mel_spectrograms:
-            # Apply normalization if requested
-            if self.use_istft_norm:
-                input_proc = M.istft_norm(input_wave, length=target_wave.shape[-1])
-                target_proc = M.istft_norm(target_wave, length=target_wave.shape[-1])
-            else:
-                input_proc, target_proc = input_wave, target_wave
-
-            x_mels = M(input_proc)
-            y_mels = M(target_proc)
-
-            loss = self.loss_fn(x_mels.squeeze(), y_mels.squeeze())
-            losses += loss * self.lambda_mel
-
-            # pitch/f0 loss
-            if self.use_pitch_loss:
-                x_pitch = normalize_unit_norm(M.compute_pitch(input_proc))
-                y_pitch = normalize_unit_norm(M.compute_pitch(target_proc))
-                f0_loss = self.loss_fn(x_pitch, y_pitch)
-                losses += f0_loss * self.lambda_pitch
-
-            # energy/rms loss
-            if self.use_rms_loss:
-                x_rms = normalize_unit_norm(M.compute_rms(input_proc, x_mels))
-                y_rms = normalize_unit_norm(M.compute_rms(target_proc, y_mels))
-                rms_loss = self.loss_fn(x_rms, y_rms)
-                losses += rms_loss * self.lambda_rms
-
-        return losses * self.weight

{lt_tensor-0.0.1a33.dist-info → lt_tensor-0.0.1a34.dist-info}/METADATA

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

@@ -1,6 +1,6 @@
-lt_tensor/__init__.py,sha256=f3wraCpbx0fV2tQgsZfKw1ifTPp87hSCOZmE0d09LYk,441
+lt_tensor/__init__.py,sha256=WAGPuMPq5c4DGAJ57x1Ykgzg3vMlLq9BiWk5EdJcUsU,441
 lt_tensor/config_templates.py,sha256=F9UvL8paAjkSvio890kp8WznpYeI50pYnm9iqQroBxk,2797
-lt_tensor/losses.py,sha256=zvkCOnE5XpF3v6ymivRIdqPTsMM5zc94ZMom7YDi3zM,4946
+lt_tensor/losses.py,sha256=fHVMqOFo3ekjORYy89R_aRjmtT6lo27Z1egzOYjQ1W8,8646
 lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
 lt_tensor/math_ops.py,sha256=ahX6Z1Mt3X-FhmwSZYZea5mB1B0S8GDuvKPfAm5e_FQ,2646
 lt_tensor/misc_utils.py,sha256=N2r3UmxC4RM2BZBQhpjDZ_BKLrzsyIlKzopTzJbnjFU,28962
@@ -27,11 +27,11 @@ lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=PDuDYN1omD1RoAXcmxH
 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=0b4ikOFy8Ubozq0Igs7X1ELQD5JrPA3jwR4dzuEa6hM,27047
+lt_tensor/model_zoo/losses/discriminators.py,sha256=ZpyByFgc7L7uV_XRBsV9vkdVItbJO3z--Y6LlvTvtwY,20765
 lt_tensor/processors/__init__.py,sha256=Pvxhh0KR65zLCgUd53_k5Z0y5JWWcO0ZBXFK9rv0o5w,109
 lt_tensor/processors/audio.py,sha256=HNr1GS-6M2q0Rda4cErf5y2Jlc9f4jD58FvpX2ua9d4,18369
-lt_tensor-0.0.1a33.dist-info/licenses/LICENSE,sha256=TbiyJWLgNqqgqhfCnrGwFIxy7EqGNrIZZcKhHrefcuU,11354
-lt_tensor-0.0.1a33.dist-info/METADATA,sha256=6xlFyxd0mYYqTi8oSS3M99mnqZUQrmtp3_AJt-rlewg,1062
-lt_tensor-0.0.1a33.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a33.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a33.dist-info/RECORD,,
+lt_tensor-0.0.1a34.dist-info/licenses/LICENSE,sha256=TbiyJWLgNqqgqhfCnrGwFIxy7EqGNrIZZcKhHrefcuU,11354
+lt_tensor-0.0.1a34.dist-info/METADATA,sha256=WkTafcY5nYZbrZ7WzUc3JXnmg9NtUAXrchx42dCok9I,1062
+lt_tensor-0.0.1a34.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a34.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a34.dist-info/RECORD,,