lt-tensor 0.0.1a11__py3-none-any.whl → 0.0.1a13__py3-none-any.whl

lt_tensor/model_zoo/{tfrms.py → transformer.py}

@@ -7,12 +7,12 @@ __all__ = [
 ]
 
 import math
-from ..torch_commons import *
-from ..model_base import Model
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
 from lt_utils.misc_utils import default
 from typing import Optional
-from .pos import *
-from .bsc import FeedForward
+from lt_tensor.model_zoo.pos_encoder import *
+from lt_tensor.model_zoo.basic import FeedForward
 
 
 def init_weights(module):

lt_tensor/noise_tools.py

@@ -14,10 +14,10 @@ __all__ = [
 
 from lt_utils.common import *
 import torch.nn.functional as F
-from .torch_commons import *
+from lt_tensor.torch_commons import *
 import math
 import random
-from .misc_utils import set_seed
+from lt_tensor.misc_utils import set_seed
 
 
 def add_gaussian_noise(x: Tensor, noise_level=0.025):

lt_tensor/processors/audio.py

@@ -5,11 +5,14 @@ from lt_utils.type_utils import is_file, is_array
 from lt_tensor.misc_utils import log_tensor
 import librosa
 import torchaudio
-
+import numpy as np
 from lt_tensor.transform import InverseTransformConfig, InverseTransform
 from lt_utils.file_ops import FileScan, get_file_name, path_to_str
-
+from torchaudio.functional import detect_pitch_frequency
 from lt_tensor.model_base import Model
+import torch.nn.functional as F
+
+DEFAULT_DEVICE = torch.tensor([0]).device
 
 
 class AudioProcessor(Model):
@@ -21,32 +24,30 @@ class AudioProcessor(Model):
         win_length: Optional[int] = None,
         hop_length: Optional[int] = None,
         f_min: float = 0,
-        f_max: float | None = None,
+        f_max: float = 12000.0,
         center: bool = True,
         mel_scale: Literal["htk", "slaney"] = "htk",
         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),
         n_iter: int = 32,
-        *__,
-        **_,
+        window: Optional[Tensor] = None,
+        normalized: bool = False,
+        onesided: Optional[bool] = None,
     ):
         super().__init__()
-        assert isinstance(inverse_transform_config, (InverseTransformConfig, dict))
         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.f_max = max(min(f_max, 12000), self.f_min + 1)
         self.n_iter = n_iter
         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(
+        self.center = center
+        self._mel_spec = torchaudio.transforms.MelSpectrogram(
             sample_rate=sample_rate,
             n_mels=n_mels,
             n_fft=n_fft,
@@ -71,14 +72,260 @@ class AudioProcessor(Model):
             win_length=win_length,
             hop_length=hop_length,
         )
-        if isinstance(inverse_transform_config, dict):
-            inverse_transform_config = InverseTransformConfig(
-                **inverse_transform_config
+        self.normalized = normalized
+        self.onesided = onesided
+
+        self.register_buffer(
+            "window",
+            (torch.hann_window(self.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 not any([device is not None, dtype is not None]):
+            return converted
+        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 not any([device is not None, dtype is not None]):
+            return stacked
+        return stacked.to(device=device, dtype=dtype)
+
+    def to_numpy_safe(self, tensor: Tensor):
+        return tensor.detach().to(DEFAULT_DEVICE).numpy(force=True)
+
+    def compute_rms(
+        self, audio: Union[Tensor, np.ndarray], mel: Optional[Tensor] = None
+    ):
+        default_dtype = audio.dtype
+        default_device = audio.device
+        assert audio.ndim in [1, 2], (
+            f"Audio should have 1D for unbatched and 2D for batched"
+            ", received instead a: {audio.ndim}D"
+        )
+        if mel is not None:
+            assert mel.ndim in [2, 3], (
+                "Mel spectogram should have 2D dim for non-batched or 3D dim for both non-batched or batched"
+                f". Received instead {mel.ndim}D."
+            )
+        if audio.ndim == 2:
+            B = audio.shape[0]
+        else:
+            B = 1
+            audio = audio.unsqueeze(0)
+
+        if mel is not None:
+            if mel.ndim == 2:
+                assert B == 1, "Batch from mel and audio must be the same!"
+                mel = mel.unsqueeze(0)
+            else:
+                assert B == mel.shape[0], "Batch from mel and audio must be the same!"
+            mel = self.to_numpy_safe(mel)
+            gt_mel = lambda idx: mel[idx, :, :]
+
+        else:
+            gt_mel = lambda idx: None
+        audio = self.to_numpy_safe(audio)
+        if B == 1:
+            _r = librosa.feature.rms(
+                y=audio, frame_length=self.n_fft, hop_length=self.hop_length
+            )[0]
+            rms = self.from_numpy(_r, default_device, default_dtype)
+
+        else:
+            rms_ = []
+            for i in range(B):
+                _r = librosa.feature.rms(
+                    y=audio[i, :],
+                    S=gt_mel(i),
+                    frame_length=self.n_fft,
+                    hop_length=self.hop_length,
+                )[0]
+                rms_.append(_r)
+            rms = self.from_numpy_batch(rms_, default_device, default_dtype)
+
+        return rms
+
+    def compute_pitch(
+        self,
+        audio: Tensor,
+    ):
+        default_dtype = audio.dtype
+        default_device = audio.device
+        assert audio.ndim in [1, 2], (
+            f"Audio should have 1D for unbatched and 2D for batched"
+            ", received instead a: {audio.ndim}D"
+        )
+        if audio.ndim == 2:
+            B = audio.shape[0]
+        else:
+            B = 1
+        fmin = max(self.f_min, 80)
+        if B == 1:
+            f0 = self.from_numpy(
+                librosa.yin(
+                    self.to_numpy_safe(audio),
+                    fmin=fmin,
+                    fmax=self.f_max,
+                    frame_length=self.n_fft,
+                    sr=self.sample_rate,
+                    hop_length=self.hop_length,
+                    center=self.center,
+                ),
+                default_device,
+                default_dtype,
+            )
+
+        else:
+            f0_ = []
+            for i in range(B):
+                r = librosa.yin(
+                    self.to_numpy_safe(audio[i, :]),
+                    fmin=fmin,
+                    fmax=self.f_max,
+                    frame_length=self.n_fft,
+                    sr=self.sample_rate,
+                    hop_length=self.hop_length,
+                    center=self.center,
+                )
+                f0_.append(r)
+            f0 = self.from_numpy_batch(f0_, default_device, default_dtype)
+
+        # librosa.pyin(self.f_min, self.f_max)
+        return f0  # dict(f0=f0, attention_mask=f0 != f_max)
+
+    def compute_pitch_torch(self, audio: Tensor):
+        return detect_pitch_frequency(
+            audio,
+            sample_rate=self.sample_rate,
+            frame_time=self.n_fft,
+            win_length=self.win_length,
+            freq_low=max(self.f_min, 35),
+            freq_high=self.f_max,
+        )
+
+    def interpolate_tensor(
+        self,
+        tensor: Tensor,
+        target_len: int,
+        mode: Literal[
+            "nearest",
+            "linear",
+            "bilinear",
+            "bicubic",
+            "trilinear",
+            "area",
+            "nearest-exact",
+        ] = "nearest",
+        align_corners: Optional[bool] = None,
+        scale_factor: Optional[list[float]] = None,
+        recompute_scale_factor: Optional[bool] = None,
+        antialias: bool = False,
+    ):
+        """
+        The modes available for upsampling are: `nearest`, `linear` (3D-only),
+        `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only)
+        """
+
+        if tensor.ndim == 2:  # [1, T]
+            tensor = tensor.unsqueeze(1)  # [1, 1, T]
+        return F.interpolate(
+            tensor,
+            size=target_len,
+            mode=mode,
+            align_corners=align_corners,
+            scale_factor=scale_factor,
+            recompute_scale_factor=recompute_scale_factor,
+            antialias=antialias,
+        )
+
+    def inverse_transform(
+        self,
+        spec: Tensor,
+        phase: Tensor,
+        n_fft: Optional[int] = None,
+        hop_length: Optional[int] = None,
+        win_length: Optional[int] = None,
+        length: Optional[int] = None,
+        *,
+        _recall: bool = False,
+    ):
+        try:
+            return torch.istft(
+                spec * torch.exp(phase * 1j),
+                n_fft=n_fft or self.n_fft,
+                hop_length=hop_length or self.hop_length,
+                win_length=win_length or self.win_length,
+                window=torch.hann_window(
+                    win_length or self.win_length, device=spec.device
+                ),
+                center=self.center,
+                normalized=self.normalized,
+                onesided=self.onesided,
+                length=length,
+                return_complex=False,
             )
-        self._inv_transform = InverseTransform(**inverse_transform_config.to_dict())
+        except RuntimeError as e:
+            if not _recall and spec.device != self.window.device:
+                self.window = self.window.to(spec.device)
+                return self.inverse_transform(
+                    spec, phase, n_fft, hop_length, win_length, length, _recall=True
+                )
+            raise e
 
-    def inverse_transform(self, spec: Tensor, phase: Tensor, *_, **kwargs):
-        return self._inv_transform(spec, phase, **kwargs)
+    def normalize_audio(
+        self,
+        wave: Tensor,
+        length: Optional[int] = None,
+        *,
+        _recall: bool = False,
+    ):
+        try:
+            spectrogram = torch.stft(
+                input=wave,
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                window=self.window,
+                center=self.center,
+                pad_mode="reflect",
+                normalized=self.normalized,
+                onesided=self.onesided,
+                return_complex=True,  # needed for the istft
+            )
+            return torch.istft(
+                spectrogram
+                * torch.full(
+                    spectrogram.size(),
+                    fill_value=1,
+                    device=spectrogram.device,
+                ),
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                window=self.window,
+                length=length,
+                center=self.center,
+                normalized=self.normalized,
+                onesided=self.onesided,
+                return_complex=False,
+            )
+        except RuntimeError as e:
+            if not _recall and wave.device != self.window.device:
+                self.window = self.window.to(wave.device)
+                return self.normalize_audio(wave, length, _recall=True)
+            raise e
 
     def compute_mel(
         self,
@@ -87,21 +334,16 @@ class AudioProcessor(Model):
         add_base: bool = True,
     ) -> Tensor:
         """Returns: [B, M, ML]"""
-        mel_tensor = self.mel_spec(wave.to(self.device))  # [M, ML]
+        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
         ).squeeze()
 
-    def reverse_mel(self, mel: Tensor, n_iter: Optional[int] = None):
+    def inverse_mel_spectogram(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.giffin_lim.n_iter = n_iter
             self.n_iter = n_iter
         return self.giffin_lim.forward(
             self.mel_rscale(mel),
@@ -111,21 +353,26 @@ class AudioProcessor(Model):
         self,
         path: PathLike,
         top_db: float = 30,
+        normalize: bool = False,
+        alpha: float = 1.0,
     ) -> 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)
-        )
+        if sr != self.sample_rate:
+            wave = librosa.resample(wave, orig_sr=sr, target_sr=self.sample_rate)
+        if normalize:
+            wave = librosa.util.normalize(wave)
+        if alpha not in [0.0, 1.0]:
+            wave = wave * alpha
+        return torch.from_numpy(wave).float().unsqueeze(0)
 
-    def find_audios(self, path: PathLike, additional_extensions: List[str] = []):
+    def find_audios(
+        self,
+        path: PathLike,
+        additional_extensions: List[str] = [],
+        maximum: int | None = None,
+    ):
         extensions = [
             "*.wav",
             "*.aac",
@@ -141,6 +388,7 @@ class AudioProcessor(Model):
         return FileScan.files(
             path,
             extensions,
+            maximum,
         )
 
     def find_audio_text_pairs(
@@ -169,57 +417,6 @@ class AudioProcessor(Model):
                     break
         return results
 
-    def slice_mismatch_outputs(
-        self,
-        tensor_1: Tensor,
-        tensor_2: Tensor,
-        smallest_size: Optional[int] = None,
-        left_to_right: bool = True,
-    ):
-        assert tensor_1.ndim == tensor_2.ndim, (
-            "Tensors must have the same dimentions to be sliced! \n"
-            f"Received instead a tensor_1 with {tensor_1.ndim}D and tensor_2 with {tensor_1.ndim}D."
-        )
-        dim = tensor_1.ndim
-        assert dim < 5, (
-            "Excpected to receive tensors with from 1D up to 4D. "
-            f"Received instead a {dim}D tensor."
-        )
-
-        if tensor_1.shape[-1] == tensor_2.shape[-1]:
-            return tensor_1, tensor_2
-
-        if smallest_size is None:
-            smallest_size = min(tensor_1.shape[-1], tensor_2.shape[-1])
-        if dim == 0:
-            tensor_1 = tensor_1.unsqueeze(0)
-            tensor_2 = tensor_2.unsqueeze(0)
-            dim = 1
-
-        if dim == 1:
-            if left_to_right:
-                return tensor_1[:smallest_size], tensor_2[:smallest_size]
-            return tensor_1[-smallest_size:], tensor_2[-smallest_size:]
-        elif dim == 2:
-            if left_to_right:
-                return tensor_1[:, :smallest_size], tensor_2[:, :smallest_size]
-            return tensor_1[:, -smallest_size:], tensor_2[:, -smallest_size:]
-        elif dim == 3:
-            if left_to_right:
-                return tensor_1[:, :, :smallest_size], tensor_2[:, :, :smallest_size]
-            return tensor_1[:, :, -smallest_size:], tensor_2[:, :, -smallest_size:]
-
-        # else:
-        if left_to_right:
-            return (
-                tensor_1[:, :, :, :smallest_size],
-                tensor_2[:, :, :, :smallest_size],
-            )
-        return (
-            tensor_1[:, :, :, -smallest_size:],
-            tensor_2[:, :, :, -smallest_size:],
-        )
-
     def stft_loss(
         self,
         signal: Tensor,
@@ -231,11 +428,23 @@ class AudioProcessor(Model):
             smallest = min(signal.shape[-1], ground.shape[-1])
             signal = signal[:, -smallest:]
             ground = ground[:, -smallest:]
-        sig_mel = self.compute_mel(signal, base, True).detach().cpu()
-        gnd_mel = self.compute_mel(ground, base, True).detach().cpu()
+        sig_mel = self.compute_mel(signal, base, True)
+        gnd_mel = self.compute_mel(ground, base, True)
         return torch.norm(gnd_mel - sig_mel, p=1) / torch.norm(gnd_mel, p=1)
 
-    # def forward(self, wave: Tensor, base: Optional[float] = None):
+    @staticmethod
+    def plot_spectrogram(spectrogram, ax_):
+        import matplotlib.pylab as plt
+
+        fig, ax = plt.subplots(figsize=(10, 2))
+        im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+        plt.colorbar(im, ax=ax)
+
+        fig.canvas.draw()
+        plt.close()
+
+        return fig
+
     def forward(
         self,
         *inputs: Union[Tensor, float],
@@ -251,6 +460,6 @@ class AudioProcessor(Model):
         elif ap_task == "inv_transform":
             return self.inverse_transform(*inputs, **inputs_kwargs)
         elif ap_task == "revert_mel":
-            return self.reverse_mel(*inputs, **inputs_kwargs)
+            return self.inverse_mel_spectogram(*inputs, **inputs_kwargs)
         else:
             raise ValueError(f"Invalid task '{ap_task}'")

lt_tensor/transform.py

@@ -20,14 +20,14 @@ __all__ = [
     "stft_istft_rebuild",
 ]
 
-from .torch_commons import *
+from lt_tensor.torch_commons import *
 import torchaudio
 import math
-from .misc_utils import log_tensor
+from lt_tensor.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
+from lt_tensor.model_base import Model
 import warnings
 
 
@@ -316,24 +316,12 @@ def inverse_transform(
     )
 
 
-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,
+        hop_length: Optional[int] = None,
         length: Optional[int] = None,
         window: Optional[Tensor] = None,
         onesided: Optional[bool] = None,
@@ -342,8 +330,8 @@ class InverseTransformConfig:
         center: bool = True,
     ):
         self.n_fft = n_fft
-        self.hop_length = hop_length
         self.win_length = win_length
+        self.hop_length = hop_length
         self.length = length
         self.onesided = onesided
         self.return_complex = return_complex
@@ -359,8 +347,8 @@ class InverseTransform(Model):
     def __init__(
         self,
         n_fft: int = 1024,
-        hop_length: Optional[int] = None,
         win_length: Optional[int] = None,
+        hop_length: Optional[int] = None,
         length: Optional[int] = None,
         window: Optional[Tensor] = None,
         onesided: Optional[bool] = None,
@@ -378,10 +366,10 @@ class InverseTransform(Model):
         ----------
         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`.
+        hop_length : int, optional
+            Number of audio samples between STFT columns. Defaults to `n_fft`.
         length : int, optional
             Output waveform length. If not provided, length will be inferred.
         window : Tensor, optional
@@ -403,14 +391,8 @@ class InverseTransform(Model):
             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),
-            )
-        )
+        assert window is None or isinstance(window, (Tensor, nn.Module))
+        assert not bool(return_complex and onesided)
         self.n_fft = n_fft
         self.length = length
         self.win_length = win_length or n_fft // 4
@@ -419,13 +401,11 @@ class InverseTransform(Model):
         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)
+        self.register_buffer(
+            "window", torch.hann_window(self.win_length) if window is None else window
+        )
 
-    def forward(self, spec: Tensor, phase: Tensor, **kwargs):
+    def forward(self, spec: Tensor, phase: Tensor, *, _recall: bool = False):
         """
         Perform the inverse short-time Fourier transform.
 
@@ -435,24 +415,28 @@ class InverseTransform(Model):
             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`.
         """
+        try:
+            return torch.istft(
+                spec * torch.exp(phase * 1j),
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                window=self.window,
+                center=self.center,
+                normalized=self.normalized,
+                onesided=self.onesided,
+                length=self.length,
+                return_complex=self.return_complex,
+            )
+        except RuntimeError as e:
+            if not _recall and spec.device != self.window.device:
+                self.window = self.window.to(spec.device)
+                return self.forward(spec, phase, _recall=True)
+            raise e
 
-        return torch.istft(
-            spec * torch.exp(phase * 1j),
-            n_fft = self.n_fft,
-            hop_length=self.hop_length,
-            win_length=self.win_length,
-            window=self.window,
-            center=self.center,
-            normalized=self.normalized,
-            onesided=self.onesided,
-            length=self.length,
-            return_complex=self.return_complex,
-        )

{lt_tensor-0.0.1a11.dist-info → lt_tensor-0.0.1a13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a11
+Version: 0.0.1a13
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336
@@ -11,14 +11,17 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Utilities
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: torch>=2.2.0
-Requires-Dist: torchaudio>=2.2.0
+Requires-Dist: torch>=2.7.0
+Requires-Dist: torchaudio>=2.7.0
 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
-Requires-Dist: librosa>=0.11.0
+Requires-Dist: lt-utils>=0.0.2a1
+Requires-Dist: librosa==0.11.*
+Requires-Dist: einops
+Requires-Dist: plotly
+Requires-Dist: scipy
 Dynamic: author
 Dynamic: classifier
 Dynamic: description