Merge branch 'reroute-rewrite' of https://github.com/pythongosssss/ComfyUI

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,
 			})
 		);
 

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) {