From c226b2615a18c5f3921125149f0a878af0bfb27c Mon Sep 17 00:00:00 2001
From: ssit <ssit5568@gmail.com>
Date: Thu, 22 Jun 2023 16:49:49 -0400
Subject: [PATCH] Allow primitives to connect to reroutes with type

- Added a check so that reroute nodes won't connect widget inputs with the same type but different configs (e.g. seed and steps).
- The widget property of widget inputs are distributed among every reroute node in the reroute chain, allowing primitives to connect.
---
 web/extensions/core/rerouteNode.js | 31 +++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/web/extensions/core/rerouteNode.js b/web/extensions/core/rerouteNode.js
index 499a171da..83aea6361 100644
--- a/web/extensions/core/rerouteNode.js
+++ b/web/extensions/core/rerouteNode.js
@@ -21,8 +21,21 @@ app.registerExtension({
 
 					// 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 !== "*"));
+						const callback = function (l) {
+							// Map the link to a type, including the widget config if it exists
+							const link = app.graph.links[l];
+							const input = app.graph.getNodeById(link.target_id).getInputInfo(link.target_slot);
+							return [link.type, input?.widget ? JSON.stringify(input.widget.config) : ""];
+						}
+						const types = new Set(
+							this.outputs[0].links
+							// Map each link to a type
+							.map((l) => callback(l))
+							// Ignore wildcard nodes as these will be updated to real types
+							.filter((t) => t[0] !== "*")
+							// Concatenate the type and widget config into a single string
+							.map((t) => t[0] + t[1])
+						);
 						if (types.size > 1) {
 							const linksToDisconnect = [];
 							for (let i = 0; i < this.outputs[0].links.length - 1; i++) {
@@ -75,6 +88,7 @@ app.registerExtension({
 					// Find all outputs
 					const nodes = [this];
 					let outputType = null;
+					let outputWidget = null;
 					while (nodes.length) {
 						currentNode = nodes.pop();
 						const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || [];
@@ -95,8 +109,9 @@ app.registerExtension({
 								} 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;
-									if (inputType && nodeOutType !== inputType) {
-										// The output doesnt match our input so disconnect it
+									outputWidget = node.inputs[link.target_slot]?.widget ?? null;
+									if (inputType && nodeOutType !== inputType && inputType !== "*") {
+										// The output doesnt match our input and is not a wildcard so disconnect it
 										node.disconnectInput(link.target_slot);
 									} else {
 										outputType = nodeOutType;
@@ -108,7 +123,7 @@ app.registerExtension({
 						}
 					}
 
-					const displayType = inputType || outputType || "*";
+					const displayType = (inputType === "*" && outputType) || inputType || outputType || "*";
 					const color = LGraphCanvas.link_type_colors[displayType];
 
 					// Update the types of each node
@@ -118,6 +133,12 @@ app.registerExtension({
 						node.outputs[0].type = inputType || "*";
 						node.__outputType = displayType;
 						node.outputs[0].name = node.properties.showOutputText ? displayType : "";
+						if (outputWidget) {
+							// Inherit the widget from the connected input 
+							node.inputs[0].widget = outputWidget;
+						} else {
+							delete node.inputs[0].widget;
+						}
 						node.size = node.computeSize();
 						node.applyOrientation();