From b09620f89c09a4b7882f787e1be0fd79833b2575 Mon Sep 17 00:00:00 2001 From: Jacob Segal Date: Sat, 22 Jul 2023 23:02:10 -0700 Subject: [PATCH] In For loops, display feedback on original node Rather than displaying live feedback on the "End For" node, the feedback will be displayed in the UI on the original node. --- comfy/graph_utils.py | 9 ++++++++- .../execution-inversion-demo-comfyui/flow_control.py | 1 + execution.py | 11 ++++++----- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/comfy/graph_utils.py b/comfy/graph_utils.py index 9103a9622..e436840ff 100644 --- a/comfy/graph_utils.py +++ b/comfy/graph_utils.py @@ -68,6 +68,7 @@ class Node: self.id = id self.class_type = class_type self.inputs = inputs + self.override_parent_id = None def out(self, index): return [self.id, index] @@ -82,11 +83,17 @@ class Node: def get_input(self, key): return self.inputs.get(key) + def set_override_parent_id(self, override_parent_id): + self.override_parent_id = override_parent_id + def serialize(self): - return { + serialized = { "class_type": self.class_type, "inputs": self.inputs } + if self.override_parent_id is not None: + serialized["override_parent_id"] = self.override_parent_id + return serialized def add_graph_prefix(graph, outputs, prefix): # Change the node IDs and any internal links diff --git a/custom_nodes/execution-inversion-demo-comfyui/flow_control.py b/custom_nodes/execution-inversion-demo-comfyui/flow_control.py index 626a82f39..38ddea8e1 100644 --- a/custom_nodes/execution-inversion-demo-comfyui/flow_control.py +++ b/custom_nodes/execution-inversion-demo-comfyui/flow_control.py @@ -103,6 +103,7 @@ class WhileLoopClose: for node_id in contained: original_node = dynprompt.get_node(node_id) node = graph.node(original_node["class_type"], node_id) + node.set_override_parent_id(node_id) for node_id in contained: original_node = dynprompt.get_node(node_id) node = graph.lookup_node(node_id) diff --git a/execution.py b/execution.py index 5c31b1a7d..c537dc313 100644 --- a/execution.py +++ b/execution.py @@ -146,13 +146,13 @@ class DynamicPrompt: return self.original_prompt[node_id] return None - def add_ephemeral_node(self, real_parent_id, node_id, node_info): + def add_ephemeral_node(self, parent_id, node_id, node_info): self.ephemeral_prompt[node_id] = node_info - self.ephemeral_parents[node_id] = real_parent_id + self.ephemeral_parents[node_id] = parent_id def get_real_node_id(self, node_id): - if node_id in self.ephemeral_parents: - return self.ephemeral_parents[node_id] + while node_id in self.ephemeral_parents: + node_id = self.ephemeral_parents[node_id] return node_id def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, dynprompt=None, extra_data={}): @@ -340,7 +340,8 @@ def non_recursive_execute(server, dynprompt, outputs, current_item, extra_data, break new_output_ids = [] for node_id, node_info in new_graph.items(): - dynprompt.add_ephemeral_node(real_node_id, node_id, node_info) + parent_id = node_info.get("override_parent_id", real_node_id) + dynprompt.add_ephemeral_node(parent_id, node_id, node_info) # Figure out if the newly created node is an output node class_type = node_info["class_type"] class_def = nodes.NODE_CLASS_MAPPINGS[class_type]