feat: add additional skip connection

Author: flavioschneider

README.md

@@ -125,7 +125,7 @@ from audio_diffusion_pytorch import UNet1d
 unet = UNet1d(
     in_channels=1,
     channels=128,
-    patch_size=16,
+    patch_blocks=4,
     kernel_sizes_init=[1, 3, 7],
     multipliers=[1, 2, 4, 4, 4, 4, 4],
     factors=[4, 4, 4, 2, 2, 2],

audio_diffusion_pytorch/modules.py

@@ -1079,8 +1079,8 @@ def forward(
         mapping = self.get_mapping(time, features)
 
         x = self.to_in(x, mapping)
+        skips_list = [x]
 
-        skips_list = []
         for i, downsample in enumerate(self.downsamples):
             channels = self.get_channels(channels_list, layer=i + 1)
             x, skips = downsample(
@@ -1094,6 +1094,7 @@ def forward(
             skips = skips_list.pop()
             x = upsample(x, skips, mapping=mapping, embedding=embedding)
 
+        x += skips_list.pop()
         x = self.to_out(x, mapping)
 
         return x

setup.py

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