feat: add convmean tail

Author: flavioschneider

audio_diffusion_pytorch/diffusion.py

@@ -153,7 +153,6 @@ def step(self, x: Tensor, fn: Callable, sigma: float, sigma_next: float) -> Tens
         x_next = x + d * (sigma_down - sigma)
         # Add randomness
         x_next = x_next + torch.randn_like(x) * sigma_up
-        print(sigma_up)
         return x_next
 
     def forward(

audio_diffusion_pytorch/model.py

@@ -23,6 +23,7 @@ def __init__(
         channels: int,
         patch_size: int,
         kernel_sizes_init: Sequence[int],
+        out_means: int,
         multipliers: Sequence[int],
         factors: Sequence[int],
         num_blocks: Sequence[int],
@@ -50,6 +51,7 @@ def __init__(
             resnet_groups=resnet_groups,
             kernel_multiplier_downsample=kernel_multiplier_downsample,
             kernel_sizes_init=kernel_sizes_init,
+            out_means=out_means,
             multipliers=multipliers,
             factors=factors,
             num_blocks=num_blocks,
@@ -98,6 +100,7 @@ def __init__(self, *args, **kwargs):
             channels=128,
             patch_size=16,
             kernel_sizes_init=[1, 3, 7],
+            out_means=4,
             multipliers=[1, 2, 4, 4, 4, 4, 4],
             factors=[4, 4, 4, 2, 2, 2],
             num_blocks=[2, 2, 2, 2, 2, 2],
@@ -133,6 +136,7 @@ def __init__(self, factor: int, in_channels: int = 1, *args, **kwargs):
             in_channels=in_channels,
             channels=128,
             patch_size=16,
+            out_means=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

@@ -3,7 +3,7 @@
 
 import torch
 import torch.nn as nn
-from einops import rearrange
+from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
 from einops_exts import rearrange_many, repeat_many
 from einops_exts.torch import EinopsToAndFrom
@@ -25,6 +25,17 @@ def ConvTranspose1d(*args, **kwargs) -> nn.Module:
     return nn.ConvTranspose1d(*args, **kwargs)
 
 
+class ConvMean1d(nn.Module):
+    def __init__(self, num_means: int, *args, **kwargs):
+        super().__init__()
+        self.convs = nn.ModuleList([Conv1d(*args, **kwargs) for _ in range(num_means)])
+
+    def forward(self, x: Tensor) -> Tensor:
+        xs = torch.stack([conv(x) for conv in self.convs])
+        x = reduce(xs, "n b c t -> b c t", "mean")
+        return x
+
+
 def Downsample1d(
     in_channels: int, out_channels: int, factor: int, kernel_multiplier: int = 2
 ) -> nn.Module:
@@ -713,6 +724,7 @@ def __init__(
         use_skip_scale: bool,
         use_attention_bottleneck: bool,
         out_channels: Optional[int] = None,
+        out_means: int = 1,
         context_channels: Optional[Sequence[int]] = None,
     ):
         super().__init__()
@@ -821,7 +833,8 @@ def __init__(
                 num_groups=resnet_groups,
                 time_context_features=time_context_features,
             ),
-            Conv1d(
+            ConvMean1d(
+                num_means=out_means,
                 in_channels=channels,
                 out_channels=out_channels * patch_size,
                 kernel_size=1,

setup.py

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