Allow connecting primitivenode to reroutes

This commit is contained in:
pythongosssss 2023-12-05 20:27:13 +00:00
parent be3468ddd5
commit 44265e0810
2 changed files with 112 additions and 30 deletions

View File

@ -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;
if (inputType && nodeOutType !== inputType) {
const nodeOutType =
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
// each slot in the inputs and outputs are layed out horizontally,
// 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: () => {
@ -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];

View File

@ -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;
if (widgetName) {
const 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, {});
}
let widget;
if (input.widget[TARGET]) {
widget = input.widget[TARGET];
} else {
const widgetName = input.widget.name;
if (widgetName) {
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() {
const values = this.widgets.map((w) => w.value);
recreateWidget() {
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];
return this.widgets[0];
if (values?.length) {
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 = "*";