TI: Bring back old embedding decay

1 files changed, 7 insertions, 8 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index a356434..63a141f 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -37,7 +37,7 @@ def resize_embedding(old_embedding: nn.Embedding, new_num_embeddings: int, initi
 
 
 class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
-    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings, alpha: float = 1.0):
+    def __init__(self, config: CLIPTextConfig, embeddings: CLIPTextEmbeddings):
         super().__init__(config)
 
         self.token_embedding = embeddings.token_embedding
@@ -49,7 +49,6 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
             device=self.token_embedding.weight.device,
             dtype=self.token_embedding.weight.dtype,
         )
-        self.alpha = alpha
 
     def resize(self, size: int):
         self.token_override_embedding.resize(size)
@@ -87,7 +86,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         token_ids = torch.tensor(token_ids, dtype=torch.long)
 
         self.token_embedding.weight.data[token_ids] = initializer
-        self.token_override_embedding.set(token_ids)
+        self.token_override_embedding.set(token_ids, initializer)
 
     def load_embed(self, input_ids: list[int], filename: Path):
         with safe_open(filename, framework="pt", device="cpu") as file:
@@ -101,8 +100,8 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         embs, mask = self.token_override_embedding(input_ids)
         if embs is not None:
             input_ids = input_ids[mask]
-            self.token_embedding.weight.data[input_ids] += self.alpha * embs
-        self.token_override_embedding.unset(input_ids)
+            self.token_embedding.weight.data[input_ids] = embs
+            self.token_override_embedding.unset(input_ids)
 
     def get_embed(self, input_ids: Union[list[int], torch.LongTensor]):
         if isinstance(input_ids, list):
@@ -111,7 +110,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         embs = self.token_embedding(input_ids)
         embs_override, mask = self.token_override_embedding(input_ids)
         if embs_override is not None:
-            embs[mask] += self.alpha * embs_override
+            embs[mask] = embs_override
 
         return embs
 
@@ -135,7 +134,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         return 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)
+def patch_managed_embeddings(text_encoder: CLIPTextModel) -> ManagedCLIPTextEmbeddings:
+    text_embeddings = ManagedCLIPTextEmbeddings(text_encoder.config, text_encoder.text_model.embeddings)
     text_encoder.text_model.embeddings = text_embeddings
     return text_embeddings