Add support for Adafactor, add TI initializer noise

1 files changed, 9 insertions, 1 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 4166dc6..9abd1bb 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -52,7 +52,12 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         self.temp_token_embedding = resize_embedding(self.temp_token_embedding, size, self.initializer_factor)
         self.token_embedding = resize_embedding(self.token_embedding, size, self.initializer_factor)
 
-    def add_embed(self, token_ids: Union[int, list[int]], initializer: Optional[Union[int, list[int], torch.FloatTensor]] = None):
+    def add_embed(
+        self,
+        token_ids: Union[int, list[int]],
+        initializer: Optional[Union[int, list[int], torch.FloatTensor]] = None,
+        initializer_noise: float = 0.0,
+    ):
         if isinstance(token_ids, int):
             token_ids = [token_ids]
 
@@ -73,6 +78,9 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
             dtype=self.temp_token_embedding.weight.dtype,
         )
 
+        if initializer_noise != 0:
+            initializer += torch.randn_like(initializer) * initializer_noise
+
         token_ids = torch.tensor(token_ids, dtype=torch.long)
 
         self.temp_token_ids = torch.cat([self.temp_token_ids, token_ids])