From 230d22a62894705a62ea701d7d75a14b2c129209 Mon Sep 17 00:00:00 2001 From: nkchocoai Date: Thu, 15 Feb 2024 21:58:21 +0900 Subject: [PATCH] Refactoring --- py/capture.py | 79 +++++++++++++++++++++++++++ py/defs/{hooks.py => captures.py} | 2 +- py/defs/samplers.py | 2 +- py/defs/validators.py | 6 +-- py/nodes/node.py | 90 ++++--------------------------- 5 files changed, 93 insertions(+), 86 deletions(-) create mode 100644 py/capture.py rename py/defs/{hooks.py => captures.py} (97%) diff --git a/py/capture.py b/py/capture.py new file mode 100644 index 0000000..45ddf41 --- /dev/null +++ b/py/capture.py @@ -0,0 +1,79 @@ +from . import hook +from .defs.captures import CAPTURE_FIELD_LIST +from .defs.meta import MetaField + +from nodes import NODE_CLASS_MAPPINGS +from execution import get_input_data + + +class Capture: + @classmethod + def get_inputs(cls): + inputs = {} + prompt = hook.current_prompt + extra_data = hook.current_extra_data + outputs = hook.prompt_executer.outputs + + for node_id, obj in prompt.items(): + class_type = obj["class_type"] + obj_class = NODE_CLASS_MAPPINGS[class_type] + node_inputs = prompt[node_id]["inputs"] + input_data_all = get_input_data( + node_inputs, obj_class, node_id, outputs, prompt, extra_data + ) + for node_class, metas in CAPTURE_FIELD_LIST.items(): + if class_type == node_class: + for meta, field_data in metas.items(): + validate = field_data.get("validate") + if validate is not None and not validate( + node_id, obj, prompt, extra_data, outputs, input_data_all + ): + print("validate failed", meta, node_id, obj) + continue + + if meta not in inputs: + inputs[meta] = [] + field_name = field_data["field_name"] + value = input_data_all.get(field_name) + if value is not None: + format = field_data.get("format") + v = value + if format is not None: + v = format(v) + if isinstance(value, list) and len(value) > 0: + v = value[0] + inputs[meta].append((node_id, v)) + return inputs + + @classmethod + def gen_pnginfo_dict(cls, inputs): + pnginfo_dict = {} + + def update_pnginfo_dict(metafield, key): + x = inputs.get(metafield, []) + if len(x) > 0: + pnginfo_dict[key] = x[-1][1] + + update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "") + update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt") + update_pnginfo_dict(MetaField.STEPS, "Steps") + + sampler_names = inputs.get(MetaField.SAMPLER_NAME, []) + if len(sampler_names) > 0: + pnginfo_dict["Sampler"] = sampler_names[0][1] + + schedulers = inputs.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") + + image_widths = inputs.get(MetaField.IMAGE_WIDTH, []) + image_heights = inputs.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]}" + + return pnginfo_dict diff --git a/py/defs/hooks.py b/py/defs/captures.py similarity index 97% rename from py/defs/hooks.py rename to py/defs/captures.py index 9374f10..4f2f1ed 100644 --- a/py/defs/hooks.py +++ b/py/defs/captures.py @@ -2,7 +2,7 @@ from .meta import MetaField from .validators import is_positive_prompt, is_negative_prompt -hooks = { +CAPTURE_FIELD_LIST = { "EmptyLatentImage": { MetaField.IMAGE_WIDTH: {"field_name": "width"}, MetaField.IMAGE_HEIGHT: {"field_name": "height"}, diff --git a/py/defs/samplers.py b/py/defs/samplers.py index 34cb447..6a84790 100644 --- a/py/defs/samplers.py +++ b/py/defs/samplers.py @@ -1,4 +1,4 @@ -samplers = { +SAMPLERS = { "KSampler": { "positive": "positive", "negative": "negative", diff --git a/py/defs/validators.py b/py/defs/validators.py index f48558f..2d5d151 100644 --- a/py/defs/validators.py +++ b/py/defs/validators.py @@ -1,10 +1,10 @@ -from .samplers import samplers +from .samplers import SAMPLERS def is_positive_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all): positive_node_id_list = {} for nid, node in prompt.items(): - for sampler_type, field_map in samplers.items(): + for sampler_type, field_map in SAMPLERS.items(): if node["class_type"] == sampler_type: positive_node_id = node["inputs"][field_map["positive"]][0] positive_node_id_list[nid] = positive_node_id @@ -14,7 +14,7 @@ def is_positive_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all def is_negative_prompt(node_id, obj, prompt, extra_data, outputs, input_data_all): negative_node_id_list = {} for nid, node in prompt.items(): - for sampler_type, field_map in samplers.items(): + for sampler_type, field_map in SAMPLERS.items(): if node["class_type"] == sampler_type: negative_node_id = node["inputs"][field_map["negative"]][0] negative_node_id_list[nid] = negative_node_id diff --git a/py/nodes/node.py b/py/nodes/node.py index 4ec4fa4..b8d5d56 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -8,13 +8,8 @@ import folder_paths from comfy.cli_args import args from .base import BaseNode -from .. import hook -from ..defs.hooks import hooks -from ..defs.samplers import samplers -from ..defs.meta import MetaField -from nodes import NODE_CLASS_MAPPINGS -from execution import get_input_data +from ..capture import Capture # refer. https://github.com/comfyanonymous/ComfyUI/blob/38b7ac6e269e6ecc5bdd6fefdfb2fb1185b09c9d/nodes.py#L1411 @@ -57,6 +52,14 @@ class SaveImageWithMetaData(BaseNode): for image in images: i = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + inputs = Capture.get_inputs() + print("[get_inputs]") + print(inputs) + pnginfo_dict = Capture.gen_pnginfo_dict(inputs) + print("[PNGInfo Dict]") + print(pnginfo_dict) + metadata = None if not args.disable_metadata: metadata = PngInfo() @@ -73,84 +76,9 @@ class SaveImageWithMetaData(BaseNode): # pnginfo=metadata, # compress_level=self.compress_level, # ) - inputs = self.get_inputs() - print("[get_inputs]") - print(inputs) - pnginfo_dict = self.gen_pnginfo_dict(inputs) - print("[PNGInfo Dict]") - print(pnginfo_dict) results.append( {"filename": file, "subfolder": subfolder, "type": self.type} ) counter += 1 return {"ui": {"images": results}} - - def get_inputs(self): - inputs = {} - prompt = hook.current_prompt - extra_data = hook.current_extra_data - outputs = hook.prompt_executer.outputs - - for node_id, obj in prompt.items(): - class_type = obj["class_type"] - obj_class = NODE_CLASS_MAPPINGS[class_type] - node_inputs = prompt[node_id]["inputs"] - input_data_all = get_input_data( - node_inputs, obj_class, node_id, outputs, prompt, extra_data - ) - for node_class, metas in hooks.items(): - if class_type == node_class: - for meta, field_data in metas.items(): - validate = field_data.get("validate") - if validate is not None and not validate( - node_id, obj, prompt, extra_data, outputs, input_data_all - ): - print("validate failed", meta, node_id, obj) - continue - - if meta not in inputs: - inputs[meta] = [] - field_name = field_data["field_name"] - value = input_data_all.get(field_name) - if value is not None: - format = field_data.get("format") - v = value - if format is not None: - v = format(v) - if isinstance(value, list) and len(value) > 0: - v = value[0] - inputs[meta].append((node_id, v)) - return inputs - - def gen_pnginfo_dict(self, inputs): - pnginfo_dict = {} - - def update_pnginfo_dict(metafield, key): - x = inputs.get(metafield, []) - if len(x) > 0: - pnginfo_dict[key] = x[-1][1] - - update_pnginfo_dict(MetaField.POSITIVE_PROMPT, "") - update_pnginfo_dict(MetaField.NEGATIVE_PROMPT, "Negative prompt") - update_pnginfo_dict(MetaField.STEPS, "Steps") - - sampler_names = inputs.get(MetaField.SAMPLER_NAME, []) - if len(sampler_names) > 0: - pnginfo_dict["Sampler"] = sampler_names[0][1] - - schedulers = inputs.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") - - image_widths = inputs.get(MetaField.IMAGE_WIDTH, []) - image_heights = inputs.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]}" - - return pnginfo_dict