fix: cross attention skip connection

Author: flavioschneider

audio_diffusion_pytorch/modules.py

@@ -421,7 +421,7 @@ def __init__(
     def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
         x = self.attention(x) + x
         if self.use_cross_attention:
-            x = self.cross_attention(x, context=context)
+            x = self.cross_attention(x, context=context) + x
         x = self.feed_forward(x) + x
         return x
 

setup.py

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