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

lt_tensor/datasets/audio.py

@@ -20,6 +20,7 @@ class WaveMelDataset(Dataset):
     randomize_ranges: bool = False
     alpha_wv: float = 1.0
     limit_files: Optional[int] = None
+    min_frame_length: Optional[int] = None
     max_frame_length: Optional[int] = None
 
     def __init__(
@@ -27,12 +28,13 @@ class WaveMelDataset(Dataset):
         audio_processor: AudioProcessor,
         dataset_path: PathLike,
         limit_files: Optional[int] = None,
+        min_frame_length: Optional[int] = None,
         max_frame_length: Optional[int] = None,
         randomize_ranges: Optional[bool] = None,
         pre_load: bool = False,
         normalize_waves: Optional[bool] = None,
         alpha_wv: Optional[float] = None,
-        n_noises: int = 0,  # TODO: Implement the random noises into the dataset
+        lib_norm: bool = True,
     ):
         super().__init__()
         assert max_frame_length is None or max_frame_length >= (
@@ -52,11 +54,15 @@ class WaveMelDataset(Dataset):
             self.randomize_ranges = randomize_ranges
 
         self.post_n_fft = (audio_processor.n_fft // 2) + 1
-
+        self.lib_norm = lib_norm
         if max_frame_length is not None:
             max_frame_length = max(self.post_n_fft + 1, max_frame_length)
             self.r_range = max(self.post_n_fft + 1, max_frame_length // 3)
             self.max_frame_length = max_frame_length
+        if min_frame_length is not None:
+            self.min_frame_length = max(
+                self.post_n_fft + 1, min(min_frame_length, max_frame_length)
+            )
 
         self.files = self.ap.find_audios(dataset_path, maximum=None)
         if limit_files:
@@ -96,21 +102,26 @@ class WaveMelDataset(Dataset):
         }
 
     def load_data(self, file: PathLike):
-        initial_audio = self.ap.normalize_audio(
-            self.ap.load_audio(
-                file, normalize=self.normalize_waves, alpha=self.alpha_wv
-            )
+        initial_audio = self.ap.load_audio(
+            file, normalize=self.lib_norm, alpha=self.alpha_wv
         )
+        if self.normalize_waves:
+            initial_audio = self.ap.normalize_audio(initial_audio)
         if initial_audio.shape[-1] < self.post_n_fft:
             return None
 
+        if self.min_frame_length is not None:
+            if self.min_frame_length > initial_audio.shape[-1]:
+                return None
         if (
             not self.max_frame_length
             or initial_audio.shape[-1] <= self.max_frame_length
         ):
+
             audio_rms = self.ap.compute_rms(initial_audio)
             audio_pitch = self.ap.compute_pitch(initial_audio)
             audio_mel = self.ap.compute_mel(initial_audio, add_base=True)
+
             return [
                 self._add_dict(initial_audio, audio_mel, audio_pitch, audio_rms, file)
             ]
@@ -129,6 +140,12 @@ class WaveMelDataset(Dataset):
             if fragment.shape[-1] < self.post_n_fft:
                 # Too small
                 continue
+            if (
+                self.min_frame_length is not None
+                and self.min_frame_length > fragment.shape[-1]
+            ):
+                continue
+
             audio_rms = self.ap.compute_rms(fragment)
             audio_pitch = self.ap.compute_pitch(fragment)
             audio_mel = self.ap.compute_mel(fragment, add_base=True)

lt_tensor/misc_utils.py

@@ -240,7 +240,7 @@ class LogTensor:
     stored_items: List[
         Dict[str, Union[str, Number, Tensor, List[Union[Tensor, Number, str]]]]
     ] = []
-    max_stored_items: int = 32
+    max_stored_items: int = 8
 
     def _setup_message(self, title: str, t: Union[Tensor, str, int]):
         try:

lt_tensor/model_base.py

@@ -179,6 +179,20 @@ class Model(_Devices_Base, ABC):
     """
 
     _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 = (
+        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
+    _low_rank_lambda: Optional[Callable[[], nn.Module]] = (
+        None  # Example: lambda: nn.Linear(32, 32, True)
+    )
+    low_rank_adapter: Union[nn.Identity, nn.Module, nn.Sequential] = nn.Identity()
+
+    # dont save list:
+    _dont_save_items: List[str] = []
 
     def _apply_device_to(self):
         """Add here components that are needed to have device applied to them,
@@ -192,116 +206,67 @@ class Model(_Devices_Base, ABC):
         """
         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():
+        no_exclusions = not exclude
+        results = []
+        for name, module in self.named_modules():
+            if name not in self.registered_freezable_modules:
+                results.append(
+                    (
+                        name,
+                        "Unregistered module, to freeze/unfreeze it add its name into 'registered_freezable_modules'.",
+                    )
+                )
+                continue
             if no_exclusions:
-                try:
-                    if param.requires_grad:
-                        param.requires_grad_(False)
-                        frozen.append(name)
-                    else:
-                        not_frozen.append((name, "was_frozen"))
-                except Exception as e:
-                    not_frozen.append((name, str(e)))
-            elif any(layer in name for layer in exclude):
-                try:
-                    if param.requires_grad:
-                        param.requires_grad_(False)
-                        frozen.append(name)
-                    else:
-                        not_frozen.append((name, "was_frozen"))
-                except Exception as e:
-                    not_frozen.append((name, str(e)))
+                self.change_frozen_state(True, module)
+            elif not any(exclusion in name for exclusion in exclude):
+                results.append((name, self.change_frozen_state(True, module)))
             else:
-                not_frozen.append((name, "excluded"))
-        return dict(frozen=frozen, not_frozen=not_frozen)
+                results.append((name, "excluded"))
+        return results
 
     def unfreeze_all(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():
+        results = []
+        for name, module in self.named_modules():
+            if name not in self.registered_freezable_modules:
+                results.append(
+                    (
+                        name,
+                        "Unregistered module, to freeze/unfreeze it add it into 'registered_freezable_modules'.",
+                    )
+                )
+                continue
             if no_exclusions:
-                try:
-                    if not param.requires_grad:
-                        param.requires_grad_(True)
-                        unfrozen.append(name)
-                    else:
-                        not_unfrozen.append((name, "was_unfrozen"))
-                except Exception as e:
-                    not_unfrozen.append((name, str(e)))
-            elif any(layer in name for layer in exclude):
-                try:
-                    if not param.requires_grad:
-                        param.requires_grad_(True)
-                        unfrozen.append(name)
-                    else:
-                        not_unfrozen.append((name, "was_unfrozen"))
-                except Exception as e:
-                    not_unfrozen.append((name, str(e)))
+                self.change_frozen_state(False, module)
+            elif not any(exclusion in name for exclusion in exclude):
+                results.append((name, self.change_frozen_state(False, module)))
             else:
-                not_unfrozen.append((name, "excluded"))
-        return dict(unfrozen=unfrozen, not_unfrozen=not_unfrozen)
+                results.append((name, "excluded"))
+        return results
+
+    def change_frozen_state(self, freeze: bool, module: nn.Module):
+        try:
+            if isinstance(module, Model):
+                if module._is_gradient_freezable:
+                    if freeze:
+                        return module.freeze_all()
+                    return module.unfreeze_all()
+                else:
+                    return "Not Allowed"
+            elif isinstance(module, nn.Module):
+                module.requires_grad_(not freeze)
+                return not freeze
+        except Exception as e:
+            return e
 
-    def count_trainable_parameters(self, module_name: Optional[str] = None):
+    def trainable_parameters(self, module_name: Optional[str] = None):
         """Gets the number of trainable parameters from either the entire model or from a specific module."""
         if module_name is not None:
-            assert hasattr(self, module_name), f"Module {module_name} does not exits"
+            assert hasattr(self, module_name), f"Module '{module_name}' not found."
             module = getattr(self, module_name)
             return sum(
                 [
@@ -318,10 +283,10 @@ class Model(_Devices_Base, ABC):
             ]
         )
 
-    def count_non_trainable_parameters(self, module_name: Optional[str] = None):
+    def non_trainable_parameters(self, module_name: Optional[str] = None):
         """Gets the number of non-trainable parameters from either the entire model or from a specific module."""
         if module_name is not None:
-            assert hasattr(self, module_name), f"Module {module_name} does not exits"
+            assert hasattr(self, module_name), f"Module '{module_name}' not found."
             module = getattr(self, module_name)
             return sum(
                 [
@@ -338,10 +303,10 @@ class Model(_Devices_Base, ABC):
             ]
         )
 
-    def get_weights(self, module_name: Optional[str] = None) -> List[Tensor]:
+    def extract_weights(self, module_name: Optional[str] = None) -> List[Tensor]:
         """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"
+            assert hasattr(self, module_name), f"Module '{module_name}' not found."
             module = getattr(self, module_name)
             params = []
             if isinstance(module, nn.Module):
@@ -351,27 +316,29 @@ class Model(_Devices_Base, ABC):
             raise (f"{module_name} is has no weights")
         return [x.data.detach() for x in self.parameters()]
 
-    def print_trainable_parameters(
-        self, module_name: Optional[str] = None
-    ) -> List[Tensor]:
-        params = format(self.count_trainable_parameters(module_name), ",").replace(
+    def format_trainable_parameters(self, module_name: Optional[str] = None) -> str:
+        params = format(self.trainable_parameters(module_name), ",").replace(",", ".")
+        return params
+
+    def format_non_trainable_parameters(self, module_name: Optional[str] = None) -> str:
+        params = format(self.non_trainable_parameters(module_name), ",").replace(
             ",", "."
         )
-        if module_name:
-            print(f'Trainable Parameters from "{module_name}": {params}')
+        return params
+
+    def print_trainable_parameters(self, module_name: Optional[str] = None) -> str:
+        fmt = self.format_trainable_parameters(module_name)
+        if module_name is not None:
+            print(f"Trainable parameter(s) for module '{module_name}': {fmt}")
         else:
-            print(f"Trainable Parameters: {params}")
+            print(f"Trainable parameter(s): {fmt}")
 
-    def print_non_trainable_parameters(
-        self, module_name: Optional[str] = None
-    ) -> List[Tensor]:
-        params = format(self.count_non_trainable_parameters(module_name), ",").replace(
-            ",", "."
-        )
-        if module_name:
-            print(f'Non-Trainable Parameters from "{module_name}": {params}')
+    def print_non_trainable_parameters(self, module_name: Optional[str] = None) -> str:
+        fmt = self.format_non_trainable_parameters(module_name)
+        if module_name is not None:
+            print(f"Non-Trainable parameter(s) for module '{module_name}': {fmt}")
         else:
-            print(f"Non-Trainable Parameters: {params}")
+            print(f"Non-Trainable parameter(s): {fmt}")
 
     @classmethod
     def from_pretrained(
@@ -390,6 +357,7 @@ class Model(_Devices_Base, ABC):
         self,
         path: Union[Path, str],
         replace: bool = False,
+        save_with_adapters: bool = False,
     ):
         path = Path(path)
         model_dir = path
@@ -405,17 +373,86 @@ class Model(_Devices_Base, ABC):
             if not "." in str(path):
                 model_dir = Path(path, f"model_{get_current_time()}.pt")
         path.parent.mkdir(exist_ok=True, parents=True)
-        torch.save(obj=self.state_dict(), f=str(model_dir))
+
+        state_dict = self.state_dict()
+        if not save_with_adapters or isinstance(self.low_rank_adapter, nn.Identity):
+            state_dict.pop("low_rank_adapter", None)
+        torch.save(obj=state_dict, f=str(model_dir))
+
+    def save_lora(
+        self,
+        path: Union[Path, str],
+        replace: bool = False,
+    ):
+        assert not isinstance(
+            self.low_rank_adapter, nn.Identity
+        ), "The adapter is empty!"
+        path = Path(path)
+        model_dir = path
+        if path.exists():
+            if path.is_dir():
+                model_dir = Path(path, f"adapter_{get_current_time()}.pt")
+            elif path.is_file():
+                if replace:
+                    path.unlink()
+                else:
+                    model_dir = Path(path.parent, f"adapter_{get_current_time()}.pt")
+        else:
+            if not "." in str(path):
+                model_dir = Path(path, f"adapter_{get_current_time()}.pt")
+
+        state_dict = self.low_rank_adapter.state_dict()
+        torch.save(obj=state_dict, f=str(model_dir))
+
+    def load_lora(
+        self,
+        path: Union[Path, str],
+        raise_if_not_exists: bool = False,
+        strict: bool = False,
+        assign: bool = False,
+        weights_only: bool = True,
+        mmap: Optional[bool] = None,
+        **pickle_load_args,
+    ):
+        assert (
+            self._low_rank_lambda is not None
+        ), "Lora not implemented! '_low_rank_lambda' must be setup to deploy a proper module"
+        path = Path(path)
+        if not path.exists():
+            assert not raise_if_not_exists, "Path does not exists!"
+            return None
+
+        if path.is_dir():
+            possible_files = list(Path(path).rglob("adapter_*.pt"))
+            assert (
+                possible_files or not raise_if_not_exists
+            ), "No model could be found in the given path!"
+            if not possible_files:
+                return None
+            path = sorted(possible_files)[-1]
+
+        state_dict = torch.load(
+            str(path), weights_only=weights_only, mmap=mmap, **pickle_load_args
+        )
+        self.low_rank_adapter = None
+        gc.collect()
+        self.low_rank_adapter = self._low_rank_lambda()
+        incompatible_keys = self.low_rank_adapter.load_state_dict(
+            state_dict,
+            strict=strict,
+            assign=assign,
+        )
+        return incompatible_keys
 
     def load_weights(
         self,
         path: Union[Path, str],
         raise_if_not_exists: bool = False,
-        strict: bool = True,
+        strict: bool = False,
         assign: bool = False,
-        weights_only: bool = False,
+        weights_only: bool = True,
         mmap: Optional[bool] = None,
-        **torch_loader_kwargs,
+        **pickle_load_args,
     ):
         path = Path(path)
         if not path.exists():
@@ -430,7 +467,7 @@ class Model(_Devices_Base, ABC):
                 return None
             path = sorted(possible_files)[-1]
         state_dict = torch.load(
-            str(path), weights_only=weights_only, mmap=mmap, **torch_loader_kwargs
+            str(path), weights_only=weights_only, mmap=mmap, **pickle_load_args
         )
         incompatible_keys = self.load_state_dict(
             state_dict,
@@ -439,6 +476,9 @@ class Model(_Devices_Base, ABC):
         )
         return incompatible_keys
 
+    def lora_step(self, *arg, **kwargs):
+        raise NotImplementedError("Not implemented for this model")
+
     @torch.no_grad()
     def inference(self, *args, **kwargs):
         if self.training:
@@ -524,7 +564,7 @@ class Model(_Devices_Base, ABC):
             bool: True when its frozen and false when its trainable.
         """
         if losses is not None:
-            self.add_loss(losses)
+            self.add_loss(losses, "train")
 
         if isinstance(trigger_loss, bool):
             if trigger_loss:

lt_tensor/model_zoo/diffwave/model.py

@@ -0,0 +1,200 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from math import sqrt
+
+
+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
+
+
+params = AttrDict(
+    # Training params
+    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.
+    # 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
+)
+
+
+def Conv1d(*args, **kwargs):
+    layer = nn.Conv1d(*args, **kwargs)
+    nn.init.kaiming_normal_(layer.weight)
+    return layer
+
+
+class DiffusionEmbedding(nn.Module):
+    def __init__(self, max_steps):
+        super().__init__()
+        self.register_buffer(
+            "embedding", self._build_embedding(max_steps), persistent=False
+        )
+        self.projection1 = nn.Linear(128, 512)
+        self.projection2 = nn.Linear(512, 512)
+        self.activation = nn.SiLU()
+
+    def forward(self, diffusion_step):
+        if diffusion_step.dtype in [torch.int32, torch.int64]:
+            x = self.embedding[diffusion_step]
+        else:
+            x = self._lerp_embedding(diffusion_step)
+        x = self.projection1(x)
+        x = self.activation(x)
+        x = self.projection2(x)
+        x = self.activation(x)
+        return x
+
+    def _lerp_embedding(self, t):
+        low_idx = torch.floor(t).long()
+        high_idx = torch.ceil(t).long()
+        low = self.embedding[low_idx]
+        high = self.embedding[high_idx]
+        return low + (high - low) * (t - low_idx)
+
+    def _build_embedding(self, max_steps):
+        steps = torch.arange(max_steps).unsqueeze(1)  # [T,1]
+        dims = torch.arange(64).unsqueeze(0)  # [1,64]
+        table = steps * 10.0 ** (dims * 4.0 / 63.0)  # [T,64]
+        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)
+        return table
+
+
+class SpectrogramUpsampler(nn.Module):
+    def __init__(self, n_mels):
+        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])
+
+    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 = torch.squeeze(x, 1)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, n_mels, residual_channels, dilation, uncond=False):
+        """
+        :param n_mels: inplanes of conv1x1 for spectrogram conditional
+        :param residual_channels: audio conv
+        :param dilation: audio conv dilation
+        :param uncond: disable spectrogram conditional
+        """
+        super().__init__()
+        self.dilated_conv = Conv1d(
+            residual_channels,
+            2 * residual_channels,
+            3,
+            padding=dilation,
+            dilation=dilation,
+        )
+        self.diffusion_projection = nn.Linear(512, residual_channels)
+        if not uncond:  # conditional model
+            self.conditioner_projection = Conv1d(n_mels, 2 * residual_channels, 1)
+        else:  # unconditional model
+            self.conditioner_projection = None
+
+        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
+        )
+
+        diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
+        y = x + diffusion_step
+        if self.conditioner_projection is None:  # using a unconditional model
+            y = self.dilated_conv(y)
+        else:
+            conditioner = self.conditioner_projection(conditioner)
+            y = self.dilated_conv(y) + conditioner
+
+        gate, filter = torch.chunk(y, 2, dim=1)
+        y = torch.sigmoid(gate) * torch.tanh(filter)
+
+        y = self.output_projection(y)
+        residual, skip = torch.chunk(y, 2, dim=1)
+        return (x + residual) / sqrt(2.0), skip
+
+
+class DiffWave(nn.Module):
+    def __init__(self, params):
+        super().__init__()
+        self.params = params
+        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
+        else:
+            self.spectrogram_upsampler = SpectrogramUpsampler(params.n_mels)
+
+        self.residual_layers = nn.ModuleList(
+            [
+                ResidualBlock(
+                    params.n_mels,
+                    params.residual_channels,
+                    2 ** (i % params.dilation_cycle_length),
+                    uncond=params.unconditional,
+                )
+                for i in range(params.residual_layers)
+            ]
+        )
+        self.skip_projection = Conv1d(
+            params.residual_channels, params.residual_channels, 1
+        )
+        self.output_projection = Conv1d(params.residual_channels, 1, 1)
+        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)
+
+        diffusion_step = self.diffusion_embedding(diffusion_step)
+        if self.spectrogram_upsampler:  # use conditional model
+            spectrogram = self.spectrogram_upsampler(spectrogram)
+
+        skip = None
+        for layer in 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)
+        x = self.output_projection(x)
+        return x