add checkpoint info to saved embeddings

1 files changed, 12 insertions, 1 deletions

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 9d2241ce..1183aab7 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -7,7 +7,7 @@ import tqdm
 import html

 import datetime

 

-from modules import shared, devices, sd_hijack, processing

+from modules import shared, devices, sd_hijack, processing, sd_models

 import modules.textual_inversion.dataset

 

 

@@ -17,6 +17,8 @@ class Embedding:
         self.name = name

         self.step = step

         self.cached_checksum = None

+        self.sd_checkpoint = None

+        self.sd_checkpoint_name = None

 

     def save(self, filename):

         embedding_data = {

@@ -24,6 +26,8 @@ class Embedding:
             "string_to_param": {"*": self.vec},

             "name": self.name,

             "step": self.step,

+            "sd_checkpoint": self.sd_checkpoint,

+            "sd_checkpoint_name": self.sd_checkpoint_name,

         }

 

         torch.save(embedding_data, filename)

@@ -41,6 +45,7 @@ class Embedding:
         self.cached_checksum = f'{const_hash(self.vec.reshape(-1) * 100) & 0xffff:04x}'

         return self.cached_checksum

 

+

 class EmbeddingDatabase:

     def __init__(self, embeddings_dir):

         self.ids_lookup = {}

@@ -96,6 +101,8 @@ class EmbeddingDatabase:
             vec = emb.detach().to(devices.device, dtype=torch.float32)

             embedding = Embedding(vec, name)

             embedding.step = data.get('step', None)

+            embedding.sd_checkpoint = data.get('hash', None)

+            embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)

             self.register_embedding(embedding, shared.sd_model)

 

         for fn in os.listdir(self.embeddings_dir):

@@ -249,6 +256,10 @@ Last saved image: {html.escape(last_saved_image)}<br/>
 </p>

 """

 

+    checkpoint = sd_models.select_checkpoint()

+

+    embedding.sd_checkpoint = checkpoint.hash

+    embedding.sd_checkpoint_name = checkpoint.model_name

     embedding.cached_checksum = None

     embedding.save(filename)