diff --git a/web/extensions/core/rerouteNode.js b/web/extensions/core/rerouteNode.js index bdc7306b..7188dfd2 100644 --- a/web/extensions/core/rerouteNode.js +++ b/web/extensions/core/rerouteNode.js @@ -14,11 +14,111 @@ app.registerExtension({ this.addInput("", "*"); this.addOutput(this.properties.showOutputText ? "*" : "", "*"); - this.onConnectInput = function (_, type) { - if (type !== this.outputs[0].type) { - this.removeOutput(0); - this.addOutput(this.properties.showOutputText ? type : "", type); - this.size = this.computeSize(); + + this.onConnectionsChange = function (type, index, connected, link_info) { + // Prevent multiple connections to different types when we have no input + if (connected && type === LiteGraph.OUTPUT) { + // Ignore wildcard nodes as these will be updated to real types + const types = new Set(this.outputs[0].links.map((l) => app.graph.links[l].type).filter((t) => t !== "*")); + if (types.size > 1) { + for (let i = 0; i < this.outputs[0].links.length - 1; i++) { + const linkId = this.outputs[0].links[i]; + const link = app.graph.links[linkId]; + const node = app.graph.getNodeById(link.target_id); + node.disconnectInput(link.target_slot); + } + } + } + + // Find root input + let currentNode = this; + let updateNodes = []; + let inputType = null; + let inputNode = null; + while (currentNode) { + updateNodes.unshift(currentNode); + const linkId = currentNode.inputs[0].link; + if (linkId !== null) { + const link = app.graph.links[linkId]; + const node = app.graph.getNodeById(link.origin_id); + const type = node.constructor.type; + if (type === "Reroute") { + // Move the previous node + currentNode = node; + } else { + // We've found the end + inputNode = currentNode; + inputType = node.outputs[link.origin_slot].type; + break; + } + } else { + // This path has no input node + currentNode = null; + break; + } + } + + // Find all outputs + const nodes = [this]; + let outputType = null; + while (nodes.length) { + currentNode = nodes.pop(); + const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || []; + if (outputs.length) { + for (const linkId of outputs) { + const link = app.graph.links[linkId]; + + // When disconnecting sometimes the link is still registered + if (!link) continue; + + const node = app.graph.getNodeById(link.target_id); + const type = node.constructor.type; + + if (type === "Reroute") { + // Follow reroute nodes + nodes.push(node); + updateNodes.push(node); + } else { + // We've found an output + const nodeOutType = node.inputs[link.target_slot].type; + if (inputType && nodeOutType !== inputType) { + // The output doesnt match our input so disconnect it + node.disconnectInput(link.target_slot); + } else { + outputType = nodeOutType; + } + } + } + } else { + // No more outputs for this path + } + } + + const displayType = inputType || outputType || "*"; + const color = LGraphCanvas.link_type_colors[displayType]; + + // Update the types of each node + for (const node of updateNodes) { + // If we dont have an input type we are always wildcard but we'll show the output type + // This lets you change the output link to a different type and all nodes will update + node.outputs[0].type = inputType || "*"; + node.__outputType = displayType; + node.outputs[0].name = node.properties.showOutputText ? displayType : ""; + node.size = node.computeSize(); + + for (const l of node.outputs[0].links || []) { + const link = app.graph.links[l]; + if (link) { + link.color = color; + } + } + } + + if (inputNode) { + const link = app.graph.links[inputNode.inputs[0].link]; + if (link) { + link.color = color; + } } }; @@ -41,12 +141,12 @@ app.registerExtension({ callback: () => { this.properties.showOutputText = !this.properties.showOutputText; if (this.properties.showOutputText) { - this.outputs[0].name = this.outputs[0].type; + this.outputs[0].name = this.__outputType || this.outputs[0].type; } else { this.outputs[0].name = ""; } this.size = this.computeSize(); - app.graph.setDirtyCanvas(true); + app.graph.setDirtyCanvas(true, true); }, }, { @@ -61,8 +161,8 @@ app.registerExtension({ computeSize() { return [ this.properties.showOutputText && this.outputs && this.outputs.length - ? Math.max(55, LiteGraph.NODE_TEXT_SIZE * this.outputs[0].name.length * 0.6 + 40) - : 55, + ? Math.max(75, LiteGraph.NODE_TEXT_SIZE * this.outputs[0].name.length * 0.6 + 40) + : 75, 26, ]; } @@ -85,6 +185,7 @@ app.registerExtension({ Object.assign(RerouteNode, { title_mode: LiteGraph.NO_TITLE, title: "Reroute", + collapsable: false, }) ); diff --git a/web/scripts/app.js b/web/scripts/app.js index 5e5f4f56..86edcdcf 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -678,24 +678,10 @@ class ComfyApp { for (let i in node.inputs) { let parent = node.getInputNode(i); if (parent) { - let link; - if (parent.isVirtualNode) { - // Follow the path of virtual nodes until we reach the first real one - while (parent != null) { - link = parent.getInputLink(0); - if (link) { - const from = graph.getNodeById(link.origin_id); - if (from.isVirtualNode) { - parent = from; - } else { - parent = null; - } - } else { - parent = null; - } - } - } else { - link = node.getInputLink(i); + let link = node.getInputLink(i); + while (parent && parent.isVirtualNode) { + link = parent.getInputLink(link.origin_slot); + parent = parent.getInputNode(link.origin_slot); } if (link) {