add missing position_embeddings

Author: damian0815

ldm/modules/textual_inversion_manager.py

@@ -21,7 +21,7 @@ def embedding_vector_length(self) -> int:
         return self.embedding.shape[0]
 
 class TextualInversionManager():
-    def __init__(self, clip_embedder: FrozenCLIPEmbedder, full_precision: bool):
+    def __init__(self, clip_embedder: FrozenCLIPEmbedder, full_precision: bool=True):
         self.clip_embedder = clip_embedder
         self.full_precision = full_precision
         self.hf_concepts_library = HuggingFaceConceptsLibrary()
@@ -169,17 +169,17 @@ def overwrite_textual_inversion_embeddings(self, prompt_token_ids: Union[torch.T
         textual_inversion_token_ids = [ti.token_id for ti in self.textual_inversions]
         pad_token_id = self.clip_embedder.tokenizer.pad_token_id
         overwritten_prompt_embeddings = prompt_embeddings.clone()
-        for i, token_id in enumerate(prompt_token_ids):
-            if token_id == pad_token_id:
-                continue
-            if token_id in textual_inversion_token_ids:
-                textual_inversion = next(ti for ti in self.textual_inversions if ti.token_id == token_id)
-                end_index = min(i + textual_inversion.embedding_vector_length, self.clip_embedder.max_length-1)
-                count_to_overwrite = end_index - i
-                for j in range(0, count_to_overwrite):
-                    # only overwrite the textual inversion token id or the padding token id
-                    if prompt_token_ids[i+j] != pad_token_id and prompt_token_ids[i+j] != token_id:
-                        break
-                    overwritten_prompt_embeddings[i+j] = textual_inversion.embedding[j]
+
+        indices_of_textual_inversion_tokens_in_prompt = [index for index in range(0, len(prompt_token_ids)) if prompt_token_ids[index] in textual_inversion_token_ids]
+        eos_marker_index = self.clip_embedder.max_length-1
+        for i in indices_of_textual_inversion_tokens_in_prompt:
+            token_id = prompt_token_ids[i]
+            textual_inversion = next(ti for ti in self.textual_inversions if ti.token_id == token_id)
+            # don't overwrite the eos marker
+            after_end_index = min(i + textual_inversion.embedding_vector_length, eos_marker_index)
+            actual_count_to_overwrite = after_end_index - i
+            position_embeddings = self.clip_embedder.position_embedding(torch.arange(i, after_end_index, dtype=int))
+            embeddings_to_write = position_embeddings + textual_inversion.embedding[0:actual_count_to_overwrite]
+            overwritten_prompt_embeddings[i:i+actual_count_to_overwrite] = embeddings_to_write
 
         return overwritten_prompt_embeddings

tests/test_textual_inversion.py

@@ -1,9 +1,10 @@
 
 import unittest
+from typing import Union
 
 import torch
 
-from ldm.modules.embedding_manager import TextualInversionManager
+from ldm.modules.textual_inversion_manager import TextualInversionManager
 
 
 KNOWN_WORDS = ['a', 'b', 'c']
@@ -53,7 +54,16 @@ def __init__(self):
         self.max_length = 77
         self.transformer = DummyTransformer()
         self.tokenizer = DummyTokenizer()
+        self.position_embeddings_tensor = torch.randn([77,768], dtype=torch.float32)
 
+    def position_embedding(self, indices: Union[list,torch.Tensor]):
+        if type(indices) is list:
+            indices = torch.tensor(indices, dtype=int)
+        return torch.index_select(self.position_embeddings_tensor, 0, indices)
+
+
+def was_embedding_overwritten_correctly(tim: TextualInversionManager, overwritten_embedding: torch.Tensor, ti_indices: list, ti_embedding: torch.Tensor) -> bool:
+    return torch.allclose(overwritten_embedding[ti_indices], ti_embedding + tim.clip_embedder.position_embedding(ti_indices))
 
 class TextualInversionManagerTestCase(unittest.TestCase):
 
@@ -270,7 +280,7 @@ def test_overwrite_textual_inversion_1v_single(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:4], default_prompt_embeddings[0:4]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[4], test_embedding_1v[0]))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [4], test_embedding_1v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[5:77], default_prompt_embeddings[5:77]))
 
         # at the start
@@ -283,7 +293,7 @@ def test_overwrite_textual_inversion_1v_single(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:1], default_prompt_embeddings[0:1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[1], test_embedding_1v[0]))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [1], test_embedding_1v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[2:77], default_prompt_embeddings[2:77]))
 
         # in the middle
@@ -296,7 +306,7 @@ def test_overwrite_textual_inversion_1v_single(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:2], default_prompt_embeddings[0:2]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[2], test_embedding_1v[0]))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [2], test_embedding_1v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[3:77], default_prompt_embeddings[3:77]))
 
 
@@ -326,8 +336,8 @@ def test_overwrite_textual_inversion_1v_multiple(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:4], default_prompt_embeddings[0:4]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[4], test_embedding_1v_1[0]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[5], test_embedding_1v_2[0]))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [4], test_embedding_1v_1))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [5], test_embedding_1v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[6:77], default_prompt_embeddings[6:77]))
 
         # at the start
@@ -340,8 +350,10 @@ def test_overwrite_textual_inversion_1v_multiple(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:1], default_prompt_embeddings[0:1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[1], test_embedding_1v_1[0]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[2], test_embedding_1v_2[0]))
+        self.assertTrue(
+            was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [1], test_embedding_1v_1))
+        self.assertTrue(
+            was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [2], test_embedding_1v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[3:77], default_prompt_embeddings[3:77]))
 
         # clumped in the middle
@@ -354,8 +366,10 @@ def test_overwrite_textual_inversion_1v_multiple(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:2], default_prompt_embeddings[0:2]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[2], test_embedding_1v_1[0]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[3], test_embedding_1v_2[0]))
+        self.assertTrue(
+            was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [2], test_embedding_1v_1))
+        self.assertTrue(
+            was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [3], test_embedding_1v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[4:77], default_prompt_embeddings[4:77]))
 
         # scattered
@@ -368,9 +382,11 @@ def test_overwrite_textual_inversion_1v_multiple(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:2], default_prompt_embeddings[0:2]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[2], test_embedding_1v_1[0]))
+        self.assertTrue(
+            was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [2], test_embedding_1v_1))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[3], default_prompt_embeddings[3]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[4], test_embedding_1v_2[0]))
+        self.assertTrue(
+            was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings, [4], test_embedding_1v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[5:77], default_prompt_embeddings[5:77]))
 
     def test_overwrite_textual_inversion_4v_single(self):
@@ -393,7 +409,9 @@ def test_overwrite_textual_inversion_4v_single(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:4], default_prompt_embeddings[0:4]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[4:8], test_embedding_4v))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(4,8)),
+                                                            test_embedding_4v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[8:77], default_prompt_embeddings[8:77]))
 
         # at the start
@@ -406,7 +424,9 @@ def test_overwrite_textual_inversion_4v_single(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:1], default_prompt_embeddings[0:1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[1:5], test_embedding_4v))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(1,5)),
+                                                            test_embedding_4v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[5:77], default_prompt_embeddings[5:77]))
 
         # in the middle
@@ -419,7 +439,9 @@ def test_overwrite_textual_inversion_4v_single(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:2], default_prompt_embeddings[0:2]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[2:6], test_embedding_4v))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(2,6)),
+                                                            test_embedding_4v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[6:77], default_prompt_embeddings[6:77]))
 
     def test_overwrite_textual_inversion_4v_overflow(self):
@@ -445,8 +467,11 @@ def test_overwrite_textual_inversion_4v_overflow(self):
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         base_prompt_length = len(base_prompt)
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:base_prompt_length+1], default_prompt_embeddings[0:base_prompt_length+1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[base_prompt_length+1:base_prompt_length+1+3], test_embedding_4v[0:3]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[base_prompt_length+1+3:77], default_prompt_embeddings[base_prompt_length+1+3:77]))
+        truncated_overflowed_overwrite_count = min(75 - len(base_prompt), test_embedding_4v.shape[0])
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(base_prompt_length+1,base_prompt_length+1+truncated_overflowed_overwrite_count)),
+                                                            test_embedding_4v[0:truncated_overflowed_overwrite_count]))
+        self.assertTrue(torch.equal(overwritten_prompt_embeddings[base_prompt_length+1+4:77], default_prompt_embeddings[base_prompt_length+1+4:77]))
 
         # at the start
         prompt_token_ids = [test_embedding_4v_token_id] + base_prompt
@@ -459,7 +484,9 @@ def test_overwrite_textual_inversion_4v_overflow(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:1], default_prompt_embeddings[0:1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[1:5], test_embedding_4v))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(1,5)),
+                                                            test_embedding_4v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[5:77], default_prompt_embeddings[5:77]))
 
         # in the middle
@@ -472,7 +499,9 @@ def test_overwrite_textual_inversion_4v_overflow(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:21], default_prompt_embeddings[0:21]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[21:25], test_embedding_4v))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(21,25)),
+                                                            test_embedding_4v))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[25:77], default_prompt_embeddings[25:77]))
 
 
@@ -504,8 +533,12 @@ def test_overwrite_textual_inversion_4v_multiple(self):
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         base_prompt_length = len(base_prompt)
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:base_prompt_length+1], default_prompt_embeddings[0:base_prompt_length+1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[base_prompt_length+1:base_prompt_length+1+4], test_embedding_4v_1))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[base_prompt_length+1+4:base_prompt_length+1+4+4], test_embedding_4v_2))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(base_prompt_length+1, base_prompt_length+1+4)),
+                                                            test_embedding_4v_1))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(base_prompt_length+1+4, base_prompt_length+1+4+4)),
+                                                            test_embedding_4v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[base_prompt_length+1+4+4:77], default_prompt_embeddings[base_prompt_length+1+4+4:77]))
 
         # at the start
@@ -519,8 +552,12 @@ def test_overwrite_textual_inversion_4v_multiple(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:1], default_prompt_embeddings[0:1]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[1:5], test_embedding_4v_1))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[5:9], test_embedding_4v_2))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(1,5)),
+                                                            test_embedding_4v_1))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(5,9)),
+                                                            test_embedding_4v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[9:77], default_prompt_embeddings[9:77]))
 
         # in the middle
@@ -533,7 +570,11 @@ def test_overwrite_textual_inversion_4v_multiple(self):
         overwritten_prompt_embeddings = tim.overwrite_textual_inversion_embeddings(padded_prompt_token_ids, default_prompt_embeddings)
         self.assertFalse(torch.equal(default_prompt_embeddings, overwritten_prompt_embeddings))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[0:11], default_prompt_embeddings[0:11]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[11:15], test_embedding_4v_1))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(11,15)),
+                                                            test_embedding_4v_1))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[15:25], default_prompt_embeddings[15:25]))
-        self.assertTrue(torch.equal(overwritten_prompt_embeddings[25:29], test_embedding_4v_2))
+        self.assertTrue(was_embedding_overwritten_correctly(tim, overwritten_prompt_embeddings,
+                                                            list(range(25,29)),
+                                                            test_embedding_4v_2))
         self.assertTrue(torch.equal(overwritten_prompt_embeddings[29:77], default_prompt_embeddings[29:77]))