lt-tensor 0.0.1a6__tar.gz → 0.0.1a7__tar.gz

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a6
+Version: 0.0.1a7
 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.1.a3
+Requires-Dist: lt-utils==0.0.1
 Requires-Dist: librosa>=0.11.0
 Dynamic: author
 Dynamic: classifier

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor/__init__.py

@@ -10,6 +10,8 @@ from . import (
     transform,
     noise_tools,
     losses,
+    processors,
+    datasets,
 )
 
 __all__ = [
@@ -22,4 +24,6 @@ __all__ = [
     "lr_schedulers",
     "noise_tools",
     "losses",
+    "processors",
+    "datasets",
 ]

lt_tensor-0.0.1a7/lt_tensor/datasets/audio.py

@@ -0,0 +1,109 @@
+__all__ = ["WaveMelDatasets"]
+from ..torch_commons import *
+from lt_utils.common import *
+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 WaveMelDataset(Dataset):
+    """Untested!"""
+
+    data: Union[list[dict[str, Tensor]], Tuple[Tensor, Tensor]] = []
+
+    def __init__(
+        self,
+        audio_processor: AudioProcessor,
+        path: PathLike,
+        limit_files: Optional[int] = None,
+        max_frame_length: Optional[int] = None,
+    ):
+        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 = []
+
+        for file in self.files:
+            results = self.load_data(file, max_frame_length)
+            self.data.extend(results)
+
+    def _add_dict(self, audio_raw: Tensor, audio_mel: Tensor, file: PathLike):
+        return {"mel": audio_mel, "raw": audio_raw, "file": file}
+
+    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 get_data_loader(
+        self,
+        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-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor/model_base.py

@@ -41,15 +41,15 @@ class Model(nn.Module, ABC):
     def device(self, device: Union[torch.device, str]):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
-        self.tp_apply_device_to()
+        self._apply_device_to()
 
-    def tp_apply_device_to(self):
+    def _apply_device_to(self):
         """Add here components that are needed to have device applied to them,
         that usually the '.to()' function fails to apply
 
         example:
         ```
-        def tp_apply_device_to(self):
+        def _apply_device_to(self):
             self.my_tensor = self.my_tensor.to(device=self.device)
         ```
         """

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor/noise_tools.py

@@ -128,7 +128,7 @@ def apply_noise(
             on_error != "raise"
         ), f"Noise '{noise_type}' is not supported for {x.ndim}D input."
         if on_error == "return_unchanged":
-            return x
+            return x, None
         elif on_error == "try_others":
             remaining = [
                 n
@@ -136,27 +136,33 @@ def apply_noise(
                 if n not in last_tries and x.ndim in _NOISE_DIM_SUPPORT[n]
             ]
             if not remaining:
-                return x
+                return x, None
             new_type = random.choice(remaining)
             last_tries.append(new_type)
-            return apply_noise(
-                x, new_type, noise_level, seed, on_error, last_tries.copy()
+            return (
+                apply_noise(
+                    x, new_type, noise_level, seed, on_error, last_tries.copy()
+                ),
+                noise_type,
             )
     try:
         if isinstance(seed, int):
             set_seed(seed)
-        return _NOISE_MAP[noise_type](x, noise_level)
+        return _NOISE_MAP[noise_type](x, noise_level), noise_type
     except Exception as e:
         if on_error == "raise":
             raise e
         elif on_error == "return_unchanged":
-            return x
+            return x, None
         if len(last_tries) == len(_VALID_NOISES):
-            return x
+            return x, None
         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())
+        return (
+            apply_noise(x, new_type, noise_level, seed, on_error, last_tries.copy()),
+            noise_type,
+        )
 
 
 class NoiseSchedulerA(nn.Module):

lt_tensor-0.0.1a7/lt_tensor/processors/__init__.py

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

{lt_tensor-0.0.1a6/lt_tensor/datasets → lt_tensor-0.0.1a7/lt_tensor/processors}/audio.py

@@ -1,33 +1,40 @@
 __all__ = ["AudioProcessor"]
 from ..torch_commons import *
-import torchaudio
 from lt_utils.common import *
-import librosa
 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
+from ..misc_utils import log_tensor
+import librosa
+import torchaudio
 
+from ..transform import InverseTransformConfig, InverseTransform
+from lt_utils.file_ops import FileScan, get_file_name, path_to_str
 
-class AudioProcessor:
+from lt_tensor.model_base import Model
 
+
+class AudioProcessor(Model):
     def __init__(
         self,
         sample_rate: int = 24000,
         n_mels: int = 80,
         n_fft: int = 1024,
-        win_length: int = 1024,
-        hop_length: int = 256,
+        win_length: Optional[int] = None,
+        hop_length: Optional[int] = None,
         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,
+        inverse_transform_config: Union[
+            Dict[str, Union[Number, Tensor, bool]], InverseTransformConfig
+        ] = dict(n_fft=16, hop_length=4, win_length=16, center=True),
+        *__,
+        **_,
     ):
+        super().__init__()
+        assert isinstance(inverse_transform_config, (InverseTransformConfig, dict))
         self.mean = mean
         self.std = std
         self.n_mels = n_mels
@@ -36,7 +43,8 @@ class AudioProcessor:
         self.f_min = f_min
         self.f_max = f_max
         self.n_iter = n_iter
-        self.hop_length = hop_length
+        self.hop_length = hop_length or n_fft // 4
+        self.win_length = win_length or n_fft
         self.sample_rate = sample_rate
         self.mel_spec = torchaudio.transforms.MelSpectrogram(
             sample_rate=sample_rate,
@@ -63,21 +71,26 @@ class AudioProcessor:
             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
-        )
+        if isinstance(inverse_transform_config, dict):
+            inverse_transform_config = InverseTransformConfig(
+                **inverse_transform_config
+            )
+        self._inv_transform = InverseTransform(**inverse_transform_config.to_dict())
 
-    def inverse_transform(self, spec: Tensor, phase: Tensor):
-        return self._inverse_transform(spec, phase)
+    def inverse_transform(self, spec: Tensor, phase: Tensor, *_, **kwargs):
+        return self._inv_transform(spec, phase, **kwargs)
 
     def compute_mel(
-        self,
-        wave: Tensor,
+        self, wave: Tensor, base: float = 1e-6, add_base: bool = False
     ) -> 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]
+        wave_device = wave.device
+        mel_tensor = self.mel_spec(wave.to(self.device))  # [M, ML]
+        if not add_base:
+            return (mel_tensor - self.mean) / self.std
+        return ((torch.log(base + mel_tensor.unsqueeze(0)) - self.mean) / self.std).to(
+            device=wave_device
+        )
 
     def reverse_mel(self, mel: Tensor, n_iter: Optional[int] = None):
         if isinstance(n_iter, int) and n_iter != self.n_iter:
@@ -143,12 +156,38 @@ class AudioProcessor:
             and not any(list(map(lambda y: y in x), text_file_patterns))
         ]
         audio_files = self.find_audios(path, additional_extensions)
-        text_files = []
+        results = []
         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
+                    results.append((audio, path_to_str(possible_txt_file)))
+                    break
+        return results
+
+    def stft_loss(self, signal: Tensor, ground: Tensor, base: float = 1e-5):
+        sig_mel = self(signal, base)
+        gnd_mel = self(ground, base)
+        return torch.norm(gnd_mel - sig_mel, p=1) / torch.norm(gnd_mel, p=1)
+
+    # def forward(self, wave: Tensor, base: Optional[float] = None):
+    def forward(
+        self,
+        *inputs: Union[Tensor, float],
+        ap_task: Literal[
+            "get_mel", "get_loss", "inv_transform", "revert_mel"
+        ] = "get_mel",
+        **inputs_kwargs,
+    ):
+        if ap_task == "get_mel":
+            return self.compute_mel(*inputs, **inputs_kwargs)
+        elif ap_task == "get_loss":
+            return self.stft_loss(*inputs, **inputs_kwargs)
+        elif ap_task == "inv_transform":
+            return self.inverse_transform(*inputs, **inputs_kwargs)
+        elif ap_task == "revert_mel":
+            return self.reverse_mel(*inputs, **inputs_kwargs)
+        else:
+            raise ValueError(f"Invalid task '{ap_task}'")

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor/transform.py

@@ -15,6 +15,8 @@ __all__ = [
     "normalize",
     "window_sumsquare",
     "inverse_transform",
+    "InverseTransformConfig",
+    "InverseTransform",
     "stft_istft_rebuild",
 ]
 
@@ -23,12 +25,15 @@ import torchaudio
 import math
 from .misc_utils import log_tensor
 from lt_utils.common import *
+from lt_utils.misc_utils import cache_wrapper, default
 import torch.nn.functional as F
+from .model_base import Model
+import warnings
 
 
 def to_mel_spectrogram(
     waveform: torch.Tensor,
-    sample_rate: int = 22050,
+    sample_rate: int = 24000,
     n_fft: int = 1024,
     hop_length: Optional[int] = None,
     win_length: Optional[int] = None,
@@ -198,8 +203,7 @@ def generate_window(
         return torch.ones(1, device=device)
 
     n = torch.arange(M, dtype=torch.float32, device=device)
-    window = alpha - (1.0 - alpha) * torch.cos(2.0 * math.pi * n / (M - 1))
-    return window
+    return alpha - (1.0 - alpha) * torch.cos(2.0 * math.pi * n / (M - 1))
 
 
 def pad_center(tensor: torch.Tensor, size: int, axis: int = -1) -> torch.Tensor:
@@ -260,8 +264,8 @@ def normalize(
 def window_sumsquare(
     window_spec: Union[str, int, float, Callable, List[Any], Tuple[Any, ...]],
     n_frames: int,
-    hop_length: int = 300,
-    win_length: int = 1200,
+    hop_length: int = 256,
+    win_length: int = 1024,
     n_fft: int = 2048,
     dtype: torch.dtype = torch.float32,
     norm: Optional[Union[int, float]] = None,
@@ -294,9 +298,9 @@ def window_sumsquare(
 def inverse_transform(
     spec: Tensor,
     phase: Tensor,
-    n_fft: int = 2048,
-    hop_length: int = 300,
-    win_length: int = 1200,
+    n_fft: int = 1024,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
     length: Optional[Any] = None,
     window: Optional[Tensor] = None,
 ):
@@ -310,3 +314,168 @@ def inverse_transform(
         window=window,
         length=length,
     )
+
+
+def is_nand(a: bool, b: bool):
+    """[a -> b = result]
+    ```
+    False -> False = True
+    False -> True = True
+    True -> False = True
+    True -> True = False
+    ```
+    """
+    return not (a and b)
+
+
+class InverseTransformConfig:
+    def __init__(
+        self,
+        n_fft: int = 1024,
+        hop_length: Optional[int] = None,
+        win_length: Optional[int] = None,
+        length: Optional[int] = None,
+        window: Optional[Tensor] = None,
+        onesided: Optional[bool] = None,
+        return_complex: bool = False,
+        normalized: bool = False,
+        center: bool = True,
+    ):
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.length = length
+        self.onesided = onesided
+        self.return_complex = return_complex
+        self.normalized = normalized
+        self.center = center
+        self.window = window
+
+    def to_dict(self):
+        return self.__dict__.copy()
+
+
+class InverseTransform(Model):
+    def __init__(
+        self,
+        n_fft: int = 1024,
+        hop_length: Optional[int] = None,
+        win_length: Optional[int] = None,
+        length: Optional[int] = None,
+        window: Optional[Tensor] = None,
+        onesided: Optional[bool] = None,
+        return_complex: bool = False,
+        normalized: bool = False,
+        center: bool = True,
+    ):
+        """
+        Module for inverting a magnitude + phase spectrogram to a waveform using ISTFT.
+
+        This class encapsulates common ISTFT parameters at initialization and applies
+        the inverse transformation in the `forward()` method with minimal per-call overhead.
+
+        Parameters
+        ----------
+        n_fft : int, optional
+            Size of FFT to use during inversion. Default is 1024.
+        hop_length : int, optional
+            Number of audio samples between STFT columns. Defaults to `n_fft`.
+        win_length : int, optional
+            Size of the window function. Defaults to `n_fft // 4`.
+        length : int, optional
+            Output waveform length. If not provided, length will be inferred.
+        window : Tensor, optional
+            Custom window tensor. If None, a Hann window is used.
+        onesided : bool, optional
+            Whether the input STFT was onesided. Required only for consistency checks.
+        return_complex : bool, default=False
+            Must be False if `onesided` is True. Not used internally.
+        normalized : bool, default=False
+            Whether the STFT was normalized.
+        center : bool, default=True
+            Whether the signal was padded during STFT.
+
+        Methods
+        -------
+        forward(spec, phase)
+            Applies ISTFT using stored settings on the given magnitude and phase tensors.
+        update_settings(...)
+            Updates ISTFT parameters dynamically (used internally during forward).
+        """
+        super().__init__()
+        assert window is None or isinstance(window, Tensor)
+        assert any(
+            (
+                (not return_complex and not onesided),
+                (not onesided and return_complex),
+                (not return_complex and onesided),
+            )
+        )
+        self.n_fft = n_fft
+        self.length = length
+        self.win_length = win_length or n_fft // 4
+        self.hop_length = hop_length or n_fft
+        self.center = center // 4
+        self.return_complex = return_complex
+        self.onesided = onesided
+        self.normalized = normalized
+        self.window = torch.hann_window(win_length) if window is None else window
+
+    def _apply_device_to(self):
+        """Applies to device while used with module `Model`"""
+        self.window = self.window.to(device=self.device)
+
+    def update_settings(
+        self,
+        *,
+        n_fft: Optional[int] = None,
+        hop_length: Optional[int] = None,
+        win_length: Optional[int] = None,
+        length: Optional[int] = None,
+        window: Optional[Tensor] = None,
+        onesided: Optional[bool] = None,
+        return_complex: Optional[bool] = None,
+        center: Optional[bool] = None,
+        normalized: Optional[bool] = None,
+        **_,
+    ):
+
+        self.kwargs = dict(
+            n_fft=default(n_fft, self.n_fft),
+            hop_length=default(hop_length, self.hop_length),
+            win_length=default(win_length, self.win_length),
+            length=default(length, self.length),
+            window=default(window, self.window),
+            onesided=default(onesided, self.onesided),
+            return_complex=default(return_complex, self.return_complex),
+            center=default(center, self.center),
+            normalized=default(normalized, self.normalized),
+        )
+        if self.kwargs["onesided"] and self.kwargs["return_complex"]:
+            warnings.warn(
+                "You cannot use return_complex with `onesided` enabled. `return_complex` is set to False."
+            )
+            self.kwargs["return_complex"] = False
+
+    def forward(self, spec: Tensor, phase: Tensor, **kwargs):
+        """
+        Perform the inverse short-time Fourier transform.
+
+        Parameters
+        ----------
+        spec : Tensor
+            Magnitude spectrogram of shape (batch, freq, time).
+        phase : Tensor
+            Phase angles tensor, same shape as `spec`, in radians.
+        **kwargs : dict, optional
+            Optional ISTFT override parameters (same as in `update_settings`).
+
+        Returns
+        -------
+        Tensor
+            Time-domain waveform reconstructed from `spec` and `phase`.
+        """
+        if kwargs:
+            self.update_settings(**kwargs)
+
+        return torch.istft(spec * torch.exp(phase * 1j), **self.kwargs)

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a6
+Version: 0.0.1a7
 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.1.a3
+Requires-Dist: lt-utils==0.0.1
 Requires-Dist: librosa>=0.11.0
 Dynamic: author
 Dynamic: classifier

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor.egg-info/SOURCES.txt

@@ -26,4 +26,6 @@ lt_tensor/model_zoo/gns.py
 lt_tensor/model_zoo/istft.py
 lt_tensor/model_zoo/pos.py
 lt_tensor/model_zoo/rsd.py
-lt_tensor/model_zoo/tfrms.py
+lt_tensor/model_zoo/tfrms.py
+lt_tensor/processors/__init__.py
+lt_tensor/processors/audio.py

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/lt_tensor.egg-info/requires.txt

@@ -4,5 +4,5 @@ numpy>=1.26.4
 tokenizers
 pyyaml>=6.0.0
 numba>0.60.0
-lt-utils==0.0.1.a3
+lt-utils==0.0.1
 librosa>=0.11.0

{lt_tensor-0.0.1a6 → lt_tensor-0.0.1a7}/setup.py

@@ -4,7 +4,7 @@ with open("README.md", "r", encoding="utf-8") as f:
     long_description = f.read()
 
 setup(
-    version="0.0.1a6",
+    version="0.0.1a7",
     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.1.a3",
+        "lt-utils==0.0.1",
         "librosa>=0.11.0",
     ],
     author="gr1336",