lt-tensor 0.0.1a16__py3-none-any.whl → 0.0.1a17__py3-none-any.whl

lt_tensor/config_templates.py

@@ -7,16 +7,14 @@ from lt_tensor.misc_utils import updateDict
 
 class ModelConfig(ABC, OrderedDict):
     _default_settings: Dict[str, Any] = {}
-    _forbidden_list: List[str] = [
-        "_settings",
-    ]
+    _forbidden_list: List[str] = ["_default_settings", "_forbidden_list" "path_name"]
 
     def __init__(
         self,
-        settings: Dict[str, Any] = None,
+        settings: Dict[str, Any] = {},
         path_name: Optional[Union[str, PathLike]] = None,
     ):
-        assert is_dict(settings)
+        assert is_dict(settings, False)
         self._default_settings = settings
         if path_name is not None and is_pathlike(path_name):
             if not str(path_name).endswith(".json"):
@@ -37,31 +35,41 @@ class ModelConfig(ABC, OrderedDict):
                 self.path_name += ".json"
 
     def reset_settings(self):
+        dk_keys = self.__dict__.keys()
         for s_name, setting in self._default_settings.items():
-            if s_name in self._forbidden_list:
+            if s_name in self._forbidden_list or s_name not in dk_keys:
                 continue
             updateDict(self, {s_name: setting})
 
     def save_config(
         self,
-        path_name: Union[PathLike, str],
+        path_name: Optional[Union[PathLike, str]] = None,
     ):
-        assert is_pathlike(
-            path_name, True
-        ), f"path_name should be a non-empty string or pathlike object! received instead: {path_name}"
-        self._setup_path_name(path_name)
-        base = {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
-        save_json(self.path_name, base, indent=2)
+        if not is_pathlike(path_name, True):
+            assert (
+                path_name is None
+            ), f"path_name should be a non-empty string or pathlike object! received instead: {path_name}."
+            path_name = self.path_name
+        else:
+            self._setup_path_name(path_name)
 
-    def to_dict(self):
-        return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
+        base = self.get_state_dict()
+        save_json(self.path_name, base, indent=2)
 
     def set_value(self, var_name: str, value: str) -> None:
+        assert var_name in self.__dict__, "Value not registered!"
+        assert var_name not in self._forbidden_list, "Not allowed!"
         updateDict(self, {var_name: value})
 
     def get_value(self, var_name: str) -> Any:
         return self.__dict__.get(var_name)
 
+    def __getattribute__(self, name):
+        return self.__dict__.get(name)
+
+    def get_state_dict(self):
+        return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
+
     @classmethod
     def from_dict(
         cls, dictionary: Dict[str, Any], path: Optional[Union[str, PathLike]] = None

lt_tensor/model_base.py

@@ -70,16 +70,6 @@ class LossTracker:
 
 class _Devices_Base(nn.Module):
     _device: torch.device = ROOT_DEVICE
-    _autocast: bool = False
-    _loss_history: LossTracker = LossTracker(100_000)
-
-    @property
-    def autocast(self):
-        return self._autocast
-
-    @autocast.setter
-    def autocast(self, value: bool):
-        self._autocast = value
 
     @property
     def device(self):
@@ -90,6 +80,30 @@ class _Devices_Base(nn.Module):
         assert isinstance(device, (str, torch.device))
         self._device = torch.device(device) if isinstance(device, str) else device
 
+    def _apply_device(self):
+        """Add here components that are needed to have device applied to them,
+        that usually the '.to()' function fails to apply
+
+        example:
+        ```
+        def _apply_device_to(self):
+            self.my_tensor = self.my_tensor.to(device=self.device)
+        ```
+        """
+        pass
+
+    def _to_dvc(
+        self, device_name: str, device_id: Optional[Union[int, torch.device]] = None
+    ):
+        device = device_name
+        if device_id is not None:
+            if isinstance(device_id, Number):
+                device += ":" + str(int(device_id))
+            elif hasattr(device_id, "index"):
+                device += ":" + str(device_id.index)
+        self.device = device
+        self._apply_device()
+
     def to(self, *args, **kwargs):
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(
             *args, **kwargs
@@ -133,20 +147,9 @@ class _Devices_Base(nn.Module):
                     raise
 
         self._apply(convert)
-        self.device = device
+        self._to_dvc(device)
         return self
 
-    def _to_dvc(
-        self, device_name: str, device_id: Optional[Union[int, torch.device]] = None
-    ):
-        device = device_name
-        if device_id is not None:
-            if isinstance(device_id, Number):
-                device += ":" + str(int(device_id))
-            elif hasattr(device_id, "index"):
-                device += ":" + str(device_id.index)
-        self.device = device
-
     def ipu(self, device: Optional[Union[int, torch.device]] = None) -> T:
         super().ipu(device)
         self._to_dvc("ipu", device)
@@ -178,11 +181,12 @@ class Model(_Devices_Base, ABC):
     This makes it easier to assign a device and retrieves it later
     """
 
+    _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
-    _is_gradient_freezable: bool = (
+    _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
@@ -193,18 +197,15 @@ class Model(_Devices_Base, ABC):
 
     # dont save list:
     _dont_save_items: List[str] = []
+    _loss_history: LossTracker = LossTracker(20_000)
 
-    def _apply_device_to(self):
-        """Add here components that are needed to have device applied to them,
-        that usually the '.to()' function fails to apply
+    @property
+    def autocast(self):
+        return self._autocast
 
-        example:
-        ```
-        def _apply_device_to(self):
-            self.my_tensor = self.my_tensor.to(device=self.device)
-        ```
-        """
-        pass
+    @autocast.setter
+    def autocast(self, value: bool):
+        self._autocast = value
 
     def freeze_all(self, exclude: Optional[List[str]] = None):
         no_exclusions = not exclude
@@ -251,7 +252,7 @@ class Model(_Devices_Base, ABC):
     def change_frozen_state(self, freeze: bool, module: nn.Module):
         try:
             if isinstance(module, Model):
-                if module._is_gradient_freezable:
+                if module._can_be_frozen:
                     if freeze:
                         return module.freeze_all()
                     return module.unfreeze_all()
@@ -496,10 +497,7 @@ class Model(_Devices_Base, ABC):
         return self(*inputs, **kwargs)
 
     def __call__(self, *args, **kwds) -> POSSIBLE_OUTPUT_TYPES:
-        if self.autocast and not self.training:
-            with torch.autocast(device_type=self.device.type):
-                return super().__call__(*args, **kwds)
-        else:
+        with torch.autocast(device_type=self.device.type, enabled=self.autocast):
             return super().__call__(*args, **kwds)
 
     @abstractmethod
@@ -541,52 +539,3 @@ class Model(_Devices_Base, ABC):
         if quantity > 0:
             t_list = t_list[-quantity:]
         return sum(t_list) / len(t_list)
-
-    def freeze_unfreeze_loss(
-        self,
-        losses: Optional[Union[float, List[float]]] = None,
-        trigger_loss: Union[float, bool] = 0.1,
-        excluded_modules: Optional[List[str]] = None,
-        max_items: int = 1000,
-        loss_name: str = "train",
-    ):
-        """If a certain threshold is reached the weights will freeze or unfreeze the modules.
-        the biggest use-case for this function is when training GANs where the balance
-        from the discriminator and generator must be kept.
-
-        Args:
-            losses (Union[float, List[float]], Optional): The loss value or a list of losses that will be used to determine if it has reached or not the threshold. Defaults to None.
-            trigger_loss (float, bool, optional): The value where the weights will be either freeze or unfreeze. If set to a boolean it will freeze or unfreeze immediately according to the value (True = Freeze, False = Unfreeze). Defaults to 0.1.
-            excluded_modules (list[str], optional): The list of modules (names) that is not to be changed by either freezing nor unfreezing. Defaults to None.
-            max_items (float, optional): The number of previous losses to be locked behind to calculate the current average. Default to 1000.
-            loss_name (str, optional): Responsible to define with key to recover the loss.
-        returns:
-            bool: True when its frozen and false when its trainable.
-        """
-        if losses is not None:
-            self.add_loss(losses, "train")
-
-        if isinstance(trigger_loss, bool):
-            if trigger_loss:
-                if self._is_unfrozen:
-                    self.freeze_all(excluded_modules)
-                    self._is_unfrozen = False
-                return True
-            # else
-            if not self._is_unfrozen:
-                self.unfreeze_all(excluded_modules)
-                self._is_unfrozen = True
-            return False
-
-        value = self.get_loss_avg(loss_name, max_items)
-
-        if value <= trigger_loss:
-            if self._is_unfrozen:
-                self.freeze_all(excluded_modules)
-                self._is_unfrozen = False
-            return True
-        else:
-            if not self._is_unfrozen:
-                self.unfreeze_all(excluded_modules)
-                self._is_unfrozen = True
-            return False

lt_tensor/model_zoo/audio_models/__init__.py

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

lt_tensor/model_zoo/audio_models/diffwave/__init__.py

@@ -1,3 +1,219 @@
 __all__ = ["DiffWave", "SpectrogramUpsampler", "DiffusionEmbedding"]
 
-from .model import DiffWave, SpectrogramUpsampler, DiffusionEmbedding
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from lt_tensor.config_templates import ModelConfig
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from math import sqrt
+from lt_utils.common import *
+
+F.t
+
+
+class DiffWaveConfig(ModelConfig):
+    # Training params
+    batch_size = 16
+    learning_rate = 2e-4
+    max_grad_norm = None
+    # Data params
+    sample_rate = 24000
+    n_mels = 80
+    n_fft = 1024
+    hop_samples = 256
+    # Model params
+    residual_layers = 30
+    residual_channels = 64
+    dilation_cycle_length = 10
+    unconditional = False
+    noise_schedule: list[int] = np.linspace(1e-4, 0.05, 50).tolist()
+    # settings for auto-fixes
+    interpolate = False
+    interpolation_mode: Literal[
+        "nearest", "linear", "bilinear", "bicubic", "trilinear", "area", "nearest-exact"
+    ] = "nearest"
+
+    def __init__(
+        self,
+        settings: Dict[str, Any] = {},
+        path_name: Optional[Union[str, PathLike]] = None,
+    ):
+        self._forbidden_list.extend()
+        super().__init__(settings, path_name)
+
+
+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__()
+        self.register_buffer(
+            "embedding", self._build_embedding(max_steps), persistent=False
+        )
+        self.projection1 = nn.Linear(128, 512)
+        self.projection2 = nn.Linear(512, 512)
+        self.activation = nn.SiLU()
+
+    def forward(self, diffusion_step):
+        if diffusion_step.dtype in [torch.int32, torch.int64]:
+            x = self.embedding[diffusion_step]
+        else:
+            x = self._lerp_embedding(diffusion_step)
+        x = self.projection1(x)
+        x = self.activation(x)
+        x = self.projection2(x)
+        x = self.activation(x)
+        return x
+
+    def _lerp_embedding(self, t):
+        low_idx = torch.floor(t).long()
+        high_idx = torch.ceil(t).long()
+        low = self.embedding[low_idx]
+        high = self.embedding[high_idx]
+        return low + (high - low) * (t - low_idx)
+
+    def _build_embedding(self, max_steps):
+        steps = torch.arange(max_steps).unsqueeze(1)  # [T,1]
+        dims = torch.arange(64).unsqueeze(0)  # [1,64]
+        table = steps * 10.0 ** (dims * 4.0 / 63.0)  # [T,64]
+        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)
+        return table
+
+
+class SpectrogramUpsampler(Model):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.ConvTranspose2d(1, 1, [3, 32], stride=[1, 16], padding=[1, 8])
+        self.conv2 = nn.ConvTranspose2d(1, 1, [3, 32], stride=[1, 16], padding=[1, 8])
+        self.activation = nn.LeakyReLU(0.4)
+
+    def forward(self, x):
+        x = torch.unsqueeze(x, 1)
+        x = self.activation(self.conv1(x))
+        x = self.activation(self.conv2(x))
+        x = torch.squeeze(x, 1)
+        return x
+
+
+class ResidualBlock(Model):
+    def __init__(self, n_mels, residual_channels, dilation, uncond=False):
+        """
+        :param n_mels: inplanes of conv1x1 for spectrogram conditional
+        :param residual_channels: audio conv
+        :param dilation: audio conv dilation
+        :param uncond: disable spectrogram conditional
+        """
+        super().__init__()
+        self.dilated_conv = Conv1d(
+            residual_channels,
+            2 * residual_channels,
+            3,
+            padding=dilation,
+            dilation=dilation,
+        )
+        self.diffusion_projection = nn.Linear(512, residual_channels)
+        if not uncond:  # conditional model
+            self.conditioner_projection = Conv1d(n_mels, 2 * residual_channels, 1)
+        else:  # unconditional model
+            self.conditioner_projection = None
+
+        self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1)
+
+    def forward(
+        self,
+        x: Tensor,
+        diffusion_step: Tensor,
+        conditioner: Optional[Tensor] = None,
+    ):
+
+        diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
+        y = x + diffusion_step
+        if (
+            conditioner is None or self.conditioner_projection is None
+        ):  # using a unconditional model
+            y = self.dilated_conv(y)
+        else:
+            conditioner = self.conditioner_projection(conditioner)
+            y = self.dilated_conv(y) + conditioner
+
+        gate, filter = torch.chunk(y, 2, dim=1)
+        y = torch.sigmoid(gate) * torch.tanh(filter)
+
+        y = self.output_projection(y)
+        residual, skip = torch.chunk(y, 2, dim=1)
+        return (x + residual) / sqrt(2.0), skip
+
+
+class DiffWave(Model):
+    def __init__(self, params: DiffWaveConfig = DiffWaveConfig()):
+        super().__init__()
+        self.params = params
+        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.diffusion_embedding = DiffusionEmbedding(len(params.noise_schedule))
+        if self.params.unconditional:  # use unconditional model
+            self.spectrogram_upsample = None
+        else:
+            self.spectrogram_upsample = SpectrogramUpsampler()
+
+        self.residual_layers = nn.ModuleList(
+            [
+                ResidualBlock(
+                    params.n_mels,
+                    params.residual_channels,
+                    2 ** (i % params.dilation_cycle_length),
+                    uncond=params.unconditional,
+                )
+                for i in range(params.residual_layers)
+            ]
+        )
+        self.skip_projection = Conv1d(
+            params.residual_channels, params.residual_channels, 1
+        )
+        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)
+
+    def forward(
+        self,
+        audio: Tensor,
+        diffusion_step: Tensor,
+        spectrogram: Optional[Tensor] = None,
+    ):
+        T = x.shape[-1]
+        if x.ndim == 2:
+            x = audio.unsqueeze(1)
+        x = self.activation(self.input_projection(x))
+
+        diffusion_step = self.diffusion_embedding(diffusion_step)
+        if spectrogram is not None and self.spectrogram_upsample is not None:
+            if self.auto_interpolate:
+                # a little heavy, but helps a lot to fix mismatched shapes,
+                # not always recommended due to data loss
+                spectrogram = F.interpolate(
+                    input=spectrogram,
+                    size=int(T * self.n_hop),
+                    mode=self.interpolate_mode,
+                )
+            spectrogram = self.spectrogram_upsample(spectrogram)
+
+        skip = None
+        for i, layer in enumerate(self.residual_layers):
+            x, skip_connection = layer(x, diffusion_step, spectrogram)
+            if i == 0:
+                skip = skip_connection
+            else:
+                skip = skip_connection + skip
+        x = skip / self.r_sqrt
+        x = self.activation(self.skip_projection(x))
+        x = self.output_projection(x)
+        return x

{lt_tensor-0.0.1a16.dist-info → lt_tensor-0.0.1a17.dist-info}/METADATA

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

@@ -1,10 +1,10 @@
 lt_tensor/__init__.py,sha256=XxNCGcVL-haJyMpifr-GRaamo32R6jmqe3iOuS4ecfs,469
-lt_tensor/config_templates.py,sha256=FRN4-i1amoqMh_wyp4gNsw61ABWTIhGC62Uc3l3SNss,3515
+lt_tensor/config_templates.py,sha256=xWZhktYVlkwvJVreqyACpWo-lJ5htG9vTZyqZ6OexzA,3899
 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=S57M5XuGsIuaOKnEGZJsY3B2dTmggpdhsqQr51CQsYo,28754
-lt_tensor/model_base.py,sha256=qqqIVpYz6nv01MnZuuAj1dxq4_NN-zSivP1GaegA9TI,21597
+lt_tensor/model_base.py,sha256=J-f-iQ9qGyYD4NkLljyAEkwtHKKbUKIrIpunMiSmh90,19155
 lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,2239
 lt_tensor/noise_tools.py,sha256=wFeAsHhLhSlEc5XU5LbFKaXoHeVxrWjiMeljjGdIKyM,11363
 lt_tensor/torch_commons.py,sha256=fntsEU8lhBQo0ebonI1iXBkMbWMN3HpBsG13EWlP5s8,718
@@ -18,15 +18,14 @@ lt_tensor/model_zoo/fusion.py,sha256=usC1bcjQRNivDc8xzkIS5T1glm78OLcs2V_tPqfp-eI
 lt_tensor/model_zoo/pos_encoder.py,sha256=3d1EYLinCU9UAy-WuEWeYMGhMqaGknCiQ5qEmhw_UYM,4487
 lt_tensor/model_zoo/residual.py,sha256=i5V4ju7DB3WesKBVm6KH_LyPoKGDUOyo2Usfs-PyP58,9394
 lt_tensor/model_zoo/transformer.py,sha256=HUFoFFh7EQJErxdd9XIxhssdjvNVx2tNGDJOTUfwG2A,4301
-lt_tensor/model_zoo/audio_models/__init__.py,sha256=CLoLqvbA_ltqm3KOg5AH3A0co0HtsLfFPUBsxxLSCgI,39
-lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=aFSmr8PYpmOfbe15lhNoj-ZzP5ChrZcikovKLZKg7nw,140
-lt_tensor/model_zoo/audio_models/diffwave/model.py,sha256=kHo76bxLJtTBn1m0gq5KKrUsjm9ASsCCwf8MvWaB1R8,6901
+lt_tensor/model_zoo/audio_models/__init__.py,sha256=CmoakfBLoxqtJuYc1NYrB_z0x1kS2WQNaYQRmCaC5ko,40
+lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=8DbKJpQ44s9iPlajfs7_A2N1diYGXzkhet_wS4hX6mU,7421
 lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=BOBZSK2HFOdMcFyjrzwZi_TeAtBGIcpb8pQxiGlwLEE,12302
 lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=o7Ie1qI22u_g9t1252PX4vl4uF6JHynAJryuz2lAZE0,12920
 lt_tensor/processors/__init__.py,sha256=4b9MxAJolXiJfSm20ZEspQTDm1tgLazwlPWA_jB1yLM,63
 lt_tensor/processors/audio.py,sha256=SMqNSl4Den-x1awTCQ8-TcR-0jPiv5lDaUpU93SRRaw,14749
-lt_tensor-0.0.1a16.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
-lt_tensor-0.0.1a16.dist-info/METADATA,sha256=uxk1cMeQkLniYUIgEjHD2eJ8_JGwAKS2minrCmAJfMo,1033
-lt_tensor-0.0.1a16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a16.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a16.dist-info/RECORD,,
+lt_tensor-0.0.1a17.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
+lt_tensor-0.0.1a17.dist-info/METADATA,sha256=nHIAMKShjCnhr2KdEiAhi8IIqP9PoTWnGHgsSfNcTDs,1033
+lt_tensor-0.0.1a17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a17.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a17.dist-info/RECORD,,

lt_tensor/model_zoo/audio_models/diffwave/model.py

@@ -1,201 +0,0 @@
-__all__ = ["DiffWave", "SpectrogramUpsampler", "DiffusionEmbedding"]
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from math import sqrt
-
-
-class AttrDict(dict):
-    def __init__(self, *args, **kwargs):
-        super(AttrDict, self).__init__(*args, **kwargs)
-        self.__dict__ = self
-
-    def override(self, attrs):
-        if isinstance(attrs, dict):
-            self.__dict__.update(**attrs)
-        elif isinstance(attrs, (list, tuple, set)):
-            for attr in attrs:
-                self.override(attr)
-        elif attrs is not None:
-            raise NotImplementedError
-        return self
-
-
-params = AttrDict(
-    # Training params
-    batch_size=16,
-    learning_rate=2e-4,
-    max_grad_norm=None,
-    # Data params
-    sample_rate=22050,
-    n_mels=80,
-    n_fft=1024,
-    hop_samples=256,
-    crop_mel_frames=62,  # Probably an error in paper.
-    # Model params
-    residual_layers=30,
-    residual_channels=64,
-    dilation_cycle_length=10,
-    unconditional=False,
-    noise_schedule=np.linspace(1e-4, 0.05, 50).tolist(),
-    inference_noise_schedule=[0.0001, 0.001, 0.01, 0.05, 0.2, 0.5],
-    # unconditional sample len
-    audio_len=22050 * 5,  # unconditional_synthesis_samples
-)
-
-
-def Conv1d(*args, **kwargs):
-    layer = nn.Conv1d(*args, **kwargs)
-    nn.init.kaiming_normal_(layer.weight)
-    return layer
-
-
-class DiffusionEmbedding(nn.Module):
-    def __init__(self, max_steps):
-        super().__init__()
-        self.register_buffer(
-            "embedding", self._build_embedding(max_steps), persistent=False
-        )
-        self.projection1 = nn.Linear(128, 512)
-        self.projection2 = nn.Linear(512, 512)
-        self.activation = nn.SiLU()
-
-    def forward(self, diffusion_step):
-        if diffusion_step.dtype in [torch.int32, torch.int64]:
-            x = self.embedding[diffusion_step]
-        else:
-            x = self._lerp_embedding(diffusion_step)
-        x = self.projection1(x)
-        x = self.activation(x)
-        x = self.projection2(x)
-        x = self.activation(x)
-        return x
-
-    def _lerp_embedding(self, t):
-        low_idx = torch.floor(t).long()
-        high_idx = torch.ceil(t).long()
-        low = self.embedding[low_idx]
-        high = self.embedding[high_idx]
-        return low + (high - low) * (t - low_idx)
-
-    def _build_embedding(self, max_steps):
-        steps = torch.arange(max_steps).unsqueeze(1)  # [T,1]
-        dims = torch.arange(64).unsqueeze(0)  # [1,64]
-        table = steps * 10.0 ** (dims * 4.0 / 63.0)  # [T,64]
-        table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)
-        return table
-
-
-class SpectrogramUpsampler(nn.Module):
-    def __init__(self, n_mels):
-        super().__init__()
-        self.conv1 = nn.ConvTranspose2d(1, 1, [3, 32], stride=[1, 16], padding=[1, 8])
-        self.conv2 = nn.ConvTranspose2d(1, 1, [3, 32], stride=[1, 16], padding=[1, 8])
-
-    def forward(self, x):
-        x = torch.unsqueeze(x, 1)
-        x = self.conv1(x)
-        x = F.leaky_relu(x, 0.4)
-        x = self.conv2(x)
-        x = F.leaky_relu(x, 0.4)
-        x = torch.squeeze(x, 1)
-        return x
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, n_mels, residual_channels, dilation, uncond=False):
-        """
-        :param n_mels: inplanes of conv1x1 for spectrogram conditional
-        :param residual_channels: audio conv
-        :param dilation: audio conv dilation
-        :param uncond: disable spectrogram conditional
-        """
-        super().__init__()
-        self.dilated_conv = Conv1d(
-            residual_channels,
-            2 * residual_channels,
-            3,
-            padding=dilation,
-            dilation=dilation,
-        )
-        self.diffusion_projection = nn.Linear(512, residual_channels)
-        if not uncond:  # conditional model
-            self.conditioner_projection = Conv1d(n_mels, 2 * residual_channels, 1)
-        else:  # unconditional model
-            self.conditioner_projection = None
-
-        self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1)
-
-    def forward(self, x, diffusion_step, conditioner=None):
-        assert (conditioner is None and self.conditioner_projection is None) or (
-            conditioner is not None and self.conditioner_projection is not None
-        )
-
-        diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
-        y = x + diffusion_step
-        if self.conditioner_projection is None:  # using a unconditional model
-            y = self.dilated_conv(y)
-        else:
-            conditioner = self.conditioner_projection(conditioner)
-            y = self.dilated_conv(y) + conditioner
-
-        gate, filter = torch.chunk(y, 2, dim=1)
-        y = torch.sigmoid(gate) * torch.tanh(filter)
-
-        y = self.output_projection(y)
-        residual, skip = torch.chunk(y, 2, dim=1)
-        return (x + residual) / sqrt(2.0), skip
-
-
-class DiffWave(nn.Module):
-    def __init__(self, params):
-        super().__init__()
-        self.params = params
-        self.input_projection = Conv1d(1, params.residual_channels, 1)
-        self.diffusion_embedding = DiffusionEmbedding(len(params.noise_schedule))
-        if self.params.unconditional:  # use unconditional model
-            self.spectrogram_upsampler = None
-        else:
-            self.spectrogram_upsampler = SpectrogramUpsampler(params.n_mels)
-
-        self.residual_layers = nn.ModuleList(
-            [
-                ResidualBlock(
-                    params.n_mels,
-                    params.residual_channels,
-                    2 ** (i % params.dilation_cycle_length),
-                    uncond=params.unconditional,
-                )
-                for i in range(params.residual_layers)
-            ]
-        )
-        self.skip_projection = Conv1d(
-            params.residual_channels, params.residual_channels, 1
-        )
-        self.output_projection = Conv1d(params.residual_channels, 1, 1)
-        nn.init.zeros_(self.output_projection.weight)
-
-    def forward(self, audio, diffusion_step, spectrogram=None):
-        assert (spectrogram is None and self.spectrogram_upsampler is None) or (
-            spectrogram is not None and self.spectrogram_upsampler is not None
-        )
-        x = audio.unsqueeze(1)
-        x = self.input_projection(x)
-        x = F.relu(x)
-
-        diffusion_step = self.diffusion_embedding(diffusion_step)
-        if self.spectrogram_upsampler:  # use conditional model
-            spectrogram = self.spectrogram_upsampler(spectrogram)
-
-        skip = None
-        for layer in self.residual_layers:
-            x, skip_connection = layer(x, diffusion_step, spectrogram)
-            skip = skip_connection if skip is None else skip_connection + skip
-
-        x = skip / sqrt(len(self.residual_layers))
-        x = self.skip_projection(x)
-        x = F.relu(x)
-        x = self.output_projection(x)
-        return x