feat: set default parameters to unet

Author: flavioschneider

audio_diffusion_pytorch/model.py

@@ -350,7 +350,6 @@ def forward(self, x: Tensor, **kwargs) -> Tensor:
 def get_default_model_kwargs():
     return dict(
         channels=128,
-        patch_blocks=1,
         patch_factor=16,
         multipliers=[1, 2, 4, 4, 4, 4, 4],
         factors=[4, 4, 4, 2, 2, 2],
@@ -360,11 +359,6 @@ def get_default_model_kwargs():
         attention_features=64,
         attention_multiplier=2,
         attention_use_rel_pos=False,
-        resnet_groups=8,
-        kernel_multiplier_downsample=2,
-        use_nearest_upsample=False,
-        use_skip_scale=True,
-        use_context_time=True,
         diffusion_type="v",
         diffusion_sigma_distribution=UniformDistribution(),
     )
@@ -416,13 +410,6 @@ def decode(self, *args, **kwargs):
 class AudioDiffusionMAE(DiffusionMAE1d):
     def __init__(self, *args, **kwargs):
         default_kwargs = dict(
-            patch_blocks=1,
-            patch_factor=1,
-            resnet_groups=8,
-            kernel_multiplier_downsample=2,
-            use_nearest_upsample=False,
-            use_skip_scale=True,
-            use_context_time=True,
             diffusion_type="v",
             diffusion_sigma_distribution=UniformDistribution(),
             stft_num_fft=1023,
@@ -470,8 +457,6 @@ def __init__(self, in_channels: int, **kwargs):
             stft_num_fft=1023,
             stft_hop_length=256,
             channels=512,
-            patch_blocks=1,
-            patch_factor=1,
             multipliers=[3, 2, 1, 1, 1, 1, 1, 1],
             factors=[1, 2, 2, 2, 2, 2, 2],
             num_blocks=[1, 1, 1, 1, 1, 1, 1],
@@ -480,11 +465,6 @@ def __init__(self, in_channels: int, **kwargs):
             attention_features=64,
             attention_multiplier=2,
             attention_use_rel_pos=False,
-            resnet_groups=8,
-            kernel_multiplier_downsample=2,
-            use_nearest_upsample=False,
-            use_skip_scale=True,
-            use_context_time=True,
             diffusion_type="v",
             diffusion_sigma_distribution=UniformDistribution(),
         )

audio_diffusion_pytorch/modules.py

@@ -923,17 +923,17 @@ def __init__(
         self,
         in_channels: int,
         channels: int,
-        patch_blocks: int,
-        patch_factor: int,
         multipliers: Sequence[int],
         factors: Sequence[int],
         num_blocks: Sequence[int],
         attentions: Sequence[int],
-        resnet_groups: int,
-        kernel_multiplier_downsample: int,
-        use_nearest_upsample: bool,
-        use_skip_scale: bool,
-        use_context_time: bool,
+        patch_blocks: int = 1,
+        patch_factor: int = 1,
+        resnet_groups: int = 8,
+        use_context_time: bool = True,
+        kernel_multiplier_downsample: int = 2,
+        use_nearest_upsample: bool = False,
+        use_skip_scale: bool = True,
         use_stft: bool = False,
         use_stft_context: bool = False,
         out_channels: Optional[int] = None,

setup.py

@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.88",
+    version="0.0.89",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",