diff --git a/server.py b/server.py index fab33be3e..045d71204 100644 --- a/server.py +++ b/server.py @@ -384,6 +384,23 @@ class PromptServer(): else: info['output_node'] = False + # In litegraph horizontal controls the orientation of the inputs/outputs + # In comfy horizontal controls the orientation of the node, so we need to invert it + if hasattr(obj_class, 'DEFAULT_ORIENTATION') and obj_class.DEFAULT_ORIENTATION is not None: + if obj_class.DEFAULT_ORIENTATION in ['vertical', 'horizontal']: + info['default_orientation'] = obj_class.DEFAULT_ORIENTATION + else: + raise Exception("Invalid DEFAULT_ORIENTATION - must be 'vertical' or 'horizontal'") + + if hasattr(obj_class, 'FORCED_ORIENTATION'): + if obj_class.FORCED_ORIENTATION in ['vertical', 'horizontal']: + info['forced_orientation'] = obj_class.FORCED_ORIENTATION + else: + raise Exception("Invalid FORCED_ORIENTATION - must be 'vertical' or 'horizontal'") + + if 'default_orientation' in info and 'forced_orientation' in info and info['default_orientation'] != info['forced_orientation']: + raise Exception("FORCED_ORIENTATION and DEFAULT_ORIENTATION must be the same if both are set") + if hasattr(obj_class, 'CATEGORY'): info['category'] = obj_class.CATEGORY return info diff --git a/web/extensions/core/nodeOrientation.js b/web/extensions/core/nodeOrientation.js new file mode 100644 index 000000000..6c29f2c70 --- /dev/null +++ b/web/extensions/core/nodeOrientation.js @@ -0,0 +1,69 @@ +import {app} from "../../scripts/app.js"; + +app.registerExtension({ + name: "Comfy.NodeOrientation", + + // On node creation, set the horizontal property if it's set in the node properties + // This is needed as the horizontal status is not serialized + loadedGraphNode(node, _) { + if (node?.forced_orientation !== undefined) { + // Orientation is already handled by nodeCreated + return; + } + + if (node?.properties?.orientation !== undefined) { + node.horizontal = node.properties.orientation !== 'horizontal'; + } + }, + nodeCreated(node, _) { + if (node?.forced_orientation !== undefined) { + // forced_orientation handles node orientation + // Litegraph uses the opposite of horizontal for the orientation + node.horizontal = node.forced_orientation !== 'horizontal'; + return; + } + + if (node?.default_orientation !== undefined) { + node.horizontal = node.default_orientation !== 'horizontal'; + } + }, + async beforeRegisterNodeDef(nodeType, nodeData, app) { + const hasForcedOrientation = Object.hasOwn(nodeData, "forced_orientation"); + const hasDefaultOrientation = Object.hasOwn(nodeData, "default_orientation"); + if (hasForcedOrientation) { + nodeType.prototype.forced_orientation = nodeData.forced_orientation; + } + if (hasDefaultOrientation) { + nodeType.prototype.default_orientation = nodeData.default_orientation; + } + + // nodeType.prototype.forced_orientation = false; + // Add menu option to change node orientation + const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + const origOnConfigure = nodeType.prototype.onConfigure; + if (!Object.hasOwn(nodeData, "forced_orientation")) { + nodeType.prototype.onConfigure = function (settings) { + origOnConfigure.apply(this, arguments); + if (settings?.properties?.orientation !== undefined) { + this.horizontal = settings.properties.orientation !== 'horizontal'; + } + console.log("onConfigure", this, settings); + } + nodeType.prototype.getExtraMenuOptions = function (_, options) { + const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, arguments) : undefined; + options.push({ + content: `Set ${this.horizontal ? "horizontal" : "vertical"}`, + callback: () => { + // properties get serialized, so we need to set the property on the node + this.properties.orientation = this.horizontal ? "horizontal" : "vertical"; + this.horizontal = !this.horizontal; + app.graph.setDirtyCanvas(false, true); + }, + }) + return r; + } + } + }, + } +) +