From d2f7a1b993c19b7b08e834eca8d53c010696747b Mon Sep 17 00:00:00 2001 From: nkchocoai Date: Sun, 18 Feb 2024 23:06:49 +0900 Subject: [PATCH] Support VAE --- py/capture.py | 67 ++++++++++++++++++++++++++++--------------- py/defs/captures.py | 5 ++++ py/defs/formatters.py | 5 ++++ py/defs/meta.py | 2 ++ py/nodes/node.py | 9 ++++-- 5 files changed, 63 insertions(+), 25 deletions(-) diff --git a/py/capture.py b/py/capture.py index 858dffb..795420b 100644 --- a/py/capture.py +++ b/py/capture.py @@ -58,46 +58,59 @@ class Capture: 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" + ) - pnginfo_dict.update(cls.gen_loras(inputs)) - pnginfo_dict.update(cls.gen_embeddings(inputs)) + update_pnginfo_dict(inputs_before_this_node, MetaField.VAE_NAME, "VAE") + update_pnginfo_dict(inputs_before_this_node, MetaField.VAE_HASH, "VAE hash") - hashes_for_civitai = cls.get_hashes_for_civitai(inputs) + 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) @@ -121,22 +134,30 @@ class Capture: return result + ", ".join(s_list) @classmethod - def get_hashes_for_civitai(cls, inputs): + def get_hashes_for_civitai( + cls, inputs_before_sampler_node, inputs_before_this_node + ): resource_hashes = {} - model_hashes = inputs.get(MetaField.MODEL_HASH, []) + model_hashes = inputs_before_sampler_node.get(MetaField.MODEL_HASH, []) if len(model_hashes) > 0: resource_hashes["model"] = model_hashes[0][1] - lora_model_names = inputs.get(MetaField.LORA_MODEL_NAME, []) - lora_model_hashes = inputs.get(MetaField.LORA_MODEL_HASH, []) + 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.get(MetaField.EMBEDDING_NAME, []) - embedding_hashes = inputs.get(MetaField.EMBEDDING_HASH, []) + 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] diff --git a/py/defs/captures.py b/py/defs/captures.py index 16baa01..d842adb 100644 --- a/py/defs/captures.py +++ b/py/defs/captures.py @@ -2,6 +2,7 @@ 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, @@ -22,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"}, diff --git a/py/defs/formatters.py b/py/defs/formatters.py index 0ce598c..0c86031 100644 --- a/py/defs/formatters.py +++ b/py/defs/formatters.py @@ -18,6 +18,11 @@ def calc_model_hash(model_name, input_data): return calc_hash(filename) +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) diff --git a/py/defs/meta.py b/py/defs/meta.py index 773ee0c..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 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