lt-tensor 0.0.1a12__py3-none-any.whl → 0.0.1a14__py3-none-any.whl

lt_tensor/model_zoo/residual.py

@@ -0,0 +1,252 @@
+__all__ = [
+    "spectral_norm_select",
+    "get_weight_norm",
+    "ResBlock1D",
+    "ResBlock2D",
+    "ResBlock1DShuffled",
+    "AdaResBlock1D",
+    "ResBlocks",
+]
+import math
+from lt_utils.common import *
+import torch.nn.functional as F
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from lt_tensor.misc_utils import log_tensor
+from lt_tensor.model_zoo.fusion import AdaFusion1D, AdaIN1D
+
+
+def spectral_norm_select(module: nn.Module, enabled: bool):
+    if enabled:
+        return spectral_norm(module)
+    return module
+
+
+def get_weight_norm(norm_type: Optional[Literal["weight", "spectral"]] = None):
+    if not norm_type:
+        return lambda x: x
+    if norm_type == "weight":
+        return lambda x: weight_norm(x)
+    return lambda x: spectral_norm(x)
+
+
+class ConvNets(Model):
+    def remove_weight_norm(self):
+        for module in self.modules():
+            try:
+                remove_weight_norm(module)
+            except ValueError:
+                pass
+
+    @staticmethod
+    def init_weights(m, mean=0.0, std=0.01):
+        classname = m.__class__.__name__
+        if "Conv" in classname:
+            m.weight.data.normal_(mean, std)
+
+
+class ResBlocks(ConvNets):
+    def __init__(
+        self,
+        channels: int,
+        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
+        resblock_dilation_sizes: List[Union[int, List[int]]] = [
+            [1, 3, 5],
+            [1, 3, 5],
+            [1, 3, 5],
+        ],
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.rb = nn.ModuleList()
+        self.activation = activation
+
+        for k, j in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+            self.rb.append(ResBlock1D(channels, k, j, activation))
+
+        self.rb.apply(self.init_weights)
+
+    def forward(self, x: torch.Tensor):
+        xs = None
+        for i, block in enumerate(self.rb):
+            if i == 0:
+                xs = block(x)
+            else:
+                xs += block(x)
+        x = xs / self.num_kernels
+        return x
+
+
+
+class ResBlock1D(ConvNets):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+
+        self.conv_nets = nn.ModuleList(
+            [
+                self._get_conv_layer(i, channels, kernel_size, 1, dilation, activation)
+                for i in range(3)
+            ]
+        )
+        self.conv_nets.apply(self.init_weights)
+        self.last_index = len(self.conv_nets) - 1
+
+    def _get_conv_layer(self, id, ch, k, stride, d, actv):
+        get_padding = lambda ks, d: int((ks * d - d) / 2)
+        return nn.Sequential(
+            actv,  # 1
+            weight_norm(
+                nn.Conv1d(
+                    ch, ch, k, stride, dilation=d[id], padding=get_padding(k, d[id])
+                )
+            ),  # 2
+            actv,  # 3
+            weight_norm(
+                nn.Conv1d(ch, ch, k, stride, dilation=1, padding=get_padding(k, 1))
+            ),  # 4
+        )
+
+    def forward(self, x: Tensor):
+        for cnn in self.conv_nets:
+            x = cnn(x) + x
+        return x
+
+
+class ResBlock1DShuffled(ConvNets):
+    def __init__(
+        self,
+        channels,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        add_channel_shuffle: bool = False,  # requires pytorch 2.7.0 +
+        channel_shuffle_groups=1,
+    ):
+        super().__init__()
+
+        self.channel_shuffle = (
+            nn.ChannelShuffle(channel_shuffle_groups)
+            if add_channel_shuffle
+            else nn.Identity()
+        )
+
+        self.conv_nets = nn.ModuleList(
+            [
+                self._get_conv_layer(i, channels, kernel_size, 1, dilation, activation)
+                for i in range(3)
+            ]
+        )
+        self.conv_nets.apply(self.init_weights)
+        self.last_index = len(self.conv_nets) - 1
+
+    def _get_conv_layer(self, id, ch, k, stride, d, actv):
+        get_padding = lambda ks, d: int((ks * d - d) / 2)
+        return nn.Sequential(
+            actv,  # 1
+            weight_norm(
+                nn.Conv1d(
+                    ch, ch, k, stride, dilation=d[id], padding=get_padding(k, d[id])
+                )
+            ),  # 2
+            actv,  # 3
+            weight_norm(
+                nn.Conv1d(ch, ch, k, stride, dilation=1, padding=get_padding(k, 1))
+            ),  # 4
+        )
+
+    def forward(self, x: Tensor):
+        b = x.clone() * 0.5
+        for cnn in self.conv_nets:
+            x = cnn(self.channel_shuffle(x)) + b
+        return x
+
+
+class ResBlock2D(Model):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        downsample=False,
+    ):
+        super().__init__()
+        stride = 2 if downsample else 1
+
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, stride, 1),
+            nn.LeakyReLU(0.2),
+            nn.Conv2d(out_channels, out_channels, 3, 1, 1),
+        )
+
+        self.skip = nn.Identity()
+        if downsample or in_channels != out_channels:
+            self.skip = spectral_norm_select(
+                nn.Conv2d(in_channels, out_channels, 1, stride)
+            )
+        # on less to be handled every cicle
+        self.sqrt_2 = math.sqrt(2)
+
+    def forward(self, x: Tensor):
+        return (self.block(x) + self.skip(x)) / self.sqrt_2
+
+
+class AdaResBlock1D(ConvNets):
+    def __init__(
+        self,
+        res_block_channels: int,
+        ada_channel_in: int,
+        kernel_size=3,
+        dilation=(1, 3, 5),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        super().__init__()
+
+        self.conv_nets = nn.ModuleList(
+            [
+                self._get_conv_layer(
+                    i,
+                    res_block_channels,
+                    ada_channel_in,
+                    kernel_size,
+                    1,
+                    dilation,
+                )
+                for i in range(3)
+            ]
+        )
+        self.conv_nets.apply(self.init_weights)
+        self.last_index = len(self.conv_nets) - 1
+        self.activation = activation
+
+    def _get_conv_layer(self, id, ch, ada_ch, k, stride, d):
+        get_padding = lambda ks, d: int((ks * d - d) / 2)
+        return nn.ModuleDict(
+            dict(
+                norm1=AdaFusion1D(ada_ch, ch),
+                norm2=AdaFusion1D(ada_ch, ch),
+                alpha1=nn.Parameter(torch.ones(1, ada_ch, 1)),
+                alpha2=nn.Parameter(torch.ones(1, ada_ch, 1)),
+                conv1=weight_norm(
+                    nn.Conv1d(
+                        ch, ch, k, stride, dilation=d[id], padding=get_padding(k, d[id])
+                    )
+                ),  # 2
+                conv2=weight_norm(
+                    nn.Conv1d(ch, ch, k, stride, dilation=1, padding=get_padding(k, 1))
+                ),  # 4
+            )
+        )
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        for cnn in self.conv_nets:
+            xt = self.activation(cnn["norm1"](x, y, cnn["alpha1"]))
+            xt = cnn["conv1"](xt)
+            xt = self.activation(cnn["norm2"](xt, y, cnn["alpha2"]))
+            x = cnn["conv2"](xt) + x
+        return x

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

@@ -11,8 +11,8 @@ 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 lt_tensor.model_zoo.pos import *
-from lt_tensor.model_zoo.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/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,14 +24,14 @@ 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,
         n_iter: int = 32,
         window: Optional[Tensor] = None,
-        normalized: bool =False,
+        normalized: bool = False,
         onesided: Optional[bool] = None,
     ):
         super().__init__()
@@ -38,7 +41,7 @@ class AudioProcessor(Model):
         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
@@ -76,7 +79,165 @@ class AudioProcessor(Model):
             "window",
             (torch.hann_window(self.win_length) if window is None else window),
         )
-        # self._inv_transform = InverseTransform(**inverse_transform_config.to_dict())
+
+    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
+        if audio.ndim > 1:
+            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
+        if audio.ndim > 1:
+            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,
@@ -95,7 +256,9 @@ class AudioProcessor(Model):
                 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),
+                window=torch.hann_window(
+                    win_length or self.win_length, device=spec.device
+                ),
                 center=self.center,
                 normalized=self.normalized,
                 onesided=self.onesided,
@@ -105,10 +268,12 @@ class AudioProcessor(Model):
         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)
+                return self.inverse_transform(
+                    spec, phase, n_fft, hop_length, win_length, length, _recall=True
+                )
             raise e
 
-    def rebuild_spectrogram(
+    def normalize_audio(
         self,
         wave: Tensor,
         length: Optional[int] = None,
@@ -148,7 +313,7 @@ class AudioProcessor(Model):
         except RuntimeError as e:
             if not _recall and wave.device != self.window.device:
                 self.window = self.window.to(wave.device)
-                return self.rebuild_spectrogram(wave, length, _recall=True)
+                return self.normalize_audio(wave, length, _recall=True)
             raise e
 
     def compute_mel(
@@ -167,12 +332,7 @@ class AudioProcessor(Model):
 
     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),
@@ -182,21 +342,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",
@@ -212,6 +377,7 @@ class AudioProcessor(Model):
         return FileScan.files(
             path,
             extensions,
+            maximum,
         )
 
     def find_audio_text_pairs(
@@ -240,57 +406,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,
@@ -302,11 +417,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],

lt_tensor/transform.py

@@ -316,18 +316,6 @@ 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,
@@ -413,9 +401,11 @@ class InverseTransform(Model):
         self.return_complex = return_complex
         self.onesided = onesided
         self.normalized = normalized
-        self.register_buffer('window', torch.hann_window(self.win_length) if window is None else window)
+        self.register_buffer(
+            "window", torch.hann_window(self.win_length) if window is None else window
+        )
 
-    def forward(self, spec: Tensor, phase: Tensor, *, _recall:bool = False):
+    def forward(self, spec: Tensor, phase: Tensor, *, _recall: bool = False):
         """
         Perform the inverse short-time Fourier transform.
 
@@ -434,7 +424,7 @@ class InverseTransform(Model):
         try:
             return torch.istft(
                 spec * torch.exp(phase * 1j),
-                n_fft = self.n_fft,
+                n_fft=self.n_fft,
                 hop_length=self.hop_length,
                 win_length=self.win_length,
                 window=self.window,
@@ -448,4 +438,5 @@ class InverseTransform(Model):
             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
+            raise e
+