Allow connecting primitivenode to reroutes

web/extensions/core/rerouteNode.js

@@ -1,10 +1,11 @@
 import { app } from "../../scripts/app.js";
+import { mergeIfValid, getWidgetConfig, setWidgetConfig } from "./widgetInputs.js";
 
 // Node that allows you to redirect connections for cleaner graphs
 
 app.registerExtension({
 	name: "Comfy.RerouteNode",
-	registerCustomNodes() {
+	registerCustomNodes(app) {
 		class RerouteNode {
 			constructor() {
 				if (!this.properties) {
@@ -16,6 +17,12 @@ app.registerExtension({
 				this.addInput("", "*");
 				this.addOutput(this.properties.showOutputText ? "*" : "", "*");
 
+				this.onAfterGraphConfigured = function () {
+					requestAnimationFrame(() => {
+						this.onConnectionsChange(LiteGraph.INPUT, null, true, null);
+					});
+				};
+
 				this.onConnectionsChange = function (type, index, connected, link_info) {
 					this.applyOrientation();
 
@@ -54,8 +61,7 @@ app.registerExtension({
 									// We've found a circle
 									currentNode.disconnectInput(link.target_slot);
 									currentNode = null;
-								}
+								} else {
-								else {
 									// Move the previous node
 									currentNode = node;
 								}
@@ -94,8 +100,11 @@ app.registerExtension({
 									updateNodes.push(node);
 								} else {
 									// We've found an output
-									const nodeOutType = node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type ? node.inputs[link.target_slot].type : null;
+									const nodeOutType =
-									if (inputType && nodeOutType !== inputType) {
+										node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type
+											? node.inputs[link.target_slot].type
+											: null;
+									if (inputType && inputType !== "*" && nodeOutType !== inputType) {
 										// The output doesnt match our input so disconnect it
 										node.disconnectInput(link.target_slot);
 									} else {
@@ -111,6 +120,9 @@ app.registerExtension({
 					const displayType = inputType || outputType || "*";
 					const color = LGraphCanvas.link_type_colors[displayType];
 
+					let widgetConfig;
+					let targetWidget;
+					let widgetType;
 					// 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
@@ -125,10 +137,38 @@ app.registerExtension({
 							const link = app.graph.links[l];
 							if (link) {
 								link.color = color;
+
+								if (app.configuringGraph) continue;
+								const targetNode = app.graph.getNodeById(link.target_id);
+								const targetInput = targetNode.inputs?.[link.target_slot];
+								if (targetInput?.widget) {
+									const config = getWidgetConfig(targetInput);
+									if (!widgetConfig) {
+										widgetConfig = config[1] ?? {};
+										widgetType = config[0];
+									}
+									if (!targetWidget) {
+										targetWidget = targetNode.widgets?.find((w) => w.name === targetInput.widget.name);
+									}
+
+									const merged = mergeIfValid(targetInput, [config[0], widgetConfig]);
+									if (merged.customConfig) {
+										widgetConfig = merged.customConfig;
+									}
+								}
 							}
 						}
 					}
 
+					for (const node of updateNodes) {
+						if (widgetConfig && outputType) {
+							node.inputs[0].widget = { name: "value" };
+							setWidgetConfig(node.inputs[0], [widgetType ?? displayType, widgetConfig], targetWidget);
+						} else {
+							setWidgetConfig(node.inputs[0], null);
+						}
+					}
+
 					if (inputNode) {
 						const link = app.graph.links[inputNode.inputs[0].link];
 						if (link) {
@@ -173,8 +213,8 @@ app.registerExtension({
 					},
 					{
 						// naming is inverted with respect to LiteGraphNode.horizontal
-						// LiteGraphNode.horizontal == true means that 
+						// LiteGraphNode.horizontal == true means that
-						// each slot in the inputs and outputs are layed out horizontally, 
+						// each slot in the inputs and outputs are layed out horizontally,
 						// which is the opposite of the visual orientation of the inputs and outputs as a node
 						content: "Set " + (this.properties.horizontal ? "Horizontal" : "Vertical"),
 						callback: () => {
@@ -187,7 +227,7 @@ app.registerExtension({
 			applyOrientation() {
 				this.horizontal = this.properties.horizontal;
 				if (this.horizontal) {
-					// we correct the input position, because LiteGraphNode.horizontal 
+					// we correct the input position, because LiteGraphNode.horizontal
 					// doesn't account for title presence
 					// which reroute nodes don't have
 					this.inputs[0].pos = [this.size[0] / 2, 0];

web/extensions/core/widgetInputs.js

@@ -5,6 +5,11 @@ const CONVERTED_TYPE = "converted-widget";
 const VALID_TYPES = ["STRING", "combo", "number", "BOOLEAN"];
 const CONFIG = Symbol();
 const GET_CONFIG = Symbol();
+const TARGET = Symbol(); // Used for reroutes to specify the real target widget
+
+export function getWidgetConfig(slot) {
+	return slot.widget[CONFIG] ?? slot.widget[GET_CONFIG]();
+}
 
 function getConfig(widgetName) {
 	const { nodeData } = this.constructor;
@@ -100,7 +105,6 @@ function getWidgetType(config) {
 	return { type };
 }
 
-
 function isValidCombo(combo, obj) {
 	// New input isnt a combo
 	if (!(obj instanceof Array)) {
@@ -121,6 +125,31 @@ function isValidCombo(combo, obj) {
 	return true;
 }
 
+export function setWidgetConfig(slot, config, target) {
+	if (!slot.widget) return;
+	if (config) {
+		slot.widget[GET_CONFIG] = () => config;
+		slot.widget[TARGET] = target;
+	} else {
+		delete slot.widget;
+	}
+
+	if (slot.link) {
+		const link = app.graph.links[slot.link];
+		if (link) {
+			const originNode = app.graph.getNodeById(link.origin_id);
+			if (originNode.type === "PrimitiveNode") {
+				if (config) {
+					originNode.recreateWidget();
+				} else if(!app.configuringGraph) {
+					originNode.disconnectOutput(0);
+					originNode.onLastDisconnect();
+				}
+			}
+		}
+	}
+}
+
 export function mergeIfValid(output, config2, forceUpdate, recreateWidget, config1) {
 	if (!config1) {
 		config1 = output.widget[CONFIG] ?? output.widget[GET_CONFIG]();
@@ -434,14 +463,20 @@ app.registerExtension({
 				for (const linkInfo of links) {
 					const node = this.graph.getNodeById(linkInfo.target_id);
 					const input = node.inputs[linkInfo.target_slot];
-					const widgetName = input.widget.name;
+					let widget;
-					if (widgetName) {
+					if (input.widget[TARGET]) {
-						const widget = node.widgets.find((w) => w.name === widgetName);
+						widget = input.widget[TARGET];
-						if (widget) {
+					} else {
-							widget.value = this.widgets[0].value;
+						const widgetName = input.widget.name;
-							if (widget.callback) {
+						if (widgetName) {
-								widget.callback(widget.value, app.canvas, node, app.canvas.graph_mouse, {});
+							widget = node.widgets.find((w) => w.name === widgetName);
-							}
+						}
+					}
+
+					if (widget) {
+						widget.value = this.widgets[0].value;
+						if (widget.callback) {
+							widget.callback(widget.value, app.canvas, node, app.canvas.graph_mouse, {});
 						}
 					}
 				}
@@ -494,14 +529,13 @@ app.registerExtension({
 					this.#mergeWidgetConfig();
 
 					if (!links?.length) {
-						this.#onLastDisconnect();
+						this.onLastDisconnect();
 					}
 				}
 			}
 
 			onConnectOutput(slot, type, input, target_node, target_slot) {
 				// Fires before the link is made allowing us to reject it if it isn't valid
-
 				// No widget, we cant connect
 				if (!input.widget) {
 					if (!(input.type in ComfyWidgets)) return false;
@@ -519,6 +553,10 @@ app.registerExtension({
 
 			#onFirstConnection(recreating) {
 				// First connection can fire before the graph is ready on initial load so random things can be missing
+				if (!this.outputs[0].links) {
+					this.onLastDisconnect();
+					return;
+				}
 				const linkId = this.outputs[0].links[0];
 				const link = this.graph.links[linkId];
 				if (!link) return;
@@ -546,10 +584,10 @@ app.registerExtension({
 				this.outputs[0].name = type;
 				this.outputs[0].widget = widget;
 
-				this.#createWidget(widget[CONFIG] ?? config, theirNode, widget.name, recreating);
+				this.#createWidget(widget[CONFIG] ?? config, theirNode, widget.name, recreating, widget[TARGET]);
 			}
 
-			#createWidget(inputData, node, widgetName, recreating) {
+			#createWidget(inputData, node, widgetName, recreating, targetWidget) {
 				let type = inputData[0];
 
 				if (type instanceof Array) {
@@ -563,7 +601,9 @@ app.registerExtension({
 					widget = this.addWidget(type, "value", null, () => {}, {});
 				}
 
-				if (node?.widgets && widget) {
+				if (targetWidget) {
+					widget.value = targetWidget.value;
+				} else if (node?.widgets && widget) {
 					const theirWidget = node.widgets.find((w) => w.name === widgetName);
 					if (theirWidget) {
 						widget.value = theirWidget.value;
@@ -577,7 +617,7 @@ app.registerExtension({
 					}
 					addValueControlWidgets(this, widget, control_value, undefined, inputData);
 					let filter = this.widgets_values?.[2];
-					if(filter && this.widgets.length === 3) {
+					if (filter && this.widgets.length === 3) {
 						this.widgets[2].value = filter;
 					}
 				}
@@ -610,12 +650,14 @@ app.registerExtension({
 				}
 			}
 
-			#recreateWidget() {
+			recreateWidget() {
-				const values = this.widgets.map((w) => w.value);
+				const values = this.widgets?.map((w) => w.value);
 				this.#removeWidgets();
 				this.#onFirstConnection(true);
-				for (let i = 0; i < this.widgets?.length; i++) this.widgets[i].value = values[i];
+				if (values?.length) {
-				return this.widgets[0];
+					for (let i = 0; i < this.widgets?.length; i++) this.widgets[i].value = values[i];
+				}
+				return this.widgets?.[0];
 			}
 
 			#mergeWidgetConfig() {
@@ -631,7 +673,7 @@ app.registerExtension({
 				if (links?.length < 2 && hasConfig) {
 					// Copy the widget options from the source
 					if (links.length) {
-						this.#recreateWidget();
+						this.recreateWidget();
 					}
 
 					return;
@@ -657,7 +699,7 @@ app.registerExtension({
 				// Only allow connections where the configs match
 				const output = this.outputs[0];
 				const config2 = input.widget[GET_CONFIG]();
-				return !!mergeIfValid.call(this, output, config2, forceUpdate, this.#recreateWidget);
+				return !!mergeIfValid.call(this, output, config2, forceUpdate, this.recreateWidget);
 			}
 
 			#removeWidgets() {
@@ -672,7 +714,7 @@ app.registerExtension({
 				}
 			}
 
-			#onLastDisconnect() {
+			onLastDisconnect() {
 				// We cant remove + re-add the output here as if you drag a link over the same link
 				// it removes, then re-adds, causing it to break
 				this.outputs[0].type = "*";