lt-tensor 0.0.1a22__tar.gz → 0.0.1a26__tar.gz

{lt_tensor-0.0.1a22 → lt_tensor-0.0.1a26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a22
+Version: 0.0.1a26
 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.1a22 → lt_tensor-0.0.1a26}/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-0.0.1a22 → lt_tensor-0.0.1a26}/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-0.0.1a22 → lt_tensor-0.0.1a26}/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-0.0.1a22 → lt_tensor-0.0.1a26}/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
 
 

{lt_tensor-0.0.1a22 → lt_tensor-0.0.1a26}/lt_tensor/model_zoo/audio_models/istft/__init__.py

@@ -1,11 +1,11 @@
 __all__ = ["iSTFTNetGenerator", "iSTFTNetConfig"]
 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_tensor.config_templates import ModelConfig
+from lt_tensor.misc_utils import get_config, get_weights
 from lt_utils.file_ops import load_json, is_file, is_dir, is_path_valid
-from huggingface_hub import hf_hub_download
 
 
 class iSTFTNetConfig(ModelConfig):
@@ -39,7 +39,7 @@ class iSTFTNetConfig(ModelConfig):
             [1, 3, 5],
         ],
         activation: nn.Module = nn.LeakyReLU(0.1),
-        resblock: int = 0,
+        resblock: Union[int, str] = "1",
         gen_istft_n_fft: int = 16,
         sampling_rate: Number = 24000,
         *args,
@@ -59,7 +59,11 @@ class iSTFTNetConfig(ModelConfig):
         }
         super().__init__(**settings)
 
-
+    def post_process(self):
+        if isinstance(self.resblock, str):
+            self.resblock = 0 if self.resblock == "1" else 1
+            
+            
 def get_padding(ks, d):
     return int((ks * d - d) / 2)
 
@@ -271,15 +275,15 @@ class iSTFTNetGenerator(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,
@@ -288,6 +292,18 @@ class iSTFTNetGenerator(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(
@@ -303,95 +319,37 @@ class iSTFTNetGenerator(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[iSTFTNetConfig, 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(iSTFTNetConfig(**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, iSTFTNetConfig):
+            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 = iSTFTNetConfig(**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