2 files changed, 33 insertions, 17 deletions

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 2d60c28..e8cc865 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -38,18 +38,18 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
         self.token_embedding = embeddings.token_embedding
         self.position_embedding = embeddings.position_embedding
         self.initializer_factor = config.initializer_factor
+        self.init_temp_embeddings()
 
+    def init_temp_embeddings(self):
         self.temp_token_embedding = nn.Embedding(
-            self.token_embedding.num_embeddings,
+            0,
             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_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)
         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):
@@ -74,9 +74,17 @@ 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.temp_token_embedding = resize_embedding(
+            self.temp_token_embedding,
+            self.temp_token_ids.shape[0],
+            self.initializer_factor
+        )
+
+        mask = torch.nonzero(torch.isin(self.temp_token_ids, token_ids)).squeeze(1)
+        self.temp_token_embedding.weight.data[mask] = initializer
+        self.token_embedding.weight.data[token_ids] = initializer
 
     def load_embed(self, input_ids: list[int], filename: Path):
         with safe_open(filename, framework="pt", device="cpu") as file:
@@ -86,17 +94,25 @@ 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]
-        self.temp_token_ids = torch.tensor([], dtype=torch.long)
+        self.token_embedding.weight.data[self.temp_token_ids] = self.temp_token_embedding.weight.data[:]
+        self.init_temp_embeddings()
 
     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]
+        embeds_mask = torch.isin(input_ids, all_temp_token_ids)
+        temp_token_ids = input_ids[embeds_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()
+
+        embeds[embeds_mask] = self.temp_token_embedding(temp_token_ids)
 
         return embeds
 
diff --git a/train_ti.py b/train_ti.py
index 9ae8d1b..e4fd464 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -588,14 +588,6 @@ def main():
         unet.enable_gradient_checkpointing()
         text_encoder.gradient_checkpointing_enable()
 
-    if args.embeddings_dir is not None:
-        embeddings_dir = Path(args.embeddings_dir)
-        if not embeddings_dir.exists() or not embeddings_dir.is_dir():
-            raise ValueError("--embeddings_dir must point to an existing directory")
-
-        added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
-        print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
-
     if len(args.alias_tokens) != 0:
         alias_placeholder_tokens = args.alias_tokens[::2]
         alias_initializer_tokens = args.alias_tokens[1::2]
@@ -609,6 +601,14 @@ def main():
         embeddings.persist()
         print(f"Added {len(added_tokens)} aliases: {list(zip(alias_placeholder_tokens, added_tokens, alias_initializer_tokens, added_ids))}")
 
+    if args.embeddings_dir is not None:
+        embeddings_dir = Path(args.embeddings_dir)
+        if not embeddings_dir.exists() or not embeddings_dir.is_dir():
+            raise ValueError("--embeddings_dir must point to an existing directory")
+
+        added_tokens, added_ids = load_embeddings_from_dir(tokenizer, embeddings, embeddings_dir)
+        print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
+
     if args.scale_lr:
         args.learning_rate = (
             args.learning_rate * args.gradient_accumulation_steps *