120 lines
4.9 KiB
TypeScript
120 lines
4.9 KiB
TypeScript
import LGraphCanvas from "@litegraph-ts/core/src/LGraphCanvas";
|
|
import ComfyGraphNode from "./ComfyGraphNode";
|
|
import ComfyWidgets from "$lib/widgets"
|
|
import type { ComfyWidgetNode } from "$lib/nodes/widgets";
|
|
import { BuiltInSlotShape, BuiltInSlotType, LiteGraph, type SerializedLGraphNode } from "@litegraph-ts/core";
|
|
import type IComfyInputSlot from "$lib/IComfyInputSlot";
|
|
import type { ComfyInputConfig } from "$lib/IComfyInputSlot";
|
|
import { iterateNodeDefOutputs, type ComfyNodeDef, iterateNodeDefInputs } from "$lib/ComfyNodeDef";
|
|
import type { SerializedPromptOutput } from "$lib/utils";
|
|
|
|
/*
|
|
* Base class for any node with configuration sent by the backend.
|
|
*/
|
|
export class ComfyBackendNode extends ComfyGraphNode {
|
|
comfyClass: string;
|
|
comfyNodeDef: ComfyNodeDef;
|
|
displayName: string | null;
|
|
|
|
constructor(title: string, comfyClass: string, nodeDef: ComfyNodeDef) {
|
|
super(title)
|
|
this.type = comfyClass; // XXX: workaround dependency in LGraphNode.addInput()
|
|
this.displayName = nodeDef.display_name;
|
|
this.comfyNodeDef = nodeDef;
|
|
this.comfyClass = comfyClass;
|
|
this.isBackendNode = true;
|
|
|
|
const color = LGraphCanvas.node_colors["yellow"];
|
|
if (this.color == null)
|
|
this.color = color.color
|
|
if (this.bgColor == null)
|
|
this.bgColor = color.bgColor
|
|
|
|
this.setup(nodeDef)
|
|
|
|
if (nodeDef.output_node) {
|
|
this.addOutput("OUTPUT", BuiltInSlotType.EVENT, { color_off: "rebeccapurple", color_on: "rebeccapurple" });
|
|
}
|
|
}
|
|
|
|
// comfy class -> input name -> input config
|
|
private static defaultInputConfigs: Record<string, Record<string, ComfyInputConfig>> = {}
|
|
|
|
private setup(nodeDef: ComfyNodeDef) {
|
|
ComfyBackendNode.defaultInputConfigs[this.type] = {}
|
|
|
|
for (const [inputName, inputData] of iterateNodeDefInputs(nodeDef)) {
|
|
const config: Partial<IComfyInputSlot> = {};
|
|
|
|
const [type, opts] = inputData;
|
|
|
|
if (opts?.forceInput) {
|
|
if (Array.isArray(type)) {
|
|
throw new Error(`Can't have forceInput set to true for an enum type! ${type}`)
|
|
}
|
|
this.addInput(inputName, type as string);
|
|
} else {
|
|
if (Array.isArray(type)) {
|
|
// Enums
|
|
Object.assign(config, ComfyWidgets.COMBO.callback(this, inputName, inputData) || {});
|
|
} else if (`${type}:${inputName}` in ComfyWidgets) {
|
|
// Support custom ComfyWidgets by Type:Name
|
|
Object.assign(config, ComfyWidgets[`${type}:${inputName}`].callback(this, inputName, inputData) || {});
|
|
} else if (type in ComfyWidgets) {
|
|
// Standard type ComfyWidgets
|
|
Object.assign(config, ComfyWidgets[type].callback(this, inputName, inputData) || {});
|
|
} else {
|
|
// Node connection inputs (backend)
|
|
this.addInput(inputName, type);
|
|
}
|
|
}
|
|
|
|
if ("widgetNodeType" in config)
|
|
ComfyBackendNode.defaultInputConfigs[this.type][inputName] = (config as IComfyInputSlot).config
|
|
}
|
|
|
|
for (const output of iterateNodeDefOutputs(nodeDef)) {
|
|
const outputShape = output.is_list ? BuiltInSlotShape.GRID_SHAPE : BuiltInSlotShape.CIRCLE_SHAPE;
|
|
this.addOutput(output.name, output.type, { shape: outputShape });
|
|
}
|
|
|
|
this.serialize_widgets = false;
|
|
// app.#invokeExtensionsAsync("nodeCreated", this);
|
|
}
|
|
|
|
override onSerialize(o: SerializedLGraphNode) {
|
|
super.onSerialize(o);
|
|
for (const input of o.inputs) {
|
|
// strip user-identifying data, it will be reinstantiated later
|
|
if ((input as any).config != null) {
|
|
(input as any).config = {};
|
|
}
|
|
}
|
|
}
|
|
|
|
override onConfigure(o: SerializedLGraphNode) {
|
|
super.onConfigure(o);
|
|
|
|
const configs = ComfyBackendNode.defaultInputConfigs[o.type]
|
|
for (let index = 0; index < this.inputs.length; index++) {
|
|
const input = this.inputs[index] as IComfyInputSlot
|
|
const config = configs[input.name]
|
|
if (config != null && index >= 0 && index < this.inputs.length) {
|
|
if (input.config == null || Object.keys(input.config).length !== Object.keys(config).length) {
|
|
console.debug("[ComfyBackendNode] restore input config", input, config)
|
|
input.config = config
|
|
}
|
|
}
|
|
else {
|
|
console.debug("[ComfyBackendNode] Missing input config in onConfigure()", input, configs)
|
|
input.config = {}
|
|
}
|
|
}
|
|
}
|
|
|
|
override onExecuted(outputData: SerializedPromptOutput) {
|
|
console.warn("onExecuted outputs", outputData)
|
|
this.triggerSlot(0, outputData)
|
|
}
|
|
}
|