feat: add audio autoencoder head option

Author: flavioschneider

README.md

@@ -146,7 +146,7 @@ y_long = composer(y, keep_start=True) # [1, 1, 98304]
 - [x] Add elucidated diffusion.
 - [x] Add ancestral DPM2 sampler.
 - [x] Add dynamic thresholding.
-- [ ] Add support with (variational) autoencoder to compress audio before diffusion.
+- [x] Add (variational) autoencoder option to compress audio before diffusion.
 - [ ] Fix inpainting and make it work with ADPM2 sampler.
 
 ## Appreciation

audio_diffusion_pytorch/__init__.py

@@ -11,5 +11,5 @@
     Schedule,
     SpanBySpanComposer,
 )
-from .model import AudioDiffusionModel, Model1d
+from .model import AudioAutoEncoderModel, AudioDiffusionModel, Model1d
 from .modules import AutoEncoder1d, UNet1d

audio_diffusion_pytorch/model.py

@@ -12,7 +12,8 @@
     Sampler,
     Schedule,
 )
-from .modules import UNet1d
+from .modules import AutoEncoder1d, UNet1d
+from .utils import exists
 
 
 class Model1d(nn.Module):
@@ -39,9 +40,19 @@ def __init__(
         diffusion_sigma_data: int,
         diffusion_dynamic_threshold: float,
         out_channels: Optional[int] = None,
+        use_autoencoder: bool = False,
+        autoencoder: Optional[AutoEncoder1d] = None,
+        autoencoder_scale: float = 1.0,
     ):
         super().__init__()
 
+        self.use_autoencoder = use_autoencoder
+
+        if use_autoencoder:
+            assert exists(autoencoder)
+            self.autoencoder_scale = autoencoder_scale
+            self.autoencoder = autoencoder
+
         self.unet = UNet1d(
             in_channels=in_channels,
             channels=channels,
@@ -71,6 +82,8 @@ def __init__(
         )
 
     def forward(self, x: Tensor) -> Tensor:
+        if self.use_autoencoder:
+            x = self.autoencoder_scale * self.autoencoder.encode(x)  # type: ignore
         return self.diffusion(x)
 
     def sample(
@@ -82,19 +95,35 @@ def sample(
             sigma_schedule=sigma_schedule,
             num_steps=num_steps,
         )
-        return diffusion_sampler(noise)
+        x = diffusion_sampler(noise)
 
+        if self.use_autoencoder:
+            x = (1.0 / self.autoencoder_scale) * self.autoencoder.decode(x)
 
-class AudioDiffusionModel(Model1d):
+        return x
+
+
+class AudioAutoEncoderModel(AutoEncoder1d):
     def __init__(self, *args, **kwargs):
         default_kwargs = dict(
             in_channels=1,
+            bottleneck_channels=128,
             channels=128,
             patch_size=16,
-            multipliers=[1, 2, 4, 4, 4, 4, 4],
-            factors=[4, 4, 4, 2, 2, 2],
-            num_blocks=[2, 2, 2, 2, 2, 2],
-            attentions=[False, False, False, True, True, True],
+            multipliers=[1, 1, 1, 1, 1],
+            factors=[1, 4, 4, 4],
+            num_blocks=[2, 2, 2, 2],
+            resnet_groups=8,
+            kernel_multiplier_downsample=2,
+            loss_kl_weight=1e-8,
+        )
+        super().__init__(*args, **{**default_kwargs, **kwargs})
+
+
+class AudioDiffusionModel(Model1d):
+    def __init__(self, *args, **kwargs):
+        default_kwargs = dict(
+            channels=128,
             attention_heads=8,
             attention_features=64,
             attention_multiplier=2,
@@ -106,14 +135,41 @@ def __init__(self, *args, **kwargs):
             use_attention_bottleneck=True,
             use_learned_time_embedding=True,
             diffusion_sigma_distribution=LogNormalDistribution(mean=-3.0, std=1.0),
-            diffusion_sigma_data=0.1,
-            diffusion_dynamic_threshold=0.95,
         )
-        super().__init__(*args, **{**default_kwargs, **kwargs})
+
+        model_kwargs = None
+
+        if "autoencoder" in kwargs:
+            sigma_data = 0.2
+            model_kwargs = dict(
+                in_channels=128,
+                patch_size=1,
+                multipliers=[1, 4, 4, 4],
+                factors=[2, 2, 2],
+                num_blocks=[2, 2, 2],
+                attentions=[True, True, True],
+                diffusion_sigma_data=sigma_data,
+                diffusion_dynamic_threshold=0.0,
+                use_autoencoder=True,
+                autoencoder_scale=sigma_data,
+            )
+        else:
+            model_kwargs = dict(
+                in_channels=1,
+                patch_size=16,
+                multipliers=[1, 2, 4, 4, 4, 4, 4],
+                factors=[4, 4, 4, 2, 2, 2],
+                num_blocks=[2, 2, 2, 2, 2, 2],
+                attentions=[False, False, False, True, True, True],
+                diffusion_sigma_data=0.1,
+                diffusion_dynamic_threshold=0.95,
+                use_autoencoder=False,
+            )
+        super().__init__(*args, **{**default_kwargs, **model_kwargs, **kwargs})
 
     def sample(self, *args, **kwargs):
         default_kwargs = dict(
-            sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3, rho=9.0),
-            sampler=ADPM2Sampler(rho=1),
+            sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
+            sampler=ADPM2Sampler(rho=1.0),
         )
         return super().sample(*args, **{**default_kwargs, **kwargs})

setup.py

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