File size: 3,344 Bytes
4450790
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import { app } from "scripts/app.js";
import type {
  LGraphCanvas as TLGraphCanvas,
  ContextMenuItem,
  LGraphNode,
} from "typings/litegraph.js";
import type { ComfyNodeConstructor, ComfyObjectInfo } from "typings/comfy.js";
import { rgthree } from "./rgthree.js";
import { getOutputNodes } from "./utils.js";
import { SERVICE as CONFIG_SERVICE } from "./services/config_service.js";


function showQueueNodesMenuIfOutputNodesAreSelected(existingOptions: ContextMenuItem[]) {
  if (CONFIG_SERVICE.getConfigValue("features.menu_queue_selected_nodes") === false) {
    return;
  }
  const outputNodes = getOutputNodes(Object.values(app.canvas.selected_nodes));
  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);
}

function showQueueGroupNodesMenuIfGroupIsSelected(existingOptions: ContextMenuItem[]) {
  if (CONFIG_SERVICE.getConfigValue("features.menu_queue_selected_nodes") === false) {
    return;
  }
  const group =
    rgthree.lastAdjustedMouseEvent &&
    app.graph.getGroupOnPos(
      rgthree.lastAdjustedMouseEvent.canvasX,
      rgthree.lastAdjustedMouseEvent.canvasY,
    );

  const outputNodes = group && getOutputNodes(group._nodes);
  const menuItem = {
    content: `Queue Group Output Nodes (rgthree)  `,
    className: "rgthree-contextmenu-item",
    callback: () => {
      outputNodes && rgthree.queueOutputNodes(outputNodes.map((n) => n.id));
    },
    disabled: !outputNodes?.length,
  };

  let idx = existingOptions.findIndex((o) => o?.content?.startsWith("Queue Selected ")) + 1;
  idx = 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);
}

/**

 * Adds a "Queue Node" menu item to all output nodes, working with `rgthree.queueOutputNode` to

 * execute only a single node's path.

 */
app.registerExtension({
  name: "rgthree.QueueNode",
  async beforeRegisterNodeDef(nodeType: ComfyNodeConstructor, nodeData: ComfyObjectInfo) {
    const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
    nodeType.prototype.getExtraMenuOptions = function (

      canvas: TLGraphCanvas,

      options: ContextMenuItem[],

    ) {
      getExtraMenuOptions ? getExtraMenuOptions.apply(this, arguments) : undefined;
      showQueueNodesMenuIfOutputNodesAreSelected(options);
      showQueueGroupNodesMenuIfGroupIsSelected(options);
    };
  },

  async setup() {
    const getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions;
    LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) {
      const options = getCanvasMenuOptions.apply(this, [...args] as any);
      showQueueNodesMenuIfOutputNodesAreSelected(options);
      showQueueGroupNodesMenuIfGroupIsSelected(options);
      return options;
    };
  },
});