lt-tensor 0.0.1a29__py3-none-any.whl → 0.0.1a31__py3-none-any.whl

lt_tensor/model_base.py

@@ -70,6 +70,7 @@ class LossTracker:
 
 class _Devices_Base(nn.Module):
     _device: torch.device = ROOT_DEVICE
+    _setting_device: bool = False
 
     @property
     def device(self):
@@ -136,25 +137,35 @@ class _Devices_Base(nn.Module):
                 f"Item '{module_or_name}' is not a valid module, parameter or tensor."
             )
 
-    def _apply_device(self):
-        """This may be seem as overkill, but its necessary"""
-        for modules in self.modules():
-            try:
-                modules.to(self.device)
-            except:
-                pass
-
-        for buffer in self.buffers():
-            try:
-                buffer.to(self.device)
-            except:
-                pass
-
-        for tensor in self.parameters():
-            try:
-                tensor.to(self.device)
-            except:
-                pass
+    def apply_device(self):
+        """Helps to apply devices towards all the internal components"""
+        if self._setting_device:
+            return
+
+        self._setting_device = True
+
+        try:
+            for modules in self.modules():
+                try:
+                    modules.to(self.device)
+                except:
+                    pass
+
+            for buffer in self.buffers():
+                try:
+                    buffer.to(self.device)
+                except:
+                    pass
+
+            for tensor in self.parameters():
+                try:
+                    tensor.to(self.device)
+                except:
+                    pass
+        except:
+            pass
+        finally:
+            self._setting_device = False
 
     def _to_dvc(
         self, device_name: str, device_id: Optional[Union[int, torch.device]] = None
@@ -166,7 +177,11 @@ class _Devices_Base(nn.Module):
             elif hasattr(device_id, "index"):
                 device += ":" + str(device_id.index)
         self.device = device
-        self._apply_device()
+        if not self._setting_device:
+            self.apply_device()
+
+    def _to(self, *args, **kwargs):
+        self.to(*args, _internal=True, **kwargs)
 
     def to(self, *args, **kwargs):
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
@@ -378,6 +393,7 @@ class Model(_Devices_Base, ABC):
         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)
+        state_dict.pop("_setting_device", None)
         torch.save(obj=state_dict, f=str(model_dir))
 
     def save_lora(

lt_tensor/model_zoo/losses/discriminators.py

@@ -20,6 +20,7 @@ def get_padding(kernel_size, dilation=1):
 
 class MultiDiscriminatorWrapper(ConvNets):
     """Base for all multi-steps type of discriminators"""
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.leaky_relu = nn.LeakyReLU(kwargs.get("negative_slope", 0.1))
@@ -207,7 +208,7 @@ class MultiPeriodDiscriminator(MultiDiscriminatorWrapper):
         return y_d_rs, y_d_gs, fmap_rs, fmap_gs
 
 
-class EnvelopeExtractor(nn.Module):
+class EnvelopeExtractor(Model):
     """Extracts the amplitude envelope of the audio signal."""
 
     def __init__(self, kernel_size=101):
@@ -216,7 +217,7 @@ class EnvelopeExtractor(nn.Module):
         self.kernel_size = kernel_size
         self.register_buffer("kernel", torch.ones(1, 1, kernel_size) / kernel_size)
 
-    def forward(self, x):
+    def forward(self, x: Tensor):
         # x: (B, 1, T) -> abs(x)
         envelope = torch.abs(x)
         # Apply low-pass smoothing (via conv1d)
@@ -274,7 +275,6 @@ class MultiEnvelopeDiscriminator(MultiDiscriminatorWrapper):
     def forward(self, y, y_hat):
         y_d_rs, y_d_gs = [], []
         fmap_rs, fmap_gs = [], []
-
         for i, d in enumerate(self.discriminators):
             if i != 0:
                 y = self.meanpools[i - 1](y)
@@ -555,7 +555,19 @@ class MultiResolutionDiscriminator(MultiDiscriminatorWrapper):
 
 
 class MultiDiscriminatorStep(Model):
-    def __init__(self, list_discriminator: List[MultiDiscriminatorWrapper]):
+    def __init__(
+        self, list_discriminator: List[MultiDiscriminatorWrapper]
+    ):
+        """Setup example:
+        model_d = MultiDiscriminatorStep(
+            [
+                MultiEnvelopeDiscriminator(),
+                MultiBandDiscriminator(),
+                MultiResolutionDiscriminator(),
+                MultiPeriodDiscriminator(0.5),
+            ]
+        )
+        """
         super().__init__()
         self.disc: Sequence[MultiDiscriminatorWrapper] = nn.ModuleList(
             list_discriminator

{lt_tensor-0.0.1a29.dist-info → lt_tensor-0.0.1a31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a29
+Version: 0.0.1a31
 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.1a29.dist-info → lt_tensor-0.0.1a31.dist-info}/RECORD

@@ -4,7 +4,7 @@ lt_tensor/losses.py,sha256=zvkCOnE5XpF3v6ymivRIdqPTsMM5zc94ZMom7YDi3zM,4946
 lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
 lt_tensor/math_ops.py,sha256=TkD4WQG42KsQ9Fg7FXOjf8f-ixtW0apf2XjaooecVx4,2257
 lt_tensor/misc_utils.py,sha256=N2r3UmxC4RM2BZBQhpjDZ_BKLrzsyIlKzopTzJbnjFU,28962
-lt_tensor/model_base.py,sha256=8qcFXe0_y8f1_tAwt18gwQjyyapbnVEKcjCMrKnQatw,17614
+lt_tensor/model_base.py,sha256=5T4dbAh4MXbQmPRpihGtMYwTY8sJTQOhY6An3VboM58,18086
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
 lt_tensor/noise_tools.py,sha256=wFeAsHhLhSlEc5XU5LbFKaXoHeVxrWjiMeljjGdIKyM,11363
 lt_tensor/torch_commons.py,sha256=8l0bxmrAzwvyqjivCIVISXlbvKarlg4DdE0BOGSnMuQ,812
@@ -27,11 +27,11 @@ lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=PDuDYN1omD1RoAXcmxH
 lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=7GJqKLw7-juXpfp5IFzjASLut0uouDhjZ1CQknf3H68,16533
 lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=ltIuD9t1gmS3bTmCqZIwJHKrhC6DYya3OaXlskWX9kw,17606
 lt_tensor/model_zoo/losses/__init__.py,sha256=B9RAUxBiOZwooztnij1oLeRwZ7_MjnN3mPoum7saD6s,59
-lt_tensor/model_zoo/losses/discriminators.py,sha256=yYh7HzRTUtr0RVTG7cWpcYsJZsRCz6yzg6Loq8FtyOk,20405
+lt_tensor/model_zoo/losses/discriminators.py,sha256=ZA7Qqrhe8kELrI1-IITadGSl8JCgpgPKFCW6qvSOk1E,20724
 lt_tensor/processors/__init__.py,sha256=Pvxhh0KR65zLCgUd53_k5Z0y5JWWcO0ZBXFK9rv0o5w,109
 lt_tensor/processors/audio.py,sha256=1JuxxexfUsXkLjVjWUk-oTRU-QNnCCwvKX3eP0m7LGE,16452
-lt_tensor-0.0.1a29.dist-info/licenses/LICENSE,sha256=tQHc38scHOba4kDBNG4U0U6PpObaloiZG-FvKSgv2b0,11336
-lt_tensor-0.0.1a29.dist-info/METADATA,sha256=F03dNMnEydcKjjZF3IntNaIj34FwLdoy-L0pBB_yz0E,1062
-lt_tensor-0.0.1a29.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a29.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a29.dist-info/RECORD,,
+lt_tensor-0.0.1a31.dist-info/licenses/LICENSE,sha256=tQHc38scHOba4kDBNG4U0U6PpObaloiZG-FvKSgv2b0,11336
+lt_tensor-0.0.1a31.dist-info/METADATA,sha256=qhs6RI_KE0LNjKiD7kky9d_4wcolq6PHyInLPVePzNo,1062
+lt_tensor-0.0.1a31.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a31.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a31.dist-info/RECORD,,