feat: add option to change encoder num_blocks

Author: flavioschneider

audio_diffusion_pytorch/model.py

@@ -152,6 +152,7 @@ def __init__(
         encoder_channels: int,
         context_channels: int,
         bottleneck: Optional[Bottleneck] = None,
+        encoder_num_blocks: Optional[Sequence[int]] = None,
         **kwargs
     ):
         super().__init__(
@@ -172,14 +173,18 @@ def __init__(
         self.encoder_factor = patch_size * prod(factors[0:encoder_depth])
         self.bottleneck = bottleneck
 
+        encoder_num_blocks = default(encoder_num_blocks, num_blocks)
+        assert_message = "The number of encoder_num_blocks must match encoder_depth"
+        assert len(encoder_num_blocks) >= encoder_depth, assert_message
+
         self.encoder = Encoder1d(
             in_channels=in_channels,
             channels=channels,
             patch_size=patch_size,
             kernel_sizes_init=kernel_sizes_init,
             multipliers=multipliers,
             factors=factors,
-            num_blocks=num_blocks,
+            num_blocks=encoder_num_blocks,
             resnet_groups=resnet_groups,
             kernel_multiplier_downsample=kernel_multiplier_downsample,
             extract_channels=[0] * (encoder_depth - 1) + [encoder_channels],

setup.py

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