add a time token

README.md

@@ -73,3 +73,14 @@ sampled = e2tts.sample(mel[:, :5], text = text)
     url     = {https://api.semanticscholar.org/CorpusID:263134283}
 }
 ```
+
+```bibtex
+@article{Bao2022AllAW,
+    title   = {All are Worth Words: A ViT Backbone for Diffusion Models},
+    author  = {Fan Bao and Shen Nie and Kaiwen Xue and Yue Cao and Chongxuan Li and Hang Su and Jun Zhu},
+    journal = {2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+    year    = {2022},
+    pages   = {22669-22679},
+    url     = {https://api.semanticscholar.org/CorpusID:253581703}
+}
+```

e2_tts_pytorch/e2_tts.py

@@ -289,6 +289,8 @@ def __init__(
                 nn.SiLU()
             )
 
+            self.to_time_token = nn.Linear(dim, dim, bias = False)
+
         for ind in range(depth):
             is_later_half = ind >= (depth // 2)
 
@@ -343,6 +345,14 @@ def forward(
             times = self.time_cond_mlp(times)
             norm_kwargs.update(condition = times)
 
+            # u-vit paper claims using a time token helps better condition https://arxiv.org/abs/2209.12152
+
+            time_token = self.to_time_token(times)
+            x, time_packed_shape = pack((time_token, x), 'b * d')
+
+            if exists(mask):
+                mask = F.pad(mask, (1, 0), value = True)
+
         # register tokens
 
         registers = repeat(self.registers, 'r d -> b r d', b = batch)
@@ -394,6 +404,9 @@ def forward(
 
         _, x = unpack(x, registers_packed_shape, 'b * d')
 
+        if exists(times):
+            _, x = unpack(x, time_packed_shape, 'b * d')
+
         return self.final_norm(x, **norm_kwargs)
 
 # main classes

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "e2-tts-pytorch"
-version = "0.3.1"
+version = "0.4.0"
 description = "E2-TTS in Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }