From d8c8bd160339ef7cc6357898a90b15fa6b686eb7 Mon Sep 17 00:00:00 2001 From: Davemane42 Date: Sat, 8 Apr 2023 03:05:09 -0400 Subject: [PATCH] Allow reroute to have in/out in different orientation --- web/extensions/core/rerouteNode.js | 178 ++++++++++++++++++++++++----- 1 file changed, 147 insertions(+), 31 deletions(-) diff --git a/web/extensions/core/rerouteNode.js b/web/extensions/core/rerouteNode.js index c31f63cd0..8ca9f01dd 100644 --- a/web/extensions/core/rerouteNode.js +++ b/web/extensions/core/rerouteNode.js @@ -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; + }; + } });