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

lt_tensor/__init__.py

@@ -1,13 +1,17 @@
 __version__ = "0.0.1a"
 
 from . import (
+    lr_schedulers,
     model_zoo,
     model_base,
     math_ops,
     misc_utils,
     monotonic_align,
     transform,
-    lr_schedulers,
+    noise_tools,
+    losses,
+    processors,
+    datasets,
 )
 
 __all__ = [
@@ -18,4 +22,8 @@ __all__ = [
     "monotonic_align",
     "transform",
     "lr_schedulers",
+    "noise_tools",
+    "losses",
+    "processors",
+    "datasets",
 ]

lt_tensor/datasets/audio.py

@@ -1,110 +1,109 @@
-__all__ = ["AudioProcessor"]
+__all__ = ["WaveMelDatasets"]
 from ..torch_commons import *
-import torchaudio
 from lt_utils.common import *
-import librosa
-from lt_utils.type_utils import is_file
-from torchaudio.functional import resample
-from ..transform import inverse_transform
+import random
+from torch.utils.data import Dataset, DataLoader, Sampler
+from ..processors import AudioProcessor
+import torch.nn.functional as FT
+from ..misc_utils import log_tensor
 
 
-class AudioProcessor:
+class WaveMelDataset(Dataset):
+    """Untested!"""
+
+    data: Union[list[dict[str, Tensor]], Tuple[Tensor, Tensor]] = []
 
     def __init__(
         self,
-        sample_rate: int = 24000,
-        n_mels: int = 80,
-        n_fft: int = 1024,
-        win_length: int = 1024,
-        hop_length: int = 256,
-        f_min: float = 0,
-        f_max: float | None = None,
-        n_iter: int = 32,
-        center: bool = True,
-        mel_scale: Literal["htk", "slaney"] = "htk",
-        inv_n_fft: int = 16,
-        inv_hop: int = 4,
-        std: int = 4,
-        mean: int = -4,
+        audio_processor: AudioProcessor,
+        path: PathLike,
+        limit_files: Optional[int] = None,
+        max_frame_length: Optional[int] = None,
     ):
-        self.mean = mean
-        self.std = std
-        self.n_mels = n_mels
-        self.n_fft = n_fft
-        self.n_stft = n_fft // 2 + 1
-        self.f_min = f_min
-        self.f_max = f_max
-        self.n_iter = n_iter
-        self.hop_length = hop_length
-        self.sample_rate = sample_rate
-        self.mel_spec = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate,
-            n_mels=n_mels,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            center=center,
-            f_min=f_min,
-            f_max=f_max,
-            mel_scale=mel_scale,
-        )
-        self.mel_rscale = torchaudio.transforms.InverseMelScale(
-            n_stft=self.n_stft,
-            n_mels=n_mels,
-            sample_rate=sample_rate,
-            f_min=f_min,
-            f_max=f_max,
-            mel_scale=mel_scale,
-        )
-        self.giffin_lim = torchaudio.transforms.GriffinLim(
-            n_fft=n_fft,
-            n_iter=n_iter,
-            win_length=win_length,
-            hop_length=hop_length,
-        )
-        self._inverse_transform = lambda x, y: inverse_transform(
-            x, y, inv_n_fft, inv_hop, inv_n_fft
+        super().__init__()
+        assert max_frame_length is None or max_frame_length >= (
+            (audio_processor.n_fft // 2) + 1
         )
+        self.post_n_fft = (audio_processor.n_fft // 2) + 1
+        self.ap = audio_processor
+        self.files = self.ap.find_audios(path)
+        if limit_files:
+            random.shuffle(self.files)
+            self.files = self.files[:limit_files]
+        self.data = []
 
-    def inverse_transform(self, spec: Tensor, phase: Tensor):
-        return self._inverse_transform(spec, phase)
+        for file in self.files:
+            results = self.load_data(file, max_frame_length)
+            self.data.extend(results)
 
-    def compute_mel(
-        self,
-        wave: Tensor,
-    ) -> Tensor:
-        """Returns: [B, M, ML]"""
-        mel_tensor = self.mel_spec(wave)  # [M, ML]
-        mel_tensor = (mel_tensor - self.mean) / self.std
-        return mel_tensor  # [B, M, ML]
+    def _add_dict(self, audio_raw: Tensor, audio_mel: Tensor, file: PathLike):
+        return {"mel": audio_mel, "raw": audio_raw, "file": file}
 
-    def reverse_mel(self, mel: Tensor, n_iter: Optional[int] = None):
-        if isinstance(n_iter, int) and n_iter != self.n_iter:
-            self.giffin_lim = torchaudio.transforms.GriffinLim(
-                n_fft=self.n_fft,
-                n_iter=n_iter,
-                win_length=self.win_length,
-                hop_length=self.hop_length,
-            )
-            self.n_iter = n_iter
-        return self.giffin_lim.forward(
-            self.mel_rscale(mel),
-        )
+    def load_data(self, file: PathLike, audio_frames_limit: Optional[int] = None):
+        initial_audio = self.ap.load_audio(file)
+        if not audio_frames_limit or initial_audio.shape[-1] <= audio_frames_limit:
+            audio_mel = self.ap.compute_mel(initial_audio, add_base=True)
+            return [self._add_dict(initial_audio, audio_mel, file)]
+        results = []
+        for fragment in torch.split(
+            initial_audio, split_size_or_sections=audio_frames_limit, dim=-1
+        ):
+            if fragment.shape[-1] < self.post_n_fft:
+                # sometimes the tensor will be too small to be able to pass on mel
+                continue
+            audio_mel = self.ap.compute_mel(fragment, add_base=True)
+            results.append(self._add_dict(fragment, audio_mel, file))
+        return results
 
-    def load_audio(
+    def get_data_loader(
         self,
-        path: PathLike,
-        top_db: float = 30,
-    ) -> Tensor:
-        is_file(path, True)
-        wave, sr = librosa.load(str(path), sr=self.sample_rate)
-        wave, _ = librosa.effects.trim(wave, top_db=top_db)
-        return (
-            torch.from_numpy(
-                librosa.resample(wave, orig_sr=sr, target_sr=self.sample_rate)
-                if sr != self.sample_rate
-                else wave
-            )
-            .float()
-            .unsqueeze(0)
+        batch_size: int = 1,
+        shuffle: Optional[bool] = None,
+        sampler: Optional[Union[Sampler, Iterable]] = None,
+        batch_sampler: Optional[Union[Sampler[list], Iterable[list]]] = None,
+        num_workers: int = 0,
+        pin_memory: bool = False,
+        drop_last: bool = False,
+        timeout: float = 0,
+    ):
+        return DataLoader(
+            self,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            sampler=sampler,
+            batch_sampler=batch_sampler,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            drop_last=drop_last,
+            timeout=timeout,
+            collate_fn=self.collate_fn,
         )
+
+    @staticmethod
+    def collate_fn(batch: Sequence[Dict[str, Tensor]]):
+        mels = []
+        audios = []
+        files = []
+        for x in batch:
+            mels.append(x["mel"])
+            audios.append(x["raw"])
+            files.append(x["file"])
+        # Find max time in mel (dim -1), and max audio length
+        max_mel_len = max([m.shape[-1] for m in mels])
+        max_audio_len = max([a.shape[-1] for a in audios])
+
+        padded_mels = torch.stack(
+            [FT.pad(m, (0, max_mel_len - m.shape[-1])) for m in mels]
+        )  # shape: [B, 80, T_max]
+
+        padded_audios = torch.stack(
+            [FT.pad(a, (0, max_audio_len - a.shape[-1])) for a in audios]
+        )  # shape: [B, L_max]
+
+        return padded_mels, padded_audios, files
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]

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