diff --git a/rgthree_config.json.default b/rgthree_config.json.default index 5d5b17b..b8e7289 100644 --- a/rgthree_config.json.default +++ b/rgthree_config.json.default @@ -4,6 +4,7 @@ "patch_recursive_execution": true, "show_alerts_for_corrupt_workflows": false, "monitor_for_corrupt_links": false, + "menu_queue_selected_nodes": true, "progress_bar": { "enabled": true, "height": 16, diff --git a/src_web/comfyui/config.ts b/src_web/comfyui/config.ts index e093b87..2d91abc 100644 --- a/src_web/comfyui/config.ts +++ b/src_web/comfyui/config.ts @@ -25,7 +25,7 @@ type ConfigurationSchema = { key: string; type: ConfigType; label: string; - options?: string[]|number[]; + options?: string[] | number[]; description?: string; subconfig?: ConfigurationSchema[]; }; @@ -62,6 +62,14 @@ const CONFIGURABLE: { features: ConfigurationSchema[] } = { }, ], }, + { + key: "features.menu_queue_selected_nodes", + type: ConfigType.BOOLEAN, + label: "Show 'Queue Selected Output Nodes' menu item", + description: + "Will show a menu item in the right-click context menus to queue (only) the selected " + + "output nodes.", + }, { key: "features.show_alerts_for_corrupt_workflows", type: ConfigType.BOOLEAN, @@ -73,7 +81,6 @@ const CONFIGURABLE: { features: ConfigurationSchema[] } = { ], }; - /** * Creates a new fieldrow for main or sub configuration items. */ @@ -99,12 +106,12 @@ function fieldrow(item: ConfigurationSchema) { if (item.options?.length) { input = $el(`select[id="${item.key}"]`, { parent: container, - children: item.options.map(o => { + children: item.options.map((o) => { return $el(`option[value="${String(o)}"]`, { text: String(o), - selected: o === initialValue + selected: o === initialValue, }); - }) + }), }); } else if (item.type === ConfigType.BOOLEAN) { container.classList.toggle("-checked", initialValue); @@ -112,7 +119,6 @@ function fieldrow(item: ConfigurationSchema) { parent: container, checked: initialValue, }); - } else { input = $el(`input[id="${item.key}"]`, { parent: container, diff --git a/src_web/comfyui/menu_queue_node.ts b/src_web/comfyui/menu_queue_node.ts index 526ab0b..17d776c 100644 --- a/src_web/comfyui/menu_queue_node.ts +++ b/src_web/comfyui/menu_queue_node.ts @@ -1,8 +1,54 @@ // @ts-ignore import { app } from "../../scripts/app.js"; -import type { LGraphCanvas, ContextMenuItem } from "typings/litegraph.js"; +import type { + LGraphCanvas as TLGraphCanvas, + ContextMenuItem, + LGraphNode, + LiteGraph as TLiteGraph, +} from "typings/litegraph.js"; import type { ComfyNodeConstructor, ComfyObjectInfo } from "typings/comfy.js"; import { rgthree } from "./rgthree.js"; +import { SERVICE as CONFIG_SERVICE } from "./config_service.js"; + +declare const LiteGraph: typeof TLiteGraph; +declare const LGraphCanvas: typeof TLGraphCanvas; + +function getOutputNodesFromSelected(canvas: TLGraphCanvas) { + return ( + (canvas.selected_nodes && + Object.values(canvas.selected_nodes).filter((n) => { + return ( + n.mode != LiteGraph.NEVER && + ((n.constructor as any).nodeData as ComfyObjectInfo)?.output_node + ); + })) || + [] + ); +} + +function showQueueNodesMenuIfOutputNodesAreSelected( + existingOptions: ContextMenuItem[], + node?: LGraphNode, +) { + if (CONFIG_SERVICE.getConfigValue("features.menu_queue_selected_nodes") != false) { + const canvas = app.canvas as TLGraphCanvas; + const outputNodes = getOutputNodesFromSelected(canvas); + const menuItem = { + content: `Queue Selected Output Nodes (rgthree)  `, + className: "rgthree-contextmenu-item", + callback: () => { + rgthree.queueOutputNodes(outputNodes.map((n) => n.id)); + }, + disabled: !outputNodes.length, + }; + + let idx = existingOptions.findIndex((o) => o?.content === "Outputs") + 1; + idx = idx || existingOptions.findIndex((o) => o?.content === "Align") + 1; + idx = idx || 3; + existingOptions.splice(idx, 0, menuItem); + } + return existingOptions; +} /** * Adds a "Queue Node" menu item to all output nodes, working with `rgthree.queueOutputNode` to @@ -11,24 +57,23 @@ import { rgthree } from "./rgthree.js"; app.registerExtension({ name: "rgthree.QueueNode", async beforeRegisterNodeDef(nodeType: ComfyNodeConstructor, nodeData: ComfyObjectInfo) { - if (nodeData.output_node == true) { - const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; - nodeType.prototype.getExtraMenuOptions = function ( - canvas: LGraphCanvas, - options: ContextMenuItem[], - ) { - getExtraMenuOptions ? getExtraMenuOptions.apply(this, arguments) : undefined; + const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function ( + canvas: TLGraphCanvas, + options: ContextMenuItem[], + ) { + showQueueNodesMenuIfOutputNodesAreSelected(options, this); + }; + }, - const menuItem: ContextMenuItem = { - content: `Queue Node (rgthree)`, - className: "rgthree-contextmenu-item", - callback: () => { - rgthree.queueOutputNode(this.id); - }, - }; - const idx = options.findIndex(o => o?.content === "Outputs") + 1 || options.length - 1; - options.splice(idx, 0, menuItem); - }; - } + async setup() { + console.log("rgthree-menu-setup"); + + const getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions; + LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { + const options = getCanvasMenuOptions.apply(this, [...args] as any); + showQueueNodesMenuIfOutputNodesAreSelected(options); + return options; + }; }, }); diff --git a/src_web/comfyui/rgthree.ts b/src_web/comfyui/rgthree.ts index 1cbdfed..6e2e6c5 100644 --- a/src_web/comfyui/rgthree.ts +++ b/src_web/comfyui/rgthree.ts @@ -145,7 +145,7 @@ class Rgthree extends EventTarget { private progressBarEl: RgthreeProgressBar | null = null; /** Stores a node id that we will use to queu only that output node (with `queueOutputNode`). */ - private queueNodeId: number | null = null; + private queueNodeIds: number[] | null = null; /** Are any functional keys pressed in this given moment? */ ctrlKey = false; @@ -402,15 +402,15 @@ class Rgthree extends EventTarget { * Wraps an `app.queuePrompt` call setting a specific node id that we will inspect and change the * serialized graph when set (below, in our `graphToPrompt` override). */ - async queueOutputNode(nodeId: number) { + async queueOutputNodes(nodeIds: number[]) { try { - this.queueNodeId = nodeId; + this.queueNodeIds = nodeIds; await app.queuePrompt(); } catch(e) { - const [n, v] = this.logParts(LogLevel.ERROR, `There was an error queuing node #${nodeId}`, e); + const [n, v] = this.logParts(LogLevel.ERROR, `There was an error queuing nodes ${nodeIds}`, e); console[n]?.(...v); } finally { - this.queueNodeId = null; + this.queueNodeIds = null; } } @@ -472,11 +472,15 @@ class Rgthree extends EventTarget { rgthree.dispatchEvent(new CustomEvent("graph-to-prompt")); let promise = graphToPrompt.apply(app, [...arguments]); const result = (await promise as ComfyApiPrompt); - // If queueNodeId is set, then we wonly want to queue one node. We'll capture that and rewrite - // the api format, 'output' field. - if (rgthree.queueNodeId != null) { + // If queueNodeIds is set, then we only want to queue those nodes. We'll capture that and + // rewrite the api format, 'output' field, so only those are evaluated. + if (rgthree.queueNodeIds?.length) { const oldOutput = result.output; - const newOutput = rgthree.recursiveAddNodes(String(rgthree.queueNodeId), oldOutput, {}); + let newOutput = {}; + for (const queueNodeId of rgthree.queueNodeIds) { + rgthree.recursiveAddNodes(String(queueNodeId), oldOutput, newOutput); + } + console.log('newOutput', newOutput); result.output = newOutput; } rgthree.dispatchEvent(new CustomEvent("graph-to-prompt-end")); diff --git a/web/comfyui/config.js b/web/comfyui/config.js index 0f8d3a8..7df110a 100644 --- a/web/comfyui/config.js +++ b/web/comfyui/config.js @@ -45,6 +45,13 @@ const CONFIGURABLE = { }, ], }, + { + key: "features.menu_queue_selected_nodes", + type: ConfigType.BOOLEAN, + label: "Show 'Queue Selected Output Nodes' menu item", + description: "Will show a menu item in the right-click context menus to queue (only) the selected " + + "output nodes.", + }, { key: "features.show_alerts_for_corrupt_workflows", type: ConfigType.BOOLEAN, @@ -75,12 +82,12 @@ function fieldrow(item) { if ((_a = item.options) === null || _a === void 0 ? void 0 : _a.length) { input = $el(`select[id="${item.key}"]`, { parent: container, - children: item.options.map(o => { + children: item.options.map((o) => { return $el(`option[value="${String(o)}"]`, { text: String(o), - selected: o === initialValue + selected: o === initialValue, }); - }) + }), }); } else if (item.type === ConfigType.BOOLEAN) { diff --git a/web/comfyui/menu_queue_node.js b/web/comfyui/menu_queue_node.js index c1f34a9..e1f3705 100644 --- a/web/comfyui/menu_queue_node.js +++ b/web/comfyui/menu_queue_node.js @@ -1,22 +1,49 @@ import { app } from "../../scripts/app.js"; import { rgthree } from "./rgthree.js"; +import { SERVICE as CONFIG_SERVICE } from "./config_service.js"; +function getOutputNodesFromSelected(canvas) { + return ((canvas.selected_nodes && + Object.values(canvas.selected_nodes).filter((n) => { + var _a; + return (n.mode != LiteGraph.NEVER && + ((_a = n.constructor.nodeData) === null || _a === void 0 ? void 0 : _a.output_node)); + })) || + []); +} +function showQueueNodesMenuIfOutputNodesAreSelected(existingOptions, node) { + if (CONFIG_SERVICE.getConfigValue("features.menu_queue_selected_nodes") != false) { + const canvas = app.canvas; + const outputNodes = getOutputNodesFromSelected(canvas); + const menuItem = { + content: `Queue Selected Output Nodes (rgthree)  `, + className: "rgthree-contextmenu-item", + callback: () => { + rgthree.queueOutputNodes(outputNodes.map((n) => n.id)); + }, + disabled: !outputNodes.length, + }; + let idx = existingOptions.findIndex((o) => (o === null || o === void 0 ? void 0 : o.content) === "Outputs") + 1; + idx = idx || existingOptions.findIndex((o) => (o === null || o === void 0 ? void 0 : o.content) === "Align") + 1; + idx = idx || 3; + existingOptions.splice(idx, 0, menuItem); + } + return existingOptions; +} app.registerExtension({ name: "rgthree.QueueNode", async beforeRegisterNodeDef(nodeType, nodeData) { - if (nodeData.output_node == true) { - const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; - nodeType.prototype.getExtraMenuOptions = function (canvas, options) { - getExtraMenuOptions ? getExtraMenuOptions.apply(this, arguments) : undefined; - const menuItem = { - content: `Queue Node (rgthree)`, - className: "rgthree-contextmenu-item", - callback: () => { - rgthree.queueOutputNode(this.id); - }, - }; - const idx = options.findIndex(o => (o === null || o === void 0 ? void 0 : o.content) === "Outputs") + 1 || options.length - 1; - options.splice(idx, 0, menuItem); - }; - } + const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (canvas, options) { + showQueueNodesMenuIfOutputNodesAreSelected(options, this); + }; + }, + async setup() { + console.log("rgthree-menu-setup"); + const getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions; + LGraphCanvas.prototype.getCanvasMenuOptions = function (...args) { + const options = getCanvasMenuOptions.apply(this, [...args]); + showQueueNodesMenuIfOutputNodesAreSelected(options); + return options; + }; }, }); diff --git a/web/comfyui/rgthree.js b/web/comfyui/rgthree.js index 63ed7b6..fcfc823 100644 --- a/web/comfyui/rgthree.js +++ b/web/comfyui/rgthree.js @@ -82,7 +82,7 @@ class Rgthree extends EventTarget { this.api = api; this.settingsDialog = null; this.progressBarEl = null; - this.queueNodeId = null; + this.queueNodeIds = null; this.ctrlKey = false; this.altKey = false; this.metaKey = false; @@ -278,18 +278,18 @@ class Rgthree extends EventTarget { }, ]; } - async queueOutputNode(nodeId) { + async queueOutputNodes(nodeIds) { var _a; try { - this.queueNodeId = nodeId; + this.queueNodeIds = nodeIds; await app.queuePrompt(); } catch (e) { - const [n, v] = this.logParts(LogLevel.ERROR, `There was an error queuing node #${nodeId}`, e); + const [n, v] = this.logParts(LogLevel.ERROR, `There was an error queuing nodes ${nodeIds}`, e); (_a = console[n]) === null || _a === void 0 ? void 0 : _a.call(console, ...v); } finally { - this.queueNodeId = null; + this.queueNodeIds = null; } } recursiveAddNodes(nodeId, oldOutput, newOutput) { @@ -331,12 +331,17 @@ class Rgthree extends EventTarget { }; const graphToPrompt = app.graphToPrompt; app.graphToPrompt = async function () { + var _a; rgthree.dispatchEvent(new CustomEvent("graph-to-prompt")); let promise = graphToPrompt.apply(app, [...arguments]); const result = await promise; - if (rgthree.queueNodeId != null) { + if ((_a = rgthree.queueNodeIds) === null || _a === void 0 ? void 0 : _a.length) { const oldOutput = result.output; - const newOutput = rgthree.recursiveAddNodes(String(rgthree.queueNodeId), oldOutput, {}); + let newOutput = {}; + for (const queueNodeId of rgthree.queueNodeIds) { + rgthree.recursiveAddNodes(String(queueNodeId), oldOutput, newOutput); + } + console.log('newOutput', newOutput); result.output = newOutput; } rgthree.dispatchEvent(new CustomEvent("graph-to-prompt-end"));