lt-tensor 0.0.1a22__py3-none-any.whl → 0.0.1a27__py3-none-any.whl

lt_tensor/config_templates.py

@@ -6,9 +6,7 @@ from lt_tensor.misc_utils import updateDict
 
 
 class ModelConfig(ABC, OrderedDict):
-    _default_settings: Dict[str, Any] = {}
     _forbidden_list: List[str] = [
-        "_default_settings",
         "_forbidden_list",
     ]
 
@@ -16,12 +14,15 @@ class ModelConfig(ABC, OrderedDict):
         self,
         **settings,
     ):
-        self._default_settings = settings
-        self.set_state_dict(self._default_settings)
+        self.set_state_dict(settings)
 
     def reset_settings(self):
         raise NotImplementedError("Not implemented")
-
+    
+    def post_process(self):
+        """Implement the post process, to do a final check to the input data"""
+        pass
+    
     def save_config(
         self,
         path: str,
@@ -48,6 +49,7 @@ class ModelConfig(ABC, OrderedDict):
         }
         updateDict(self, new_state)
         self.update(**new_state)
+        self.post_process()
 
     def state_dict(self):
         return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
@@ -89,3 +91,5 @@ class ModelConfig(ABC, OrderedDict):
         settings.pop("path_name", None)
 
         return ModelConfig(**settings)
+
+

lt_tensor/misc_utils.py

@@ -111,10 +111,10 @@ def get_weights(directory: Union[str, PathLike]):
     directory = Path(directory)
     if is_file(directory):
         if directory.name.endswith((".pt", ".ckpt", ".pth")):
-            return directory
+            return [directory]
         directory = directory.parent
     res = sorted(find_files(directory, ["*.pt", "*.ckpt", "*.pth"]))
-    return res[-1] if res else None
+    return res
 
 
 def get_config(directory: Union[str, PathLike], default: Optional[Any] = None):
@@ -128,7 +128,19 @@ def get_config(directory: Union[str, PathLike], default: Optional[Any] = None):
                 return load_json(directory, default)
             return load_yaml(directory, default)
         directory = directory.parent
-    res = sorted(find_files(directory, ["*.pt", "*.ckpt", "*.pth"]))
+    res = sorted(
+        find_files(
+            directory,
+            [
+                "config*.json",
+                "*config.json",
+                "config*.yml",
+                "*config.yml",
+                "*config.yaml",
+                "config*.yaml",
+            ],
+        )
+    )
     if res:
         res = res[-1]
         if Path(res).name.endswith(".json"):

lt_tensor/model_base.py

@@ -80,6 +80,62 @@ class _Devices_Base(nn.Module):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
 
+    def freeze_all(self, exclude: list[str] = []):
+        for name, module in self.named_modules():
+            if name in exclude or not hasattr(module, "requires_grad"):
+                continue
+            try:
+                self.freeze_module(module)
+            except:
+                pass
+
+    def unfreeze_all(self, exclude: list[str] = []):
+        for name, module in self.named_modules():
+            if name in exclude or not hasattr(module, "requires_grad"):
+                continue
+            try:
+                self.unfreeze_module(module)
+            except:
+                pass
+
+    def freeze_module(
+        self, module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor]
+    ):
+        self._change_gradient_state(module_or_name, False)
+
+    def unfreeze_module(
+        self, module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor]
+    ):
+        self._change_gradient_state(module_or_name, True)
+
+    def _change_gradient_state(
+        self,
+        module_or_name: Union[str, nn.Module, nn.Parameter, "Model", Tensor],
+        new_state: bool,  # True = unfreeze
+    ):
+        assert isinstance(
+            module_or_name, (str, nn.Module, nn.Parameter, Model, Tensor)
+        ), f"Item '{module_or_name}' is not a valid module, parameter, tensor or a string."
+        if isinstance(module_or_name, (nn.Module, nn.Parameter, Model, Tensor)):
+            target = module_or_name
+        else:
+            target = getattr(self, module_or_name)
+
+        if isinstance(target, Tensor):
+            target.requires_grad = new_state
+        elif isinstance(target, nn.Parameter):
+            target.requires_grad = new_state
+        elif isinstance(target, Model):
+            target.freeze_all()
+        elif isinstance(target, nn.Module):
+            for param in target.parameters():
+                if hasattr(param, "requires_grad"):
+                    param.requires_grad = new_state
+        else:
+            raise ValueError(
+                f"Item '{module_or_name}' is not a valid module, parameter or tensor."
+            )
+
     def _apply_device(self):
         """Add here components that are needed to have device applied to them,
         that usually the '.to()' function fails to apply
@@ -182,20 +238,12 @@ class Model(_Devices_Base, ABC):
     """
 
     _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
-    _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
     _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()
-    # never freeze:
-    _never_freeze_modules: List[str] = ["low_rank_adapter"]
 
     # dont save list:
     _dont_save_items: List[str] = []
@@ -208,75 +256,6 @@ class Model(_Devices_Base, ABC):
     def autocast(self, value: bool):
         self._autocast = value
 
-    def freeze_all(self, exclude: Optional[List[str]] = None, force: bool = False):
-        no_exclusions = not exclude
-        no_exclusions = not exclude
-        results = []
-        for name, module in self.named_modules():
-            if (
-                name in self._never_freeze_modules
-                or not force
-                and 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:
-                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:
-                results.append((name, "excluded"))
-        return results
-
-    def unfreeze_all(self, exclude: Optional[list[str]] = None, force: bool = False):
-        """Unfreezes all model parameters except specified layers."""
-        no_exclusions = not exclude
-        results = []
-        for name, module in self.named_modules():
-            if (
-                name in self._never_freeze_modules
-                or not force
-                and 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:
-                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:
-                results.append((name, "excluded"))
-        return results
-
-    def change_frozen_state(self, freeze: bool, module: nn.Module):
-        assert isinstance(module, nn.Module)
-        if module.__class__.__name__ in self._never_freeze_modules:
-            return "Not Allowed"
-        try:
-            if isinstance(module, Model):
-                if module._can_be_frozen:
-                    if freeze:
-                        return module.freeze_all()
-                    return module.unfreeze_all()
-                else:
-                    return "Not Allowed"
-            else:
-                module.requires_grad_(not freeze)
-                return not freeze
-        except Exception as e:
-            return e
-
     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:

lt_tensor/model_zoo/audio_models/diffwave/__init__.py

@@ -5,6 +5,7 @@ from lt_tensor.torch_commons import *
 from torch.nn import functional as F
 from lt_tensor.config_templates import ModelConfig
 from lt_tensor.torch_commons import *
+from lt_tensor.model_zoo.convs import ConvNets, Conv1dEXT
 from lt_tensor.model_base import Model
 from math import sqrt
 from lt_utils.common import *
@@ -18,6 +19,8 @@ class DiffWaveConfig(ModelConfig):
     residual_channels = 64
     dilation_cycle_length = 10
     unconditional = False
+    apply_norm: Optional[Literal["weight", "spectral"]] = None
+    apply_norm_resblock: Optional[Literal["weight", "spectral"]] = None
     noise_schedule: list[int] = np.linspace(1e-4, 0.05, 50).tolist()
     # settings for auto-fixes
     interpolate = False
@@ -44,6 +47,8 @@ class DiffWaveConfig(ModelConfig):
             "area",
             "nearest-exact",
         ] = "nearest",
+        apply_norm: Optional[Literal["weight", "spectral"]] = None,
+        apply_norm_resblock: Optional[Literal["weight", "spectral"]] = None,
     ):
         settings = {
             "n_mels": n_mels,
@@ -55,16 +60,12 @@ class DiffWaveConfig(ModelConfig):
             "noise_schedule": noise_schedule,
             "interpolate": interpolate_cond,
             "interpolation_mode": interpolation_mode,
+            "apply_norm": apply_norm,
+            "apply_norm_resblock": apply_norm_resblock,
         }
         super().__init__(**settings)
 
 
-def Conv1d(*args, **kwargs):
-    layer = nn.Conv1d(*args, **kwargs)
-    nn.init.kaiming_normal_(layer.weight)
-    return layer
-
-
 class DiffusionEmbedding(Model):
     def __init__(self, max_steps: int):
         super().__init__()
@@ -117,7 +118,14 @@ class SpectrogramUpsample(Model):
 
 
 class ResidualBlock(Model):
-    def __init__(self, n_mels, residual_channels, dilation, uncond=False):
+    def __init__(
+        self,
+        n_mels,
+        residual_channels,
+        dilation,
+        uncond=False,
+        apply_norm: Optional[Literal["weight", "spectral"]] = None,
+    ):
         """
         :param n_mels: inplanes of conv1x1 for spectrogram conditional
         :param residual_channels: audio conv
@@ -125,20 +133,28 @@ class ResidualBlock(Model):
         :param uncond: disable spectrogram conditional
         """
         super().__init__()
-        self.dilated_conv = Conv1d(
+        self.dilated_conv = Conv1dEXT(
             residual_channels,
             2 * residual_channels,
             3,
             padding=dilation,
             dilation=dilation,
+            apply_norm=apply_norm,
         )
         self.diffusion_projection = nn.Linear(512, residual_channels)
         if not uncond:  # conditional model
-            self.conditioner_projection = Conv1d(n_mels, 2 * residual_channels, 1)
+            self.conditioner_projection = Conv1dEXT(
+                n_mels,
+                2 * residual_channels,
+                1,
+                apply_norm=apply_norm,
+            )
         else:  # unconditional model
             self.conditioner_projection = None
 
-        self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1)
+        self.output_projection = Conv1dEXT(
+            residual_channels, 2 * residual_channels, 1, apply_norm == apply_norm
+        )
 
     def forward(
         self,
@@ -172,7 +188,12 @@ class DiffWave(Model):
         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.input_projection = Conv1dEXT(
+            in_channels=1,
+            out_channels=params.residual_channels,
+            kernel_size=1,
+            apply_norm=self.params.apply_norm,
+        )
         self.diffusion_embedding = DiffusionEmbedding(len(params.noise_schedule))
         if self.params.unconditional:  # use unconditional model
             self.spectrogram_upsample = None
@@ -186,14 +207,20 @@ class DiffWave(Model):
                     params.residual_channels,
                     2 ** (i % params.dilation_cycle_length),
                     uncond=params.unconditional,
+                    apply_norm=self.params.apply_norm_resblock,
                 )
                 for i in range(params.residual_layers)
             ]
         )
-        self.skip_projection = Conv1d(
-            params.residual_channels, params.residual_channels, 1
+        self.skip_projection = Conv1dEXT(
+            in_channels=params.residual_channels,
+            out_channels=params.residual_channels,
+            kernel_size=1,
+            apply_norm=self.params.apply_norm,
+        )
+        self.output_projection = Conv1dEXT(
+            params.residual_channels, 1, 1, apply_norm=self.params.apply_norm
         )
-        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)

lt_tensor/model_zoo/audio_models/hifigan/__init__.py

@@ -1,10 +1,10 @@
 __all__ = ["HifiganGenerator", "HifiganConfig"]
 from lt_utils.common import *
 from lt_tensor.torch_commons import *
-from lt_tensor.model_zoo.residual import ConvNets
+from lt_tensor.model_zoo.convs import ConvNets
 from torch.nn import functional as F
 from lt_utils.file_ops import load_json, is_file, is_dir, is_path_valid
-from huggingface_hub import hf_hub_download
+from lt_tensor.misc_utils import get_config, get_weights
 
 
 def get_padding(kernel_size, dilation=1):
@@ -43,7 +43,7 @@ class HifiganConfig(ModelConfig):
             [1, 3, 5],
         ],
         activation: nn.Module = nn.LeakyReLU(0.1),
-        resblock: int = 0,
+        resblock: Union[int, str] = "1",
         *args,
         **kwargs,
     ):
@@ -59,6 +59,10 @@ class HifiganConfig(ModelConfig):
         }
         super().__init__(**settings)
 
+    def post_process(self):
+        if isinstance(self.resblock, str):
+            self.resblock = 0 if self.resblock == "1" else 1
+
 
 class ResBlock1(ConvNets):
     def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
@@ -240,15 +244,15 @@ class HifiganGenerator(ConvNets):
     def load_weights(
         self,
         path,
-        raise_if_not_exists=False,
         strict=False,
         assign=False,
-        weights_only=True,
+        weights_only=False,
         mmap=None,
+        raise_if_not_exists=False,
         **pickle_load_args,
     ):
         try:
-            return super().load_weights(
+            incompatible_keys = super().load_weights(
                 path,
                 raise_if_not_exists,
                 strict,
@@ -257,6 +261,18 @@ class HifiganGenerator(ConvNets):
                 mmap,
                 **pickle_load_args,
             )
+            if incompatible_keys:
+                self.remove_norms()
+                incompatible_keys = super().load_weights(
+                    path,
+                    raise_if_not_exists,
+                    strict,
+                    assign,
+                    weights_only,
+                    mmap,
+                    **pickle_load_args,
+                )
+            return incompatible_keys
         except RuntimeError:
             self.remove_norms()
             return super().load_weights(
@@ -272,95 +288,37 @@ class HifiganGenerator(ConvNets):
     @classmethod
     def from_pretrained(
         cls,
-        model_id: str,
-        map_location: str = "cpu",
-        local_files_only: bool = False,
-        strict: bool = False,
+        model_file: PathLike,
+        model_config: Union[HifiganConfig, Dict[str, Any]],
         *,
-        subfolder: str | None = None,
-        repo_type: str | None = None,
-        revision: str | None = None,
-        cache_dir: str | Path | None = None,
-        force_download: bool = False,
-        proxies: Dict | None = None,
-        token: bool | str | None = None,
-        resume_download: bool | None = None,
-        local_dir_use_symlinks: bool | Literal["auto"] = "auto",
+        strict: bool = False,
+        map_location: str = "cpu",
+        weights_only: bool = False,
         **kwargs,
     ):
-        """Load Pytorch pretrained weights and return the loaded model."""
-        hub_kwargs = dict(
-            repo_id=model_id,
-            subfolder=subfolder,
-            repo_type=repo_type,
-            revision=revision,
-            cache_dir=cache_dir,
-            force_download=force_download,
-            proxies=proxies,
-            resume_download=resume_download,
-            token=token,
-            local_files_only=local_files_only,
-            local_dir_use_symlinks=local_dir_use_symlinks,
-        )
-
-        # Download and load hyperparameters (h) used by BigVGAN
-        _model_path = Path(model_id)
-        config_file = None
-        if is_path_valid(model_id):
-            if is_file(model_id):
-                _p_conf = _model_path.parent / "config.json"
-            else:
-                _p_conf = _model_path / "config.json"
-
-            if is_file(_p_conf):
-                print("Loading config.json from local directory")
-                config_file = Path(model_id, "config.json")
-            else:
-                if not local_files_only:
-                    print(f"Loading config from {model_id}")
-                    config_file = hf_hub_download(filename="config.json", **hub_kwargs)
-        else:
-            if not local_files_only:
-                print(f"Loading config from {model_id}")
-                config_file = hf_hub_download(filename="config.json", **hub_kwargs)
-
-        if config_file is not None:
-            model = cls(HifiganConfig(**load_json(config_file)))
-        else:
-            model = cls()
 
-        # Download and load pretrained generator weight
-        _retrieve_kwargs = dict(
-            **hub_kwargs,
-            filename="generator.pt",
+        is_file(model_file, validate=True)
+        model_state_dict = torch.load(
+            model_file, weights_only=weights_only, map_location=map_location
         )
-        path = Path(model_id)
-        if path.exists():
-            if path.is_dir():
-                path = path / "generator.pt"
-                if path.exists():
-                    print("Loading weights from local directory")
-                    model_file = str(path)
-                else:
-                    print(f"Loading weights from {model_id}")
-                    model_file = hf_hub_download(**_retrieve_kwargs)
-            else:
-                print("Loading weights from local directory")
-                model_file = str(path)
+
+        if isinstance(model_config, HifiganConfig):
+            h = model_config
         else:
-            print(f"Loading weights from {model_id}")
-            model_file = hf_hub_download(**_retrieve_kwargs)
-        checkpoint_dict = torch.load(model_file, map_location=map_location)
+            h = HifiganConfig(**model_config)
 
+        model = cls(h)
         try:
-            model.load_state_dict(checkpoint_dict, strict=strict)
+            incompatible_keys = model.load_state_dict(model_state_dict, strict=strict)
+            if incompatible_keys:
+                model.remove_norms()
+                model.load_state_dict(model_state_dict, strict=strict)
         except RuntimeError:
             print(
                 f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!"
             )
             model.remove_norms()
-            model.load_state_dict(checkpoint_dict, strict=strict)
-
+            model.load_state_dict(model_state_dict, strict=strict)
         return model