lt-tensor 0.0.1a14__py3-none-any.whl → 0.0.1a16__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/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/__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/model_zoo/audio_models/__init__.py

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

lt_tensor/model_zoo/audio_models/diffwave/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["DiffWave", "SpectrogramUpsampler", "DiffusionEmbedding"]
+
+from .model import DiffWave, SpectrogramUpsampler, DiffusionEmbedding