get pnginfo from trace tree

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
+)

py/defs/combo.py

@@ -0,0 +1 @@
+SAMPLER_FIND_METHOD = ["Farthest", "Nearest", "By node ID"]

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

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

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