lt-tensor 0.0.1a38__py3-none-any.whl → 0.0.1a40__py3-none-any.whl

lt_tensor/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.1a38"
+__version__ = "0.0.1a40"
 
 from . import (
     lr_schedulers,

lt_tensor/model_zoo/audio_models/bemaganv2/__init__.py

@@ -0,0 +1,205 @@
+from lt_utils.common import *
+from lt_tensor.torch_commons import *
+from lt_tensor.model_zoo.convs import ConvNets
+from lt_tensor.config_templates import ModelConfig
+from lt_utils.file_ops import is_file, load_json
+from lt_tensor.model_zoo.audio_models.resblocks import ResBlock1, ResBlock2
+from lt_tensor.model_zoo.activations import snake, alias_free
+from lt_tensor.model_zoo.audio_models.resblocks import AMPBlock1, AMPBlock2, get_snake
+
+
+class BemaGANv2Config(ModelConfig):
+    # Training params
+    in_channels: int = 80
+    upsample_rates: List[Union[int, List[int]]] = [8, 8, 2, 2]
+    upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16, 4, 4]
+    upsample_initial_channel: int = 1536
+    resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11]
+    resblock_dilation_sizes: List[Union[int, List[int]]] = [
+        [1, 3, 5],
+        [1, 3, 5],
+        [1, 3, 5],
+    ]
+
+    activation: Literal["snake", "snakebeta"] = "snakebeta"
+    resblock_activation: Literal["snake", "snakebeta"] = "snakebeta"
+    resblock: int = 0
+    use_bias_at_final: bool = True
+    use_tanh_at_final: bool = True
+    snake_logscale: bool = True
+
+    def __init__(
+        self,
+        in_channels: int = 80,
+        upsample_rates: List[Union[int, List[int]]] = [8, 8, 2, 2],
+        upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16, 4, 4],
+        upsample_initial_channel: int = 1536,
+        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
+        resblock_dilation_sizes: List[Union[int, List[int]]] = [
+            [1, 3, 5],
+            [1, 3, 5],
+            [1, 3, 5],
+        ],
+        activation: Literal["snake", "snakebeta"] = "snakebeta",
+        resblock_activation: Literal["snake", "snakebeta"] = "snakebeta",
+        resblock: Union[int, str] = "1",
+        use_bias_at_final: bool = False,
+        use_tanh_at_final: bool = False,
+        *args,
+        **kwargs,
+    ):
+        settings = {
+            "in_channels": in_channels,
+            "upsample_rates": upsample_rates,
+            "upsample_kernel_sizes": upsample_kernel_sizes,
+            "upsample_initial_channel": upsample_initial_channel,
+            "resblock_kernel_sizes": resblock_kernel_sizes,
+            "resblock_dilation_sizes": resblock_dilation_sizes,
+            "activation": activation,
+            "resblock_activation": resblock_activation,
+            "resblock": resblock,
+            "use_bias_at_final": use_bias_at_final,
+            "use_tanh_at_final": use_tanh_at_final,
+        }
+        super().__init__(**settings)
+
+    def post_process(self):
+        if isinstance(self.resblock, str):
+            self.resblock = 0 if self.resblock == "1" else 1
+
+
+class BemaGANv2Generator(ConvNets):
+
+    def __init__(
+        self, cfg: Union[BemaGANv2Config, Dict[str, object]] = BemaGANv2Config()
+    ):
+        super().__init__()
+        cfg = cfg if isinstance(cfg, BemaGANv2Config) else BemaGANv2Config(**cfg)
+        self.cfg = cfg
+
+        actv = get_snake(self.cfg.activation)
+
+        self.num_kernels = len(cfg.resblock_kernel_sizes)
+        self.num_upsamples = len(cfg.upsample_rates)
+
+        self.conv_pre = weight_norm(
+            nn.Conv1d(cfg.num_mels, cfg.upsample_initial_channel, 7, 1, padding=3)
+        )
+
+        # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        resblock = AMPBlock1 if cfg.resblock == 0 else AMPBlock2
+
+        # transposed conv-based upsamplers. does not apply anti-aliasing
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(cfg.upsample_rates, cfg.upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            nn.ConvTranspose1d(
+                                cfg.upsample_initial_channel // (2**i),
+                                cfg.upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+        # residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = cfg.upsample_initial_channel // (2 ** (i + 1))
+            for k, d in zip(cfg.resblock_kernel_sizes, cfg.resblock_dilation_sizes):
+                self.resblocks.append(
+                    resblock(
+                        ch,
+                        k,
+                        d,
+                        snake_logscale=cfg.snake_logscale,
+                        activation=cfg.resblock_activation,
+                    )
+                )
+
+        self.activation_post = actv(ch, alpha_logscale=cfg.snake_logscale)
+        # post conv
+
+        self.conv_post = weight_norm(
+            nn.Conv1d(ch, 1, 7, 1, padding=3, bias=self.cfg.use_bias_at_final)
+        )
+        self._use_tanh = self.cfg.use_tanh_at_final
+
+        # weight initialization
+        for i in range(len(self.ups)):
+            self.ups[i].apply(self.init_weights)
+        self.conv_post.apply(self.init_weights)
+
+    def forward(self, x: Tensor):
+        # pre conv
+        x = self.conv_pre(x)
+
+        for i in range(self.num_upsamples):
+            # upsampling
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            # AMP blocks
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+        # post conv
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        if self._use_tanh:
+            return x.tanh()
+        return x
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_file: PathLike,
+        model_config: Union[
+            BemaGANv2Config, Dict[str, Any], PathLike
+        ] = BemaGANv2Config(),
+        *,
+        remove_norms: bool = False,
+        strict: bool = True,
+        map_location: str = "cpu",
+        weights_only: bool = False,
+        **kwargs,
+    ):
+
+        is_file(model_file, validate=True)
+        model_state_dict = torch.load(
+            model_file,
+            weights_only=weights_only,
+            map_location=map_location,
+        )
+
+        if isinstance(model_config, BemaGANv2Config):
+            h = model_config
+        elif isinstance(model_config, dict):
+            h = BemaGANv2Config(**model_config)
+        elif isinstance(model_config, (str, Path, bytes)):
+            h = BemaGANv2Config(
+                **load_json(model_config, BemaGANv2Config().state_dict())
+            )
+
+        model = cls(h)
+        if remove_norms:
+            model.remove_norms()
+        try:
+            model.load_state_dict(model_state_dict, strict=strict)
+            return model
+        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(model_state_dict, strict=strict)
+        return model

lt_tensor/model_zoo/audio_models/bigvgan/__init__.py

@@ -2,9 +2,9 @@ from lt_utils.common import *
 from lt_tensor.torch_commons import *
 from lt_tensor.model_zoo.convs import ConvNets
 from lt_tensor.config_templates import ModelConfig
-from lt_tensor.model_zoo.activations import snake, alias_free
+from lt_tensor.model_zoo.activations import alias_free
 from lt_tensor.model_zoo.audio_models.resblocks import AMPBlock1, AMPBlock2, get_snake
-from lt_utils.file_ops import load_json, is_file, is_dir, is_path_valid
+from lt_utils.file_ops import load_json, is_file
 
 
 class BigVGANConfig(ModelConfig):
@@ -78,8 +78,9 @@ class BigVGAN(ConvNets):
 
     """
 
-    def __init__(self, cfg: BigVGANConfig):
+    def __init__(self, cfg: Union[BigVGANConfig, Dict[str, object]] = BigVGANConfig()):
         super().__init__()
+        cfg = cfg if isinstance(cfg, BigVGANConfig) else BigVGANConfig(**cfg)
         self.cfg = cfg
         actv = get_snake(self.cfg.activation)
 
@@ -173,46 +174,16 @@ class BigVGAN(ConvNets):
             return x.tanh()
         return x.clamp(min=-1.0, max=1.0)
 
-    def load_weights(
-        self,
-        path,
-        strict=False,
-        assign=False,
-        weights_only=False,
-        mmap=None,
-        raise_if_not_exists=False,
-        **pickle_load_args,
-    ):
-        try:
-            return super().load_weights(
-                path,
-                raise_if_not_exists,
-                strict,
-                assign,
-                weights_only,
-                mmap,
-                **pickle_load_args,
-            )
-        except RuntimeError:
-            self.remove_norms()
-            return super().load_weights(
-                path,
-                raise_if_not_exists,
-                strict,
-                assign,
-                weights_only,
-                mmap,
-                **pickle_load_args,
-            )
-
     @classmethod
     def from_pretrained(
         cls,
         model_file: PathLike,
-        model_config: Union[BigVGANConfig, Dict[str, Any]],
+        model_config: Union[
+            BigVGANConfig, Dict[str, Any], Dict[str, Any], PathLike
+        ] = BigVGANConfig(),
         *,
         remove_norms: bool = False,
-        strict: bool = False,
+        strict: bool = True,
         map_location: str = "cpu",
         weights_only: bool = False,
         **kwargs,
@@ -220,13 +191,17 @@ class BigVGAN(ConvNets):
 
         is_file(model_file, validate=True)
         model_state_dict = torch.load(
-            model_file, weights_only=weights_only, map_location=map_location
+            model_file,
+            weights_only=weights_only,
+            map_location=map_location,
         )
 
         if isinstance(model_config, BigVGANConfig):
             h = model_config
-        else:
+        elif isinstance(model_config, dict):
             h = BigVGANConfig(**model_config)
+        elif isinstance(model_config, (str, Path, bytes)):
+            h = BigVGANConfig(**load_json(model_config, BigVGANConfig().state_dict()))
 
         model = cls(h)
         if remove_norms:

lt_tensor/model_zoo/audio_models/diffwave/__init__.py

@@ -177,43 +177,44 @@ class ResidualBlock(Model):
 
 
 class DiffWave(Model):
-    def __init__(self, params: DiffWaveConfig = DiffWaveConfig()):
+    def __init__(self, cfg: Union[DiffWaveConfig, dict[str, object]] = DiffWaveConfig()):
         super().__init__()
-        self.params = params
-        self.n_hop = self.params.hop_samples
+        cfg = cfg if isinstance(cfg, DiffWaveConfig) else DiffWaveConfig(**cfg)
+        self.cfg = cfg
+        self.n_hop = self.cfg.hop_samples
         self.input_projection = ConvEXT(
             in_channels=1,
-            out_channels=params.residual_channels,
+            out_channels=cfg.residual_channels,
             kernel_size=1,
-            apply_norm=self.params.apply_norm,
+            apply_norm=self.cfg.apply_norm,
             activation_out=nn.LeakyReLU(0.1),
         )
-        self.diffusion_embedding = DiffusionEmbedding(len(params.noise_schedule))
+        self.diffusion_embedding = DiffusionEmbedding(len(cfg.noise_schedule))
         self.spectrogram_upsampler = (
-            SpectrogramUpsampler() if not self.params.unconditional else None
+            SpectrogramUpsampler() if not self.cfg.unconditional else None
         )
 
         self.residual_layers = nn.ModuleList(
             [
                 ResidualBlock(
-                    params.n_mels,
-                    params.residual_channels,
-                    2 ** (i % params.dilation_cycle_length),
-                    uncond=params.unconditional,
-                    apply_norm=self.params.apply_norm_resblock,
+                    cfg.n_mels,
+                    cfg.residual_channels,
+                    2 ** (i % cfg.dilation_cycle_length),
+                    uncond=cfg.unconditional,
+                    apply_norm=self.cfg.apply_norm_resblock,
                 )
-                for i in range(params.residual_layers)
+                for i in range(cfg.residual_layers)
             ]
         )
         self.skip_projection = ConvEXT(
-            in_channels=params.residual_channels,
-            out_channels=params.residual_channels,
+            in_channels=cfg.residual_channels,
+            out_channels=cfg.residual_channels,
             kernel_size=1,
-            apply_norm=self.params.apply_norm,
+            apply_norm=self.cfg.apply_norm,
             activation_out=nn.LeakyReLU(0.1),
         )
         self.output_projection = ConvEXT(
-            params.residual_channels, 1, 1, apply_norm=self.params.apply_norm, init_weights=True,
+            cfg.residual_channels, 1, 1, apply_norm=self.cfg.apply_norm, init_weights=True,
         )
         self.activation = nn.LeakyReLU(0.1)
         self._res_d = sqrt(len(self.residual_layers))
@@ -224,7 +225,7 @@ class DiffWave(Model):
         diffusion_step: Tensor,
         spectrogram: Optional[Tensor] = None,
     ):
-        if not self.params.unconditional:
+        if not self.cfg.unconditional:
             assert spectrogram is not None
         if audio.ndim < 3:
             if audio.ndim == 2:
@@ -234,7 +235,7 @@ class DiffWave(Model):
 
         x = self.input_projection(audio)
         diffusion_step = self.diffusion_embedding(diffusion_step)
-        if not self.params.unconditional:  # use conditional model
+        if not self.cfg.unconditional:  # use conditional model
             spectrogram = self.spectrogram_upsampler(spectrogram)
 
         skip = torch.zeros_like(x, device=x.device)

lt_tensor/model_zoo/audio_models/hifigan/__init__.py

@@ -5,7 +5,7 @@ from lt_utils.common import *
 from lt_tensor.torch_commons import *
 from lt_tensor.model_zoo.convs import ConvNets
 from lt_tensor.config_templates import ModelConfig
-from lt_utils.file_ops import is_file
+from lt_utils.file_ops import is_file, load_json
 from lt_tensor.model_zoo.audio_models.resblocks import ResBlock1, ResBlock2
 
 
@@ -16,11 +16,15 @@ def get_padding(kernel_size, dilation=1):
 class HifiganConfig(ModelConfig):
     # Training params
     in_channels: int = 80
-    upsample_rates: List[Union[int, List[int]]] = [8,8,2,2]
-    upsample_kernel_sizes: List[Union[int, List[int]]] = [16,16,4,4]
+    upsample_rates: List[Union[int, List[int]]] = [8, 8, 2, 2]
+    upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16, 4, 4]
     upsample_initial_channel: int = 512
     resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11]
-    resblock_dilation_sizes: List[Union[int, List[int]]] = [[1,3,5], [1,3,5], [1,3,5]]
+    resblock_dilation_sizes: List[Union[int, List[int]]] = [
+        [1, 3, 5],
+        [1, 3, 5],
+        [1, 3, 5],
+    ]
 
     activation: nn.Module = nn.LeakyReLU(0.1)
     resblock_activation: nn.Module = nn.LeakyReLU(0.1)
@@ -29,10 +33,10 @@ class HifiganConfig(ModelConfig):
     def __init__(
         self,
         in_channels: int = 80,
-        upsample_rates: List[Union[int, List[int]]] = [8,8,2,2],
-        upsample_kernel_sizes: List[Union[int, List[int]]] = [16,16,4,4],
+        upsample_rates: List[Union[int, List[int]]] = [8, 8, 2, 2],
+        upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16, 4, 4],
         upsample_initial_channel: int = 512,
-        resblock_kernel_sizes: List[Union[int, List[int]]] = [3,7,11],
+        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
         resblock_dilation_sizes: List[Union[int, List[int]]] = [
             [1, 3, 5],
             [1, 3, 5],
@@ -63,9 +67,11 @@ class HifiganConfig(ModelConfig):
 
 
 class HifiganGenerator(ConvNets):
-    def __init__(self, cfg: HifiganConfig = HifiganConfig()):
+    def __init__(self, cfg: Union[HifiganConfig, Dict[str, object]] = HifiganConfig()):
         super().__init__()
+        cfg = cfg if isinstance(cfg, HifiganConfig) else HifiganConfig(**cfg)
         self.cfg = cfg
+
         self.num_kernels = len(cfg.resblock_kernel_sizes)
         self.num_upsamples = len(cfg.upsample_rates)
         self.conv_pre = weight_norm(
@@ -115,46 +121,16 @@ class HifiganGenerator(ConvNets):
         x = self.conv_post(self.activation(x))
         return x.tanh()
 
-    def load_weights(
-        self,
-        path,
-        strict=False,
-        assign=False,
-        weights_only=False,
-        mmap=None,
-        raise_if_not_exists=False,
-        **pickle_load_args,
-    ):
-        try:
-            return super().load_weights(
-                path,
-                raise_if_not_exists,
-                strict,
-                assign,
-                weights_only,
-                mmap,
-                **pickle_load_args,
-            )
-        except RuntimeError:
-            self.remove_norms()
-            return super().load_weights(
-                path,
-                raise_if_not_exists,
-                strict,
-                assign,
-                weights_only,
-                mmap,
-                **pickle_load_args,
-            )
-
     @classmethod
     def from_pretrained(
         cls,
         model_file: PathLike,
-        model_config: Union[HifiganConfig, Dict[str, Any]],
+        model_config: Union[
+            HifiganConfig, Dict[str, Any], Dict[str, Any], PathLike
+        ] = HifiganConfig(),
         *,
         remove_norms: bool = False,
-        strict: bool = False,
+        strict: bool = True,
         map_location: str = "cpu",
         weights_only: bool = False,
         **kwargs,
@@ -162,13 +138,17 @@ class HifiganGenerator(ConvNets):
 
         is_file(model_file, validate=True)
         model_state_dict = torch.load(
-            model_file, weights_only=weights_only, map_location=map_location
+            model_file,
+            weights_only=weights_only,
+            map_location=map_location,
         )
 
         if isinstance(model_config, HifiganConfig):
             h = model_config
-        else:
+        elif isinstance(model_config, dict):
             h = HifiganConfig(**model_config)
+        elif isinstance(model_config, (str, Path, bytes)):
+            h = HifiganConfig(**load_json(model_config, HifiganConfig().state_dict()))
 
         model = cls(h)
         if remove_norms:

lt_tensor/model_zoo/audio_models/istft/__init__.py

@@ -3,7 +3,7 @@ from lt_utils.common import *
 from lt_tensor.torch_commons import *
 from lt_tensor.model_zoo.convs import ConvNets
 from lt_tensor.config_templates import ModelConfig
-from lt_utils.file_ops import is_file
+from lt_utils.file_ops import is_file, load_json
 from lt_tensor.model_zoo.audio_models.resblocks import ResBlock1, ResBlock2
 
 
@@ -67,8 +67,11 @@ class iSTFTNetConfig(ModelConfig):
 
 
 class iSTFTNetGenerator(ConvNets):
-    def __init__(self, cfg: iSTFTNetConfig = iSTFTNetConfig()):
+    def __init__(
+        self, cfg: Union[iSTFTNetConfig, Dict[str, object]] = iSTFTNetConfig()
+    ):
         super().__init__()
+        cfg = cfg if isinstance(cfg, iSTFTNetConfig) else iSTFTNetConfig(**cfg)
         self.cfg = cfg
         self.num_kernels = len(cfg.resblock_kernel_sizes)
         self.num_upsamples = len(cfg.upsample_rates)
@@ -146,46 +149,16 @@ class iSTFTNetGenerator(ConvNets):
 
         return spec, phase
 
-    def load_weights(
-        self,
-        path,
-        strict=False,
-        assign=False,
-        weights_only=False,
-        mmap=None,
-        raise_if_not_exists=False,
-        **pickle_load_args,
-    ):
-        try:
-            return super().load_weights(
-                path,
-                raise_if_not_exists,
-                strict,
-                assign,
-                weights_only,
-                mmap,
-                **pickle_load_args,
-            )
-        except RuntimeError:
-            self.remove_norms()
-            return super().load_weights(
-                path,
-                raise_if_not_exists,
-                strict,
-                assign,
-                weights_only,
-                mmap,
-                **pickle_load_args,
-            )
-
     @classmethod
     def from_pretrained(
         cls,
         model_file: PathLike,
-        model_config: Union[iSTFTNetConfig, Dict[str, Any]],
+        model_config: Union[
+            iSTFTNetConfig, Dict[str, Any], Dict[str, Any], PathLike
+        ] = iSTFTNetConfig(),
         *,
         remove_norms: bool = False,
-        strict: bool = False,
+        strict: bool = True,
         map_location: str = "cpu",
         weights_only: bool = False,
         **kwargs,
@@ -193,14 +166,17 @@ class iSTFTNetGenerator(ConvNets):
 
         is_file(model_file, validate=True)
         model_state_dict = torch.load(
-            model_file, weights_only=weights_only, map_location=map_location
+            model_file,
+            weights_only=weights_only,
+            map_location=map_location,
         )
 
         if isinstance(model_config, iSTFTNetConfig):
             h = model_config
-        else:
+        elif isinstance(model_config, dict):
             h = iSTFTNetConfig(**model_config)
-
+        elif isinstance(model_config, (str, Path, bytes)):
+            h = iSTFTNetConfig(**load_json(model_config, iSTFTNetConfig().state_dict()))
         model = cls(h)
         if remove_norms:
             model.remove_norms()

lt_tensor/model_zoo/convs.py

@@ -1,11 +1,7 @@
 __all__ = ["ConvNets", "ConvEXT"]
-import math
 from lt_utils.common import *
-import torch.nn.functional as F
 from lt_tensor.torch_commons import *
 from lt_tensor.model_base import Model
-from lt_tensor.misc_utils import log_tensor
-from lt_tensor.model_zoo.fusion import AdaFusion1D, AdaIN1D
 from lt_utils.misc_utils import default
 
 
@@ -52,6 +48,41 @@ class ConvNets(Model):
         if "Conv" in m.__class__.__name__:
             m.weight.data.normal_(mean, std)
 
+    def load_weights(
+        self,
+        path,
+        strict=False,
+        assign=False,
+        weights_only=False,
+        mmap=None,
+        raise_if_not_exists=False,
+        **pickle_load_args,
+    ):
+        try:
+            return super().load_weights(
+                path,
+                raise_if_not_exists,
+                strict,
+                assign,
+                weights_only,
+                mmap,
+                **pickle_load_args,
+            )
+        except RuntimeError as e:
+            try:
+                self.remove_norms()
+                return super().load_weights(
+                    path,
+                    raise_if_not_exists,
+                    strict,
+                    assign,
+                    weights_only,
+                    mmap,
+                    **pickle_load_args,
+                )
+            except:
+                raise e
+
 
 class ConvEXT(ConvNets):
     def __init__(