mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-11 22:12:33 +08:00
5d72965863
Execution blocking can be done by returning an `ExecutionBlocker` (imported from graph_utils) either in place of results or as a specific output. Any node that uses an `ExecutionBlocker` as input will be skipped. This operates on a per-entry basis when inputs are lists. If the `ExecutionBlocker` is initialized with an error message, that message will be displayed on the first node it's used on (and further downstream nodes will be silently skipped).
128 lines
3.9 KiB
Python
128 lines
3.9 KiB
Python
import json
|
|
import random
|
|
|
|
def is_link(obj):
|
|
if not isinstance(obj, list):
|
|
return False
|
|
if len(obj) != 2:
|
|
return False
|
|
if not isinstance(obj[0], str):
|
|
return False
|
|
if not isinstance(obj[1], int) and not isinstance(obj[1], float):
|
|
return False
|
|
return True
|
|
|
|
# The GraphBuilder is just a utility class that outputs graphs in the form expected by the ComfyUI back-end
|
|
class GraphBuilder:
|
|
def __init__(self, prefix = True):
|
|
if isinstance(prefix, str):
|
|
self.prefix = prefix
|
|
elif prefix:
|
|
self.prefix = "%d.%d." % (random.randint(0, 0xffffffffffffffff), random.randint(0, 0xffffffffffffffff))
|
|
else:
|
|
self.prefix = ""
|
|
self.nodes = {}
|
|
self.id_gen = 1
|
|
|
|
def node(self, class_type, id=None, **kwargs):
|
|
if id is None:
|
|
id = str(self.id_gen)
|
|
self.id_gen += 1
|
|
id = self.prefix + id
|
|
if id in self.nodes:
|
|
return self.nodes[id]
|
|
|
|
node = Node(id, class_type, kwargs)
|
|
self.nodes[id] = node
|
|
return node
|
|
|
|
def lookup_node(self, id):
|
|
id = self.prefix + id
|
|
return self.nodes.get(id)
|
|
|
|
def finalize(self):
|
|
output = {}
|
|
for node_id, node in self.nodes.items():
|
|
output[node_id] = node.serialize()
|
|
return output
|
|
|
|
def replace_node_output(self, node_id, index, new_value):
|
|
node_id = self.prefix + node_id
|
|
to_remove = []
|
|
for node in self.nodes.values():
|
|
for key, value in node.inputs.items():
|
|
if is_link(value) and value[0] == node_id and value[1] == index:
|
|
if new_value is None:
|
|
to_remove.append((node, key))
|
|
else:
|
|
node.inputs[key] = new_value
|
|
for node, key in to_remove:
|
|
del node.inputs[key]
|
|
|
|
def remove_node(self, id):
|
|
id = self.prefix + id
|
|
del self.nodes[id]
|
|
|
|
# Return this from a node and any users will be blocked with the given error message.
|
|
class ExecutionBlocker:
|
|
def __init__(self, message):
|
|
self.message = message
|
|
|
|
class Node:
|
|
def __init__(self, id, class_type, inputs):
|
|
self.id = id
|
|
self.class_type = class_type
|
|
self.inputs = inputs
|
|
self.override_parent_id = None
|
|
|
|
def out(self, index):
|
|
return [self.id, index]
|
|
|
|
def set_input(self, key, value):
|
|
if value is None:
|
|
if key in self.inputs:
|
|
del self.inputs[key]
|
|
else:
|
|
self.inputs[key] = value
|
|
|
|
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):
|
|
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
|
|
new_graph = {}
|
|
for node_id, node_info in graph.items():
|
|
# Make sure the added nodes have unique IDs
|
|
new_node_id = prefix + node_id
|
|
new_node = { "class_type": node_info["class_type"], "inputs": {} }
|
|
for input_name, input_value in node_info.get("inputs", {}).items():
|
|
if is_link(input_value):
|
|
new_node["inputs"][input_name] = [prefix + input_value[0], input_value[1]]
|
|
else:
|
|
new_node["inputs"][input_name] = input_value
|
|
new_graph[new_node_id] = new_node
|
|
|
|
# Change the node IDs in the outputs
|
|
new_outputs = []
|
|
for n in range(len(outputs)):
|
|
output = outputs[n]
|
|
if is_link(output):
|
|
new_outputs.append([prefix + output[0], output[1]])
|
|
else:
|
|
new_outputs.append(output)
|
|
|
|
return new_graph, tuple(new_outputs)
|
|
|