fix: forward context_features, big bug that prevented conditioning from working

Author: flavioschneider

audio_diffusion_pytorch/modules.py

@@ -414,7 +414,10 @@ def __init__(
 
         if self.use_cross_attention:
             self.cross_attention = Attention(
-                features=features, num_heads=num_heads, head_features=head_features
+                features=features,
+                num_heads=num_heads,
+                head_features=head_features,
+                context_features=context_features,
             )
 
         self.feed_forward = FeedForward(features=features, multiplier=multiplier)
@@ -592,6 +595,7 @@ def __init__(
                 num_heads=attention_heads,
                 head_features=attention_features,
                 multiplier=attention_multiplier,
+                context_features=context_embedding_features,
             )
 
         if self.use_extract:
@@ -692,6 +696,7 @@ def __init__(
                 num_heads=attention_heads,
                 head_features=attention_features,
                 multiplier=attention_multiplier,
+                context_features=context_embedding_features,
             )
 
         self.upsample = Upsample1d(
@@ -776,6 +781,7 @@ def __init__(
                 num_heads=attention_heads,
                 head_features=attention_features,
                 multiplier=attention_multiplier,
+                context_features=context_embedding_features,
             )
 
         self.post_block = ResnetBlock1d(

setup.py

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