make it possible for extensions/scripts to add their own embedding directories

1 files changed, 4 insertions, 3 deletions

diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py
index cfdb09d6..6b0d95af 100644
--- a/modules/sd_hijack.py
+++ b/modules/sd_hijack.py
@@ -83,10 +83,12 @@ class StableDiffusionModelHijack:
     clip = None

     optimization_method = None

 

-    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase(cmd_opts.embeddings_dir)

+    embedding_db = modules.textual_inversion.textual_inversion.EmbeddingDatabase()

 

-    def hijack(self, m):

+    def __init__(self):

+        self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir)

 

+    def hijack(self, m):

         if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

             model_embeddings = m.cond_stage_model.roberta.embeddings

             model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.word_embeddings, self)

@@ -117,7 +119,6 @@ class StableDiffusionModelHijack:
         self.layers = flatten(m)

 

     def undo_hijack(self, m):

-

         if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation:

             m.cond_stage_model = m.cond_stage_model.wrapped