nef7/my-comfyui-workflow

1def is_link(obj):2    if not isinstance(obj, list):3        return False4    if len(obj) != 2:5        return False6    if not isinstance(obj[0], str):7        return False8    if not isinstance(obj[1], int) and not isinstance(obj[1], float):9        return False10    return True11 12# The GraphBuilder is just a utility class that outputs graphs in the form expected by the ComfyUI back-end13class GraphBuilder:14    _default_prefix_root = ""15    _default_prefix_call_index = 016    _default_prefix_graph_index = 017 18    def __init__(self, prefix = None):19        if prefix is None:20            self.prefix = GraphBuilder.alloc_prefix()21        else:22            self.prefix = prefix23        self.nodes = {}24        self.id_gen = 125 26    @classmethod27    def set_default_prefix(cls, prefix_root, call_index, graph_index = 0):28        cls._default_prefix_root = prefix_root29        cls._default_prefix_call_index = call_index30        cls._default_prefix_graph_index = graph_index31 32    @classmethod33    def alloc_prefix(cls, root=None, call_index=None, graph_index=None):34        if root is None:35            root = GraphBuilder._default_prefix_root36        if call_index is None:37            call_index = GraphBuilder._default_prefix_call_index38        if graph_index is None:39            graph_index = GraphBuilder._default_prefix_graph_index40        result = f"{root}.{call_index}.{graph_index}."41        GraphBuilder._default_prefix_graph_index += 142        return result43 44    def node(self, class_type, id=None, **kwargs):45        if id is None:46            id = str(self.id_gen)47            self.id_gen += 148        id = self.prefix + id49        if id in self.nodes:50            return self.nodes[id]51 52        node = Node(id, class_type, kwargs)53        self.nodes[id] = node54        return node55 56    def lookup_node(self, id):57        id = self.prefix + id58        return self.nodes.get(id)59 60    def finalize(self):61        output = {}62        for node_id, node in self.nodes.items():63            output[node_id] = node.serialize()64        return output65 66    def replace_node_output(self, node_id, index, new_value):67        node_id = self.prefix + node_id68        to_remove = []69        for node in self.nodes.values():70            for key, value in node.inputs.items():71                if is_link(value) and value[0] == node_id and value[1] == index:72                    if new_value is None:73                        to_remove.append((node, key))74                    else:75                        node.inputs[key] = new_value76        for node, key in to_remove:77            del node.inputs[key]78 79    def remove_node(self, id):80        id = self.prefix + id81        del self.nodes[id]82 83class Node:84    def __init__(self, id, class_type, inputs):85        self.id = id86        self.class_type = class_type87        self.inputs = inputs88        self.override_display_id = None89 90    def out(self, index):91        return [self.id, index]92 93    def set_input(self, key, value):94        if value is None:95            if key in self.inputs:96                del self.inputs[key]97        else:98            self.inputs[key] = value99 100    def get_input(self, key):101        return self.inputs.get(key)102 103    def set_override_display_id(self, override_display_id):104        self.override_display_id = override_display_id105 106    def serialize(self):107        serialized = {108            "class_type": self.class_type,109            "inputs": self.inputs110        }111        if self.override_display_id is not None:112            serialized["override_display_id"] = self.override_display_id113        return serialized114 115def add_graph_prefix(graph, outputs, prefix):116    # Change the node IDs and any internal links117    new_graph = {}118    for node_id, node_info in graph.items():119        # Make sure the added nodes have unique IDs120        new_node_id = prefix + node_id121        new_node = { "class_type": node_info["class_type"], "inputs": {} }122        for input_name, input_value in node_info.get("inputs", {}).items():123            if is_link(input_value):124                new_node["inputs"][input_name] = [prefix + input_value[0], input_value[1]]125            else:126                new_node["inputs"][input_name] = input_value127        new_graph[new_node_id] = new_node128 129    # Change the node IDs in the outputs130    new_outputs = []131    for n in range(len(outputs)):132        output = outputs[n]133        if is_link(output):134            new_outputs.append([prefix + output[0], output[1]])135        else:136            new_outputs.append(output)137 138    return new_graph, tuple(new_outputs)139 140