fix: bottleneck channels computed twice

Author: flavioschneider

audio_diffusion_pytorch/model.py

@@ -7,7 +7,7 @@
 from torch import Tensor, nn
 
 from .diffusion import LinearSchedule, UniformDistribution, VSampler, XDiffusion
-from .modules import STFT, Conv1d, SinusoidalEmbedding, UNet1d, UNetConditional1d
+from .modules import STFT, SinusoidalEmbedding, UNet1d, UNetConditional1d
 from .utils import (
     default,
     downsample,
@@ -153,13 +153,6 @@ def __init__(
             **encoder_kwargs,
         )
 
-        if exists(bottleneck_channels):
-            self.to_bottleneck = Conv1d(
-                in_channels=encoder_channels * encoder_multipliers[-1],
-                out_channels=bottleneck_channels,
-                kernel_size=1,
-            )
-
         self.encoder_downsample_factor = encoder_patch_size * prod(encoder_factors)
         self.bottleneck_channels = bottleneck_channels
         self.bottlenecks = nn.ModuleList(to_list(bottleneck))
@@ -168,9 +161,6 @@ def encode(
         self, x: Tensor, with_info: bool = False
     ) -> Union[Tensor, Tuple[Tensor, Any]]:
         latent, info = self.encoder(x, with_info=True)
-        # Convert latent channels
-        if exists(self.bottleneck_channels):
-            latent = self.to_bottleneck(latent)
         # Apply bottlenecks if present
         for bottleneck in self.bottlenecks:
             latent, info_bottleneck = bottleneck(latent, with_info=True)

setup.py

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