Add menu item for queueing only the selected nodes whenever there are selected output nodes. Fixes #132

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;
+    };
   },
 });