"scripts/vscode:/vscode.git/clone" did not exist on "f65c13b5595a8437b23b716a3d0f76657e15df14"
rerouteNode.js 7.26 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
import { app } from "../../scripts/app.js";

// Node that allows you to redirect connections for cleaner graphs

app.registerExtension({
	name: "Comfy.RerouteNode",
	registerCustomNodes() {
		class RerouteNode {
			constructor() {
				if (!this.properties) {
					this.properties = {};
				}
				this.properties.showOutputText = RerouteNode.defaultVisibility;
14
				this.properties.horizontal = false;
15
16
17

				this.addInput("", "*");
				this.addOutput(this.properties.showOutputText ? "*" : "", "*");
18
19

				this.onConnectionsChange = function (type, index, connected, link_info) {
20
21
					this.applyOrientation();

22
23
24
25
26
					// 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) {
27
							const linksToDisconnect = [];
28
29
30
							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];
31
32
33
								linksToDisconnect.push(link);
							}
							for (const link of linksToDisconnect) {
34
35
36
37
38
39
40
41
42
43
								const node = app.graph.getNodeById(link.target_id);
								node.disconnectInput(link.target_slot);
							}
						}
					}

					// Find root input
					let currentNode = this;
					let updateNodes = [];
					let inputType = null;
pythongosssss's avatar
pythongosssss committed
44
					let inputNode = null;
45
46
47
48
49
50
51
52
					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") {
53
54
55
56
57
58
								if (node === this) {
									// We've found a circle
									currentNode.disconnectInput(link.target_slot);
									currentNode = null;
								}
								else {
59
60
									// Move the previous node
									currentNode = node;
61
								}
62
63
							} else {
								// We've found the end
pythongosssss's avatar
pythongosssss committed
64
								inputNode = currentNode;
65
								inputType = node.outputs[link.origin_slot]?.type ?? null;
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
								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];
pythongosssss's avatar
pythongosssss committed
84

85
86
87
88
89
90
91
92
93
94
95
96
								// 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
97
									const nodeOutType = node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type ? node.inputs[link.target_slot].type : null;
98
99
100
101
102
103
104
105
106
107
108
109
110
									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
						}
					}

pythongosssss's avatar
pythongosssss committed
111
					const displayType = inputType || outputType || "*";
pythongosssss's avatar
pythongosssss committed
112
					const color = LGraphCanvas.link_type_colors[displayType];
pythongosssss's avatar
pythongosssss committed
113

114
115
116
117
118
					// 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 || "*";
pythongosssss's avatar
pythongosssss committed
119
120
						node.__outputType = displayType;
						node.outputs[0].name = node.properties.showOutputText ? displayType : "";
121
						node.size = node.computeSize();
122
						node.applyOrientation();
pythongosssss's avatar
pythongosssss committed
123
124

						for (const l of node.outputs[0].links || []) {
pythongosssss's avatar
pythongosssss committed
125
126
127
128
							const link = app.graph.links[l];
							if (link) {
								link.color = color;
							}
pythongosssss's avatar
pythongosssss committed
129
						}
130
					}
pythongosssss's avatar
pythongosssss committed
131
132

					if (inputNode) {
pythongosssss's avatar
pythongosssss committed
133
134
135
136
						const link = app.graph.links[inputNode.inputs[0].link];
						if (link) {
							link.color = color;
						}
pythongosssss's avatar
pythongosssss committed
137
					}
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
				};

				this.clone = function () {
					const cloned = RerouteNode.prototype.clone.apply(this);
					cloned.removeOutput(0);
					cloned.addOutput(this.properties.showOutputText ? "*" : "", "*");
					cloned.size = cloned.computeSize();
					return cloned;
				};

				// This node is purely frontend and does not impact the resulting prompt so should not be serialized
				this.isVirtualNode = true;
			}

			getExtraMenuOptions(_, options) {
				options.unshift(
					{
						content: (this.properties.showOutputText ? "Hide" : "Show") + " Type",
						callback: () => {
							this.properties.showOutputText = !this.properties.showOutputText;
							if (this.properties.showOutputText) {
159
								this.outputs[0].name = this.__outputType || this.outputs[0].type;
160
161
162
163
							} else {
								this.outputs[0].name = "";
							}
							this.size = this.computeSize();
164
							this.applyOrientation();
165
							app.graph.setDirtyCanvas(true, true);
166
167
168
169
170
171
172
						},
					},
					{
						content: (RerouteNode.defaultVisibility ? "Hide" : "Show") + " Type By Default",
						callback: () => {
							RerouteNode.setDefaultTextVisibility(!RerouteNode.defaultVisibility);
						},
173
174
175
176
177
178
179
180
181
182
183
					},
					{
						// naming is inverted with respect to LiteGraphNode.horizontal
						// LiteGraphNode.horizontal == true means that 
						// 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: () => {
							this.properties.horizontal = !this.properties.horizontal;
							this.applyOrientation();
						},
184
185
186
					}
				);
			}
187
188
189
190
191
192
193
194
195
196
197
198
			applyOrientation() {
				this.horizontal = this.properties.horizontal;
				if (this.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];
				} else {
					delete this.inputs[0].pos;
				}
				app.graph.setDirtyCanvas(true, true);
			}
199
200
201
202

			computeSize() {
				return [
					this.properties.showOutputText && this.outputs && this.outputs.length
pythongosssss's avatar
pythongosssss committed
203
204
						? Math.max(75, LiteGraph.NODE_TEXT_SIZE * this.outputs[0].name.length * 0.6 + 40)
						: 75,
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
					26,
				];
			}

			static setDefaultTextVisibility(visible) {
				RerouteNode.defaultVisibility = visible;
				if (visible) {
					localStorage["Comfy.RerouteNode.DefaultVisibility"] = "true";
				} else {
					delete localStorage["Comfy.RerouteNode.DefaultVisibility"];
				}
			}
		}

		// Load default visibility
		RerouteNode.setDefaultTextVisibility(!!localStorage["Comfy.RerouteNode.DefaultVisibility"]);

		LiteGraph.registerNodeType(
			"Reroute",
			Object.assign(RerouteNode, {
				title_mode: LiteGraph.NO_TITLE,
				title: "Reroute",
pythongosssss's avatar
pythongosssss committed
227
				collapsable: false,
228
229
230
231
232
233
			})
		);

		RerouteNode.category = "utils";
	},
});