lt-tensor 0.0.1a32__tar.gz → 0.0.1a34__tar.gz

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright 2025 gr1336
+   Copyright 2025 gr1336 (Gabriel Ribeiro)
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a32
+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
@@ -17,7 +17,7 @@ Requires-Dist: numpy>=1.26.4
 Requires-Dist: tokenizers
 Requires-Dist: pyyaml>=6.0.0
 Requires-Dist: numba>0.60.0
-Requires-Dist: lt-utils>=0.0.3
+Requires-Dist: lt-utils>=0.0.4
 Requires-Dist: librosa==0.11.*
 Requires-Dist: einops
 Requires-Dist: plotly

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/lt_tensor/__init__.py

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

lt_tensor-0.0.1a34/lt_tensor/losses.py

@@ -0,0 +1,277 @@
+__all__ = [
+    "masked_cross_entropy",
+    "adaptive_l1_loss",
+    "contrastive_loss",
+    "smooth_l1_loss",
+    "hybrid_loss",
+    "diff_loss",
+    "cosine_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)
+
+
+def adaptive_l1_loss(
+    inp: Tensor,
+    tgt: Tensor,
+    weight: Optional[Tensor] = None,
+    scale: float = 1.0,
+    inverted: bool = False,
+):
+
+    if weight is not None:
+        loss = torch.mean(torch.abs((inp - tgt) + weight.mean()))
+    else:
+        loss = torch.mean(torch.abs(inp - tgt))
+    loss *= scale
+    if inverted:
+        return -loss
+    return loss
+
+
+def smooth_l1_loss(inp: Tensor, tgt: Tensor, beta=1.0, weight=None):
+    diff = torch.abs(inp - tgt)
+    loss = torch.where(diff < beta, 0.5 * diff**2 / beta, diff - 0.5 * beta)
+    if weight is not None:
+        loss *= weight
+    return loss.mean()
+
+
+def contrastive_loss(x1: Tensor, x2: Tensor, label: Tensor, margin: float = 1.0):
+    # label == 1: similar, label == 0: dissimilar
+    dist = torch.nn.functional.pairwise_distance(x1, x2)
+    loss = label * dist**2 + (1 - label) * torch.clamp(margin - dist, min=0.0) ** 2
+    return loss.mean()
+
+
+def cosine_loss(inp, tgt):
+    cos = torch.nn.functional.cosine_similarity(inp, tgt, dim=-1)
+    return 1 - cos.mean()  # Lower is better
+
+
+def masked_cross_entropy(
+    logits: torch.Tensor,  # [B, T, V]
+    targets: torch.Tensor,  # [B, T]
+    lengths: torch.Tensor,  # [B]
+    reduction: str = "mean",
+) -> torch.Tensor:
+    """
+    CrossEntropyLoss with masking for variable-length sequences.
+    - logits: unnormalized scores [B, T, V]
+    - targets: ground truth indices [B, T]
+    - lengths: actual sequence lengths [B]
+    """
+    B, T, V = logits.size()
+    logits = logits.view(-1, V)
+    targets = targets.view(-1)
+
+    # Create mask
+    mask = torch.arange(T, device=lengths.device).expand(B, T) < lengths.unsqueeze(1)
+    mask = mask.reshape(-1)
+
+    # Apply CE only where mask == True
+    loss = F.cross_entropy(
+        logits[mask], targets[mask], reduction="mean" if reduction == "mean" else "none"
+    )
+    if reduction == "none":
+        return loss
+    return loss
+
+
+def diff_loss(pred_noise, true_noise, mask=None):
+    """Standard diffusion noise-prediction loss (e.g., DDPM)"""
+    if mask is not None:
+        return F.mse_loss(pred_noise * mask, true_noise * mask)
+    return F.mse_loss(pred_noise, true_noise)
+
+
+def hybrid_diff_loss(pred_noise, true_noise, alpha=0.5):
+    """Combines L1 and L2"""
+    l1 = F.l1_loss(pred_noise, true_noise)
+    l2 = F.mse_loss(pred_noise, true_noise)
+    return alpha * l1 + (1 - alpha) * l2
+
+
+def gan_d_loss(real_preds, fake_preds, use_lsgan=True):
+    loss = 0
+    for real, fake in zip(real_preds, fake_preds):
+        if use_lsgan:
+            loss += F.mse_loss(real, torch.ones_like(real)) + F.mse_loss(
+                fake, torch.zeros_like(fake)
+            )
+        else:
+            loss += -torch.mean(torch.log(real + 1e-7)) - torch.mean(
+                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-0.0.1a32 → lt_tensor-0.0.1a34}/lt_tensor/math_ops.py

@@ -6,10 +6,12 @@ __all__ = [
     "apply_window",
     "shift_ring",
     "dot_product",
-    "normalize_tensor",
     "log_magnitude",
     "shift_time",
     "phase",
+    "normalize_unit_norm",
+    "normalize_minmax",
+    "normalize_zscore",
 ]
 
 from lt_tensor.torch_commons import *
@@ -61,11 +63,6 @@ def dot_product(x: Tensor, y: Tensor, dim: int = -1) -> Tensor:
     return torch.sum(x * y, dim=dim)
 
 
-def normalize_tensor(x: Tensor, eps: float = 1e-8) -> Tensor:
-    """Normalizes a tensor to unit norm (L2)."""
-    return x / (torch.norm(x, dim=-1, keepdim=True) + eps)
-
-
 def log_magnitude(stft_complex: Tensor, eps: float = 1e-5) -> Tensor:
     """Returns log magnitude from complex STFT."""
     magnitude = torch.abs(stft_complex)
@@ -76,3 +73,19 @@ def phase(stft_complex: Tensor) -> Tensor:
     """Returns phase from complex STFT."""
     return torch.angle(stft_complex)
 
+
+def normalize_unit_norm(x: torch.Tensor, eps: float = 1e-6):
+    norm = torch.norm(x, dim=-1, keepdim=True)
+    return x / (norm + eps)
+
+
+def normalize_minmax(x: torch.Tensor, eps: float = 1e-6):
+    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-6):
+    mean = x.mean(dim=-1, keepdim=True)
+    std = x.std(dim=-1, keepdim=True)
+    return (x - mean) / (std + eps)

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/lt_tensor/model_zoo/losses/discriminators.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from lt_tensor.model_zoo.audio_models.hifigan import ConvNets
 from lt_utils.common import *
 from lt_tensor.torch_commons import *
@@ -5,6 +7,8 @@ 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
+from lt_tensor.processors import AudioProcessor, AudioProcessorConfig
+
 
 MULTI_DISC_OUT_TYPE: TypeAlias = Tuple[
     List[Tensor],
@@ -14,11 +18,75 @@ MULTI_DISC_OUT_TYPE: TypeAlias = Tuple[
 ]
 
 
+class MultiDiscriminatorWrapper(Model):
+    def __init__(self, list_discriminator: List["_MultiDiscriminatorT"]):
+        """Setup example:
+        model_d = MultiDiscriminatorStep(
+            [
+                MultiEnvelopeDiscriminator(),
+                MultiBandDiscriminator(),
+                MultiResolutionDiscriminator(),
+                MultiPeriodDiscriminator(0.5),
+            ]
+        )
+        """
+        super().__init__()
+        self.disc: Sequence[_MultiDiscriminatorT] = 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
+
+
 def get_padding(kernel_size, dilation=1):
     return int((kernel_size * dilation - dilation) / 2)
 
 
-class MultiDiscriminatorWrapper(ConvNets):
+class _MultiDiscriminatorT(ConvNets):
     """Base for all multi-steps type of discriminators"""
 
     def __init__(self, *args, **kwargs):
@@ -171,7 +239,7 @@ class DiscriminatorP(ConvNets):
         return x.flatten(1, -1), fmap
 
 
-class MultiPeriodDiscriminator(MultiDiscriminatorWrapper):
+class MultiPeriodDiscriminator(_MultiDiscriminatorT):
     def __init__(
         self,
         discriminator_channel_mult: Number = 1,
@@ -258,7 +326,7 @@ class DiscriminatorEnvelope(ConvNets):
         return x.flatten(1), fmap
 
 
-class MultiEnvelopeDiscriminator(MultiDiscriminatorWrapper):
+class MultiEnvelopeDiscriminator(_MultiDiscriminatorT):
     def __init__(self, use_spectral_norm: bool = False):
         super().__init__()
         self.discriminators = nn.ModuleList(
@@ -375,7 +443,7 @@ class DiscriminatorB(ConvNets):
         return x, fmap
 
 
-class MultiBandDiscriminator(MultiDiscriminatorWrapper):
+class MultiBandDiscriminator(_MultiDiscriminatorT):
     """
     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.
@@ -514,7 +582,7 @@ class DiscriminatorR(ConvNets):
         return mag
 
 
-class MultiResolutionDiscriminator(MultiDiscriminatorWrapper):
+class MultiResolutionDiscriminator(_MultiDiscriminatorT):
     def __init__(
         self,
         use_spectral_norm: bool = False,
@@ -552,71 +620,3 @@ class MultiResolutionDiscriminator(MultiDiscriminatorWrapper):
             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]
-    ):
-        """Setup example:
-        model_d = MultiDiscriminatorStep(
-            [
-                MultiEnvelopeDiscriminator(),
-                MultiBandDiscriminator(),
-                MultiResolutionDiscriminator(),
-                MultiPeriodDiscriminator(0.5),
-            ]
-        )
-        """
-        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-0.0.1a32 → lt_tensor-0.0.1a34}/lt_tensor/processors/audio.py

@@ -23,7 +23,7 @@ class AudioProcessorConfig(ModelConfig):
     win_length: int = 1024
     hop_length: int = 256
     f_min: float = 0
-    f_max: float = 8000.0
+    f_max: Optional[float] = None
     center: bool = True
     mel_scale: Literal["htk" "slaney"] = "htk"
     std: int = 4
@@ -41,8 +41,8 @@ class AudioProcessorConfig(ModelConfig):
         n_fft: int = 1024,
         win_length: Optional[int] = None,
         hop_length: Optional[int] = None,
-        f_min: float = 1,
-        f_max: float = 12000.0,
+        f_min: float = 0,
+        f_max: Optional[float] = None,
         center: bool = True,
         mel_scale: Literal["htk", "slaney"] = "htk",
         std: int = 4,
@@ -71,9 +71,12 @@ class AudioProcessorConfig(ModelConfig):
         self.post_process()
 
     def post_process(self):
-        self.f_min = max(self.f_min, 1)
-        self.f_max = max(min(self.f_max, self.n_fft // 2), self.f_min + 1)
         self.n_stft = self.n_fft // 2 + 1
+        # some functions needs this to be a non-zero or not None value.
+        self.f_min = max(self.f_min, (self.sample_rate / (self.n_fft - 1))  * 2)
+        self.default_f_max = min(
+            default(self.f_max, self.sample_rate // 2), self.sample_rate // 2
+        )
         self.hop_length = default(self.hop_length, self.n_fft // 4)
         self.win_length = default(self.win_length, self.n_fft)
 
@@ -105,7 +108,7 @@ class AudioProcessor(Model):
             onesided=self.cfg.onesided,
             normalized=self.cfg.normalized,
         )
-        self.mel_rscale = torchaudio.transforms.InverseMelScale(
+        self._mel_rscale = torchaudio.transforms.InverseMelScale(
             n_stft=self.cfg.n_stft,
             n_mels=self.cfg.n_mels,
             sample_rate=self.cfg.sample_rate,
@@ -119,32 +122,39 @@ class AudioProcessor(Model):
             (torch.hann_window(self.cfg.win_length) if window is None else window),
         )
 
-    def from_numpy(
-        self,
-        array: np.ndarray,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        converted = torch.from_numpy(array)
-        if device is None:
-            device = self.device
-        return converted.to(device=device, dtype=dtype)
 
-    def from_numpy_batch(
+    
+    def compute_mel(
         self,
-        arrays: List[np.ndarray],
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        stacked = torch.stack([torch.from_numpy(x) for x in arrays])
-        if device is None:
-            device = self.device
-        return stacked.to(device=device, dtype=dtype)
+        wave: Tensor,
+        raw_mel_only: bool = False,
+        eps: float = 1e-5,
+        *,
+        _recall: bool = False,
+    ) -> Tensor:
+        """Returns: [B, M, T]"""
+        try:
+            mel_tensor = self._mel_spec(wave.to(self.device))  # [M, T]
+            if not raw_mel_only:
+                mel_tensor = (
+                    torch.log(eps + mel_tensor.unsqueeze(0)) - self.cfg.mean
+                ) / self.cfg.std
+            return mel_tensor.squeeze()
 
-    def to_numpy_safe(self, tensor: Union[Tensor, np.ndarray]):
-        if isinstance(tensor, np.ndarray):
-            return tensor
-        return tensor.detach().to(DEFAULT_DEVICE).numpy(force=True)
+        except RuntimeError as e:
+            if not _recall:
+                self._mel_spec.to(self.device)
+                return self.compute_mel(wave, raw_mel_only, eps, _recall=True)
+            raise e
+
+    def compute_inverse_mel(self, melspec: Tensor, *, _recall=False):
+        try:
+            return self._mel_rscale.forward(melspec.to(self.device)).squeeze()
+        except RuntimeError as e:
+            if not _recall:
+                self._mel_rscale.to(self.device)
+                return self.compute_inverse_mel(melspec, _recall=True)
+            raise e
 
     def compute_rms(
         self,
@@ -192,12 +202,44 @@ class AudioProcessor(Model):
         else:
             rms_ = []
             for i in range(B):
-                _r = librosa.feature.rms(_comp_rms_helper(i, audio, mel), **rms_kwargs)[
+                _t = _comp_rms_helper(i, audio, mel)
+                _r = librosa.feature.rms(**_t, **rms_kwargs)[
                     0
                 ]
                 rms_.append(_r)
             return self.from_numpy_batch(rms_, default_device, default_dtype).squeeze()
 
+    def pitch_shift(self, audio: torch.Tensor, sample_rate: Optional[int] = None, n_steps: float = 2.0):
+        """
+        Shifts the pitch of an audio tensor by `n_steps` semitones.
+
+        Args:
+            audio (torch.Tensor): Tensor of shape (B, T) or (T,)
+            sample_rate (int, optional): Sample rate of the audio. Will use the class sample rate if unset.
+            n_steps (float): Number of semitones to shift. Can be negative.
+
+        Returns:
+            torch.Tensor: Pitch-shifted audio.
+        """
+        src_device = audio.device
+        src_dtype = audio.dtype
+        audio = audio.squeeze()
+        sample_rate = default(sample_rate, self.cfg.sample_rate)
+        def _shift_one(wav):
+            wav_np = self.to_numpy_safe(wav)
+            shifted_np = librosa.effects.pitch_shift(wav_np, sr=sample_rate, n_steps=n_steps)
+            return torch.from_numpy(shifted_np)
+
+        if audio.ndim == 1:
+            return _shift_one(audio).to(device=src_device, dtype=src_dtype)
+        return torch.stack([_shift_one(a) for a in audio]).to(device=src_device, dtype=src_dtype)
+
+
+    @staticmethod
+    def calc_pitch_fmin(sr:int, frame_length:float):
+        """For pitch f_min"""
+        return (sr / (frame_length - 1)) * 2
+    
     def compute_pitch(
         self,
         audio: Tensor,
@@ -218,9 +260,9 @@ class AudioProcessor(Model):
         else:
             B = 1
         sr = default(sr, self.cfg.sample_rate)
-        fmin = max(default(fmin, self.cfg.f_min), 65)
-        fmax = min(default(fmax, self.cfg.f_max), sr // 2)
         frame_length = default(frame_length, self.cfg.n_fft)
+        fmin = max(default(fmin, self.cfg.f_min), self.calc_pitch_fmin(sr, frame_length))
+        fmax = min(max(default(fmax, self.cfg.default_f_max), fmin+1), sr // 2)
         hop_length = default(hop_length, self.cfg.hop_length)
         center = default(center, self.cfg.center)
         yn_kwargs = dict(
@@ -257,10 +299,10 @@ class AudioProcessor(Model):
         frame_length: Optional[Number] = None,
     ):
         sr = default(sr, self.sample_rate)
-        fmin = max(default(fmin, self.f_min), 1)
-        fmax = min(default(fmax, self.f_max), sr // 2)
-        win_length = default(win_length, self.win_length)
-        frame_length = default(frame_length, self.n_fft)
+        win_length = default(win_length, self.cfg.win_length)
+        frame_length = default(frame_length, self.cfg.n_fft)
+        fmin = default(fmin, self.calc_pitch_fmin(sr, frame_length))
+        fmax = default(fmax, self.cfg.default_f_max)
         return detect_pitch_frequency(
             audio,
             sample_rate=sr,
@@ -270,6 +312,33 @@ class AudioProcessor(Model):
             freq_high=fmax,
         ).squeeze()
 
+    def from_numpy(
+        self,
+        array: np.ndarray,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        converted = torch.from_numpy(array)
+        if device is None:
+            device = self.device
+        return converted.to(device=device, dtype=dtype)
+
+    def from_numpy_batch(
+        self,
+        arrays: List[np.ndarray],
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        stacked = torch.stack([torch.from_numpy(x) for x in arrays])
+        if device is None:
+            device = self.device
+        return stacked.to(device=device, dtype=dtype)
+
+    def to_numpy_safe(self, tensor: Union[Tensor, np.ndarray]):
+        if isinstance(tensor, np.ndarray):
+            return tensor
+        return tensor.detach().to(DEFAULT_DEVICE).numpy(force=True)
+
     def interpolate(
         self,
         tensor: Tensor,
@@ -391,29 +460,6 @@ class AudioProcessor(Model):
                 return self.istft_norm(wave, length, _recall=True)
             raise e
 
-    def compute_mel(
-        self,
-        wave: Tensor,
-        raw_mel_only: bool = False,
-        eps: float = 1e-5,
-        *,
-        _recall: bool = False,
-    ) -> Tensor:
-        """Returns: [B, M, T]"""
-        try:
-            mel_tensor = self._mel_spec(wave.to(self.device))  # [M, T]
-            if not raw_mel_only:
-                mel_tensor = (
-                    torch.log(eps + mel_tensor.unsqueeze(0)) - self.cfg.mean
-                ) / self.cfg.std
-            return mel_tensor.squeeze()
-
-        except RuntimeError as e:
-            if not _recall:
-                self._mel_spec.to(self.device)
-                return self.compute_mel(wave, raw_mel_only, eps, _recall=True)
-            raise e
-
     def load_audio(
         self,
         path: PathLike,

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/lt_tensor.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a32
+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
@@ -17,7 +17,7 @@ Requires-Dist: numpy>=1.26.4
 Requires-Dist: tokenizers
 Requires-Dist: pyyaml>=6.0.0
 Requires-Dist: numba>0.60.0
-Requires-Dist: lt-utils>=0.0.3
+Requires-Dist: lt-utils>=0.0.4
 Requires-Dist: librosa==0.11.*
 Requires-Dist: einops
 Requires-Dist: plotly

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/lt_tensor.egg-info/requires.txt

@@ -4,7 +4,7 @@ numpy>=1.26.4
 tokenizers
 pyyaml>=6.0.0
 numba>0.60.0
-lt-utils>=0.0.3
+lt-utils>=0.0.4
 librosa==0.11.*
 einops
 plotly

{lt_tensor-0.0.1a32 → lt_tensor-0.0.1a34}/setup.py

@@ -4,7 +4,7 @@ with open("README.md", "r", encoding="utf-8") as f:
     long_description = f.read()
 
 setup(
-    version="0.0.1a32",
+    version="0.0.1a34",
     name="lt-tensor",
     description="General utilities for PyTorch and others. Built for general use.",
     long_description=long_description,
@@ -17,7 +17,7 @@ setup(
         "tokenizers",
         "pyyaml>=6.0.0",
         "numba>0.60.0",
-        "lt-utils>=0.0.3",
+        "lt-utils>=0.0.4",
         "librosa==0.11.*",
         "einops",
         "plotly",

lt_tensor-0.0.1a32/lt_tensor/losses.py

@@ -1,159 +0,0 @@
-__all__ = [
-    "masked_cross_entropy",
-    "adaptive_l1_loss",
-    "contrastive_loss",
-    "smooth_l1_loss",
-    "hybrid_loss",
-    "diff_loss",
-    "cosine_loss",
-    "gan_loss",
-    "ft_n_loss",
-]
-import math
-import random
-from lt_tensor.torch_commons import *
-from lt_utils.common import *
-import torch.nn.functional as F
-
-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)
-
-def adaptive_l1_loss(
-    inp: Tensor,
-    tgt: Tensor,
-    weight: Optional[Tensor] = None,
-    scale: float = 1.0,
-    inverted: bool = False,
-):
-
-    if weight is not None:
-        loss = torch.mean(torch.abs((inp - tgt) + weight.mean()))
-    else:
-        loss = torch.mean(torch.abs(inp - tgt))
-    loss *= scale
-    if inverted:
-        return -loss
-    return loss
-
-
-def smooth_l1_loss(inp: Tensor, tgt: Tensor, beta=1.0, weight=None):
-    diff = torch.abs(inp - tgt)
-    loss = torch.where(diff < beta, 0.5 * diff**2 / beta, diff - 0.5 * beta)
-    if weight is not None:
-        loss *= weight
-    return loss.mean()
-
-
-def contrastive_loss(x1: Tensor, x2: Tensor, label: Tensor, margin: float = 1.0):
-    # label == 1: similar, label == 0: dissimilar
-    dist = torch.nn.functional.pairwise_distance(x1, x2)
-    loss = label * dist**2 + (1 - label) * torch.clamp(margin - dist, min=0.0) ** 2
-    return loss.mean()
-
-
-def cosine_loss(inp, tgt):
-    cos = torch.nn.functional.cosine_similarity(inp, tgt, dim=-1)
-    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]
-    lengths: torch.Tensor,  # [B]
-    reduction: str = "mean",
-) -> torch.Tensor:
-    """
-    CrossEntropyLoss with masking for variable-length sequences.
-    - logits: unnormalized scores [B, T, V]
-    - targets: ground truth indices [B, T]
-    - lengths: actual sequence lengths [B]
-    """
-    B, T, V = logits.size()
-    logits = logits.view(-1, V)
-    targets = targets.view(-1)
-
-    # Create mask
-    mask = torch.arange(T, device=lengths.device).expand(B, T) < lengths.unsqueeze(1)
-    mask = mask.reshape(-1)
-
-    # Apply CE only where mask == True
-    loss = F.cross_entropy(
-        logits[mask], targets[mask], reduction="mean" if reduction == "mean" else "none"
-    )
-    if reduction == "none":
-        return loss
-    return loss
-
-
-def diff_loss(pred_noise, true_noise, mask=None):
-    """Standard diffusion noise-prediction loss (e.g., DDPM)"""
-    if mask is not None:
-        return F.mse_loss(pred_noise * mask, true_noise * mask)
-    return F.mse_loss(pred_noise, true_noise)
-
-
-def hybrid_diff_loss(pred_noise, true_noise, alpha=0.5):
-    """Combines L1 and L2"""
-    l1 = F.l1_loss(pred_noise, true_noise)
-    l2 = F.mse_loss(pred_noise, true_noise)
-    return alpha * l1 + (1 - alpha) * l2
-
-
-def gan_d_loss(real_preds, fake_preds, use_lsgan=True):
-    loss = 0
-    for real, fake in zip(real_preds, fake_preds):
-        if use_lsgan:
-            loss += F.mse_loss(real, torch.ones_like(real)) + F.mse_loss(
-                fake, torch.zeros_like(fake)
-            )
-        else:
-            loss += -torch.mean(torch.log(real + 1e-7)) - torch.mean(
-                torch.log(1 - fake + 1e-7)
-            )
-    return loss