handle the conditioning token

lucidrains · May 17, 2023 · 7e85f47 · 7e85f47
1 parent 1dcbc48
commit 7e85f47
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'soundstorm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.1',
+  version = '0.0.2',
   license='MIT',
   description = 'SoundStorm - Efficient Parallel Audio Generation from Google Deepmind, in Pytorch',
   author = 'Phil Wang',

diff --git a/soundstorm_pytorch/soundstorm.py b/soundstorm_pytorch/soundstorm.py
@@ -2,7 +2,7 @@
 from torch import Tensor, nn, einsum
 import torch.nn.functional as F
 
-from einops import rearrange, reduce
+from einops import rearrange, reduce, repeat
 from einops.layers.torch import Rearrange
 
 from beartype import beartype
@@ -50,9 +50,17 @@ def __init__(
 
     def forward(
         self,
-        x
+        x,
+        cond = None
     ):
         x = reduce(x, 'b (n h) d -> b n d', h = self.num_tokens_reduce)
+
+        if exists(cond):
+            if cond.ndim == 2:
+                cond = rearrange(cond, 'b d -> b 1 d')
+
+            x = x + cond
+
         logits = self.conformer(x)
         out = self.heads(logits)
         return out