sd3 TI support

modules/sd_hijack.py

@@ -359,13 +359,28 @@ class EmbeddingsWithFixes(torch.nn.Module):
                 vec = embedding.vec[self.textual_inversion_key] if isinstance(embedding.vec, dict) else embedding.vec
                 emb = devices.cond_cast_unet(vec)
                 emb_len = min(tensor.shape[0] - offset - 1, emb.shape[0])
-                tensor = torch.cat([tensor[0:offset + 1], emb[0:emb_len], tensor[offset + 1 + emb_len:]])
+                tensor = torch.cat([tensor[0:offset + 1], emb[0:emb_len], tensor[offset + 1 + emb_len:]]).to(dtype=inputs_embeds.dtype)
 
             vecs.append(tensor)
 
         return torch.stack(vecs)
 
 
+class TextualInversionEmbeddings(torch.nn.Embedding):
+    def __init__(self, num_embeddings: int, embedding_dim: int, textual_inversion_key='clip_l', **kwargs):
+        super().__init__(num_embeddings, embedding_dim, **kwargs)
+
+        self.embeddings = model_hijack
+        self.textual_inversion_key = textual_inversion_key
+
+    @property
+    def wrapped(self):
+        return super().forward
+
+    def forward(self, input_ids):
+        return EmbeddingsWithFixes.forward(self, input_ids)
+
+
 def add_circular_option_to_conv_2d():
     conv2d_constructor = torch.nn.Conv2d.__init__