321 lines
13 KiB
TypeScript
321 lines
13 KiB
TypeScript
import type { ComfyInputConfig } from "$lib/IComfyInputSlot";
|
|
import type { SerializedPrompt } from "$lib/components/ComfyApp";
|
|
import type ComfyWidget from "$lib/components/widgets/ComfyWidget";
|
|
import { LGraph, LGraphNode, LLink, LiteGraph, NodeMode, type INodeInputSlot, type SerializedLGraphNode, type Vector2, type INodeOutputSlot, LConnectionKind, type SlotType, LGraphCanvas, getStaticPropertyOnInstance, type PropertyLayout, type SlotLayout } from "@litegraph-ts/core";
|
|
import type { SvelteComponentDev } from "svelte/internal";
|
|
import type { ComfyWidgetNode } from "./ComfyWidgetNodes";
|
|
import type IComfyInputSlot from "$lib/IComfyInputSlot";
|
|
import uiState from "$lib/stores/uiState";
|
|
import { get } from "svelte/store";
|
|
import configState from "$lib/stores/configState";
|
|
|
|
export type DefaultWidgetSpec = {
|
|
defaultWidgetNode: new (name?: string) => ComfyWidgetNode,
|
|
config?: ComfyInputConfig
|
|
}
|
|
|
|
export type DefaultWidgetLayout = {
|
|
inputs?: Record<number, DefaultWidgetSpec>,
|
|
}
|
|
|
|
export interface ComfyGraphNodeProperties extends Record<string, any> {
|
|
tags: string[]
|
|
}
|
|
|
|
export default class ComfyGraphNode extends LGraphNode {
|
|
override properties: ComfyGraphNodeProperties = {
|
|
tags: []
|
|
}
|
|
|
|
isBackendNode?: boolean;
|
|
|
|
/*
|
|
* Triggered when the user presses the global "Queue Prompt" button in the fixed toolbar.
|
|
*/
|
|
onDefaultQueueAction?(): void;
|
|
|
|
/*
|
|
* Triggered before a prompt containing this node is passed to the backend.
|
|
*/
|
|
beforeQueued?(subgraph: string | null): void;
|
|
|
|
/*
|
|
* Triggered after a prompt containing this node is passed to the backend.
|
|
*/
|
|
afterQueued?(prompt: SerializedPrompt, subgraph: string | null): void;
|
|
|
|
/*
|
|
* Triggered when the backend sends a finished output back with this node's ID.
|
|
* Valid for output nodes like SaveImage and PreviewImage.
|
|
*/
|
|
onExecuted?(output: any): void;
|
|
|
|
/*
|
|
* Allows you to manually specify an auto-config for certain input slot
|
|
* indices, so that when a ComfyWidgetNode is connected to the input slot it
|
|
* receives the specified min/max/values/etc.
|
|
* Otherwise the config passed from the backend is used.
|
|
*
|
|
* Use this if you're creating a frontend-only node and want some input
|
|
* slots to have auto-configs, like for connected combo box widgets.
|
|
*/
|
|
defaultWidgets?: DefaultWidgetLayout
|
|
|
|
/*
|
|
* If true, attempt to reconcile wildcard types in slots ("*")
|
|
* when a new input/output is connected
|
|
*
|
|
* Only set this to true if all output slots are wildcard typed in the
|
|
* static slotLayout property by default!
|
|
*/
|
|
canInheritSlotTypes: boolean = false;
|
|
|
|
/*
|
|
* If false, don't serialize user-set properties into the workflow.
|
|
* Useful for removing personal information from shared workflows.
|
|
*/
|
|
saveUserState: boolean = true;
|
|
|
|
/*
|
|
* Called to remove user-set properties from this node.
|
|
*/
|
|
stripUserState(o: SerializedLGraphNode) {
|
|
o.widgets_values = []
|
|
}
|
|
|
|
/*
|
|
* Traverses this node backwards in the graph in order to reach a connecting
|
|
* backend node, if any. For example, reroute nodes will simply follow their
|
|
* single input, while branching nodes have conditional logic that
|
|
* determines which link to follow backwards.
|
|
*/
|
|
getUpstreamLink(): LLink | null {
|
|
return null;
|
|
}
|
|
|
|
constructor(title?: string) {
|
|
super(title)
|
|
this.addProperty("tags", [], "array")
|
|
}
|
|
|
|
/*
|
|
* Adjusts output slot types to have the same type as the first connected
|
|
* input. Used for frontend-only nodes with inputs and outputs that act as
|
|
* wildcards, so that they can be connected to ComfyBackendNodes without
|
|
* rejection.
|
|
*/
|
|
private inheritSlotTypes(type: LConnectionKind, isConnected: boolean) {
|
|
// Prevent multiple connections to different types when we have no input
|
|
if (isConnected && type === LConnectionKind.OUTPUT) {
|
|
// Ignore wildcard nodes as these will be updated to real types
|
|
const types = new Set(this.outputs.flatMap(o => o.links.map((l) => this.graph.links[l].type).filter((t) => t !== "*")));
|
|
if (types.size > 1) {
|
|
for (let j = 0; j < this.outputs.length; j++) {
|
|
for (let i = 0; i < this.outputs[j].links.length - 1; i++) {
|
|
const linkId = this.outputs[j].links[i];
|
|
const link = this.graph.links[linkId];
|
|
const node = this.graph.getNodeById(link.target_id);
|
|
node.disconnectInput(link.target_slot);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Find root input
|
|
let currentNode: ComfyGraphNode = this;
|
|
let updateNodes: ComfyGraphNode[] = [];
|
|
let inputType: SlotType | null = null;
|
|
let inputNode = null;
|
|
|
|
while (currentNode) {
|
|
updateNodes.unshift(currentNode);
|
|
const link = currentNode.getUpstreamLink();
|
|
if (link !== null) {
|
|
const node = this.graph.getNodeById(link.origin_id) as ComfyGraphNode;
|
|
console.warn(node.type)
|
|
if (node.canInheritSlotTypes) {
|
|
console.log("REROUTE2", node)
|
|
if (node === this) {
|
|
// We've found a circle
|
|
currentNode.disconnectInput(link.target_slot);
|
|
currentNode = null;
|
|
}
|
|
else {
|
|
// Move the previous node
|
|
currentNode = node;
|
|
}
|
|
} else {
|
|
// We've found the end
|
|
inputNode = currentNode;
|
|
inputType = node.outputs[link.origin_slot]?.type ?? null;
|
|
break;
|
|
}
|
|
} else {
|
|
// This path has no input node
|
|
currentNode = null;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Find all outputs
|
|
const nodes: ComfyGraphNode[] = [this];
|
|
let outputType: SlotType | null = null;
|
|
while (nodes.length) {
|
|
currentNode = nodes.pop();
|
|
if (currentNode.outputs) {
|
|
for (let i = 0; i < currentNode.outputs.length; i++) {
|
|
const outputs = currentNode.outputs[i].links || [];
|
|
if (outputs.length) {
|
|
for (const linkId of outputs) {
|
|
const link = this.graph.links[linkId];
|
|
|
|
// When disconnecting sometimes the link is still registered
|
|
if (!link) continue;
|
|
|
|
const node = this.graph.getNodeById(link.target_id) as ComfyGraphNode;
|
|
|
|
if (node.canInheritSlotTypes) {
|
|
console.log("REROUTE", node)
|
|
// Follow reroute nodes
|
|
nodes.push(node);
|
|
updateNodes.push(node);
|
|
} else {
|
|
// We've found an output
|
|
const nodeOutType = node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type ? node.inputs[link.target_slot].type : null;
|
|
if (inputType && nodeOutType !== inputType) {
|
|
// The output doesnt match our input so disconnect it
|
|
node.disconnectInput(link.target_slot);
|
|
} else {
|
|
outputType = nodeOutType;
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
// No more outputs for this path
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
const displayType = inputType || outputType || "*";
|
|
const color = LGraphCanvas.DEFAULT_LINK_TYPE_COLORS[displayType];
|
|
|
|
// Update the types of each node
|
|
for (const node of updateNodes) {
|
|
// in lieu of static abstract properties
|
|
const slotLayout = getStaticPropertyOnInstance<SlotLayout>(node, "slotLayout");
|
|
if (!slotLayout)
|
|
continue
|
|
|
|
const layoutOutputs = slotLayout.outputs || []
|
|
|
|
for (let i = 0; i < node.outputs.length; i++) {
|
|
// Check if this output was defined as starting off as a
|
|
// wildcard. If for example it was something else like a string,
|
|
// it wouldn't make sense to change its type dynamically.
|
|
const isWildcardOutput = layoutOutputs.length > i && layoutOutputs[i].type === "*";
|
|
if (!isWildcardOutput) {
|
|
console.error("not wildcard", node.outputs[i], layoutOutputs[i])
|
|
continue;
|
|
}
|
|
|
|
// If we dont have an input type we are always wildcard but we'll show the output type
|
|
// This lets you change the output link to a different type and all nodes will update
|
|
node.outputs[i].type = inputType || "*";
|
|
(node as any).__outputType = displayType;
|
|
node.outputs[i].name = node.properties.showOutputText ? String(displayType) : "";
|
|
node.size = node.computeSize();
|
|
|
|
// TODO from ComfyReroute
|
|
if ("applyOrientation" in node && typeof node.applyOrientation === "function")
|
|
node.applyOrientation();
|
|
|
|
for (const l of node.outputs[i].links || []) {
|
|
const link = this.graph.links[l];
|
|
if (link) {
|
|
link.color = color;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (inputNode) {
|
|
for (let i = 0; i < inputNode.inputs.length; i++) {
|
|
const link = this.graph.links[inputNode.inputs[i].link];
|
|
if (link) {
|
|
link.color = color;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
override onConnectionsChange(
|
|
type: LConnectionKind,
|
|
slotIndex: number,
|
|
isConnected: boolean,
|
|
link: LLink,
|
|
ioSlot: (INodeInputSlot | INodeOutputSlot)
|
|
) {
|
|
if (this.canInheritSlotTypes) {
|
|
this.inheritSlotTypes(type, isConnected);
|
|
}
|
|
}
|
|
|
|
override onResize(size: Vector2) {
|
|
// Snap to grid if shift is held down.
|
|
if ((window as any)?.app?.shiftDown) {
|
|
const w = LiteGraph.CANVAS_GRID_SIZE * Math.round(this.size[0] / LiteGraph.CANVAS_GRID_SIZE);
|
|
const h = LiteGraph.CANVAS_GRID_SIZE * Math.round(this.size[1] / LiteGraph.CANVAS_GRID_SIZE);
|
|
this.size[0] = w;
|
|
this.size[1] = h;
|
|
}
|
|
|
|
if (super.onResize)
|
|
super.onResize(size)
|
|
}
|
|
|
|
override onSerialize(o: SerializedLGraphNode) {
|
|
// Resync the widget node types for each input.
|
|
// This is so combo widget nodes will be correctly detected by ComfyApp.refreshComboInNodes().
|
|
for (let index = 0; index < this.inputs.length; index++) {
|
|
const input = this.inputs[index]
|
|
const serInput = o.inputs[index]
|
|
if ("defaultWidgetNode" in input) {
|
|
const comfyInput = input as IComfyInputSlot
|
|
const widgetNode = comfyInput.defaultWidgetNode
|
|
const ty = Object.values(LiteGraph.registered_node_types)
|
|
.find(v => v.class === widgetNode)
|
|
if (ty)
|
|
(serInput as any).widgetNodeType = ty.type;
|
|
(serInput as any).defaultWidgetNode = null
|
|
}
|
|
}
|
|
|
|
(o as any).saveUserState = this.saveUserState
|
|
if (!this.saveUserState && (!get(uiState).isSavingToLocalStorage || get(configState).alwaysStripUserState)) {
|
|
this.stripUserState(o)
|
|
console.warn("[ComfyGraphNode] stripUserState", this, o)
|
|
}
|
|
}
|
|
|
|
override onConfigure(o: SerializedLGraphNode) {
|
|
// Save the litegraph type of the default ComfyWidgetNode for each input.
|
|
for (let index = 0; index < this.inputs.length; index++) {
|
|
const input = this.inputs[index]
|
|
const serInput = o.inputs[index]
|
|
if (serInput && "widgetNodeType" in serInput) {
|
|
const comfyInput = input as IComfyInputSlot
|
|
const ty: string = serInput.widgetNodeType as any
|
|
const widgetNode = Object.values(LiteGraph.registered_node_types)
|
|
.find(v => v.type === ty)
|
|
if (widgetNode)
|
|
comfyInput.defaultWidgetNode = widgetNode.class as any
|
|
}
|
|
}
|
|
|
|
this.saveUserState = (o as any).saveUserState;
|
|
if (this.saveUserState == null)
|
|
this.saveUserState = true
|
|
}
|
|
}
|