fix: first/last block channels prod with multiplier

archinetai · Oct 26, 2022 · 22e5d75 · 22e5d75
1 parent c458fc2
commit 22e5d75
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/audio_diffusion_pytorch/modules.py b/audio_diffusion_pytorch/modules.py
@@ -1019,7 +1019,7 @@ def __init__(
 
         self.to_in = Patcher(
             in_channels=in_channels + context_channels[0],
-            out_channels=channels,
+            out_channels=channels * multipliers[0],
             blocks=patch_blocks,
             factor=patch_factor,
             context_mapping_features=context_mapping_features,
@@ -1093,7 +1093,7 @@ def __init__(
         )
 
         self.to_out = Unpatcher(
-            in_channels=channels,
+            in_channels=channels * multipliers[0],
             out_channels=out_channels * (2 if use_magnitude_channels else 1),
             blocks=patch_blocks,
             factor=patch_factor,
@@ -1383,7 +1383,7 @@ def __init__(
 
         self.to_in = Patcher(
             in_channels=in_channels,
-            out_channels=channels,
+            out_channels=channels * multipliers[0],
             blocks=patch_blocks,
             factor=patch_factor,
         )
@@ -1474,7 +1474,7 @@ def __init__(
         )
 
         self.to_out = Unpatcher(
-            in_channels=channels,
+            in_channels=channels * multipliers[0],
             out_channels=out_channels * (2 if use_magnitude_channels else 1),
             blocks=patch_blocks,
             factor=patch_factor,

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.78",
+    version="0.0.79",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",