1 files changed, 33 insertions, 17 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 1e21965..d8343a0 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -12,7 +12,7 @@ from transformers.models.clip import CLIPTextConfig
 from transformers.models.clip.modeling_clip import CLIPTextEmbeddings
 
 
-def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initializer_factor: float = 1.0) -> nn.Embedding:
+def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initializer_factor: Optional[float] = None) -> nn.Embedding:
     old_num_embeddings, old_embedding_dim = old_embedding.weight.shape
 
     if old_num_embeddings == new_num_embeddings:
@@ -26,13 +26,16 @@ def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initi
         device=old_embedding.weight.device,
         dtype=old_embedding.weight.dtype
     )
-    new_embedding.weight.data.normal_(mean=0.0, std=initializer_factor * 0.02)
+    if initializer_factor is not None:
+        new_embedding.weight.data.normal_(mean=0.0, std=initializer_factor * 0.02)
+    else:
+        nn.init.zeros_(new_embedding.weight.data)
     new_embedding.weight.data[:n, :] = old_embedding.weight.data[:n, :]
     return new_embedding
 
 
 class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
-    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, alpha: float = 1.0, rank: int = 4):
+    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, alpha: float = 1.0):
         super().__init__(config)
 
         self.token_embedding = embeddings.token_embedding
@@ -40,17 +43,16 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         self.initializer_factor = config.initializer_factor
         self.alpha = alpha
 
-        self.temp_token_embedding = nn.Embedding(
-            self.token_embedding.num_embeddings,
-            self.token_embedding.embedding_dim,
-            device=self.token_embedding.weight.device,
-            dtype=self.token_embedding.weight.dtype
-        )
-        self.temp_token_embedding.weight.data = self.token_embedding.weight.data.clone().detach()
+        self.temp_token_embedding = nn.ParameterList()
         self.temp_token_ids = torch.tensor([], dtype=torch.long)
 
     def resize(self, size: int):
-        self.temp_token_embedding = resize_embedding(self.temp_token_embedding, size, self.initializer_factor)
+        for _ in range(len(self.temp_token_embedding), size):
+            self.temp_token_embedding.append(torch.zeros(
+                self.token_embedding.embedding_dim,
+                device=self.token_embedding.weight.device,
+                dtype=self.token_embedding.weight.dtype,
+            ))
         self.token_embedding = resize_embedding(self.token_embedding, size, self.initializer_factor)
 
     def add_embed(
@@ -85,7 +87,6 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         token_ids = torch.tensor(token_ids, dtype=torch.long)
 
         self.temp_token_ids = torch.cat([self.temp_token_ids, token_ids])
-        self.temp_token_embedding.weight.data[token_ids] = initializer
         self.token_embedding.weight.data[token_ids] = initializer
 
     def load_embed(self, input_ids: list[int], filename: Path):
@@ -96,16 +97,31 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         save_file({"embed": self.get_embed(input_ids)}, filename)
 
     def persist(self):
-        self.token_embedding.weight.data[self.temp_token_ids] = self.temp_token_embedding.weight.data[self.temp_token_ids]
+        for id, emb in zip(self.temp_token_ids, self.temp_token_embedding):
+            self.token_embedding.weight.data[id] += self.alpha * emb
+            nn.init.zeros_(emb)
         self.temp_token_ids = torch.tensor([], dtype=torch.long)
 
     def get_embed(self, input_ids: Union[list[int], torch.LongTensor]):
         if isinstance(input_ids, list):
             input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device, dtype=torch.long)
 
+        all_temp_token_ids = self.temp_token_ids.to(input_ids.device)
+
         embeds = self.token_embedding(input_ids)
-        mask = torch.isin(input_ids, self.temp_token_ids.to(input_ids.device))
-        embeds[mask] = self.temp_token_embedding(input_ids[mask])
+        mask = torch.isin(input_ids, all_temp_token_ids)
+        temp_token_ids = input_ids[mask]
+
+        temp_token_ids = temp_token_ids.unsqueeze(1)
+        all_temp_token_ids = all_temp_token_ids.unsqueeze(0)
+        temp_token_ids = torch.nonzero(temp_token_ids == all_temp_token_ids)[:, 1].squeeze()
+
+        if len(temp_token_ids):
+            embeds_override = torch.stack([
+                self.temp_token_embedding[id]
+                for id in temp_token_ids
+            ])
+            embeds[mask] += self.alpha * embeds_override
 
         return embeds
 
@@ -129,7 +145,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         return embeddings
 
 
-def patch_managed_embeddings(text_encoder: CLIPTextModel) -> ManagedCLIPTextEmbeddings:
-    text_embeddings = ManagedCLIPTextEmbeddings(text_encoder.config, text_encoder.text_model.embeddings)
+def patch_managed_embeddings(text_encoder: CLIPTextModel, alpha: float = 1.0) -> ManagedCLIPTextEmbeddings:
+    text_embeddings = ManagedCLIPTextEmbeddings(text_encoder.config, text_encoder.text_model.embeddings, alpha)
     text_encoder.text_model.embeddings = text_embeddings
     return text_embeddings