custom_nodes/rgthree-comfy/src_web/comfyui/fast_actions_button.ts

import type { RgthreeBaseVirtualNodeConstructor } from "typings/rgthree.js";
import type { ComfyApp, ComfyWidget } from "typings/comfy.js";
import type { IWidget, LGraph, LGraphNode, SerializedLGraphNode } from "typings/litegraph.js";
import type { RgthreeBaseNode } from "./base_node.js";

import { app } from "scripts/app.js";
import { BaseAnyInputConnectedNode } from "./base_any_input_connected_node.js";
import { NodeTypesString } from "./constants.js";
import { addMenuItem } from "./utils.js";
import { rgthree } from "./rgthree.js";

const MODE_ALWAYS = 0;
const MODE_MUTE = 2;
const MODE_BYPASS = 4;

/**
 * The Fast Actions Button.
 *
 * This adds a button that the user can connect any node to and then choose an action to take on
 * that node when the button is pressed. Default actions are "Mute," "Bypass," and "Enable," but
 * Nodes can expose actions additional actions that can then be called back.
 */
class FastActionsButton extends BaseAnyInputConnectedNode {
  static override type = NodeTypesString.FAST_ACTIONS_BUTTON;
  static override title = NodeTypesString.FAST_ACTIONS_BUTTON;
  override comfyClass = NodeTypesString.FAST_ACTIONS_BUTTON;

  readonly logger = rgthree.newLogSession("[FastActionsButton]");

  static "@buttonText" = { type: "string" };
  static "@shortcutModifier" = {
    type: "combo",
    values: ["ctrl", "alt", "shift"],
  };
  static "@shortcutKey" = { type: "string" };

  static collapsible = false;

  override readonly isVirtualNode = true;

  override serialize_widgets = true;

  readonly buttonWidget: IWidget;

  readonly widgetToData = new Map<IWidget, { comfy?: ComfyApp; node?: LGraphNode }>();
  readonly nodeIdtoFunctionCache = new Map<number, string>();

  readonly keypressBound;
  readonly keyupBound;

  private executingFromShortcut = false;

  constructor(title?: string) {
    super(title);
    this.properties["buttonText"] = "🎬 Action!";
    this.properties["shortcutModifier"] = "alt";
    this.properties["shortcutKey"] = "";
    this.buttonWidget = this.addWidget(
      "button",
      this.properties["buttonText"],
      null,
      () => {
        this.executeConnectedNodes();
      },
      { serialize: false },
    );

    this.keypressBound = this.onKeypress.bind(this);
    this.keyupBound = this.onKeyup.bind(this);
    this.onConstructed();
  }

  /** When we're given data to configure, like from a PNG or JSON. */
  override configure(info: SerializedLGraphNode<LGraphNode>): void {
    super.configure(info);
    // Since we add the widgets dynamically, we need to wait to set their values
    // with a short timeout.
    setTimeout(() => {
      if (info.widgets_values) {
        for (let [index, value] of info.widgets_values.entries()) {
          if (index > 0) {
            if (value.startsWith("comfy_action:")) {
              value = value.replace("comfy_action:", "");
              this.addComfyActionWidget(index, value);
            }
            if (this.widgets[index]) {
              this.widgets[index]!.value = value;
            }
          }
        }
      }
    }, 100);
  }

  override clone() {
    const cloned = super.clone();
    cloned.properties["buttonText"] = "🎬 Action!";
    cloned.properties["shortcutKey"] = "";
    return cloned;
  }

  override onAdded(graph: LGraph): void {
    window.addEventListener("keydown", this.keypressBound);
    window.addEventListener("keyup", this.keyupBound);
  }

  override onRemoved(): void {
    window.removeEventListener("keydown", this.keypressBound);
    window.removeEventListener("keyup", this.keyupBound);
  }

  async onKeypress(event: KeyboardEvent) {
    const target = (event.target as HTMLElement)!;
    if (
      this.executingFromShortcut ||
      target.localName == "input" ||
      target.localName == "textarea"
    ) {
      return;
    }
    if (
      this.properties["shortcutKey"].trim() &&
      this.properties["shortcutKey"].toLowerCase() === event.key.toLowerCase()
    ) {
      const shortcutModifier = this.properties["shortcutModifier"];
      let good = shortcutModifier === "ctrl" && event.ctrlKey;
      good = good || (shortcutModifier === "alt" && event.altKey);
      good = good || (shortcutModifier === "shift" && event.shiftKey);
      good = good || (shortcutModifier === "meta" && event.metaKey);
      if (good) {
        setTimeout(() => {
          this.executeConnectedNodes();
        }, 20);
        this.executingFromShortcut = true;
        event.preventDefault();
        event.stopImmediatePropagation();
        app.canvas.dirty_canvas = true;
        return false;
      }
    }
    return;
  }

  onKeyup(event: KeyboardEvent) {
    const target = (event.target as HTMLElement)!;
    if (target.localName == "input" || target.localName == "textarea") {
      return;
    }
    this.executingFromShortcut = false;
  }

  override onPropertyChanged(property: string, value: any, _prevValue: any): boolean | void {
    if (property == "buttonText") {
      this.buttonWidget.name = value;
    }
    if (property == "shortcutKey") {
      value = value.trim();
      this.properties["shortcutKey"] = (value && value[0].toLowerCase()) || "";
    }
  }

  override handleLinkedNodesStabilization(linkedNodes: LGraphNode[]) {
    // Remove any widgets and data for widgets that are no longer linked.
    for (const [widget, data] of this.widgetToData.entries()) {
      if (!data.node) {
        continue;
      }
      if (!linkedNodes.includes(data.node)) {
        const index = this.widgets.indexOf(widget);
        if (index > -1) {
          this.widgetToData.delete(widget);
          this.removeWidget(widget);
        } else {
          const [m, a] = this.logger.debugParts("Connected widget is not in widgets... weird.");
          console[m]?.(...a);
        }
      }
    }

    const badNodes: LGraphNode[] = []; // Nodes that are deleted elsewhere may not exist in linkedNodes.
    let indexOffset = 1; // Start with button, increment when we hit a non-node widget (like comfy)
    for (const [index, node] of linkedNodes.entries()) {
      // Sometimes linkedNodes is stale.
      if (!node) {
        const [m, a] = this.logger.debugParts("linkedNode provided that does not exist. ");
        console[m]?.(...a);
        badNodes.push(node);
        continue;
      }
      let widgetAtSlot = this.widgets[index + indexOffset];
      if (widgetAtSlot && this.widgetToData.get(widgetAtSlot)?.comfy) {
        indexOffset++;
        widgetAtSlot = this.widgets[index + indexOffset];
      }

      if (!widgetAtSlot || this.widgetToData.get(widgetAtSlot)?.node?.id !== node.id) {
        // Find the next widget that matches the node.
        let widget: IWidget | null = null;
        for (let i = index + indexOffset; i < this.widgets.length; i++) {
          if (this.widgetToData.get(this.widgets[i]!)?.node?.id === node.id) {
            widget = this.widgets.splice(i, 1)[0]!;
            this.widgets.splice(index + indexOffset, 0, widget);
            break;
          }
        }
        if (!widget) {
          // Add a widget at this spot.
          const exposedActions: string[] = (node.constructor as any).exposedActions || [];
          widget = this.addWidget("combo", node.title, "None", "", {
            values: ["None", "Mute", "Bypass", "Enable", ...exposedActions],
          });
          (widget as ComfyWidget).serializeValue = async (_node: LGraphNode, _index: number) => {
            return widget?.value;
          };
          this.widgetToData.set(widget, { node });
        }
      }
    }

    // Go backwards through widgets, and remove any that are not in out widgetToData
    for (let i = this.widgets.length - 1; i > linkedNodes.length + indexOffset - 1; i--) {
      const widgetAtSlot = this.widgets[i];
      if (widgetAtSlot && this.widgetToData.get(widgetAtSlot)?.comfy) {
        continue;
      }
      this.removeWidget(widgetAtSlot);
    }
  }

  override removeWidget(widgetOrSlot?: number | IWidget): void {
    const widget = typeof widgetOrSlot === "number" ? this.widgets[widgetOrSlot] : widgetOrSlot;
    if (widget && this.widgetToData.has(widget)) {
      this.widgetToData.delete(widget);
    }
    super.removeWidget(widgetOrSlot);
  }

  /**
   * Runs through the widgets, and executes the actions.
   */
  async executeConnectedNodes() {
    for (const widget of this.widgets) {
      if (widget == this.buttonWidget) {
        continue;
      }
      const action = widget.value;
      const { comfy, node } = this.widgetToData.get(widget) ?? {};
      if (comfy) {
        if (action === "Queue Prompt") {
          await comfy.queuePrompt(0);
        }
        continue;
      }
      if (node) {
        if (action === "Mute") {
          node.mode = MODE_MUTE;
        } else if (action === "Bypass") {
          node.mode = MODE_BYPASS;
        } else if (action === "Enable") {
          node.mode = MODE_ALWAYS;
        }
        // If there's a handleAction, always call it.
        if ((node as RgthreeBaseNode).handleAction) {
          await (node as RgthreeBaseNode).handleAction(action);
        }
        app.graph.change();
        continue;
      }
      console.warn("Fast Actions Button has a widget without correct data.");
    }
  }

  /**
   * Adds a ComfyActionWidget at the provided slot (or end).
   */
  addComfyActionWidget(slot?: number, value?: string) {
    let widget = this.addWidget(
      "combo",
      "Comfy Action",
      "None",
      () => {
        if (widget.value.startsWith("MOVE ")) {
          this.widgets.push(this.widgets.splice(this.widgets.indexOf(widget), 1)[0]!);
          widget.value = (widget as any)["lastValue_"];
        } else if (widget.value.startsWith("REMOVE ")) {
          this.removeWidget(widget);
        }
        (widget as any)["lastValue_"] = widget.value;
      },
      {
        values: ["None", "Queue Prompt", "REMOVE Comfy Action", "MOVE to end"],
      },
    );
    (widget as any)["lastValue_"] = value;

    (widget as ComfyWidget).serializeValue = async (_node: LGraphNode, _index: number) => {
      return `comfy_app:${widget?.value}`;
    };
    this.widgetToData.set(widget, { comfy: app });

    if (slot != null) {
      this.widgets.splice(slot, 0, this.widgets.splice(this.widgets.indexOf(widget), 1)[0]!);
    }
    return widget;
  }

  override onSerialize(o: SerializedLGraphNode) {
    super.onSerialize && super.onSerialize(o);
    for (let [index, value] of (o.widgets_values || []).entries()) {
      if (this.widgets[index]?.name === "Comfy Action") {
        o.widgets_values![index] = `comfy_action:${value}`;
      }
    }
  }

  static override setUp() {
    super.setUp();
    addMenuItem(this, app, {
      name: "➕ Append a Comfy Action",
      callback: (nodeArg: LGraphNode) => {
        (nodeArg as FastActionsButton).addComfyActionWidget();
      },
    });
  }
}

app.registerExtension({
  name: "rgthree.FastActionsButton",
  registerCustomNodes() {
    FastActionsButton.setUp();
  },
  loadedGraphNode(node: LGraphNode) {
    if (node.type == FastActionsButton.title) {
      (node as FastActionsButton)._tempWidth = node.size[0];
    }
  },
});