lt-tensor 0.0.1a38__py3-none-any.whl → 0.0.1a40__py3-none-any.whl

lt_tensor/model_zoo/losses/_envelope_disc/__init__.py

@@ -0,0 +1,116 @@
+""" Modified from: https://github.com/dinhoitt/BemaGANv2/blob/9560ae9df153c956f259c261c57c4f84f89e3d72/envelope.py
+MIT License
+
+Copyright (c) 2025 Taseoo Park
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+from lt_utils.common import *
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+
+
+class Envelope(Model):
+    def __init__(self, max_freq: int, sample_rate: Number = 24000, cut_off: int = 0):
+        super().__init__()
+        self.sr = sample_rate
+        self.max_freq = max_freq
+        self.setup_low_pass_fn(max_freq, cut_off)
+
+    def forward(self, x: torch.Tensor):
+        if not self.max_freq:
+            return x
+        return self.lp_fn(x)
+
+    def _ft_signal(self, signal: torch.Tensor):
+        filtered_signal = self.butterwort_lowpass_filter(signal)
+        return torch.abs(self.hilbert(filtered_signal))
+
+    def setup_low_pass_fn(self, max_freq: int, cutoff: int = 0):
+        self.max_freq = int(max_freq)
+        cutoff = self.max_freq if cutoff == 0 else cutoff
+        self.lp_fn = self.hilbert if self.max_freq in [-1, 1] else self._ft_signal
+        self.setup_butterwort_lowpass_coefficients(cutoff)
+
+    def hilbert(self, signal: Tensor) -> Tensor:
+        """Implementing the Hilbert transform manually"""
+        N = signal.shape[2]  # Signal length
+        FFT_signal = torch.fft.fft(signal, axis=2)
+        h = torch.zeros_like(
+            signal
+        )  # Generate an array with the same shape as the signal
+
+        if N % 2 == 0:
+            h[:, 0, 0] = 1
+            h[:, 0, N // 2] = 1
+            h[:, 0, 1 : N // 2] = 2
+        else:
+            h[:, 0, 0] = 1
+            h[:, 0, 1 : (N + 1) // 2] = 2
+
+        out: Tensor = torch.fft.ifft(FFT_signal * h, axis=2)
+        if self.max_freq == -1:
+            return -out.abs()
+        return -out.abs()
+
+    def butterwort_lowpass_filter(self, signal):
+        filtered_signal = torch.zeros_like(signal)
+        # Applying the filter to the signal
+        for n in range(len(signal)):
+            if n < 2:
+                filtered_signal[n] = self.lp_coef_a[0] * signal[n]
+            else:
+                filtered_signal[n] = (
+                    self.lp_coef_b[0] * signal[n]
+                    + self.lp_coef_b[1] * signal[n - 1]
+                    + self.lp_coef_b[2] * signal[n - 2]
+                    - self.lp_coef_a[1] * filtered_signal[n - 1]
+                    - self.lp_coef_a[2] * filtered_signal[n - 2]
+                )
+
+        return filtered_signal
+
+    def setup_butterwort_lowpass_coefficients(self, cutoff: int):
+        cutoff = torch.tensor([cutoff], dtype=torch.float64)
+        fs = torch.tensor([self.sr], dtype=torch.float64)
+
+        omega = torch.tan(torch.pi * cutoff / fs)
+        # Convert float 2 to tensor
+        sqrt2 = torch.tensor(2.0, dtype=torch.float64).sqrt()
+
+        sq_omega = sqrt2 * omega + omega**2
+        # Transfer function coefficients using the bilinear transform
+        a = 2 * (omega**2 - 1) / (1 + sq_omega)
+        self.register_buffer(
+            "lp_coef_a",
+            torch.tensor(
+                [1.0, a.item(), ((1 - sq_omega) / (1 + sq_omega)).item()],
+                dtype=torch.float64,
+                device=self.device,
+            ),
+        )
+        b = omega**2 / (1 + sq_omega)
+        self.register_buffer(
+            "lp_coef_b",
+            torch.tensor(
+                [b.item(), (2 * b).item(), b.item()],
+                dtype=torch.float64,
+                device=self.device,
+            ),
+        )

lt_tensor/model_zoo/losses/discriminators.py

@@ -7,6 +7,7 @@ from lt_tensor.model_base import Model
 from lt_tensor.model_zoo.convs import ConvNets
 from torch.nn import functional as F
 from torchaudio import transforms as T
+from lt_tensor.model_zoo.losses._envelope_disc import Envelope
 
 MULTI_DISC_OUT_TYPE: TypeAlias = Tuple[
     List[Tensor],
@@ -313,7 +314,7 @@ class DiscriminatorS(ConvNets):
         return x.flatten(1, -1), fmap
 
 
-class MultiScaleDiscriminator(ConvNets):
+class MultiScaleDiscriminator(_MultiDiscriminatorT):
     def __init__(
         self,
         discriminator_channel_multi: Number = 1,
@@ -352,102 +353,71 @@ class MultiScaleDiscriminator(ConvNets):
         return y_d_rs, y_d_gs, fmap_rs, fmap_gs
 
 
-class EnvelopeExtractor(Model):
-    """Extracts the amplitude envelope of the audio signal."""
+class DiscriminatorE(ConvNets):
+    """Modified from: https://github.com/dinhoitt/BemaGANv2/blob/9560ae9df153c956f259c261c57c4f84f89e3d72/models.py"""
 
-    def __init__(self, kernel_size=101):
-        super().__init__()
-        # Lowpass filter for smoothing envelope (moving average)
-        self.kernel_size = kernel_size
-        self.register_buffer("kernel", torch.ones(1, 1, kernel_size) / kernel_size)
-
-    def forward(self, x: Tensor):
-        # x: (B, 1, T) -> abs(x)
-        envelope = torch.abs(x)
-        # Apply low-pass smoothing (via conv1d)
-        envelope = F.pad(
-            envelope, (self.kernel_size // 2, self.kernel_size // 2), mode="reflect"
-        )
-        envelope = F.conv1d(envelope, self.kernel)
-        return envelope
-
-
-class DiscriminatorEnvelope(ConvNets):
     def __init__(
         self,
-        use_spectral_norm=False,
+        max_freq: int,
         discriminator_channel_multi: Number = 1,
-        kernel_size: int = 101,
+        sample_rate: int = 24000,
+        use_spectral_norm: bool = False,
     ):
+
         super().__init__()
-        norm_f = weight_norm if not use_spectral_norm else spectral_norm
-        self.extractor = EnvelopeExtractor(kernel_size=kernel_size)
+        self.max_freq = max_freq
+        norm_f = spectral_norm if use_spectral_norm else weight_norm
         dsc = lambda x: int(x * discriminator_channel_multi)
         self.convs = nn.ModuleList(
             [
-                norm_f(nn.Conv1d(1, dsc(64), 15, stride=1, padding=7)),
-                norm_f(
-                    nn.Conv1d(dsc(64), dsc(128), 41, stride=2, groups=4, padding=20)
-                ),
-                norm_f(
-                    nn.Conv1d(dsc(128), dsc(256), 41, stride=2, groups=16, padding=20)
-                ),
-                norm_f(
-                    nn.Conv1d(dsc(256), dsc(512), 41, stride=4, groups=16, padding=20)
-                ),
-                norm_f(
-                    nn.Conv1d(dsc(512), dsc(512), 41, stride=4, groups=16, padding=20)
-                ),
-                norm_f(nn.Conv1d(dsc(512), dsc(512), 5, stride=1, padding=2)),
+                norm_f(nn.Conv1d(1, dsc(128), 15, 1, padding=7)),
+                norm_f(nn.Conv1d(dsc(128), dsc(128), 41, 2, groups=4, padding=20)),
+                norm_f(nn.Conv1d(dsc(128), dsc(256), 41, 2, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(256), dsc(512), 41, 4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(512), dsc(1024), 41, 4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(1024), dsc(1024), 41, 1, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(1024), dsc(1024), 5, 1, padding=2)),
             ]
         )
-        self.conv_post = norm_f(nn.Conv1d(dsc(512), 1, 3, stride=1, padding=1))
+        self.conv_post = norm_f(nn.Conv1d(dsc(1024), 1, 3, 1, padding=1))
+        self.envelope = Envelope(max_freq=self.max_freq, sample_rate=sample_rate)
         self.activation = nn.LeakyReLU(0.1)
 
-    def forward(self, x):
-        # Input: raw audio (B, 1, T)
-        x = self.extractor(x)
+    def forward(self, x: Tensor):
         fmap = []
-        for layer in self.convs:
-            x = self.activation(layer(x))
+        for l in self.convs:
+            x = self.envelope(x)
+            x = self.activation(l(x))
             fmap.append(x)
         x = self.conv_post(x)
         fmap.append(x)
-        return x.flatten(1), fmap
+        return x.flatten(start_dim=1, end_dim=-1), fmap
 
 
 class MultiEnvelopeDiscriminator(_MultiDiscriminatorT):
+    """Modified from: https://github.com/dinhoitt/BemaGANv2/blob/9560ae9df153c956f259c261c57c4f84f89e3d72/models.py"""
+
     def __init__(
         self,
-        use_spectral_norm: bool = False,
         discriminator_channel_multi: Number = 1,
     ):
         super().__init__()
+        f_times_values = [-1, 0, 1, 300, 500]
         self.discriminators = nn.ModuleList(
-            [
-                DiscriminatorEnvelope(
-                    use_spectral_norm, discriminator_channel_multi
-                ),  # raw envelope
-                DiscriminatorEnvelope(use_spectral_norm),  # downsampled once
-                DiscriminatorEnvelope(use_spectral_norm),  # downsampled twice
-            ]
-        )
-        self.meanpools = nn.ModuleList(
-            [nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)]
+            [DiscriminatorE(f, discriminator_channel_multi) for f in f_times_values]
         )
 
     def forward(self, y, y_hat):
-        y_d_rs, y_d_gs = [], []
-        fmap_rs, fmap_gs = [], []
-        for i, d in enumerate(self.discriminators):
-            if i != 0:
-                y = self.meanpools[i - 1](y)
-                y_hat = self.meanpools[i - 1](y_hat)
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for d in self.discriminators:
             y_d_r, fmap_r = d(y)
             y_d_g, fmap_g = d(y_hat)
             y_d_rs.append(y_d_r)
-            y_d_gs.append(y_d_g)
             fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
             fmap_gs.append(fmap_g)
 
         return y_d_rs, y_d_gs, fmap_rs, fmap_gs

lt_tensor/noise_tools.py

@@ -13,6 +13,7 @@ __all__ = [
 ]
 
 from lt_utils.common import *
+from lt_tensor.model_base import Model
 import torch.nn.functional as F
 from lt_tensor.torch_commons import *
 import math
@@ -20,17 +21,17 @@ import random
 from lt_tensor.misc_utils import set_seed
 
 
-def add_gaussian_noise(x: Tensor, noise_level=0.025):
+def add_gaussian_noise(x: Tensor, noise_level: float = 0.025) -> Tensor:
     noise = torch.randn_like(x) * noise_level
     return x + noise
 
 
-def add_uniform_noise(x: Tensor, noise_level=0.025):
+def add_uniform_noise(x: Tensor, noise_level: float = 0.025) -> Tensor:
     noise = (torch.rand_like(x) - 0.5) * 2 * noise_level
     return x + noise
 
 
-def add_linear_noise(x, noise_level=0.05):
+def add_linear_noise(x, noise_level=0.05) -> Tensor:
     T = x.shape[-1]
     ramp = torch.linspace(0, noise_level, T, device=x.device)
     for _ in range(x.dim() - 1):
@@ -38,7 +39,7 @@ def add_linear_noise(x, noise_level=0.05):
     return x + ramp.expand_as(x)
 
 
-def add_impulse_noise(x: Tensor, noise_level=0.025):
+def add_impulse_noise(x: Tensor, noise_level: float = 0.025) -> Tensor:
     # For image inputs
     probs = torch.rand_like(x)
     x_clone = x.detach().clone()
@@ -47,7 +48,7 @@ def add_impulse_noise(x: Tensor, noise_level=0.025):
     return x_clone
 
 
-def add_pink_noise(x: Tensor, noise_level=0.05):
+def add_pink_noise(x: Tensor, noise_level: float = 0.05) -> Tensor:
     # 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]
@@ -66,12 +67,12 @@ def add_pink_noise(x: Tensor, noise_level=0.05):
     return x + pink_noised * noise_level
 
 
-def add_clipped_gaussian_noise(x, noise_level=0.025):
+def add_clipped_gaussian_noise(x: Tensor, noise_level: float = 0.025) -> Tensor:
     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):
+def add_multiplicative_noise(x: Tensor, noise_level: float = 0.025) -> Tensor:
     noise = 1 + torch.randn_like(x) * noise_level
     return x * noise
 
@@ -109,7 +110,15 @@ _NOISE_DIM_SUPPORT = {
 
 def apply_noise(
     x: Tensor,
-    noise_type: str = "gaussian",
+    noise_type: Literal[
+        "gaussian",
+        "uniform",
+        "linear",
+        "impulse",
+        "pink",
+        "clipped_gaussian",
+        "multiplicative",
+    ] = "gaussian",
     noise_level: float = 0.01,
     seed: Optional[int] = None,
     on_error: Literal["raise", "try_others", "return_unchanged"] = "raise",
@@ -229,11 +238,11 @@ class NoiseSchedulerA(nn.Module):
         return collected, noise_history
 
 
-class NoiseSchedulerB(nn.Module):
-    def __init__(self, timesteps: int = 512):
+class NoiseSchedulerB(Model):
+    def __init__(self, timesteps: int = 50, l_min: float = 0.0005, l_max: float = 0.05):
         super().__init__()
 
-        betas = torch.linspace(1e-4, 0.02, timesteps)
+        betas = torch.linspace(l_min, l_max, timesteps)
         alphas = 1.0 - betas
         alpha_cumprod = torch.cumprod(alphas, dim=0)
 
@@ -272,7 +281,7 @@ class NoiseSchedulerB(nn.Module):
         self, x_0: Tensor, t: int, noise: Optional[Union[Tensor, float]] = None
     ) -> Tensor:
         assert (
-             0 <= t < self.timesteps
+            0 <= t < self.timesteps
         ), f"Time step t={t} is out of bounds for scheduler with {self.timesteps} steps."
 
         if noise is None:
@@ -286,7 +295,7 @@ class NoiseSchedulerB(nn.Module):
         return alpha_term + noise_term
 
 
-class NoiseSchedulerC(nn.Module):
+class NoiseSchedulerC(Model):
     def __init__(self, timesteps: int = 512):
         super().__init__()