1 files changed, 60 insertions, 1 deletions
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 5965c5a0..7a24192e 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -7,10 +7,36 @@ import tqdm
 import html
 import datetime
 
+from PIL import Image,PngImagePlugin
+from ..images import captionImageOverlay
+import numpy as np
+import base64
+import json
 
 from modules import shared, devices, sd_hijack, processing, sd_models
 import modules.textual_inversion.dataset
 
+class EmbeddingEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, torch.Tensor):
+            return {'TORCHTENSOR':obj.cpu().detach().numpy().tolist()}
+        return json.JSONEncoder.default(self, o)
+
+class EmbeddingDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
+    def object_hook(self, d):
+        if 'TORCHTENSOR' in d:
+            return torch.from_numpy(np.array(d['TORCHTENSOR']))
+        return d
+
+def embeddingToB64(data):
+    d = json.dumps(data,cls=EmbeddingEncoder)
+    return base64.b64encode(d.encode())
+
+def embeddingFromB64(data):
+    d = base64.b64decode(data)
+    return json.loads(d,cls=EmbeddingDecoder)
 
 class Embedding:
     def __init__(self, vec, name, step=None):
@@ -80,7 +106,15 @@ class EmbeddingDatabase:
         def process_file(path, filename):
             name = os.path.splitext(filename)[0]
 
-            data = torch.load(path, map_location="cpu")
+            data = []
+
+            if filename.upper().endswith('.PNG'):
+                embed_image = Image.open(path)
+                if 'sd-ti-embedding' in embed_image.text:
+                    data = embeddingFromB64(embed_image.text['sd-ti-embedding'])
+                    name = data.get('name',name)
+            else:
+                data = torch.load(path, map_location="cpu")
 
             # textual inversion embeddings
             if 'string_to_param' in data:
@@ -178,6 +212,12 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
     else:
         images_dir = None
 
+    if create_image_every > 0 and save_image_with_stored_embedding:
+        images_embeds_dir = os.path.join(log_directory, "image_embeddings")
+        os.makedirs(images_embeds_dir, exist_ok=True)
+    else:
+        images_embeds_dir = None
+        
     cond_model = shared.sd_model.cond_stage_model
 
     shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
@@ -252,6 +292,25 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
             image = processed.images[0]
 
             shared.state.current_image = image
+
+            if save_image_with_stored_embedding and os.path.exists(last_saved_file):
+                
+                last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{embedding.step}.png')
+
+                info = PngImagePlugin.PngInfo()
+                data = torch.load(last_saved_file)
+                info.add_text("sd-ti-embedding", embeddingToB64(data))
+
+                title = "<{}>".format(data.get('name','???'))
+                checkpoint = sd_models.select_checkpoint()
+                footer_left = checkpoint.model_name
+                footer_mid = '[{}]'.format(checkpoint.hash)
+                footer_right = '{}'.format(embedding.step)
+
+                captioned_image = captionImageOverlay(image,title,footer_left,footer_mid,footer_right)
+
+                captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info)
+            
             image.save(last_saved_image)
 
             last_saved_image += f", prompt: {text}"