feat: add number embedder

Author: flavioschneider

README.md

@@ -130,6 +130,16 @@ sampled = model.sample(
 )
 ```
 
+#### Number Conditional Generation
+
+```py
+from audio_diffusion_pytorch import NumberEmbedder
+
+embedder = NumberEmbedder(features=768)
+embedding = embedder([0.1, 0.2]) # [2, 768]
+```
+
+
 ## Usage with Components
 
 ### UNet1d

audio_diffusion_pytorch/__init__.py

@@ -33,4 +33,4 @@
     DiffusionVocoder1d,
     Model1d,
 )
-from .modules import T5Embedder, UNet1d, UNetConditional1d
+from .modules import NumberEmbedder, T5Embedder, UNet1d, UNetConditional1d

audio_diffusion_pytorch/modules.py

@@ -1271,6 +1271,28 @@ def forward(self, texts: List[str]) -> Tensor:
         return embedding
 
 
+class NumberEmbedder(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        dim: int = 256,
+    ):
+        super().__init__()
+        self.features = features
+        self.embedding = TimePositionalEmbedding(dim=dim, out_features=features)
+
+    def forward(self, x: Union[List[float], Tensor]) -> Tensor:
+        if not torch.is_tensor(x):
+            device = next(self.embedding.parameters()).device
+            x = torch.tensor(x, device=device)
+        assert isinstance(x, Tensor)
+        shape = x.shape
+        x = rearrange(x, "... -> (...)")
+        embedding = self.embedding(x)
+        x = embedding.view(*shape, self.features)
+        return x  # type: ignore
+
+
 """
 Audio Transforms
 """

setup.py

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