338 lines
11 KiB
TypeScript
338 lines
11 KiB
TypeScript
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];
|
||
}
|
||
},
|
||
});
|