Allow reroute to have in/out in different orientation

This commit is contained in:
Davemane42 2023-04-08 03:05:09 -04:00
parent ebd7f9bf80
commit d8c8bd1603

View File

@ -2,6 +2,42 @@ import { app } from "../../scripts/app.js";
// Node that allows you to redirect connections for cleaner graphs
// Context menu to change input/output orientation
function getOrientationMenu(value, options, e, menu, node) {
const isInput = value.options.isInput
const takenSlot = (isInput ? node.outputs[0].dir:node.inputs[0].dir) -1
let availableDir = ["Up" ,"Down", "Left", "Right"]
let availableValue = [LiteGraph.UP, LiteGraph.DOWN, LiteGraph.LEFT, LiteGraph.RIGHT]
availableDir.splice(takenSlot, 1);
availableValue.splice(takenSlot, 1);
new LiteGraph.ContextMenu(
availableDir,
{
event: e,
parentMenu: menu,
node: node,
callback: (v, options, mouse_event, menu, node) => {
if (!node) {
return;
}
let dir = availableValue[Object.values(availableDir).indexOf(v)];
if (isInput) {
node.inputs[0].dir = dir;
} else {
node.outputs[0].dir = dir;
}
node.applyOrientation();
}
}
);
}
app.registerExtension({
name: "Comfy.RerouteNode",
registerCustomNodes() {
@ -11,10 +47,49 @@ app.registerExtension({
this.properties = {};
}
this.properties.showOutputText = RerouteNode.defaultVisibility;
this.properties.horizontal = false;
this.addInput("", "*");
this.addOutput(this.properties.showOutputText ? "*" : "", "*");
this.addInput("", "*", {nameLocked: true});
this.addOutput(this.properties.showOutputText ? "*" : "", "*", {nameLocked: true});
this.inputs[0].dir = LiteGraph.LEFT;
this.outputs[0].dir = LiteGraph.RIGHT;
this.onResize = function(_) {
this.applyOrientation();
}
this.onDrawForeground = function(ctx, graphcanvas, canvas) {
if (this.properties.showOutputText && graphcanvas.ds.scale > 0.5) {
ctx.fillStyle = LiteGraph.NODE_TEXT_COLOR;
ctx.font = graphcanvas.inner_text_font;
ctx.textAlign = "center";
ctx.fillText(this.getDisplayName(), this.size[0] / 2, this.size[1] / 2+5);
}
}
this.onConfigure = function(data) {
// update old reroute
if (!this.inputs[0].dir) { this.inputs[0].dir = LiteGraph.LEFT; }
if (!this.outputs[0].dir) { this.outputs[0].dir = LiteGraph.RIGHT; }
if (this.inputs[0].label) { this.inputs[0].label = "" }
if (this.outputs[0].label) { this.outputs[0].label = "" }
if (!this.inputs[0].nameLocked) { this.inputs[0].nameLocked = true }
if (!this.outputs[0].nameLocked) { this.outputs[0].nameLocked = true }
// handle old horizontal property
if (this.properties.horizontal) {
this.inputs[0].dir = LiteGraph.UP;
this.outputs[0].dir = LiteGraph.DOWN;
delete this.properties.horizontal;
}
this.applyOrientation();
app.graph.setDirtyCanvas(true, true);
}
this.onConnectionsChange = function (type, index, connected, link_info) {
this.applyOrientation();
@ -113,7 +188,6 @@ app.registerExtension({
// 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();
node.applyOrientation();
@ -135,14 +209,18 @@ app.registerExtension({
this.clone = function () {
const cloned = RerouteNode.prototype.clone.apply(this);
let dir = this.outputs[0].dir
cloned.removeOutput(0);
cloned.addOutput(this.properties.showOutputText ? "*" : "", "*");
cloned.addOutput(this.properties.showOutputText ? "*" : "", "*", {nameLocked: true});
cloned.outputs[0].dir = dir
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;
this.applyOrientation();
}
getExtraMenuOptions(_, options) {
@ -151,11 +229,6 @@ app.registerExtension({
content: (this.properties.showOutputText ? "Hide" : "Show") + " Type",
callback: () => {
this.properties.showOutputText = !this.properties.showOutputText;
if (this.properties.showOutputText) {
this.outputs[0].name = this.__outputType || this.outputs[0].type;
} else {
this.outputs[0].name = "";
}
this.size = this.computeSize();
this.applyOrientation();
app.graph.setDirtyCanvas(true, true);
@ -168,37 +241,55 @@ 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,
// 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();
},
}
content: "Input Orientation",
has_submenu: true,
options: {isInput: true},
callback: getOrientationMenu
},
{
content: "Output Orientation",
has_submenu: true,
options: {isInput: false},
callback: getOrientationMenu
},
);
}
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;
// Place inputs/outputs based on the direction
function processInOut(node, slot) {
if (!slot) { return; } // weird copy/paste fix
const horizontal = ([LiteGraph.UP, LiteGraph.DOWN].indexOf(slot.dir) > -1);
const reversed = ([LiteGraph.DOWN, LiteGraph.RIGHT].indexOf(slot.dir) > -1);
if (horizontal) {
slot.pos = [node.size[0] / 2, reversed ? node.size[1]:0];
} else {
slot.pos = [reversed ? node.size[0]:0, node.size[1] / 2];
}
}
processInOut(this, this.inputs[0]);
processInOut(this, this.outputs[0]);
app.graph.setDirtyCanvas(true, true);
}
getDisplayName() {
let displayName = this.__outputType;
if (this.title !== "Reroute" && this.title !== "") {
displayName = this.title;
}
return displayName;
}
computeSize() {
return [
this.properties.showOutputText && this.outputs && this.outputs.length
? Math.max(75, LiteGraph.NODE_TEXT_SIZE * this.outputs[0].name.length * 0.6 + 40)
this.properties.showOutputText && this.outputs
? Math.max(75, LiteGraph.NODE_TEXT_SIZE * this.getDisplayName().length * 0.6)
: 75,
26,
25,
];
}
@ -226,4 +317,29 @@ app.registerExtension({
RerouteNode.category = "utils";
},
setup(app) {
// adds "Add reroute" to right click canvas menu
const orig = LGraphCanvas.prototype.getCanvasMenuOptions;
LGraphCanvas.prototype.getCanvasMenuOptions = function () {
const options = orig.apply(this, arguments);
options.push(
null,
{
content: "Add Reroute",
callback: (value, options, mouse_event, menu, node) => {
let newNode = LiteGraph.createNode("Reroute")
newNode.pos = app.canvas.convertEventToCanvasOffset(mouse_event);
newNode.pos[0] -= newNode.size[0]/2;
newNode.pos[1] -= newNode.size[1]/2;
app.graph.add(newNode);
}
}
);
return options;
};
}
});