From 05274e215b42f660827ed5485cd28eccbb393c68 Mon Sep 17 00:00:00 2001 From: nkchocoai Date: Sat, 17 Feb 2024 15:51:23 +0900 Subject: [PATCH] get pnginfo from trace tree --- py/__init__.py | 7 +++++- py/defs/combo.py | 1 + py/hook.py | 17 ++++++++++++++ py/nodes/node.py | 48 ++++++++++++++++++++++++++++++------- py/trace.py | 61 ++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 125 insertions(+), 9 deletions(-) create mode 100644 py/defs/combo.py create mode 100644 py/trace.py diff --git a/py/__init__.py b/py/__init__.py index ca14bf6..0db041a 100644 --- a/py/__init__.py +++ b/py/__init__.py @@ -1,6 +1,6 @@ import functools -from .hook import pre_execute +from .hook import pre_execute, pre_recursive_execute import execution @@ -17,3 +17,8 @@ def prefix_function(function, prefunction): execution.PromptExecutor.execute = prefix_function( execution.PromptExecutor.execute, pre_execute ) + + +execution.recursive_execute = prefix_function( + execution.recursive_execute, pre_recursive_execute +) diff --git a/py/defs/combo.py b/py/defs/combo.py new file mode 100644 index 0000000..091cd6c --- /dev/null +++ b/py/defs/combo.py @@ -0,0 +1 @@ +SAMPLER_FIND_METHOD = ["Farthest", "Nearest", "By node ID"] diff --git a/py/hook.py b/py/hook.py index d45d78a..61323a8 100644 --- a/py/hook.py +++ b/py/hook.py @@ -1,6 +1,7 @@ current_prompt = {} current_extra_data = {} prompt_executer = None +current_node_id = -1 def pre_execute(self, prompt, prompt_id, extra_data, execute_outputs): @@ -11,3 +12,19 @@ def pre_execute(self, prompt, prompt_id, extra_data, execute_outputs): current_prompt = prompt current_extra_data = extra_data prompt_executer = self + + +def pre_recursive_execute( + server, + prompt, + outputs, + current_item, + extra_data, + executed, + prompt_id, + outputs_ui, + object_storage, +): + global current_node_id + + current_node_id = current_item diff --git a/py/nodes/node.py b/py/nodes/node.py index be95122..1f9a832 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -11,6 +11,10 @@ from comfy.cli_args import args from .base import BaseNode from ..capture import Capture +from .. import hook +from ..trace import Trace + +from ..defs.combo import SAMPLER_FIND_METHOD # refer. https://github.com/comfyanonymous/ComfyUI/blob/38b7ac6e269e6ecc5bdd6fefdfb2fb1185b09c9d/nodes.py#L1411 @@ -27,6 +31,11 @@ class SaveImageWithMetaData(BaseNode): "required": { "images": ("IMAGE",), "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "sampler_find_method": (SAMPLER_FIND_METHOD,), + "sampler_find_node_id": ( + "INT", + {"default": 0, "min": 0, "max": 999999999, "step": 1}, + ), }, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -37,7 +46,13 @@ class SaveImageWithMetaData(BaseNode): OUTPUT_NODE = True def save_images( - self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None + self, + images, + filename_prefix="ComfyUI", + sampler_find_method=SAMPLER_FIND_METHOD[0], + sampler_find_node_id=0, + prompt=None, + extra_pnginfo=None, ): ( full_output_folder, @@ -54,16 +69,10 @@ class SaveImageWithMetaData(BaseNode): 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) - - print("[PNGInfo Dict]") - pnginfo_dict = Capture.gen_pnginfo_dict(inputs) + pnginfo_dict = self.gen_pnginfo() if len(images) >= 2: pnginfo_dict["Batch index"] = index pnginfo_dict["Batch size"] = len(images) - print(pnginfo_dict) metadata = None if not args.disable_metadata: @@ -90,3 +99,26 @@ class SaveImageWithMetaData(BaseNode): counter += 1 return {"ui": {"images": results}} + + @classmethod + def gen_pnginfo(cls): + # get all node inputs + inputs = Capture.get_inputs() + + # get sampler node before this node + trace_tree_from_this_node = Trace.trace( + hook.current_node_id, hook.current_prompt + ) + sampler_node_id = Trace.find_sampler_node_id( + trace_tree_from_this_node, sampler_find_method, sampler_find_node_id + ) + + # 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, trace_tree_from_sampler_node + ) + + # generate PNGInfo from inputs + pnginfo_dict = Capture.gen_pnginfo_dict(filtered_inputs) + return pnginfo_dict diff --git a/py/trace.py b/py/trace.py new file mode 100644 index 0000000..1061bda --- /dev/null +++ b/py/trace.py @@ -0,0 +1,61 @@ +import json + +from collections import deque + +from .defs.samplers import SAMPLERS +from .defs.combo import SAMPLER_FIND_METHOD + + +class Trace: + @classmethod + def trace(cls, start_node_id, prompt): + class_type = prompt[start_node_id]["class_type"] + Q = deque() + Q.append((start_node_id, 0)) + trace_tree = {start_node_id: (0, class_type)} + while len(Q) > 0: + current_node_id, distance = Q.popleft() + input_fields = prompt[current_node_id]["inputs"] + for value in input_fields.values(): + if isinstance(value, list): + nid = value[0] + class_type = prompt[nid]["class_type"] + trace_tree[nid] = (distance + 1, class_type) + Q.append((nid, distance + 1)) + return trace_tree + + @classmethod + def find_sampler_node_id(cls, trace_tree, sampler_find_method, node_id): + if sampler_find_method == SAMPLER_FIND_METHOD[2]: + node_id = str(node_id) + _, class_type = trace_tree.get(node_id, (-1, None)) + if class_type in SAMPLERS.keys(): + return node_id + return -1 + + sorted_by_distance_trace_tree = sorted( + [(k, v[0], v[1]) for k, v in trace_tree.items()], + key=lambda x: x[1], + reverse=(sampler_find_method == SAMPLER_FIND_METHOD[0]), + ) + for nid, _, class_type in sorted_by_distance_trace_tree: + if class_type in SAMPLERS.keys(): + return nid + return -1 + + @classmethod + def filter_inputs_by_trace_tree(cls, inputs, trace_tree): + filtered_inputs = {} + for meta, inputs_list in inputs.items(): + for node_id, input_value in inputs_list: + trace = trace_tree.get(node_id) + if trace is not None: + distance = trace[0] + if meta not in filtered_inputs: + filtered_inputs[meta] = [] + filtered_inputs[meta].append((node_id, input_value, distance)) + + # sort by distance + for k, v in filtered_inputs.items(): + filtered_inputs[k] = sorted(v, key=lambda x: x[2]) + return filtered_inputs