lt-tensor 0.0.1a15__tar.gz → 0.0.1a17__tar.gz

{lt_tensor-0.0.1a15 → lt_tensor-0.0.1a17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a15
+Version: 0.0.1a17
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336

{lt_tensor-0.0.1a15 → lt_tensor-0.0.1a17}/lt_tensor/config_templates.py

@@ -7,16 +7,14 @@ from lt_tensor.misc_utils import updateDict
 
 class ModelConfig(ABC, OrderedDict):
     _default_settings: Dict[str, Any] = {}
-    _forbidden_list: List[str] = [
-        "_settings",
-    ]
+    _forbidden_list: List[str] = ["_default_settings", "_forbidden_list" "path_name"]
 
     def __init__(
         self,
-        settings: Dict[str, Any] = None,
+        settings: Dict[str, Any] = {},
         path_name: Optional[Union[str, PathLike]] = None,
     ):
-        assert is_dict(settings)
+        assert is_dict(settings, False)
         self._default_settings = settings
         if path_name is not None and is_pathlike(path_name):
             if not str(path_name).endswith(".json"):
@@ -37,31 +35,41 @@ class ModelConfig(ABC, OrderedDict):
                 self.path_name += ".json"
 
     def reset_settings(self):
+        dk_keys = self.__dict__.keys()
         for s_name, setting in self._default_settings.items():
-            if s_name in self._forbidden_list:
+            if s_name in self._forbidden_list or s_name not in dk_keys:
                 continue
             updateDict(self, {s_name: setting})
 
     def save_config(
         self,
-        path_name: Union[PathLike, str],
+        path_name: Optional[Union[PathLike, str]] = None,
     ):
-        assert is_pathlike(
-            path_name, True
-        ), f"path_name should be a non-empty string or pathlike object! received instead: {path_name}"
-        self._setup_path_name(path_name)
-        base = {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
-        save_json(self.path_name, base, indent=2)
+        if not is_pathlike(path_name, True):
+            assert (
+                path_name is None
+            ), f"path_name should be a non-empty string or pathlike object! received instead: {path_name}."
+            path_name = self.path_name
+        else:
+            self._setup_path_name(path_name)
 
-    def to_dict(self):
-        return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
+        base = self.get_state_dict()
+        save_json(self.path_name, base, indent=2)
 
     def set_value(self, var_name: str, value: str) -> None:
+        assert var_name in self.__dict__, "Value not registered!"
+        assert var_name not in self._forbidden_list, "Not allowed!"
         updateDict(self, {var_name: value})
 
     def get_value(self, var_name: str) -> Any:
         return self.__dict__.get(var_name)
 
+    def __getattribute__(self, name):
+        return self.__dict__.get(name)
+
+    def get_state_dict(self):
+        return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
+
     @classmethod
     def from_dict(
         cls, dictionary: Dict[str, Any], path: Optional[Union[str, PathLike]] = None

{lt_tensor-0.0.1a15 → lt_tensor-0.0.1a17}/lt_tensor/model_base.py

@@ -70,16 +70,6 @@ class LossTracker:
 
 class _Devices_Base(nn.Module):
     _device: torch.device = ROOT_DEVICE
-    _autocast: bool = False
-    _loss_history: LossTracker = LossTracker(100_000)
-
-    @property
-    def autocast(self):
-        return self._autocast
-
-    @autocast.setter
-    def autocast(self, value: bool):
-        self._autocast = value
 
     @property
     def device(self):
@@ -90,6 +80,30 @@ class _Devices_Base(nn.Module):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
 
+    def _apply_device(self):
+        """Add here components that are needed to have device applied to them,
+        that usually the '.to()' function fails to apply
+
+        example:
+        ```
+        def _apply_device_to(self):
+            self.my_tensor = self.my_tensor.to(device=self.device)
+        ```
+        """
+        pass
+
+    def _to_dvc(
+        self, device_name: str, device_id: Optional[Union[int, torch.device]] = None
+    ):
+        device = device_name
+        if device_id is not None:
+            if isinstance(device_id, Number):
+                device += ":" + str(int(device_id))
+            elif hasattr(device_id, "index"):
+                device += ":" + str(device_id.index)
+        self.device = device
+        self._apply_device()
+
     def to(self, *args, **kwargs):
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
             *args, **kwargs
@@ -133,20 +147,9 @@ class _Devices_Base(nn.Module):
                     raise
 
         self._apply(convert)
-        self.device = device
+        self._to_dvc(device)
         return self
 
-    def _to_dvc(
-        self, device_name: str, device_id: Optional[Union[int, torch.device]] = None
-    ):
-        device = device_name
-        if device_id is not None:
-            if isinstance(device_id, Number):
-                device += ":" + str(int(device_id))
-            elif hasattr(device_id, "index"):
-                device += ":" + str(device_id.index)
-        self.device = device
-
     def ipu(self, device: Optional[Union[int, torch.device]] = None) -> T:
         super().ipu(device)
         self._to_dvc("ipu", device)
@@ -178,11 +181,12 @@ class Model(_Devices_Base, ABC):
     This makes it easier to assign a device and retrieves it later
     """
 
+    _autocast: bool = False
     _is_unfrozen: bool = False
     # list with modules that can be frozen or unfrozen
     registered_freezable_modules: List[str] = []
     is_frozen: bool = False
-    _is_gradient_freezable: bool = (
+    _can_be_frozen: bool = (
         False  # to control if the module can or cannot be freezed by other modules from 'Model' class
     )
     # this is to be used on the case of they module requires low-rank adapters
@@ -193,18 +197,15 @@ class Model(_Devices_Base, ABC):
 
     # dont save list:
     _dont_save_items: List[str] = []
+    _loss_history: LossTracker = LossTracker(20_000)
 
-    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
+    @property
+    def autocast(self):
+        return self._autocast
 
-        example:
-        ```
-        def _apply_device_to(self):
-            self.my_tensor = self.my_tensor.to(device=self.device)
-        ```
-        """
-        pass
+    @autocast.setter
+    def autocast(self, value: bool):
+        self._autocast = value
 
     def freeze_all(self, exclude: Optional[List[str]] = None):
         no_exclusions = not exclude
@@ -251,7 +252,7 @@ class Model(_Devices_Base, ABC):
     def change_frozen_state(self, freeze: bool, module: nn.Module):
         try:
             if isinstance(module, Model):
-                if module._is_gradient_freezable:
+                if module._can_be_frozen:
                     if freeze:
                         return module.freeze_all()
                     return module.unfreeze_all()
@@ -496,10 +497,7 @@ class Model(_Devices_Base, ABC):
         return self(*inputs, **kwargs)
 
     def __call__(self, *args, **kwds) -> POSSIBLE_OUTPUT_TYPES:
-        if self.autocast and not self.training:
-            with torch.autocast(device_type=self.device.type):
-                return super().__call__(*args, **kwds)
-        else:
+        with torch.autocast(device_type=self.device.type, enabled=self.autocast):
             return super().__call__(*args, **kwds)
 
     @abstractmethod
@@ -541,52 +539,3 @@ class Model(_Devices_Base, ABC):
         if quantity > 0:
             t_list = t_list[-quantity:]
         return sum(t_list) / len(t_list)
-
-    def freeze_unfreeze_loss(
-        self,
-        losses: Optional[Union[float, List[float]]] = None,
-        trigger_loss: Union[float, bool] = 0.1,
-        excluded_modules: Optional[List[str]] = None,
-        max_items: int = 1000,
-        loss_name: str = "train",
-    ):
-        """If a certain threshold is reached the weights will freeze or unfreeze the modules.
-        the biggest use-case for this function is when training GANs where the balance
-        from the discriminator and generator must be kept.
-
-        Args:
-            losses (Union[float, List[float]], Optional): The loss value or a list of losses that will be used to determine if it has reached or not the threshold. Defaults to None.
-            trigger_loss (float, bool, optional): The value where the weights will be either freeze or unfreeze. If set to a boolean it will freeze or unfreeze immediately according to the value (True = Freeze, False = Unfreeze). Defaults to 0.1.
-            excluded_modules (list[str], optional): The list of modules (names) that is not to be changed by either freezing nor unfreezing. Defaults to None.
-            max_items (float, optional): The number of previous losses to be locked behind to calculate the current average. Default to 1000.
-            loss_name (str, optional): Responsible to define with key to recover the loss.
-        returns:
-            bool: True when its frozen and false when its trainable.
-        """
-        if losses is not None:
-            self.add_loss(losses, "train")
-
-        if isinstance(trigger_loss, bool):
-            if trigger_loss:
-                if self._is_unfrozen:
-                    self.freeze_all(excluded_modules)
-                    self._is_unfrozen = False
-                return True
-            # else
-            if not self._is_unfrozen:
-                self.unfreeze_all(excluded_modules)
-                self._is_unfrozen = True
-            return False
-
-        value = self.get_loss_avg(loss_name, max_items)
-
-        if value <= trigger_loss:
-            if self._is_unfrozen:
-                self.freeze_all(excluded_modules)
-                self._is_unfrozen = False
-            return True
-        else:
-            if not self._is_unfrozen:
-                self.unfreeze_all(excluded_modules)
-                self._is_unfrozen = True
-            return False

{lt_tensor-0.0.1a15 → lt_tensor-0.0.1a17}/lt_tensor/model_zoo/__init__.py

@@ -2,18 +2,20 @@ __all__ = [
     "basic",  # basic
     "residual",  # residual
     "transformer",  # transformer
-    "pos_encoder", 
-    "fusion", 
-    "features", 
-    "discriminator",  
+    "pos_encoder",
+    "fusion",
+    "features",
+    "discriminator",
+    "audio_models",
+    "hifigan",
     "istft",
 ]
+from .audio_models import hifigan, istft
 from . import (
     basic,
-    discriminator,
     features,
     fusion,
-    istft,
+    audio_models,
     pos_encoder,
     residual,
     transformer,

lt_tensor-0.0.1a17/lt_tensor/model_zoo/audio_models/__init__.py

@@ -0,0 +1,2 @@
+from . import diffwave, istft, hifigan
+

lt_tensor-0.0.1a15/lt_tensor/model_zoo/diffwave/model.py → lt_tensor-0.0.1a17/lt_tensor/model_zoo/audio_models/diffwave/__init__.py

@@ -1,48 +1,47 @@
+__all__ = ["DiffWave", "SpectrogramUpsampler", "DiffusionEmbedding"]
+
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
+from lt_tensor.config_templates import ModelConfig
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
 from math import sqrt
+from lt_utils.common import *
 
-
-class AttrDict(dict):
-    def __init__(self, *args, **kwargs):
-        super(AttrDict, self).__init__(*args, **kwargs)
-        self.__dict__ = self
-
-    def override(self, attrs):
-        if isinstance(attrs, dict):
-            self.__dict__.update(**attrs)
-        elif isinstance(attrs, (list, tuple, set)):
-            for attr in attrs:
-                self.override(attr)
-        elif attrs is not None:
-            raise NotImplementedError
-        return self
+F.t
 
 
-params = AttrDict(
+class DiffWaveConfig(ModelConfig):
     # Training params
-    batch_size=16,
-    learning_rate=2e-4,
-    max_grad_norm=None,
+    batch_size = 16
+    learning_rate = 2e-4
+    max_grad_norm = None
     # Data params
-    sample_rate=22050,
-    n_mels=80,
-    n_fft=1024,
-    hop_samples=256,
-    crop_mel_frames=62,  # Probably an error in paper.
+    sample_rate = 24000
+    n_mels = 80
+    n_fft = 1024
+    hop_samples = 256
     # Model params
-    residual_layers=30,
-    residual_channels=64,
-    dilation_cycle_length=10,
-    unconditional=False,
-    noise_schedule=np.linspace(1e-4, 0.05, 50).tolist(),
-    inference_noise_schedule=[0.0001, 0.001, 0.01, 0.05, 0.2, 0.5],
-    # unconditional sample len
-    audio_len=22050 * 5,  # unconditional_synthesis_samples
-)
+    residual_layers = 30
+    residual_channels = 64
+    dilation_cycle_length = 10
+    unconditional = False
+    noise_schedule: list[int] = np.linspace(1e-4, 0.05, 50).tolist()
+    # settings for auto-fixes
+    interpolate = False
+    interpolation_mode: Literal[
+        "nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"
+    ] = "nearest"
+
+    def __init__(
+        self,
+        settings: Dict[str, Any] = {},
+        path_name: Optional[Union[str, PathLike]] = None,
+    ):
+        self._forbidden_list.extend()
+        super().__init__(settings, path_name)
 
 
 def Conv1d(*args, **kwargs):
@@ -51,8 +50,8 @@ def Conv1d(*args, **kwargs):
     return layer
 
 
-class DiffusionEmbedding(nn.Module):
-    def __init__(self, max_steps):
+class DiffusionEmbedding(Model):
+    def __init__(self, max_steps: int):
         super().__init__()
         self.register_buffer(
             "embedding", self._build_embedding(max_steps), persistent=False
@@ -87,23 +86,22 @@ class DiffusionEmbedding(nn.Module):
         return table
 
 
-class SpectrogramUpsampler(nn.Module):
-    def __init__(self, n_mels):
+class SpectrogramUpsampler(Model):
+    def __init__(self):
         super().__init__()
         self.conv1 = nn.ConvTranspose2d(1, 1, [3, 32], stride=[1, 16], padding=[1, 8])
         self.conv2 = nn.ConvTranspose2d(1, 1, [3, 32], stride=[1, 16], padding=[1, 8])
+        self.activation = nn.LeakyReLU(0.4)
 
     def forward(self, x):
         x = torch.unsqueeze(x, 1)
-        x = self.conv1(x)
-        x = F.leaky_relu(x, 0.4)
-        x = self.conv2(x)
-        x = F.leaky_relu(x, 0.4)
+        x = self.activation(self.conv1(x))
+        x = self.activation(self.conv2(x))
         x = torch.squeeze(x, 1)
         return x
 
 
-class ResidualBlock(nn.Module):
+class ResidualBlock(Model):
     def __init__(self, n_mels, residual_channels, dilation, uncond=False):
         """
         :param n_mels: inplanes of conv1x1 for spectrogram conditional
@@ -127,14 +125,18 @@ class ResidualBlock(nn.Module):
 
         self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1)
 
-    def forward(self, x, diffusion_step, conditioner=None):
-        assert (conditioner is None and self.conditioner_projection is None) or (
-            conditioner is not None and self.conditioner_projection is not None
-        )
+    def forward(
+        self,
+        x: Tensor,
+        diffusion_step: Tensor,
+        conditioner: Optional[Tensor] = None,
+    ):
 
         diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
         y = x + diffusion_step
-        if self.conditioner_projection is None:  # using a unconditional model
+        if (
+            conditioner is None or self.conditioner_projection is None
+        ):  # using a unconditional model
             y = self.dilated_conv(y)
         else:
             conditioner = self.conditioner_projection(conditioner)
@@ -148,16 +150,19 @@ class ResidualBlock(nn.Module):
         return (x + residual) / sqrt(2.0), skip
 
 
-class DiffWave(nn.Module):
-    def __init__(self, params):
+class DiffWave(Model):
+    def __init__(self, params: DiffWaveConfig = DiffWaveConfig()):
         super().__init__()
         self.params = params
+        self.n_hop = self.params.hop_samples
+        self.interpolate = self.params.interpolate
+        self.interpolate_mode = self.params.interpolation_mode
         self.input_projection = Conv1d(1, params.residual_channels, 1)
         self.diffusion_embedding = DiffusionEmbedding(len(params.noise_schedule))
         if self.params.unconditional:  # use unconditional model
-            self.spectrogram_upsampler = None
+            self.spectrogram_upsample = None
         else:
-            self.spectrogram_upsampler = SpectrogramUpsampler(params.n_mels)
+            self.spectrogram_upsample = SpectrogramUpsampler()
 
         self.residual_layers = nn.ModuleList(
             [
@@ -174,27 +179,41 @@ class DiffWave(nn.Module):
             params.residual_channels, params.residual_channels, 1
         )
         self.output_projection = Conv1d(params.residual_channels, 1, 1)
+        self.activation = nn.LeakyReLU(0.1)
+        self.r_sqrt = sqrt(len(self.residual_layers))
         nn.init.zeros_(self.output_projection.weight)
 
-    def forward(self, audio, diffusion_step, spectrogram=None):
-        assert (spectrogram is None and self.spectrogram_upsampler is None) or (
-            spectrogram is not None and self.spectrogram_upsampler is not None
-        )
-        x = audio.unsqueeze(1)
-        x = self.input_projection(x)
-        x = F.relu(x)
+    def forward(
+        self,
+        audio: Tensor,
+        diffusion_step: Tensor,
+        spectrogram: Optional[Tensor] = None,
+    ):
+        T = x.shape[-1]
+        if x.ndim == 2:
+            x = audio.unsqueeze(1)
+        x = self.activation(self.input_projection(x))
 
         diffusion_step = self.diffusion_embedding(diffusion_step)
-        if self.spectrogram_upsampler:  # use conditional model
-            spectrogram = self.spectrogram_upsampler(spectrogram)
+        if spectrogram is not None and self.spectrogram_upsample is not None:
+            if self.auto_interpolate:
+                # a little heavy, but helps a lot to fix mismatched shapes,
+                # not always recommended due to data loss
+                spectrogram = F.interpolate(
+                    input=spectrogram,
+                    size=int(T * self.n_hop),
+                    mode=self.interpolate_mode,
+                )
+            spectrogram = self.spectrogram_upsample(spectrogram)
 
         skip = None
-        for layer in self.residual_layers:
+        for i, layer in enumerate(self.residual_layers):
             x, skip_connection = layer(x, diffusion_step, spectrogram)
-            skip = skip_connection if skip is None else skip_connection + skip
-
-        x = skip / sqrt(len(self.residual_layers))
-        x = self.skip_projection(x)
-        x = F.relu(x)
+            if i == 0:
+                skip = skip_connection
+            else:
+                skip = skip_connection + skip
+        x = skip / self.r_sqrt
+        x = self.activation(self.skip_projection(x))
         x = self.output_projection(x)
         return x