feat: added audio diffusion upsampler

Author: flavioschneider

README.md

@@ -19,6 +19,7 @@ https://colab.research.google.com/gist/flavioschneider/39c6454bfc2d03dc7d0c5c9d8
 
 ## Usage
 
+### Generation
 ```py
 from audio_diffusion_pytorch import AudioDiffusionModel
 
@@ -37,6 +38,28 @@ sampled = model.sample(
 ) # [2, 1, 262144]
 ```
 
+### Upsampling
+```py
+from audio_diffusion_pytorch import AudioDiffusionUpsampler
+
+upsampler = AudioDiffusionUpsampler(
+    factor=4,
+    in_channels=1
+)
+
+# Train on high frequency data
+x = torch.randn(2, 1, 2 ** 18) # [batch, in_channels, samples]
+loss = upsampler(x)
+loss.backward()
+
+# Given start undersampled source, samples upsampled source
+start = torch.randn(1, 1, 2 ** 16)
+sampled = upsampler.sample(
+    start=start,
+    num_steps=5 # Suggested range: 2-100
+)
+```
+
 ## Usage with Components
 
 ### UNet1d
@@ -131,6 +154,8 @@ y_long = composer(y, keep_start=True) # [1, 1, 98304]
 ```
 
 
+
+
 ## Experiments
 
 

audio_diffusion_pytorch/__init__.py

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

audio_diffusion_pytorch/model.py

@@ -1,5 +1,6 @@
 from typing import Optional, Sequence
 
+import torch
 from torch import Tensor, nn
 
 from .diffusion import (
@@ -111,7 +112,7 @@ def __init__(self, *args, **kwargs):
             use_skip_scale=True,
             diffusion_sigma_distribution=LogNormalDistribution(mean=-3.0, std=1.0),
             diffusion_sigma_data=0.1,
-            diffusion_dynamic_threshold=0.95,
+            diffusion_dynamic_threshold=0.0,
         )
 
         super().__init__(*args, **{**default_kwargs, **kwargs})
@@ -122,3 +123,50 @@ def sample(self, *args, **kwargs):
             sampler=ADPM2Sampler(rho=1.0),
         )
         return super().sample(*args, **{**default_kwargs, **kwargs})
+
+
+class AudioDiffusionUpsampler(Model1d):
+    def __init__(self, factor: int, in_channels: int = 1, *args, **kwargs):
+        self.factor = factor
+
+        default_kwargs = dict(
+            in_channels=in_channels,
+            channels=128,
+            patch_size=16,
+            kernel_sizes_init=[1, 3, 7],
+            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],
+            attention_heads=8,
+            attention_features=64,
+            attention_multiplier=2,
+            use_attention_bottleneck=True,
+            resnet_groups=8,
+            kernel_multiplier_downsample=2,
+            use_nearest_upsample=False,
+            use_skip_scale=True,
+            diffusion_sigma_distribution=LogNormalDistribution(mean=-3.0, std=1.0),
+            diffusion_sigma_data=0.1,
+            diffusion_dynamic_threshold=0.0,
+            context_channels=[in_channels],
+        )
+
+        super().__init__(*args, {**default_kwargs, **kwargs})  # type: ignore
+
+    def forward(self, x: Tensor, **kwargs) -> Tensor:
+        # Downsample by picking every `factor` item
+        downsampled = x[:, :, :: self.factor]
+        # Upsample by interleaving to get context
+        context = torch.repeat_interleave(downsampled, repeats=self.factor, dim=2)
+        return self.diffusion(x, context=[context], **kwargs)
+
+    def sample(self, start: Tensor, *args, **kwargs):  # type: ignore
+        context = torch.repeat_interleave(start, repeats=self.factor, dim=2)
+        noise = torch.randn_like(context)
+        default_kwargs = dict(
+            context=[context],
+            sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
+            sampler=ADPM2Sampler(rho=1.0),
+        )
+        return super().sample(noise, *args, **{**default_kwargs, **kwargs})  # type: ignore # noqa

setup.py

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