feat: option to choose autoencoder out_channels

Author: flavioschneider

audio_diffusion_pytorch/modules.py

@@ -1508,11 +1508,13 @@ def __init__(
         use_noisy: bool = False,
         bottleneck: Union[Bottleneck, List[Bottleneck]] = [],
         bottleneck_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
         use_magnitude_channels: bool = False,
     ):
         super().__init__()
         num_layers = len(multipliers) - 1
         self.bottlenecks = nn.ModuleList(to_list(bottleneck))
+        out_channels = default(out_channels, in_channels)
 
         assert len(factors) >= num_layers and len(num_blocks) >= num_layers
 
@@ -1530,7 +1532,7 @@ def __init__(
 
         self.decoder = Decoder1d(
             in_channels=bottleneck_channels,
-            out_channels=in_channels,
+            out_channels=out_channels,
             channels=channels,
             patch_blocks=patch_blocks,
             patch_factor=patch_factor,

setup.py

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