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

lt_tensor/model_base.py

@@ -4,6 +4,7 @@ __all__ = ["Model"]
 import warnings
 from .torch_commons import *
 from lt_utils.common import *
+from lt_utils.misc_utils import log_traceback
 
 T = TypeVar("T")
 
@@ -40,20 +41,113 @@ 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 usualy the '.to()' function fails to apply
+        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)
         ```
         """
         pass
 
+    def freeze_weight(self, weight: Union[str, nn.Module], freeze: bool):
+        assert isinstance(weight, (str, nn.Module))
+        if isinstance(weight, str):
+            if hasattr(self, weight):
+                w = getattr(self, weight)
+                if isinstance(w, nn.Module):
+                    w.requires_grad_(not freeze)
+        else:
+            weight.requires_grad_(not freeze)
+
+    def _freeze_unfreeze(
+        self,
+        weight: Union[str, nn.Module],
+        task: Literal["freeze", "unfreeze"] = "freeze",
+        _skip_except: bool = False,
+    ):
+        try:
+            assert isinstance(weight, (str, nn.Module))
+            if isinstance(weight, str):
+                w_txt = f"Failed to {task} the module '{weight}'. Reason: is not a valid attribute of {self._get_name()}"
+                if hasattr(self, weight):
+                    w_txt = f"Failed to {task} the module '{weight}'. Reason: is not a Module type."
+                    w = getattr(self, weight)
+                    if isinstance(w, nn.Module):
+                        w_txt = f"Successfully {task} the module '{weight}'."
+                        w.requires_grad_(task == "unfreeze")
+
+            else:
+                w.requires_grad_(task == "unfreeze")
+                w_txt = f"Successfully '{task}' the module '{weight}'."
+            return w_txt
+        except Exception as e:
+            if not _skip_except:
+                raise e
+            return str(e)
+
+    def freeze_weight(
+        self,
+        weight: Union[str, nn.Module],
+        _skip_except: bool = False,
+    ):
+        return self._freeze_unfreeze(weight, "freeze", _skip_except)
+
+    def unfreeze_weight(
+        self,
+        weight: Union[str, nn.Module],
+        _skip_except: bool = False,
+    ):
+        return self._freeze_unfreeze(weight, "freeze", _skip_except)
+
+    def freeze_all(self, exclude: Optional[List[str]] = None):
+        no_exclusions = not exclude
+        frozen = []
+        not_frozen = []
+        for name, param in self.named_parameters():
+            if no_exclusions:
+                try:
+                    param.requires_grad_(False)
+                    frozen.append(name)
+                except Exception as e:
+                    not_frozen.append((name, str(e)))
+            elif any(layer in name for layer in exclude):
+                try:
+                    param.requires_grad_(False)
+                    frozen.append(name)
+                except Exception as e:
+                    not_frozen.append((name, str(e)))
+            else:
+                not_frozen.append((name, "Excluded"))
+        return dict(frozen=frozen, not_frozen=not_frozen)
+
+    def unfreeze_all_except(self, exclude: Optional[list[str]] = None):
+        """Unfreezes all model parameters except specified layers."""
+        no_exclusions = not exclude
+        unfrozen = []
+        not_unfrozen = []
+        for name, param in self.named_parameters():
+            if no_exclusions:
+                try:
+                    param.requires_grad_(True)
+                    unfrozen.append(name)
+                except Exception as e:
+                    not_unfrozen.append((name, str(e)))
+            elif any(layer in name for layer in exclude):
+                try:
+                    param.requires_grad_(True)
+                    unfrozen.append(name)
+                except Exception as e:
+                    not_unfrozen.append((name, str(e)))
+            else:
+                not_unfrozen.append((name, "Excluded"))
+        return dict(unfrozen=unfrozen, not_unfrozen=not_unfrozen)
+
     def to(self, *args, **kwargs):
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
             *args, **kwargs
@@ -186,11 +280,16 @@ class Model(nn.Module, ABC):
         )
 
     def get_weights(self, module_name: Optional[str] = None) -> List[Tensor]:
-        """Returns the weights of the model entrie model or from a specified module"""
+        """Returns the weights of the model entry model or from a specified module"""
         if module_name is not None:
             assert hasattr(self, module_name), f"Module {module_name} does not exits"
             module = getattr(self, module_name)
-            return [x.data.detach() for x in module.parameters()]
+            params = []
+            if isinstance(module, nn.Module):
+                return [x.data.detach() for x in module.parameters()]
+            elif isinstance(module, (Tensor, nn.Parameter)):
+                return [module.data.detach()]
+            raise (f"{module_name} is has no weights")
         return [x.data.detach() for x in self.parameters()]
 
     def print_trainable_parameters(

lt_tensor/model_zoo/disc.py

@@ -11,37 +11,36 @@ class PeriodDiscriminator(Model):
         use_spectral_norm=False,
         kernel_size: int = 5,
         stride: int = 3,
-        initial_s: int = 32,
     ):
         super().__init__()
         self.period = period
+        self.stride = stride
+        self.kernel_size = kernel_size
         self.norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        
+        self.channels = [32, 128, 512, 1024, 1024]
         self.first_pass = nn.Sequential(
             self.norm_f(
                 nn.Conv2d(
-                    1, initial_s * 4, (kernel_size, 1), (stride, 1), padding=(2, 0)
+                    1, self.channels[0], (kernel_size, 1), (stride, 1), padding=(2, 0)
                 )
             ),
             nn.LeakyReLU(0.1),
         )
-        self._last_sz = initial_s * 4
 
-        self.convs = nn.ModuleList([self._get_next(i == 3) for i in range(4)])
+        
+        self.convs = nn.ModuleList([self._get_next(self.channels[i+1], self.channels[i], i == 3) for i in range(4)])
 
         self.post_conv = nn.Conv2d(1024, 1, (stride, 1), 1, padding=(1, 0))
-        self.kernel_size = kernel_size
-        self.stride = stride
 
-    def _get_next(self, is_last: bool = False):
-        in_dim = self._last_sz
-        self._last_sz *= 4
-        print(self._last_sz, "-----------------------")
+    def _get_next(self, out_dim:int, last_in:int, is_last: bool = False):
         stride = (self.stride, 1) if not is_last else 1
+        
         return nn.Sequential(
             self.norm_f(
                 nn.Conv2d(
-                    in_dim,
-                    self._last_sz,
+                    last_in,
+                    out_dim,
                     (self.kernel_size, 1),
                     stride,
                     padding=(2, 0),
@@ -91,6 +90,7 @@ class ScaleDiscriminator(nn.Module):
     def __init__(self, use_spectral_norm=False):
         super().__init__()
         norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.activation = nn.LeakyReLU(0.1)
         self.convs = nn.ModuleList(
             [
                 norm_f(nn.Conv1d(1, 128, 15, 1, padding=7)),
@@ -103,7 +103,6 @@ class ScaleDiscriminator(nn.Module):
             ]
         )
         self.post_conv = norm_f(nn.Conv1d(1024, 1, 3, 1, padding=1))
-        self.activation = nn.LeakyReLU(0.1)
 
     def forward(self, x: torch.Tensor):
         """
@@ -147,9 +146,10 @@ class GeneralLossDescriminator(Model):
         super().__init__()
         self.mpd = MultiPeriodDiscriminator()
         self.msd = MultiScaleDiscriminator()
+        self.print_trainable_parameters()
 
     def _get_group_(self):
         pass
 
     def forward(self, x: Tensor, y_hat: Tensor):
-        return
+        return

lt_tensor/model_zoo/istft.py

@@ -106,3 +106,44 @@ class Generator(Model):
         classname = m.__class__.__name__
         if "Conv" in classname:
             m.weight.data.normal_(mean, std)
+
+
+# Below are items found in the Rishikesh's repo that might work for this generator.
+# https://github.com/rishikksh20/iSTFTNet-pytorch/blob/781480e9563d4dff5a8cc9ef1af6c6e0cab025c8/models.py
+
+
+def feature_loss(fmap_r, fmap_g, weight=2.0):
+    """Feature matching loss between real and generated feature maps."""
+    loss = 0.0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl - gl))
+    return loss * weight
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    """LSGAN-style loss for real and fake predictions."""
+    loss = 0.0
+    r_losses, g_losses = [], []
+
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1.0 - dr) ** 2)
+        g_loss = torch.mean(dg**2)
+        loss += r_loss + g_loss
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_generated_outputs):
+    """LSGAN generator loss encouraging fake to look like real (close to 1)."""
+    loss = 0.0
+    gen_losses = []
+
+    for dg in disc_generated_outputs:
+        l = torch.mean((1.0 - dg) ** 2)
+        gen_losses.append(l.item())
+        loss += l
+
+    return loss, gen_losses

lt_tensor/noise_tools.py

@@ -0,0 +1,368 @@
+__all__ = [
+    "NoiseSchedulerA",
+    "NoiseSchedulerB",
+    "NoiseSchedulerC",
+    "add_gaussian_noise",
+    "add_uniform_noise",
+    "add_linear_noise",
+    "add_impulse_noise",
+    "add_pink_noise",
+    "add_clipped_gaussian_noise",
+    "add_multiplicative_noise",
+    "apply_noise",
+]
+
+from lt_utils.common import *
+import torch.nn.functional as F
+from .torch_commons import *
+import math
+import random
+from .misc_utils import set_seed
+
+
+def add_gaussian_noise(x: Tensor, noise_level=0.025):
+    noise = torch.randn_like(x) * noise_level
+    return x + noise
+
+
+def add_uniform_noise(x: Tensor, noise_level=0.025):
+    noise = (torch.rand_like(x) - 0.5) * 2 * noise_level
+    return x + noise
+
+
+def add_linear_noise(x, noise_level=0.05):
+    T = x.shape[-1]
+    ramp = torch.linspace(0, noise_level, T, device=x.device)
+    for _ in range(x.dim() - 1):
+        ramp = ramp.unsqueeze(0)
+    return x + ramp.expand_as(x)
+
+
+def add_impulse_noise(x: Tensor, noise_level=0.025):
+    # For image inputs
+    probs = torch.rand_like(x)
+    x_clone = x.detach().clone()
+    x_clone[probs < (noise_level / 2)] = 0.0  # salt
+    x_clone[probs > (1 - noise_level / 2)] = 1.0  # pepper
+    return x_clone
+
+
+def add_pink_noise(x: Tensor, noise_level=0.05):
+    # 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]
+    pink_noised = []
+
+    for row in x:
+        white = torch.randn_like(row)
+        f = torch.fft.rfft(white)
+        freqs = torch.fft.rfftfreq(row.numel(), d=1.0)
+        freqs[0] = 1.0  # prevent div by 0
+        f /= freqs.sqrt()
+        pink = torch.fft.irfft(f, n=row.numel())
+        pink_noised.append(pink)
+
+    pink_noised = torch.stack(pink_noised, dim=0).view_as(x)
+    return x + pink_noised * noise_level
+
+
+def add_clipped_gaussian_noise(x, noise_level=0.025):
+    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):
+    noise = 1 + torch.randn_like(x) * noise_level
+    return x * noise
+
+
+_VALID_NOISES = [
+    "gaussian",
+    "uniform",
+    "linear",
+    "impulse",
+    "pink",
+    "clipped_gaussian",
+    "multiplicative",
+]
+
+_NOISE_MAP = {
+    "gaussian": add_gaussian_noise,
+    "uniform": add_uniform_noise,
+    "linear": add_linear_noise,
+    "impulse": add_impulse_noise,
+    "pink": add_pink_noise,
+    "clipped_gaussian": add_clipped_gaussian_noise,
+    "multiplicative": add_multiplicative_noise,
+}
+
+_NOISE_DIM_SUPPORT = {
+    "gaussian": (1, 2),
+    "uniform": (1, 2),
+    "multiplicative": (1, 2, 3),
+    "clipped_gaussian": (1, 2, 3),
+    "linear": (2, 3),
+    "impulse": (2, 3),
+    "pink": (2, 3),
+}
+
+
+def apply_noise(
+    x: Tensor,
+    noise_type: str = "gaussian",
+    noise_level: float = 0.01,
+    seed: Optional[int] = None,
+    on_error: Literal["raise", "try_others", "return_unchanged"] = "raise",
+    _last_tries: list[str] = [],
+):
+    noise_type = noise_type.lower().strip()
+    last_tries = _last_tries
+
+    if noise_type not in _NOISE_MAP:
+        raise ValueError(f"Noise type '{noise_type}' not supported.")
+
+    # Check dimension compatibility
+    allowed_dims = _NOISE_DIM_SUPPORT.get(noise_type, (1, 2))
+    if x.ndim not in allowed_dims:
+        assert (
+            on_error != "raise"
+        ), f"Noise '{noise_type}' is not supported for {x.ndim}D input."
+        if on_error == "return_unchanged":
+            return x, None
+        elif on_error == "try_others":
+            remaining = [
+                n
+                for n in _VALID_NOISES
+                if n not in last_tries and x.ndim in _NOISE_DIM_SUPPORT[n]
+            ]
+            if not remaining:
+                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()
+                ),
+                noise_type,
+            )
+    try:
+        if isinstance(seed, int):
+            set_seed(seed)
+        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, None
+        if len(last_tries) == len(_VALID_NOISES):
+            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()),
+            noise_type,
+        )
+
+
+class NoiseSchedulerA(nn.Module):
+    def __init__(self, samples: int = 64):
+        super().__init__()
+        self.base_steps = samples
+
+    def plot_noise_progression(noise_seq: list[Tensor], titles: list[str] = None):
+        import matplotlib.pyplot as plt
+
+        steps = len(noise_seq)
+        plt.figure(figsize=(15, 3))
+        for i, tensor in enumerate(noise_seq):
+            plt.subplot(1, steps, i + 1)
+            plt.imshow(tensor.squeeze().cpu().numpy(), aspect="auto", origin="lower")
+            if titles:
+                plt.title(titles[i])
+            plt.axis("off")
+        plt.tight_layout()
+        plt.show()
+
+    def forward(
+        self,
+        source_item: torch.Tensor,
+        steps: Optional[int] = None,
+        noise_type: Literal[
+            "gaussian",
+            "uniform",
+            "linear",
+            "impulse",
+            "pink",
+            "clipped_gaussian",
+            "multiplicative",
+        ] = "gaussian",
+        seed: Optional[int] = None,
+        noise_level: float = 0.01,
+        shuffle_noise_types: bool = False,
+        return_dict: bool = True,
+    ):
+        if steps is None:
+            steps = self.base_steps
+        collected = [source_item.detach().clone()]
+        noise_history = []
+        for i in range(steps):
+            if i > 0 and shuffle_noise_types:
+                noise_type = random.choice(_VALID_NOISES)
+            current, noise_name = apply_noise(
+                collected[-1],
+                noise_type,
+                noise_level,
+                seed=seed,
+                on_error="try_others",
+            )
+            noise_history.append(noise_name)
+            collected.append(current)
+
+        if return_dict:
+            return {
+                "steps": collected,
+                "history": noise_history,
+                "final": collected[-1],
+                "init": collected[0],
+            }
+        return collected, noise_history
+
+
+class NoiseSchedulerB(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: Optional[int] = None,
+    ) -> float:
+        if isinstance(seed, int):
+            random.seed(seed)
+        return random.uniform(*min_max)
+
+    def set_noise(
+        self,
+        noise: Optional[Union[Tensor, Number]] = None,
+        seed: Optional[int] = None,
+        min_max: Tuple[float, float] = (-3, 3),
+        default: bool = False,
+    ):
+        if noise is not None:
+            self.default_noise = noise
+        else:
+            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:
+        apply_noise()
+        assert (
+            0 >= t < self.timesteps
+        ), f"Time step t={t} is out of bounds for scheduler with {self.timesteps} steps."
+
+        if noise is None:
+            noise = torch.randn_like(x_0) * self.default_noise
+
+        elif 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
+
+
+class NoiseSchedulerC(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_strength = math.sqrt(1.25)
+        self.default_noise_type = "gaussian"
+        self.noise_seed = None
+
+    def _get_random_uniform(self, shape, min_val=-1.0, max_val=1.0):
+        return torch.empty(shape).uniform_(min_val, max_val)
+
+    def _get_noise(self, x: Tensor, noise_type: str, noise_level: float) -> Tensor:
+        # Basic noise types
+        if noise_type == "gaussian":
+            return torch.randn_like(x) * noise_level
+        elif noise_type == "uniform":
+            return self._get_random_uniform(x.shape) * noise_level
+        elif noise_type == "multiplicative":
+            return x * (1 + (torch.randn_like(x) * noise_level))
+        elif noise_type == "clipped_gaussian":
+            noise = torch.randn_like(x) * noise_level
+            return noise.clamp(-1.0, 1.0)
+        elif noise_type == "impulse":
+            mask = torch.rand_like(x) < noise_level
+            impulses = torch.randn_like(x) * noise_level
+            return x + impulses * mask
+        else:
+            raise ValueError(f"Unsupported noise type: '{noise_type}'")
+
+    def set_noise(
+        self,
+        noise_strength: Optional[Union[Tensor, float]] = None,
+        noise_type: Optional[str] = None,
+        seed: Optional[int] = None,
+        default: bool = False,
+    ):
+        if noise_strength is not None:
+            self.default_noise_strength = noise_strength
+        elif default:
+            self.default_noise_strength = math.sqrt(1.25)
+
+        if noise_type is not None:
+            self.default_noise_type = noise_type.lower().strip()
+
+        if isinstance(seed, int):
+            self.noise_seed = seed
+            torch.manual_seed(seed)
+            random.seed(seed)
+
+    def forward(
+        self,
+        x_0: Tensor,
+        t: int,
+        noise: Optional[Union[Tensor, float]] = None,
+        noise_type: Optional[str] = None,
+    ) -> Tensor:
+        assert 0 <= t < self.timesteps, f"t={t} is out of bounds [0, {self.timesteps})"
+
+        noise_type = noise_type or self.default_noise_type
+        noise_level = self.default_noise_strength
+
+        if noise is None:
+            noise = self._get_noise(x_0, noise_type, noise_level)
+        elif isinstance(noise, (float, int)):
+            noise = self._get_noise(x_0, noise_type, 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

lt_tensor/processors/__init__.py

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