Support embedding

py/capture.py

@@ -50,7 +50,11 @@ class Capture:
                                 v = value[0]
                             if format is not None:
                                 v = format(v, input_data)
-                            inputs[meta].append((node_id, v))
+                            if isinstance(v, list):
+                                for x in v:
+                                    inputs[meta].append((node_id, x))
+                            else:
+                                inputs[meta].append((node_id, v))
         return inputs
 
     @classmethod
@@ -91,6 +95,7 @@ class Capture:
         update_pnginfo_dict(MetaField.MODEL_HASH, "Model hash")
 
         pnginfo_dict.update(cls.gen_loras(inputs))
+        pnginfo_dict.update(cls.gen_embeddings(inputs))
 
         hashes_for_civitai = cls.get_hashes_for_civitai(inputs)
         if len(hashes_for_civitai) > 0:
@@ -130,6 +135,12 @@ class Capture:
             lora_model_name = os.path.splitext(os.path.basename(lora_model_name[1]))[0]
             resource_hashes[f"lora:{lora_model_name}"] = lora_model_hash[1]
 
+        embedding_names = inputs.get(MetaField.EMBEDDING_NAME, [])
+        embedding_hashes = inputs.get(MetaField.EMBEDDING_HASH, [])
+        for embedding_name, embedding_hash in zip(embedding_names, embedding_hashes):
+            embedding_name = os.path.splitext(os.path.basename(embedding_name[1]))[0]
+            resource_hashes[f"embed:{embedding_name}"] = embedding_hash[1]
+
         return resource_hashes
 
     @classmethod
@@ -153,3 +164,19 @@ class Capture:
             index += 1
 
         return pnginfo_dict
+
+    @classmethod
+    def gen_embeddings(cls, inputs):
+        pnginfo_dict = {}
+
+        embedding_names = inputs.get(MetaField.EMBEDDING_NAME, [])
+        embedding_hashes = inputs.get(MetaField.EMBEDDING_HASH, [])
+
+        index = 0
+        for embedding_name, embedding_hashe in zip(embedding_names, embedding_hashes):
+            field_prefix = f"Embedding_{index}"
+            pnginfo_dict[f"{field_prefix} name"] = os.path.basename(embedding_name[1])
+            pnginfo_dict[f"{field_prefix} hash"] = embedding_hashe[1]
+            index += 1
+
+        return pnginfo_dict

py/defs/captures.py

@@ -6,6 +6,8 @@ from .formatters import (
     convert_skip_clip,
     get_scaled_width,
     get_scaled_height,
+    extract_embedding_names,
+    extract_embedding_hashes,
 )
 
 
@@ -33,6 +35,14 @@ CAPTURE_FIELD_LIST = {
             "field_name": "text",
             "validate": is_negative_prompt,
         },
+        MetaField.EMBEDDING_NAME: {
+            "field_name": "text",
+            "format": extract_embedding_names,
+        },
+        MetaField.EMBEDDING_HASH: {
+            "field_name": "text",
+            "format": extract_embedding_hashes,
+        },
     },
     "KSampler": {
         MetaField.SEED: {"field_name": "seed"},

py/defs/formatters.py

@@ -1,6 +1,12 @@
+import os
+
 import folder_paths
 
 from ..utils.hash import calc_hash
+from ..utils.embedding import get_embedding_file_path
+
+from comfy.sd1_clip import escape_important, token_weights, unescape_important
+from comfy.sdxl_clip import SDXLTokenizer
 
 cache_model_hash = {}
 
@@ -27,3 +33,45 @@ def get_scaled_width(scaled_by, input_data):
 def get_scaled_height(scaled_by, input_data):
     samples = input_data["samples"][0]["samples"]
     return round(samples.shape[2] * scaled_by * 8)
+
+
+def extract_embedding_names(text, input_data):
+    embedding_names, _ = _extract_embedding_names(text, input_data)
+
+    return [os.path.basename(embedding_name) for embedding_name in embedding_names]
+
+
+def extract_embedding_hashes(text, input_data):
+    embedding_names, tokenizer = _extract_embedding_names(text, input_data)
+    embedding_hashes = []
+    for embedding_name in embedding_names:
+        embedding_file_path = get_embedding_file_path(embedding_name, tokenizer)
+        embedding_hashes.append(calc_hash(embedding_file_path))
+
+    return embedding_hashes
+
+
+def _extract_embedding_names(text, input_data):
+    clip = input_data["clip"][0]
+    tokenizer = clip.tokenizer
+    if isinstance(tokenizer, SDXLTokenizer):
+        tokenizer = tokenizer.clip_l
+    text = escape_important(text)
+    parsed_weights = token_weights(text, 1.0)
+
+    # tokenize words
+    embedding_names = []
+    for weighted_segment, weight in parsed_weights:
+        to_tokenize = unescape_important(weighted_segment).replace("\n", " ").split(" ")
+        to_tokenize = [x for x in to_tokenize if x != ""]
+        for word in to_tokenize:
+            # if we find an embedding, deal with the embedding
+            if (
+                word.startswith(tokenizer.embedding_identifier)
+                and tokenizer.embedding_directory is not None
+            ):
+                embedding_name = word[len(tokenizer.embedding_identifier) :].strip("\n")
+
+                embedding_names.append(embedding_name)
+
+    return embedding_names, tokenizer

py/defs/meta.py

@@ -15,6 +15,7 @@ class MetaField(IntEnum):
     IMAGE_WIDTH = 30
     IMAGE_HEIGHT = 31
     EMBEDDING_NAME = 40
+    EMBEDDING_HASH = 41
     LORA_MODEL_NAME = 50
     LORA_MODEL_HASH = 51
     LORA_STRENGTH_MODEL = 52

py/utils/embedding.py

@@ -0,0 +1,35 @@
+import os
+
+from comfy.sd1_clip import expand_directory_list
+
+
+# refer. https://github.com/comfyanonymous/ComfyUI/blob/dccca1daa5af1954d55918f365e83a3331019549/comfy/sd1_clip.py#L284-L309
+def get_embedding_file_path(embedding_name, tokenizer):
+    embedding_directory = tokenizer.embedding_directory
+    if isinstance(embedding_directory, str):
+        embedding_directory = [embedding_directory]
+
+    embedding_directory = expand_directory_list(embedding_directory)
+
+    valid_file = None
+    for embed_dir in embedding_directory:
+        embed_path = os.path.abspath(os.path.join(embed_dir, embedding_name))
+        embed_dir = os.path.abspath(embed_dir)
+        try:
+            if os.path.commonpath((embed_dir, embed_path)) != embed_dir:
+                continue
+        except Exception:
+            continue
+        if not os.path.isfile(embed_path):
+            extensions = [".safetensors", ".pt", ".bin"]
+            for x in extensions:
+                t = embed_path + x
+                if os.path.isfile(t):
+                    valid_file = t
+                    break
+        else:
+            valid_file = embed_path
+        if valid_file is not None:
+            break
+
+    return valid_file