mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-11 22:12:33 +08:00
Rather than displaying live feedback on the "End For" node, the feedback will be displayed in the UI on the original node.
123 lines
3.8 KiB
Python
123 lines
3.8 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]
|
|
|
|
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)
|
|
|