lt-tensor 0.0.1a4__py3-none-any.whl → 0.0.1a6__py3-none-any.whl

lt_tensor/__init__.py

@@ -1,13 +1,15 @@
 __version__ = "0.0.1a"
 
 from . import (
+    lr_schedulers,
     model_zoo,
     model_base,
     math_ops,
     misc_utils,
     monotonic_align,
     transform,
-    lr_schedulers,
+    noise_tools,
+    losses,
 )
 
 __all__ = [
@@ -18,4 +20,6 @@ __all__ = [
     "monotonic_align",
     "transform",
     "lr_schedulers",
+    "noise_tools",
+    "losses",
 ]

lt_tensor/datasets/audio.py

@@ -3,9 +3,10 @@ from ..torch_commons import *
 import torchaudio
 from lt_utils.common import *
 import librosa
-from lt_utils.type_utils import is_file
+from lt_utils.type_utils import is_file, is_array
 from torchaudio.functional import resample
 from ..transform import inverse_transform
+from lt_utils.file_ops import FileScan, load_text, get_file_name
 
 
 class AudioProcessor:
@@ -108,3 +109,46 @@ class AudioProcessor:
             .float()
             .unsqueeze(0)
         )
+
+    def find_audios(self, path: PathLike, additional_extensions: List[str] = []):
+        extensions = [
+            "*.wav",
+            "*.aac",
+            "*.m4a",
+            "*.mp3",
+            "*.ogg",
+            "*.opus",
+            "*.flac",
+        ]
+        extensions.extend(
+            [x for x in additional_extensions if isinstance(x, str) and "*" in x]
+        )
+        return FileScan.files(
+            path,
+            extensions,
+        )
+
+    def find_audio_text_pairs(
+        self,
+        path,
+        additional_extensions: List[str] = [],
+        text_file_patterns: List[str] = [".normalized.txt", ".original.txt"],
+    ):
+        is_array(text_file_patterns, True, validate=True)  # Rases if empty or not valid
+        additional_extensions = [
+            x
+            for x in additional_extensions
+            if isinstance(x, str)
+            and "*" in x
+            and not any(list(map(lambda y: y in x), text_file_patterns))
+        ]
+        audio_files = self.find_audios(path, additional_extensions)
+        text_files = []
+        for audio in audio_files:
+            base_audio_dir = Path(audio).parent
+            audio_name = get_file_name(audio, False)
+            for pattern in text_file_patterns:
+                possible_txt_file = Path(base_audio_dir, audio_name + pattern)
+                if is_file(possible_txt_file):
+                    text_files.append(audio)
+        return audio_files, text_files

lt_tensor/losses.py

@@ -0,0 +1,145 @@
+__all__ = ["masked_cross_entropy"]
+import math
+import random
+from .torch_commons import *
+from lt_utils.common import *
+import torch.nn.functional as F
+
+
+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
+
+
+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
+
+
+def gan_g_loss(fake_preds, use_lsgan=True):
+    loss = 0
+    for fake in fake_preds:
+        if use_lsgan:
+            loss += F.mse_loss(fake, torch.ones_like(fake))
+        else:
+            loss += -torch.mean(torch.log(fake + 1e-7))
+    return loss
+
+
+def feature_matching_loss(real_feats, fake_feats):
+    """real_feats and fake_feats are lists of intermediate features"""
+    loss = 0
+    for real_layers, fake_layers in zip(real_feats, fake_feats):
+        for r, f in zip(real_layers, fake_layers):
+            loss += F.l1_loss(f, r.detach())
+    return loss
+
+
+def feature_loss(real_fmaps, fake_fmaps, weight=2.0):
+    loss = 0.0
+    for dr, dg in zip(real_fmaps, fake_fmaps):  # Each (layer list from a discriminator)
+        for r_feat, g_feat in zip(dr, dg):
+            loss += F.l1_loss(r_feat, g_feat)
+    return loss * weight
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0.0
+    r_losses = []
+    g_losses = []
+
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = F.mse_loss(dr, torch.ones_like(dr))
+        g_loss = F.mse_loss(dg, torch.zeros_like(dg))
+        loss += r_loss + g_loss
+        r_losses.append(r_loss)
+        g_losses.append(g_loss)
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(fake_outputs):
+    total = 0.0
+    g_losses = []
+    for out in fake_outputs:
+        loss = F.mse_loss(out, torch.ones_like(out))
+        g_losses.append(loss)
+        total += loss
+    return total, g_losses
+
+
+def multi_resolution_stft_loss(y, y_hat, fft_sizes=[512, 1024, 2048]):
+    loss = 0
+    for fft_size in fft_sizes:
+        hop = fft_size // 4
+        win = fft_size
+        y_stft = torch.stft(
+            y, n_fft=fft_size, hop_length=hop, win_length=win, return_complex=True
+        )
+        y_hat_stft = torch.stft(
+            y_hat, n_fft=fft_size, hop_length=hop, win_length=win, return_complex=True
+        )
+
+        loss += F.l1_loss(torch.abs(y_stft), torch.abs(y_hat_stft))
+    return loss

lt_tensor/math_ops.py

@@ -8,6 +8,7 @@ __all__ = [
     "dot_product",
     "normalize_tensor",
     "log_magnitude",
+    "shift_time",
     "phase",
 ]
 
@@ -50,6 +51,11 @@ def shift_ring(x: Tensor, dim: int = -1) -> Tensor:
     return torch.roll(x, shifts=1, dims=dim)
 
 
+def shift_time(x: torch.Tensor, shift: int) -> torch.Tensor:
+    """Shifts tensor along time axis (last dim)."""
+    return torch.roll(x, shifts=shift, dims=-1)
+
+
 def dot_product(x: Tensor, y: Tensor, dim: int = -1) -> Tensor:
     """Computes dot product along the specified dimension."""
     return torch.sum(x * y, dim=dim)
@@ -69,3 +75,4 @@ def log_magnitude(stft_complex: Tensor, eps: float = 1e-5) -> Tensor:
 def phase(stft_complex: Tensor) -> Tensor:
     """Returns phase from complex STFT."""
     return torch.angle(stft_complex)
+

lt_tensor/misc_utils.py

@@ -8,7 +8,6 @@ __all__ = [
     "unfreeze_selected_weights",
     "clip_gradients",
     "detach_hidden",
-    "tensor_summary",
     "one_hot",
     "safe_divide",
     "batch_pad",
@@ -18,22 +17,20 @@ __all__ = [
     "default_device",
     "Packing",
     "Padding",
-    "MaskUtils",
-    "masked_cross_entropy",
-    "NoiseScheduler",
+    "Masking",
 ]
 
 import gc
+import sys
 import random
 import numpy as np
 from lt_utils.type_utils import is_str
 from .torch_commons import *
-from lt_utils.misc_utils import log_traceback, cache_wrapper
-from lt_utils.file_ops import load_json, load_yaml, save_json, save_yaml
-import math
+from lt_utils.misc_utils import cache_wrapper
 from lt_utils.common import *
 import torch.nn.functional as F
 
+
 def log_tensor(
     item: Union[Tensor, np.ndarray],
     title: Optional[str] = None,
@@ -64,10 +61,13 @@ def log_tensor(
                 print(f"mean: {item.mean(dim=dim):.4f}")
             except:
                 pass
-        if print_tensor:
-            print(item)
+    if print_tensor:
+        print(item)
     if has_title:
         print("".join(["-"] * _b), "\n")
+    else:
+        print("\n")
+    sys.stdout.flush()
 
 
 def set_seed(seed: int):
@@ -136,11 +136,6 @@ def detach_hidden(hidden):
         return tuple(detach_hidden(h) for h in hidden)
 
 
-def tensor_summary(tensor: torch.Tensor) -> str:
-    """Prints min/max/mean/std of a tensor for debugging."""
-    return f"Shape: {tuple(tensor.shape)}, dtype: {tensor.dtype}, min: {tensor.min():.4f}, max: {tensor.max():.4f}, mean: {tensor.mean():.4f}, std: {tensor.std():.4f}"
-
-
 def one_hot(labels: torch.Tensor, num_classes: int) -> torch.Tensor:
     """One-hot encodes a tensor of labels."""
     return F.one_hot(labels, num_classes).float()
@@ -463,7 +458,7 @@ class Padding:
         return torch.stack(padded), lengths
 
 
-class MaskUtils:
+class Masking:
 
     @staticmethod
     def apply_mask(x: Tensor, mask: Tensor, fill_value: Number = 0) -> Tensor:
@@ -546,84 +541,3 @@ class MaskUtils:
         return (
             causal_mask & pad_mask.unsqueeze(1).expand(-1, pad_mask.size(1), -1).bool()
         )
-
-
-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
-
-
-class NoiseScheduler(nn.Module):
-    def __init__(self, timesteps: int = 512):
-        super().__init__()
-
-        betas = torch.linspace(1e-4, 0.02, timesteps)
-        alphas = 1.0 - betas
-        alpha_cumprod = torch.cumprod(alphas, dim=0)
-
-        self.register_buffer("sqrt_alpha_cumprod", torch.sqrt(alpha_cumprod))
-        self.register_buffer(
-            "sqrt_one_minus_alpha_cumprod", torch.sqrt(1.0 - alpha_cumprod)
-        )
-
-        self.timesteps = timesteps
-        self.default_noise = math.sqrt(1.25)
-
-    def get_random_noise(
-        self, min_max: Tuple[float, float] = (-3, 3), seed: int = 0
-    ) -> float:
-        if seed > 0:
-            random.seed(seed)
-        return random.uniform(*min_max)
-
-    def set_noise(
-        self,
-        seed: int = 0,
-        min_max: Tuple[float, float] = (-3, 3),
-        default: bool = False,
-    ):
-        self.default_noise = (
-            math.sqrt(1.25) if default else self.get_random_noise(min_max, seed)
-        )
-
-    def forward(
-        self, x_0: Tensor, t: int, noise: Optional[Union[Tensor, float]] = None
-    ) -> Tensor:
-        if t < 0 or t >= self.timesteps:
-            raise ValueError(
-                f"Time step t={t} is out of bounds for scheduler with {self.timesteps} steps."
-            )
-
-        if noise is None:
-            noise = self.default_noise
-
-        if isinstance(noise, (float, int)):
-            noise = torch.randn_like(x_0) * noise
-
-        alpha_term = self.sqrt_alpha_cumprod[t] * x_0
-        noise_term = self.sqrt_one_minus_alpha_cumprod[t] * noise
-        return alpha_term + noise_term

lt_tensor/model_base.py

@@ -4,6 +4,7 @@ __all__ = ["Model"]
 import warnings
 from .torch_commons import *
 from lt_utils.common import *
+from lt_utils.misc_utils import log_traceback
 
 T = TypeVar("T")
 
@@ -44,7 +45,7 @@ class Model(nn.Module, ABC):
 
     def tp_apply_device_to(self):
         """Add here components that are needed to have device applied to them,
-        that usualy the '.to()' function fails to apply
+        that usually the '.to()' function fails to apply
 
         example:
         ```
@@ -54,6 +55,99 @@ class Model(nn.Module, ABC):
         """
         pass
 
+    def freeze_weight(self, weight: Union[str, nn.Module], freeze: bool):
+        assert isinstance(weight, (str, nn.Module))
+        if isinstance(weight, str):
+            if hasattr(self, weight):
+                w = getattr(self, weight)
+                if isinstance(w, nn.Module):
+                    w.requires_grad_(not freeze)
+        else:
+            weight.requires_grad_(not freeze)
+
+    def _freeze_unfreeze(
+        self,
+        weight: Union[str, nn.Module],
+        task: Literal["freeze", "unfreeze"] = "freeze",
+        _skip_except: bool = False,
+    ):
+        try:
+            assert isinstance(weight, (str, nn.Module))
+            if isinstance(weight, str):
+                w_txt = f"Failed to {task} the module '{weight}'. Reason: is not a valid attribute of {self._get_name()}"
+                if hasattr(self, weight):
+                    w_txt = f"Failed to {task} the module '{weight}'. Reason: is not a Module type."
+                    w = getattr(self, weight)
+                    if isinstance(w, nn.Module):
+                        w_txt = f"Successfully {task} the module '{weight}'."
+                        w.requires_grad_(task == "unfreeze")
+
+            else:
+                w.requires_grad_(task == "unfreeze")
+                w_txt = f"Successfully '{task}' the module '{weight}'."
+            return w_txt
+        except Exception as e:
+            if not _skip_except:
+                raise e
+            return str(e)
+
+    def freeze_weight(
+        self,
+        weight: Union[str, nn.Module],
+        _skip_except: bool = False,
+    ):
+        return self._freeze_unfreeze(weight, "freeze", _skip_except)
+
+    def unfreeze_weight(
+        self,
+        weight: Union[str, nn.Module],
+        _skip_except: bool = False,
+    ):
+        return self._freeze_unfreeze(weight, "freeze", _skip_except)
+
+    def freeze_all(self, exclude: Optional[List[str]] = None):
+        no_exclusions = not exclude
+        frozen = []
+        not_frozen = []
+        for name, param in self.named_parameters():
+            if no_exclusions:
+                try:
+                    param.requires_grad_(False)
+                    frozen.append(name)
+                except Exception as e:
+                    not_frozen.append((name, str(e)))
+            elif any(layer in name for layer in exclude):
+                try:
+                    param.requires_grad_(False)
+                    frozen.append(name)
+                except Exception as e:
+                    not_frozen.append((name, str(e)))
+            else:
+                not_frozen.append((name, "Excluded"))
+        return dict(frozen=frozen, not_frozen=not_frozen)
+
+    def unfreeze_all_except(self, exclude: Optional[list[str]] = None):
+        """Unfreezes all model parameters except specified layers."""
+        no_exclusions = not exclude
+        unfrozen = []
+        not_unfrozen = []
+        for name, param in self.named_parameters():
+            if no_exclusions:
+                try:
+                    param.requires_grad_(True)
+                    unfrozen.append(name)
+                except Exception as e:
+                    not_unfrozen.append((name, str(e)))
+            elif any(layer in name for layer in exclude):
+                try:
+                    param.requires_grad_(True)
+                    unfrozen.append(name)
+                except Exception as e:
+                    not_unfrozen.append((name, str(e)))
+            else:
+                not_unfrozen.append((name, "Excluded"))
+        return dict(unfrozen=unfrozen, not_unfrozen=not_unfrozen)
+
     def to(self, *args, **kwargs):
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
             *args, **kwargs
@@ -186,11 +280,16 @@ class Model(nn.Module, ABC):
         )
 
     def get_weights(self, module_name: Optional[str] = None) -> List[Tensor]:
-        """Returns the weights of the model entrie model or from a specified module"""
+        """Returns the weights of the model entry model or from a specified module"""
         if module_name is not None:
             assert hasattr(self, module_name), f"Module {module_name} does not exits"
             module = getattr(self, module_name)
-            return [x.data.detach() for x in module.parameters()]
+            params = []
+            if isinstance(module, nn.Module):
+                return [x.data.detach() for x in module.parameters()]
+            elif isinstance(module, (Tensor, nn.Parameter)):
+                return [module.data.detach()]
+            raise (f"{module_name} is has no weights")
         return [x.data.detach() for x in self.parameters()]
 
     def print_trainable_parameters(

lt_tensor/model_zoo/disc.py

@@ -11,37 +11,36 @@ class PeriodDiscriminator(Model):
         use_spectral_norm=False,
         kernel_size: int = 5,
         stride: int = 3,
-        initial_s: int = 32,
     ):
         super().__init__()
         self.period = period
+        self.stride = stride
+        self.kernel_size = kernel_size
         self.norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        
+        self.channels = [32, 128, 512, 1024, 1024]
         self.first_pass = nn.Sequential(
             self.norm_f(
                 nn.Conv2d(
-                    1, initial_s * 4, (kernel_size, 1), (stride, 1), padding=(2, 0)
+                    1, self.channels[0], (kernel_size, 1), (stride, 1), padding=(2, 0)
                 )
             ),
             nn.LeakyReLU(0.1),
         )
-        self._last_sz = initial_s * 4
 
-        self.convs = nn.ModuleList([self._get_next(i == 3) for i in range(4)])
+        
+        self.convs = nn.ModuleList([self._get_next(self.channels[i+1], self.channels[i], i == 3) for i in range(4)])
 
         self.post_conv = nn.Conv2d(1024, 1, (stride, 1), 1, padding=(1, 0))
-        self.kernel_size = kernel_size
-        self.stride = stride
 
-    def _get_next(self, is_last: bool = False):
-        in_dim = self._last_sz
-        self._last_sz *= 4
-        print(self._last_sz, "-----------------------")
+    def _get_next(self, out_dim:int, last_in:int, is_last: bool = False):
         stride = (self.stride, 1) if not is_last else 1
+        
         return nn.Sequential(
             self.norm_f(
                 nn.Conv2d(
-                    in_dim,
-                    self._last_sz,
+                    last_in,
+                    out_dim,
                     (self.kernel_size, 1),
                     stride,
                     padding=(2, 0),
@@ -91,6 +90,7 @@ class ScaleDiscriminator(nn.Module):
     def __init__(self, use_spectral_norm=False):
         super().__init__()
         norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.activation = nn.LeakyReLU(0.1)
         self.convs = nn.ModuleList(
             [
                 norm_f(nn.Conv1d(1, 128, 15, 1, padding=7)),
@@ -103,7 +103,6 @@ class ScaleDiscriminator(nn.Module):
             ]
         )
         self.post_conv = norm_f(nn.Conv1d(1024, 1, 3, 1, padding=1))
-        self.activation = nn.LeakyReLU(0.1)
 
     def forward(self, x: torch.Tensor):
         """
@@ -147,9 +146,10 @@ class GeneralLossDescriminator(Model):
         super().__init__()
         self.mpd = MultiPeriodDiscriminator()
         self.msd = MultiScaleDiscriminator()
+        self.print_trainable_parameters()
 
     def _get_group_(self):
         pass
 
     def forward(self, x: Tensor, y_hat: Tensor):
-        return
+        return

lt_tensor/model_zoo/istft.py

@@ -106,3 +106,44 @@ class Generator(Model):
         classname = m.__class__.__name__
         if "Conv" in classname:
             m.weight.data.normal_(mean, std)
+
+
+# Below are items found in the Rishikesh's repo that might work for this generator.
+# https://github.com/rishikksh20/iSTFTNet-pytorch/blob/781480e9563d4dff5a8cc9ef1af6c6e0cab025c8/models.py
+
+
+def feature_loss(fmap_r, fmap_g, weight=2.0):
+    """Feature matching loss between real and generated feature maps."""
+    loss = 0.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 * weight
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    """LSGAN-style loss for real and fake predictions."""
+    loss = 0.0
+    r_losses, g_losses = [], []
+
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1.0 - 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
+
+
+def generator_loss(disc_generated_outputs):
+    """LSGAN generator loss encouraging fake to look like real (close to 1)."""
+    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/noise_tools.py

@@ -0,0 +1,362 @@
+__all__ = [
+    "NoiseSchedulerA",
+    "NoiseSchedulerB",
+    "NoiseSchedulerC",
+    "add_gaussian_noise",
+    "add_uniform_noise",
+    "add_linear_noise",
+    "add_impulse_noise",
+    "add_pink_noise",
+    "add_clipped_gaussian_noise",
+    "add_multiplicative_noise",
+    "apply_noise",
+]
+
+from lt_utils.common import *
+import torch.nn.functional as F
+from .torch_commons import *
+import math
+import random
+from .misc_utils import set_seed
+
+
+def add_gaussian_noise(x: Tensor, noise_level=0.025):
+    noise = torch.randn_like(x) * noise_level
+    return x + noise
+
+
+def add_uniform_noise(x: Tensor, noise_level=0.025):
+    noise = (torch.rand_like(x) - 0.5) * 2 * noise_level
+    return x + noise
+
+
+def add_linear_noise(x, noise_level=0.05):
+    T = x.shape[-1]
+    ramp = torch.linspace(0, noise_level, T, device=x.device)
+    for _ in range(x.dim() - 1):
+        ramp = ramp.unsqueeze(0)
+    return x + ramp.expand_as(x)
+
+
+def add_impulse_noise(x: Tensor, noise_level=0.025):
+    # For image inputs
+    probs = torch.rand_like(x)
+    x_clone = x.detach().clone()
+    x_clone[probs < (noise_level / 2)] = 0.0  # salt
+    x_clone[probs > (1 - noise_level / 2)] = 1.0  # pepper
+    return x_clone
+
+
+def add_pink_noise(x: Tensor, noise_level=0.05):
+    # pink noise: divide freq spectrum by sqrt(f)
+    if x.ndim == 3:
+        x = x.view(-1, x.shape[-1])  # flatten to 2D [B*M, T]
+    pink_noised = []
+
+    for row in x:
+        white = torch.randn_like(row)
+        f = torch.fft.rfft(white)
+        freqs = torch.fft.rfftfreq(row.numel(), d=1.0)
+        freqs[0] = 1.0  # prevent div by 0
+        f /= freqs.sqrt()
+        pink = torch.fft.irfft(f, n=row.numel())
+        pink_noised.append(pink)
+
+    pink_noised = torch.stack(pink_noised, dim=0).view_as(x)
+    return x + pink_noised * noise_level
+
+
+def add_clipped_gaussian_noise(x, noise_level=0.025):
+    noise = torch.randn_like(x) * noise_level
+    return torch.clamp(x + noise, 0.0, 1.0)
+
+
+def add_multiplicative_noise(x, noise_level=0.025):
+    noise = 1 + torch.randn_like(x) * noise_level
+    return x * noise
+
+
+_VALID_NOISES = [
+    "gaussian",
+    "uniform",
+    "linear",
+    "impulse",
+    "pink",
+    "clipped_gaussian",
+    "multiplicative",
+]
+
+_NOISE_MAP = {
+    "gaussian": add_gaussian_noise,
+    "uniform": add_uniform_noise,
+    "linear": add_linear_noise,
+    "impulse": add_impulse_noise,
+    "pink": add_pink_noise,
+    "clipped_gaussian": add_clipped_gaussian_noise,
+    "multiplicative": add_multiplicative_noise,
+}
+
+_NOISE_DIM_SUPPORT = {
+    "gaussian": (1, 2),
+    "uniform": (1, 2),
+    "multiplicative": (1, 2, 3),
+    "clipped_gaussian": (1, 2, 3),
+    "linear": (2, 3),
+    "impulse": (2, 3),
+    "pink": (2, 3),
+}
+
+
+def apply_noise(
+    x: Tensor,
+    noise_type: str = "gaussian",
+    noise_level: float = 0.01,
+    seed: Optional[int] = None,
+    on_error: Literal["raise", "try_others", "return_unchanged"] = "raise",
+    _last_tries: list[str] = [],
+):
+    noise_type = noise_type.lower().strip()
+    last_tries = _last_tries
+
+    if noise_type not in _NOISE_MAP:
+        raise ValueError(f"Noise type '{noise_type}' not supported.")
+
+    # Check dimension compatibility
+    allowed_dims = _NOISE_DIM_SUPPORT.get(noise_type, (1, 2))
+    if x.ndim not in allowed_dims:
+        assert (
+            on_error != "raise"
+        ), f"Noise '{noise_type}' is not supported for {x.ndim}D input."
+        if on_error == "return_unchanged":
+            return x
+        elif on_error == "try_others":
+            remaining = [
+                n
+                for n in _VALID_NOISES
+                if n not in last_tries and x.ndim in _NOISE_DIM_SUPPORT[n]
+            ]
+            if not remaining:
+                return x
+            new_type = random.choice(remaining)
+            last_tries.append(new_type)
+            return apply_noise(
+                x, new_type, noise_level, seed, on_error, last_tries.copy()
+            )
+    try:
+        if isinstance(seed, int):
+            set_seed(seed)
+        return _NOISE_MAP[noise_type](x, noise_level)
+    except Exception as e:
+        if on_error == "raise":
+            raise e
+        elif on_error == "return_unchanged":
+            return x
+        if len(last_tries) == len(_VALID_NOISES):
+            return x
+        remaining = [n for n in _VALID_NOISES if n not in last_tries]
+        new_type = random.choice(remaining)
+        last_tries.append(new_type)
+        return apply_noise(x, new_type, noise_level, seed, on_error, last_tries.copy())
+
+
+class NoiseSchedulerA(nn.Module):
+    def __init__(self, samples: int = 64):
+        super().__init__()
+        self.base_steps = samples
+
+    def plot_noise_progression(noise_seq: list[Tensor], titles: list[str] = None):
+        import matplotlib.pyplot as plt
+
+        steps = len(noise_seq)
+        plt.figure(figsize=(15, 3))
+        for i, tensor in enumerate(noise_seq):
+            plt.subplot(1, steps, i + 1)
+            plt.imshow(tensor.squeeze().cpu().numpy(), aspect="auto", origin="lower")
+            if titles:
+                plt.title(titles[i])
+            plt.axis("off")
+        plt.tight_layout()
+        plt.show()
+
+    def forward(
+        self,
+        source_item: torch.Tensor,
+        steps: Optional[int] = None,
+        noise_type: Literal[
+            "gaussian",
+            "uniform",
+            "linear",
+            "impulse",
+            "pink",
+            "clipped_gaussian",
+            "multiplicative",
+        ] = "gaussian",
+        seed: Optional[int] = None,
+        noise_level: float = 0.01,
+        shuffle_noise_types: bool = False,
+        return_dict: bool = True,
+    ):
+        if steps is None:
+            steps = self.base_steps
+        collected = [source_item.detach().clone()]
+        noise_history = []
+        for i in range(steps):
+            if i > 0 and shuffle_noise_types:
+                noise_type = random.choice(_VALID_NOISES)
+            current, noise_name = apply_noise(
+                collected[-1],
+                noise_type,
+                noise_level,
+                seed=seed,
+                on_error="try_others",
+            )
+            noise_history.append(noise_name)
+            collected.append(current)
+
+        if return_dict:
+            return {
+                "steps": collected,
+                "history": noise_history,
+                "final": collected[-1],
+                "init": collected[0],
+            }
+        return collected, noise_history
+
+
+class NoiseSchedulerB(nn.Module):
+    def __init__(self, timesteps: int = 512):
+        super().__init__()
+
+        betas = torch.linspace(1e-4, 0.02, timesteps)
+        alphas = 1.0 - betas
+        alpha_cumprod = torch.cumprod(alphas, dim=0)
+
+        self.register_buffer("sqrt_alpha_cumprod", torch.sqrt(alpha_cumprod))
+        self.register_buffer(
+            "sqrt_one_minus_alpha_cumprod", torch.sqrt(1.0 - alpha_cumprod)
+        )
+
+        self.timesteps = timesteps
+        self.default_noise = math.sqrt(1.25)
+
+    def _get_random_noise(
+        self,
+        min_max: Tuple[float, float] = (-3, 3),
+        seed: Optional[int] = None,
+    ) -> float:
+        if isinstance(seed, int):
+            random.seed(seed)
+        return random.uniform(*min_max)
+
+    def set_noise(
+        self,
+        noise: Optional[Union[Tensor, Number]] = None,
+        seed: Optional[int] = None,
+        min_max: Tuple[float, float] = (-3, 3),
+        default: bool = False,
+    ):
+        if noise is not None:
+            self.default_noise = noise
+        else:
+            self.default_noise = (
+                math.sqrt(1.25) if default else self._get_random_noise(min_max, seed)
+            )
+
+    def forward(
+        self, x_0: Tensor, t: int, noise: Optional[Union[Tensor, float]] = None
+    ) -> Tensor:
+        apply_noise()
+        assert (
+            0 >= t < self.timesteps
+        ), f"Time step t={t} is out of bounds for scheduler with {self.timesteps} steps."
+
+        if noise is None:
+            noise = torch.randn_like(x_0) * self.default_noise
+
+        elif isinstance(noise, (float, int)):
+            noise = torch.randn_like(x_0) * noise
+
+        alpha_term = self.sqrt_alpha_cumprod[t] * x_0
+        noise_term = self.sqrt_one_minus_alpha_cumprod[t] * noise
+        return alpha_term + noise_term
+
+
+class NoiseSchedulerC(nn.Module):
+    def __init__(self, timesteps: int = 512):
+        super().__init__()
+
+        betas = torch.linspace(1e-4, 0.02, timesteps)
+        alphas = 1.0 - betas
+        alpha_cumprod = torch.cumprod(alphas, dim=0)
+
+        self.register_buffer("sqrt_alpha_cumprod", torch.sqrt(alpha_cumprod))
+        self.register_buffer(
+            "sqrt_one_minus_alpha_cumprod", torch.sqrt(1.0 - alpha_cumprod)
+        )
+
+        self.timesteps = timesteps
+        self.default_noise_strength = math.sqrt(1.25)
+        self.default_noise_type = "gaussian"
+        self.noise_seed = None
+
+    def _get_random_uniform(self, shape, min_val=-1.0, max_val=1.0):
+        return torch.empty(shape).uniform_(min_val, max_val)
+
+    def _get_noise(self, x: Tensor, noise_type: str, noise_level: float) -> Tensor:
+        # Basic noise types
+        if noise_type == "gaussian":
+            return torch.randn_like(x) * noise_level
+        elif noise_type == "uniform":
+            return self._get_random_uniform(x.shape) * noise_level
+        elif noise_type == "multiplicative":
+            return x * (1 + (torch.randn_like(x) * noise_level))
+        elif noise_type == "clipped_gaussian":
+            noise = torch.randn_like(x) * noise_level
+            return noise.clamp(-1.0, 1.0)
+        elif noise_type == "impulse":
+            mask = torch.rand_like(x) < noise_level
+            impulses = torch.randn_like(x) * noise_level
+            return x + impulses * mask
+        else:
+            raise ValueError(f"Unsupported noise type: '{noise_type}'")
+
+    def set_noise(
+        self,
+        noise_strength: Optional[Union[Tensor, float]] = None,
+        noise_type: Optional[str] = None,
+        seed: Optional[int] = None,
+        default: bool = False,
+    ):
+        if noise_strength is not None:
+            self.default_noise_strength = noise_strength
+        elif default:
+            self.default_noise_strength = math.sqrt(1.25)
+
+        if noise_type is not None:
+            self.default_noise_type = noise_type.lower().strip()
+
+        if isinstance(seed, int):
+            self.noise_seed = seed
+            torch.manual_seed(seed)
+            random.seed(seed)
+
+    def forward(
+        self,
+        x_0: Tensor,
+        t: int,
+        noise: Optional[Union[Tensor, float]] = None,
+        noise_type: Optional[str] = None,
+    ) -> Tensor:
+        assert 0 <= t < self.timesteps, f"t={t} is out of bounds [0, {self.timesteps})"
+
+        noise_type = noise_type or self.default_noise_type
+        noise_level = self.default_noise_strength
+
+        if noise is None:
+            noise = self._get_noise(x_0, noise_type, noise_level)
+        elif isinstance(noise, (float, int)):
+            noise = self._get_noise(x_0, noise_type, noise)
+
+        alpha_term = self.sqrt_alpha_cumprod[t] * x_0
+        noise_term = self.sqrt_one_minus_alpha_cumprod[t] * noise
+        return alpha_term + noise_term

lt_tensor/transform.py

@@ -8,8 +8,6 @@ __all__ = [
     "normalize",
     "min_max_scale",
     "mel_to_linear",
-    "add_noise",
-    "shift_time",
     "stretch_tensor",
     "pad_tensor",
     "get_sinusoidal_embedding",
@@ -39,16 +37,19 @@ def to_mel_spectrogram(
     f_max: Optional[float] = None,
 ) -> torch.Tensor:
     """Converts waveform to mel spectrogram."""
-    mel_spectrogram = torchaudio.transforms.MelSpectrogram(
-        sample_rate=sample_rate,
-        n_fft=n_fft,
-        hop_length=hop_length,
-        win_length=win_length,
-        n_mels=n_mels,
-        f_min=f_min,
-        f_max=f_max,
+    return (
+        torchaudio.transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            n_mels=n_mels,
+            f_min=f_min,
+            f_max=f_max,
+        )
+        .to(device=waveform.device)
+        .forward(waveform)
     )
-    return mel_spectrogram(waveform)
 
 
 def stft(
@@ -151,16 +152,6 @@ def mel_to_linear(
     return torch.matmul(mel_fb_inv, mel_spec + eps)
 
 
-def add_noise(x: torch.Tensor, noise_level: float = 0.01) -> torch.Tensor:
-    """Adds Gaussian noise to tensor."""
-    return x + noise_level * torch.randn_like(x)
-
-
-def shift_time(x: torch.Tensor, shift: int) -> torch.Tensor:
-    """Shifts tensor along time axis (last dim)."""
-    return torch.roll(x, shifts=shift, dims=-1)
-
-
 def stretch_tensor(x: torch.Tensor, rate: float, mode: str = "linear") -> torch.Tensor:
     """Time-stretch tensor using interpolation."""
     B, C, T = x.shape if x.ndim == 3 else (1, 1, x.shape[0])
@@ -274,7 +265,7 @@ def window_sumsquare(
     n_fft: int = 2048,
     dtype: torch.dtype = torch.float32,
     norm: Optional[Union[int, float]] = None,
-    device: Optional[torch.device] = "cpu",
+    device: Optional[torch.device] = None,
 ):
     if win_length is None:
         win_length = n_fft

{lt_tensor-0.0.1a4.dist-info → lt_tensor-0.0.1a6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a4
+Version: 0.0.1a6
 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.1a6.dist-info/RECORD

@@ -0,0 +1,26 @@
+lt_tensor/__init__.py,sha256=D-oEjsuKWhtk1qyiADERgNO78aRCXUJJz0hs65h8LOg,365
+lt_tensor/losses.py,sha256=TinZJP2ypZ7Tdg6d9nnFWFkPyormfgQ0Z9P2ER3sqzE,4341
+lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
+lt_tensor/math_ops.py,sha256=ewIYkvxIy_Lab_9ExjFUgLs-oYLOu8IRRDo7f1pn3i8,2248
+lt_tensor/misc_utils.py,sha256=sjWUkUaHFhaCdN4rZ6X-cQDbPieimfKchKq9VtjiwEA,17029
+lt_tensor/model_base.py,sha256=2W4m6hlvMyfRx1efWJ0NIIwctzLjL4rip208vL9_n0Y,13419
+lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
+lt_tensor/noise_tools.py,sha256=O4oq5oi0jLJuQNIuxOBZa-rB0S065QXtb1gjQUXVaLs,11212
+lt_tensor/torch_commons.py,sha256=fntsEU8lhBQo0ebonI1iXBkMbWMN3HpBsG13EWlP5s8,718
+lt_tensor/transform.py,sha256=hqsP6nXRn4nqMGkN2hBi4y-kHxEQdlIUS0y89Y1mjVI,8589
+lt_tensor/datasets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lt_tensor/datasets/audio.py,sha256=frftmRYNk0eXqHEiFggC46RMuCoGyuwBAlnPxfFsS7Y,4858
+lt_tensor/model_zoo/__init__.py,sha256=jipEk50_DTMQbGg8FnDDukxmh7Bcwvl_QVRS3rkb7aY,283
+lt_tensor/model_zoo/bsc.py,sha256=muxIR7dU-Pvf-HFE-iy3zmRb1sTJlcs1vqdlnbU1Hss,6307
+lt_tensor/model_zoo/disc.py,sha256=jZPhoSV1hlrba3ohXGutYAAcSl4pWkqGYFpOlOoN3eo,4740
+lt_tensor/model_zoo/fsn.py,sha256=5ySsg2OHjvTV_coPAdZQ0f7bz4ugJB8mDYsItmd61qA,2102
+lt_tensor/model_zoo/gns.py,sha256=Tirr_grONp_FFQ_L7K-zV2lvkaC39h8mMl4QDpx9vLQ,6028
+lt_tensor/model_zoo/istft.py,sha256=0Xms2QNPAgz_ib8XTfaWl1SCHgS53oKC6-EkDkl_qe4,4863
+lt_tensor/model_zoo/pos.py,sha256=N28v-rF8CELouYxQ9r45Jbd4ri5DNydwDgg7nzmQ4Ig,4471
+lt_tensor/model_zoo/rsd.py,sha256=5bba50g1Hm5kMexuJ4SwOIJuyQ1qJd8Acrq-Ax6CqE8,6958
+lt_tensor/model_zoo/tfrms.py,sha256=kauh-A13pk08SZ5OspEE5a-gPKD4rZr6tqMKWu3KGhk,4237
+lt_tensor-0.0.1a6.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+lt_tensor-0.0.1a6.dist-info/METADATA,sha256=-89IqEHsZD3W8moDuKWR8UodkdR2pwefqrG9C7P7y_Y,968
+lt_tensor-0.0.1a6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a6.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a6.dist-info/RECORD,,

lt_tensor-0.0.1a4.dist-info/RECORD

@@ -1,24 +0,0 @@
-lt_tensor/__init__.py,sha256=bvCjaIsYjbGFbR5MNezgLyRgN4_CsyrjmVEvuClsgOU,303
-lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
-lt_tensor/math_ops.py,sha256=ZtnJ9WB-pbFQLsXuNfQl2dAaeob5BWfxmhkwpxITUZ4,2066
-lt_tensor/misc_utils.py,sha256=e44FCQbjNHP-4WOHIbtqqH0x590DzUE6CrD_4Vl_d38,19880
-lt_tensor/model_base.py,sha256=tmRu5pTcELKMFcybOiZ1thJPuJWRSPkbUUtp9Y1NJWw,9555
-lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
-lt_tensor/torch_commons.py,sha256=fntsEU8lhBQo0ebonI1iXBkMbWMN3HpBsG13EWlP5s8,718
-lt_tensor/transform.py,sha256=IVAaQlq12OvMVhX3lX4lgsTCJYJce5n5MtMy7IK_AU4,8892
-lt_tensor/datasets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lt_tensor/datasets/audio.py,sha256=BZTceP9MlmyrVioHpWLkd_ZcyawYYZUAlVWKfKwyWAg,3318
-lt_tensor/model_zoo/__init__.py,sha256=jipEk50_DTMQbGg8FnDDukxmh7Bcwvl_QVRS3rkb7aY,283
-lt_tensor/model_zoo/bsc.py,sha256=muxIR7dU-Pvf-HFE-iy3zmRb1sTJlcs1vqdlnbU1Hss,6307
-lt_tensor/model_zoo/disc.py,sha256=ND6JR_x6b2Y1VqxZejalv8Cz5_TO3H_Z-0x6UnACbBM,4740
-lt_tensor/model_zoo/fsn.py,sha256=5ySsg2OHjvTV_coPAdZQ0f7bz4ugJB8mDYsItmd61qA,2102
-lt_tensor/model_zoo/gns.py,sha256=Tirr_grONp_FFQ_L7K-zV2lvkaC39h8mMl4QDpx9vLQ,6028
-lt_tensor/model_zoo/istft.py,sha256=RV7KVY7q4CYzzsWXH4NGJQwSqrYWwHh-16Q62lKoA2k,3594
-lt_tensor/model_zoo/pos.py,sha256=N28v-rF8CELouYxQ9r45Jbd4ri5DNydwDgg7nzmQ4Ig,4471
-lt_tensor/model_zoo/rsd.py,sha256=5bba50g1Hm5kMexuJ4SwOIJuyQ1qJd8Acrq-Ax6CqE8,6958
-lt_tensor/model_zoo/tfrms.py,sha256=kauh-A13pk08SZ5OspEE5a-gPKD4rZr6tqMKWu3KGhk,4237
-lt_tensor-0.0.1a4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-lt_tensor-0.0.1a4.dist-info/METADATA,sha256=sbT9xduzE-huVvSjnak9iCo1Eyp45bsMUarc16oTD3o,968
-lt_tensor-0.0.1a4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a4.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a4.dist-info/RECORD,,