feat: add noiser bottleneck

Author: flavioschneider

audio_diffusion_pytorch/__init__.py

@@ -31,6 +31,7 @@
 from .modules import (
     AutoEncoder1d,
     MultiEncoder1d,
+    Noiser,
     T5Embedder,
     Tanh,
     UNet1d,

audio_diffusion_pytorch/modules.py

@@ -1350,6 +1350,20 @@ def forward(
         return (x, info) if with_info else x
 
 
+class Noiser(Bottleneck):
+    def __init__(self, sigma: float = 1.0):
+        super().__init__()
+        self.sigma = sigma
+
+    def forward(
+        self, x: Tensor, with_info: bool = False
+    ) -> Union[Tensor, Tuple[Tensor, Any]]:
+        if self.training:
+            x = torch.randn_like(x) * self.sigma + x
+        info: Dict = dict()
+        return (x, info) if with_info else x
+
+
 class AutoEncoder1d(nn.Module):
     def __init__(
         self,

setup.py

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