diff --git a/README.jp.md b/README.jp.md index e1babd9..f529b86 100644 --- a/README.jp.md +++ b/README.jp.md @@ -41,9 +41,24 @@ git clone https://github.com/nkchocoai/ComfyUI-SaveImageWithMetaData.git - Size - Model - Model hash +- VAE + - KSamplerノードではなくSaveImageWithMetadataノードの入力から参照されます。 +- VAE hash + - KSamplerノードではなくSaveImageWithMetadataノードの入力から参照されます。 +- Loras + - Model name + - Model hash + - Strength model + - Strength clip +- Embeddings + - Name + - Hash - batch size >= 2の場合 : - Batch index - Batch size +- Hashes + - Model, Loras, Embeddings + - [Civitai](https://civitai.com/)用 ## 対応しているノード・拡張機能 - 対応しているノードは以下のファイルをご確認ください。 diff --git a/README.md b/README.md index f3a1b70..574bfa4 100644 --- a/README.md +++ b/README.md @@ -43,9 +43,24 @@ git clone https://github.com/nkchocoai/ComfyUI-SaveImageWithMetaData.git - Size - Model - Model hash +- VAE + - It is referenced from the input of SaveImageWithMetadata node, not KSampler node. +- VAE hash + - It is referenced from the input of SaveImageWithMetadata node, not KSampler node. +- Loras + - Model name + - Model hash + - Strength model + - Strength clip +- Embeddings + - Name + - Hash - If batch size >= 2 : - Batch index - Batch size +- Hashes + - Model, Loras, Embeddings + - For [Civitai](https://civitai.com/) ## Supported nodes and extensions - Please check the following file for supported nodes. diff --git a/examples/lora_embedding_vae.json b/examples/lora_embedding_vae.json new file mode 100644 index 0000000..78ab003 --- /dev/null +++ b/examples/lora_embedding_vae.json @@ -0,0 +1,876 @@ +{ + "last_node_id": 19, + "last_link_id": 22, + "nodes": [ + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [ + 24, + 54 + ], + "size": { + "0": 315, + "1": 98 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 2 + ], + "shape": 3 + }, + { + "name": "VAE", + "type": "VAE", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "CounterfeitV30_v30.safetensors" + ] + }, + { + "id": 5, + "type": "Anything Everywhere3", + "pos": [ + 1143, + 52 + ], + "size": { + "0": 210, + "1": 66 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "MODEL", + "type": "*", + "link": 5, + "slot_index": 0, + "color_on": "#B39DDB" + }, + { + "name": "CLIP", + "type": "*", + "link": 6, + "color_on": "#FFD500" + }, + { + "name": "anything", + "type": "*", + "link": null + } + ], + "properties": { + "group_restricted": false, + "color_restricted": false, + "Node name for S&R": "Anything Everywhere3" + } + }, + { + "id": 7, + "type": "PromptUtilitiesConstStringMultiLine", + "pos": [ + 21, + 378 + ], + "size": { + "0": 287.9851379394531, + "1": 99.69133758544922 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 7 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "PromptUtilitiesConstStringMultiLine" + }, + "widgets_values": [ + "(low quality, worst quality:1.4), embedding:EasyNegative" + ] + }, + { + "id": 9, + "type": "CLIPTextEncode", + "pos": [ + 394, + 382 + ], + "size": { + "0": 210, + "1": 54 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": null, + "slot_index": 0 + }, + { + "name": "text", + "type": "STRING", + "link": 7, + "widget": { + "name": "text" + }, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 10 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 8, + "type": "CLIPTextEncode", + "pos": [ + 390, + 257 + ], + "size": { + "0": 210, + "1": 54 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": null + }, + { + "name": "text", + "type": "STRING", + "link": 8, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 9 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 10, + "type": "Prompts Everywhere", + "pos": [ + 674, + 302 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "CONDITIONING", + "type": "*", + "link": 9, + "color_on": "#FFA931" + }, + { + "name": "CONDITIONING", + "type": "*", + "link": 10, + "slot_index": 1, + "color_on": "#FFA931" + } + ], + "properties": { + "group_restricted": false, + "color_restricted": false, + "Node name for S&R": "Prompts Everywhere" + } + }, + { + "id": 12, + "type": "EmptyLatentImageFromPresetsSD15", + "pos": [ + 23, + 550 + ], + "size": { + "0": 327.6000061035156, + "1": 122 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "latent", + "type": "LATENT", + "links": [ + 11 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "w", + "type": "INT", + "links": null, + "shape": 3 + }, + { + "name": "h", + "type": "INT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImageFromPresetsSD15" + }, + "widgets_values": [ + "512 x 768", + 1 + ] + }, + { + "id": 14, + "type": "SaveImage", + "pos": [ + 1840, + 260 + ], + "size": { + "0": 315, + "1": 270 + }, + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 14, + "slot_index": 0 + } + ], + "properties": {}, + "widgets_values": [ + "without-meta" + ] + }, + { + "id": 17, + "type": "LatentUpscaleBy", + "pos": [ + 772, + 532 + ], + "size": { + "0": 262.8999938964844, + "1": 82 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 18, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 19 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LatentUpscaleBy" + }, + "widgets_values": [ + "nearest-exact", + 2 + ] + }, + { + "id": 15, + "type": "SaveImageWithMetaData", + "pos": [ + 1840, + 580 + ], + "size": { + "0": 315, + "1": 318 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 15, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "SaveImageWithMetaData" + }, + "widgets_values": [ + "with-meta", + "Farthest", + 0 + ] + }, + { + "id": 11, + "type": "KSampler", + "pos": [ + 402, + 532 + ], + "size": { + "0": 315, + "1": 262 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": null + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": null + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": null + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 11 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 18 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 123, + "fixed", + 20, + 8, + "dpmpp_2m", + "karras", + 1 + ] + }, + { + "id": 13, + "type": "VAEDecode", + "pos": [ + 1501, + 530 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 17, + "slot_index": 0 + }, + { + "name": "vae", + "type": "VAE", + "link": 13 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 14, + 15 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 16, + "type": "KSampler", + "pos": [ + 1090, + 470 + ], + "size": { + "0": 315, + "1": 262 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": null + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": null + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": null + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 19 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 17 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 123, + "fixed", + 20, + 8, + "dpmpp_2m", + "karras", + 0.6 + ] + }, + { + "id": 4, + "type": "LoraLoader", + "pos": [ + 761, + 52 + ], + "size": { + "0": 340.8315734863281, + "1": 126 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 3, + "slot_index": 0 + }, + { + "name": "clip", + "type": "CLIP", + "link": 4 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 5 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 6 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoraLoader" + }, + "widgets_values": [ + "leco\\leco_cyborg_last.safetensors", + -3, + -3 + ] + }, + { + "id": 2, + "type": "VAELoader", + "pos": [ + 1090, + 810 + ], + "size": { + "0": 323.17578125, + "1": 58 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 13 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "clearvae_nanslesstest.safetensors" + ] + }, + { + "id": 6, + "type": "PromptUtilitiesConstStringMultiLine", + "pos": [ + 18, + 225 + ], + "size": { + "0": 287.9851379394531, + "1": 99.69133758544922 + }, + "flags": {}, + "order": 4, + "mode": 0, + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 8 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "PromptUtilitiesConstStringMultiLine" + }, + "widgets_values": [ + "masterpiece, best quality, hatsune miku, upper body" + ] + }, + { + "id": 3, + "type": "LoraLoader", + "pos": [ + 397, + 50 + ], + "size": { + "0": 315, + "1": 126 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1 + }, + { + "name": "clip", + "type": "CLIP", + "link": 2, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 3 + ], + "shape": 3 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 4 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoraLoader" + }, + "widgets_values": [ + "util\\flat2.safetensors", + -1, + -1 + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 3, + 0, + "MODEL" + ], + [ + 2, + 1, + 1, + 3, + 1, + "CLIP" + ], + [ + 3, + 3, + 0, + 4, + 0, + "MODEL" + ], + [ + 4, + 3, + 1, + 4, + 1, + "CLIP" + ], + [ + 5, + 4, + 0, + 5, + 0, + "*" + ], + [ + 6, + 4, + 1, + 5, + 1, + "*" + ], + [ + 7, + 7, + 0, + 9, + 1, + "STRING" + ], + [ + 8, + 6, + 0, + 8, + 1, + "STRING" + ], + [ + 9, + 8, + 0, + 10, + 0, + "*" + ], + [ + 10, + 9, + 0, + 10, + 1, + "*" + ], + [ + 11, + 12, + 0, + 11, + 3, + "LATENT" + ], + [ + 13, + 2, + 0, + 13, + 1, + "VAE" + ], + [ + 14, + 13, + 0, + 14, + 0, + "IMAGE" + ], + [ + 15, + 13, + 0, + 15, + 0, + "IMAGE" + ], + [ + 17, + 16, + 0, + 13, + 0, + "LATENT" + ], + [ + 18, + 11, + 0, + 17, + 0, + "LATENT" + ], + [ + 19, + 17, + 0, + 16, + 3, + "LATENT" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file diff --git a/examples/lora_embedding_vae.png b/examples/lora_embedding_vae.png new file mode 100644 index 0000000..76baffc Binary files /dev/null and b/examples/lora_embedding_vae.png differ diff --git a/py/capture.py b/py/capture.py index 4956950..795420b 100644 --- a/py/capture.py +++ b/py/capture.py @@ -1,3 +1,6 @@ +import json +import os + from . import hook from .defs.captures import CAPTURE_FIELD_LIST from .defs.meta import MetaField @@ -32,6 +35,12 @@ class Capture: if meta not in inputs: inputs[meta] = [] + + value = field_data.get("value") + if value is not None: + inputs[meta].append((node_id, value)) + continue + field_name = field_data["field_name"] value = input_data.get(field_name) if value is not None: @@ -41,45 +50,69 @@ 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 - def gen_pnginfo_dict(cls, inputs): + def gen_pnginfo_dict(cls, inputs_before_sampler_node, inputs_before_this_node): pnginfo_dict = {} - def update_pnginfo_dict(metafield, key): + def update_pnginfo_dict(inputs, metafield, key): x = inputs.get(metafield, []) if len(x) > 0: pnginfo_dict[key] = x[0][1] - update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "Positive prompt") - update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt") + update_pnginfo_dict( + inputs_before_sampler_node, MetaField.POSITIVE_PROMPT, "Positive prompt" + ) + update_pnginfo_dict( + inputs_before_sampler_node, MetaField.NEGATIVE_PROMPT, "Negative prompt" + ) - update_pnginfo_dict(MetaField.STEPS, "Steps") + update_pnginfo_dict(inputs_before_sampler_node, MetaField.STEPS, "Steps") - sampler_names = inputs.get(MetaField.SAMPLER_NAME, []) + sampler_names = inputs_before_sampler_node.get(MetaField.SAMPLER_NAME, []) if len(sampler_names) > 0: pnginfo_dict["Sampler"] = sampler_names[0][1] - schedulers = inputs.get(MetaField.SCHEDULER, []) + schedulers = inputs_before_sampler_node.get(MetaField.SCHEDULER, []) if len(schedulers) > 0: scheduler = schedulers[0][1] if scheduler != "normal": pnginfo_dict["Sampler"] += "_" + scheduler - update_pnginfo_dict(MetaField.CFG, "CFG Scale") - update_pnginfo_dict(MetaField.SEED, "Seed") + update_pnginfo_dict(inputs_before_sampler_node, MetaField.CFG, "CFG Scale") + update_pnginfo_dict(inputs_before_sampler_node, MetaField.SEED, "Seed") - update_pnginfo_dict(MetaField.CLIP_SKIP, "Clip skip") + update_pnginfo_dict( + inputs_before_sampler_node, MetaField.CLIP_SKIP, "Clip skip" + ) - image_widths = inputs.get(MetaField.IMAGE_WIDTH, []) - image_heights = inputs.get(MetaField.IMAGE_HEIGHT, []) + image_widths = inputs_before_sampler_node.get(MetaField.IMAGE_WIDTH, []) + image_heights = inputs_before_sampler_node.get(MetaField.IMAGE_HEIGHT, []) if len(image_widths) > 0 and len(image_heights) > 0: pnginfo_dict["Size"] = f"{image_widths[0][1]}x{image_heights[0][1]}" - update_pnginfo_dict(MetaField.MODEL_NAME, "Model") - update_pnginfo_dict(MetaField.MODEL_HASH, "Model hash") + update_pnginfo_dict(inputs_before_sampler_node, MetaField.MODEL_NAME, "Model") + update_pnginfo_dict( + inputs_before_sampler_node, MetaField.MODEL_HASH, "Model hash" + ) + + update_pnginfo_dict(inputs_before_this_node, MetaField.VAE_NAME, "VAE") + update_pnginfo_dict(inputs_before_this_node, MetaField.VAE_HASH, "VAE hash") + + pnginfo_dict.update(cls.gen_loras(inputs_before_sampler_node)) + pnginfo_dict.update(cls.gen_embeddings(inputs_before_sampler_node)) + + hashes_for_civitai = cls.get_hashes_for_civitai( + inputs_before_sampler_node, inputs_before_this_node + ) + if len(hashes_for_civitai) > 0: + pnginfo_dict["Hashes"] = json.dumps(hashes_for_civitai) return pnginfo_dict @@ -99,3 +132,72 @@ class Capture: s_list.append(f"{k}: {s}") return result + ", ".join(s_list) + + @classmethod + def get_hashes_for_civitai( + cls, inputs_before_sampler_node, inputs_before_this_node + ): + resource_hashes = {} + model_hashes = inputs_before_sampler_node.get(MetaField.MODEL_HASH, []) + if len(model_hashes) > 0: + resource_hashes["model"] = model_hashes[0][1] + + vae_hashes = inputs_before_this_node.get(MetaField.VAE_HASH, []) + if len(vae_hashes) > 0: + resource_hashes["vae"] = vae_hashes[0][1] + + lora_model_names = inputs_before_sampler_node.get(MetaField.LORA_MODEL_NAME, []) + lora_model_hashes = inputs_before_sampler_node.get( + MetaField.LORA_MODEL_HASH, [] + ) + for lora_model_name, lora_model_hash in zip( + lora_model_names, lora_model_hashes + ): + 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_before_sampler_node.get(MetaField.EMBEDDING_NAME, []) + embedding_hashes = inputs_before_sampler_node.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 + def gen_loras(cls, inputs): + pnginfo_dict = {} + + model_names = inputs.get(MetaField.LORA_MODEL_NAME, []) + model_hashes = inputs.get(MetaField.LORA_MODEL_HASH, []) + strength_models = inputs.get(MetaField.LORA_STRENGTH_MODEL, []) + strength_clips = inputs.get(MetaField.LORA_STRENGTH_CLIP, []) + + index = 0 + for model_name, model_hashe, strength_model, strength_clip in zip( + model_names, model_hashes, strength_models, strength_clips + ): + field_prefix = f"Lora_{index}" + pnginfo_dict[f"{field_prefix} Model name"] = os.path.basename(model_name[1]) + pnginfo_dict[f"{field_prefix} Model hash"] = model_hashe[1] + pnginfo_dict[f"{field_prefix} Strength model"] = strength_model[1] + pnginfo_dict[f"{field_prefix} Strength clip"] = strength_clip[1] + 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 diff --git a/py/defs/captures.py b/py/defs/captures.py index 0cd0e0c..d842adb 100644 --- a/py/defs/captures.py +++ b/py/defs/captures.py @@ -2,9 +2,13 @@ from .meta import MetaField from .validators import is_positive_prompt, is_negative_prompt from .formatters import ( calc_model_hash, + calc_vae_hash, + calc_lora_hash, convert_skip_clip, get_scaled_width, get_scaled_height, + extract_embedding_names, + extract_embedding_hashes, ) @@ -19,6 +23,10 @@ CAPTURE_FIELD_LIST = { "format": convert_skip_clip, }, }, + "VAELoader": { + MetaField.VAE_NAME: {"field_name": "vae_name"}, + MetaField.VAE_HASH: {"field_name": "vae_name", "format": calc_vae_hash}, + }, "EmptyLatentImage": { MetaField.IMAGE_WIDTH: {"field_name": "width"}, MetaField.IMAGE_HEIGHT: {"field_name": "height"}, @@ -32,6 +40,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"}, @@ -58,4 +74,22 @@ CAPTURE_FIELD_LIST = { "format": get_scaled_height, }, }, + "LoraLoader": { + MetaField.LORA_MODEL_NAME: {"field_name": "lora_name"}, + MetaField.LORA_MODEL_HASH: { + "field_name": "lora_name", + "format": calc_lora_hash, + }, + MetaField.LORA_STRENGTH_MODEL: {"field_name": "strength_model"}, + MetaField.LORA_STRENGTH_CLIP: {"field_name": "strength_clip"}, + }, + "LoraLoaderModelOnly": { + MetaField.LORA_MODEL_NAME: {"field_name": "lora_name"}, + MetaField.LORA_MODEL_HASH: { + "field_name": "lora_name", + "format": calc_lora_hash, + }, + MetaField.LORA_STRENGTH_MODEL: {"field_name": "strength_model"}, + MetaField.LORA_STRENGTH_CLIP: {"value": 0}, + }, } diff --git a/py/defs/formatters.py b/py/defs/formatters.py index 4a183b4..0c86031 100644 --- a/py/defs/formatters.py +++ b/py/defs/formatters.py @@ -1,24 +1,31 @@ -import hashlib +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.sd1_clip import SD1Tokenizer +from comfy.sd2_clip import SD2Tokenizer +from comfy.sdxl_clip import SDXLTokenizer + cache_model_hash = {} def calc_model_hash(model_name, input_data): - if model_name in cache_model_hash: - return cache_model_hash[model_name] - filename = folder_paths.get_full_path("checkpoints", model_name) - sha256_hash = hashlib.sha256() + return calc_hash(filename) - with open(filename, "rb") as f: - for byte_block in iter(lambda: f.read(4096), b""): - sha256_hash.update(byte_block) - model_hash = sha256_hash.hexdigest()[:10] - cache_model_hash[model_name] = model_hash - return model_hash +def calc_vae_hash(model_name, input_data): + filename = folder_paths.get_full_path("vae", model_name) + return calc_hash(filename) + + +def calc_lora_hash(model_name, input_data): + filename = folder_paths.get_full_path("loras", model_name) + return calc_hash(filename) def convert_skip_clip(stop_at_clip_layer, input_data): @@ -33,3 +40,49 @@ 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, SD1Tokenizer): + tokenizer = tokenizer.clip_l + elif isinstance(tokenizer, SD2Tokenizer): + tokenizer = tokenizer.clip_h + elif 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 diff --git a/py/defs/meta.py b/py/defs/meta.py index b18e2c9..826a854 100644 --- a/py/defs/meta.py +++ b/py/defs/meta.py @@ -4,6 +4,8 @@ from enum import IntEnum class MetaField(IntEnum): MODEL_NAME = 0 MODEL_HASH = 1 + VAE_NAME = 2 + VAE_HASH = 3 POSITIVE_PROMPT = 10 NEGATIVE_PROMPT = 11 CLIP_SKIP = 12 @@ -14,3 +16,9 @@ class MetaField(IntEnum): SCHEDULER = 24 IMAGE_WIDTH = 30 IMAGE_HEIGHT = 31 + EMBEDDING_NAME = 40 + EMBEDDING_HASH = 41 + LORA_MODEL_NAME = 50 + LORA_MODEL_HASH = 51 + LORA_STRENGTH_MODEL = 52 + LORA_STRENGTH_CLIP = 53 diff --git a/py/nodes/node.py b/py/nodes/node.py index f5fc5c5..b2a3f09 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -110,6 +110,9 @@ class SaveImageWithMetaData(BaseNode): trace_tree_from_this_node = Trace.trace( hook.current_node_id, hook.current_prompt ) + inputs_before_this_node = Trace.filter_inputs_by_trace_tree( + inputs, trace_tree_from_this_node + ) sampler_node_id = Trace.find_sampler_node_id( trace_tree_from_this_node, sampler_selection_method, @@ -118,10 +121,12 @@ class SaveImageWithMetaData(BaseNode): # get inputs before sampler node trace_tree_from_sampler_node = Trace.trace(sampler_node_id, hook.current_prompt) - filtered_inputs = Trace.filter_inputs_by_trace_tree( + inputs_before_sampler_node = Trace.filter_inputs_by_trace_tree( inputs, trace_tree_from_sampler_node ) # generate PNGInfo from inputs - pnginfo_dict = Capture.gen_pnginfo_dict(filtered_inputs) + pnginfo_dict = Capture.gen_pnginfo_dict( + inputs_before_sampler_node, inputs_before_this_node + ) return pnginfo_dict diff --git a/py/utils/embedding.py b/py/utils/embedding.py new file mode 100644 index 0000000..70c1f3a --- /dev/null +++ b/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 diff --git a/py/utils/hash.py b/py/utils/hash.py new file mode 100644 index 0000000..6cae680 --- /dev/null +++ b/py/utils/hash.py @@ -0,0 +1,18 @@ +import hashlib + + +cache_model_hash = {} + + +def calc_hash(filename): + if filename in cache_model_hash: + return cache_model_hash[filename] + sha256_hash = hashlib.sha256() + + with open(filename, "rb") as f: + for byte_block in iter(lambda: f.read(4096), b""): + sha256_hash.update(byte_block) + model_hash = sha256_hash.hexdigest()[:10] + + cache_model_hash[filename] = model_hash + return model_hash