291 lines
12 KiB
TypeScript
291 lines
12 KiB
TypeScript
import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions, LGraphCanvas, type LGraphRemoveNodeOptions, Subgraph, type LGraphAddNodeMode, type SerializedLGraphNode, type Vector2, type NodeID, reassignGraphIDs, type GraphIDMapping, type SerializedLGraph } from "@litegraph-ts/core";
|
|
import GraphSync from "./GraphSync";
|
|
import EventEmitter from "events";
|
|
import type TypedEmitter from "typed-emitter";
|
|
import uiState from "./stores/uiState";
|
|
import { get } from "svelte/store";
|
|
import type ComfyGraphNode from "./nodes/ComfyGraphNode";
|
|
import type IComfyInputSlot from "./IComfyInputSlot";
|
|
import type { ComfyBackendNode } from "./nodes/ComfyBackendNode";
|
|
import type { ComfyComboNode, ComfyWidgetNode } from "./nodes/widgets";
|
|
import selectionState from "./stores/selectionState";
|
|
import type { WritableLayoutStateStore } from "./stores/layoutStates";
|
|
import layoutStates from "./stores/layoutStates";
|
|
import type { ComfyBoxWorkflow, WorkflowInstID } from "./stores/workflowState";
|
|
import workflowState from "./stores/workflowState";
|
|
import type { SerializedComfyBoxTemplate } from "./ComfyBoxTemplate";
|
|
import { v4 as uuidv4 } from "uuid"
|
|
|
|
type ComfyGraphEvents = {
|
|
configured: (graph: LGraph) => void
|
|
nodeAdded: (node: LGraphNode) => void
|
|
nodeRemoved: (node: LGraphNode) => void
|
|
nodeConnectionChanged: (kind: LConnectionKind, node: LGraphNode, slot: SlotIndex, targetNode: LGraphNode, targetSlot: SlotIndex) => void
|
|
cleared: () => void
|
|
beforeChange: (graph: LGraph, param: any) => void
|
|
afterChange: (graph: LGraph, param: any) => void
|
|
afterExecute: () => void
|
|
}
|
|
|
|
export default class ComfyGraph extends LGraph {
|
|
eventBus: TypedEmitter<ComfyGraphEvents> = new EventEmitter() as TypedEmitter<ComfyGraphEvents>;
|
|
|
|
workflowID: WorkflowInstID | null = null;
|
|
|
|
get workflow(): ComfyBoxWorkflow | null {
|
|
const workflowID = (this.getRootGraph() as ComfyGraph)?.workflowID;
|
|
if (workflowID == null)
|
|
return null;
|
|
return workflowState.getWorkflow(workflowID)
|
|
}
|
|
|
|
get layout(): WritableLayoutStateStore | null {
|
|
return this.workflow?.layout;
|
|
}
|
|
|
|
constructor(workflowID?: WorkflowInstID) {
|
|
super();
|
|
this.workflowID = workflowID;
|
|
}
|
|
|
|
override onConfigure() {
|
|
console.debug("Configured");
|
|
}
|
|
|
|
override onBeforeChange(graph: LGraph, info: any) {
|
|
if (this.workflow != null)
|
|
this.workflow.notifyModified()
|
|
|
|
console.debug("BeforeChange", info);
|
|
}
|
|
|
|
override onAfterChange(graph: LGraph, info: any) {
|
|
console.debug("AfterChange", info);
|
|
}
|
|
|
|
override onAfterExecute() {
|
|
this.eventBus.emit("afterExecute");
|
|
}
|
|
|
|
/*
|
|
* NOTE: This function will also be called by child subgraphs on their
|
|
* parent graphs. So we have to be sure the node that receives the callback
|
|
* is a root graph (this._is_subgraph is false). If a subgraph calls this
|
|
* then options.subgraphsh will have the list of subgraphs down the chain.
|
|
*/
|
|
override onNodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
|
|
// Don't add nodes in subgraphs until this callback reaches the root
|
|
// graph
|
|
// Only root graphs will have a workflow ID, so we don't mind subgraphs
|
|
// missing it
|
|
if (node.getRootGraph() != null && !this._is_subgraph && this.workflowID != null) {
|
|
const layoutState = get(layoutStates).all[this.workflowID]
|
|
if (layoutState === null) {
|
|
throw new Error(`LGraph with workflow missing layout! ${this.workflowID}`)
|
|
}
|
|
|
|
this.doAddNode(node, layoutState, options);
|
|
}
|
|
|
|
if (this.workflow != null)
|
|
this.workflow.notifyModified()
|
|
|
|
this.eventBus.emit("nodeAdded", node);
|
|
}
|
|
|
|
/*
|
|
* Add widget UI/groups for newly added nodes.
|
|
*/
|
|
private doAddNode(node: LGraphNode, layoutState: WritableLayoutStateStore, options: LGraphAddNodeOptions) {
|
|
layoutState.nodeAdded(node, options)
|
|
|
|
// All nodes whether they come from base litegraph or ComfyBox should
|
|
// have tags added to them. Can't override serialization for litegraph's
|
|
// base node types to add `tags` as a new field so putting it in
|
|
// properties is better.
|
|
if (node.properties.tags == null)
|
|
node.properties.tags = []
|
|
|
|
if ((node as any).canInheritSlotTypes && node.inputs.length > 1) {
|
|
node.color ||= LGraphCanvas.node_colors["green"].color;
|
|
node.bgColor ||= LGraphCanvas.node_colors["green"].bgColor;
|
|
}
|
|
|
|
if ("outputProperties" in node) {
|
|
const widgetNode = node as ComfyWidgetNode;
|
|
for (const propName of widgetNode.outputProperties) {
|
|
widgetNode.addPropertyAsOutput(propName.name, propName.type)
|
|
}
|
|
}
|
|
|
|
// Check if the class declared a default widget layout
|
|
if ("defaultWidgets" in node && !("svelteComponentType" in node)) {
|
|
const comfyNode = node as ComfyGraphNode;
|
|
const widgets = comfyNode.defaultWidgets;
|
|
|
|
if (widgets) {
|
|
if (widgets.inputs) {
|
|
for (const pair of Object.entries(comfyNode.defaultWidgets.inputs)) {
|
|
const [index, spec] = pair
|
|
const input = comfyNode.inputs[index] as IComfyInputSlot;
|
|
input.defaultWidgetNode = spec.defaultWidgetNode;
|
|
if (spec.config) {
|
|
input.config = spec.config
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (get(uiState).autoAddUI) {
|
|
if (!("svelteComponentType" in node) && options.addedBy == null) {
|
|
console.debug("[ComfyGraph] AutoAdd UI")
|
|
const comfyNode = node as ComfyGraphNode;
|
|
const widgetNodesAdded: ComfyWidgetNode[] = []
|
|
for (let index = 0; index < comfyNode.inputs.length; index++) {
|
|
const input = comfyNode.inputs[index];
|
|
if ("config" in input) {
|
|
const comfyInput = input as IComfyInputSlot;
|
|
if (comfyInput.defaultWidgetNode) {
|
|
const widgetNode = LiteGraph.createNode(comfyInput.defaultWidgetNode)
|
|
const inputPos = comfyNode.getConnectionPos(true, index);
|
|
node.graph.add(widgetNode)
|
|
widgetNode.connect(0, comfyNode, index);
|
|
widgetNode.collapse();
|
|
widgetNode.pos = [inputPos[0] - 140, inputPos[1] + LiteGraph.NODE_SLOT_HEIGHT / 2];
|
|
widgetNodesAdded.push(widgetNode)
|
|
|
|
// Set combo box as loaded
|
|
if (widgetNode.type === "ui/combo" && widgetNode.properties.values != null) {
|
|
(widgetNode as ComfyComboNode).formatValues(widgetNode.properties.values);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
const dragItemIDs = widgetNodesAdded.map(wn => get(layoutState).allItemsByNode[wn.id]?.dragItem?.id).filter(Boolean)
|
|
console.debug("[ComfyGraph] Group new widgets", dragItemIDs)
|
|
|
|
// Use the default node title instead of custom node title, in
|
|
// case node was cloned
|
|
const reg = LiteGraph.registered_node_types[node.type]
|
|
|
|
layoutState.groupItems(dragItemIDs, { title: reg.title, variant: "accordion", openOnStartup: true })
|
|
}
|
|
}
|
|
|
|
// Handle nodes in subgraphs being attached to this graph indirectly
|
|
// ************** RECURSION ALERT ! **************
|
|
if (node.is(Subgraph)) {
|
|
for (const child of node.subgraph.iterateNodesInOrder()) {
|
|
this.doAddNode(child, layoutState, options)
|
|
}
|
|
}
|
|
// ************** RECURSION ALERT ! **************
|
|
|
|
if (this.workflow != null)
|
|
this.workflow.notifyModified()
|
|
}
|
|
|
|
override onNodeRemoved(node: LGraphNode, options: LGraphRemoveNodeOptions) {
|
|
selectionState.clear(); // safest option
|
|
|
|
if (!this._is_subgraph && this.workflowID != null) {
|
|
const layoutState = get(layoutStates).all[this.workflowID]
|
|
if (layoutState === null) {
|
|
throw new Error(`ComfyGraph with workflow missing layout! ${this.workflowID}`)
|
|
}
|
|
|
|
layoutState.nodeRemoved(node, options);
|
|
|
|
// Handle subgraphs being removed
|
|
if (node.is(Subgraph)) {
|
|
for (const child of node.subgraph.iterateNodesInOrder()) {
|
|
this.onNodeRemoved(child, options)
|
|
}
|
|
}
|
|
}
|
|
|
|
if (this.workflow != null)
|
|
this.workflow.notifyModified()
|
|
|
|
this.eventBus.emit("nodeRemoved", node);
|
|
}
|
|
|
|
override onInputsOutputsChange() {
|
|
if (this.workflow != null)
|
|
this.workflow.notifyModified()
|
|
}
|
|
|
|
override onNodeConnectionChange(kind: LConnectionKind, node: LGraphNode, slot: SlotIndex, targetNode: LGraphNode, targetSlot: SlotIndex) {
|
|
if (this.workflow != null)
|
|
this.workflow.notifyModified()
|
|
|
|
// console.debug("ConnectionChange", node);
|
|
this.eventBus.emit("nodeConnectionChanged", kind, node, slot, targetNode, targetSlot);
|
|
}
|
|
|
|
/*
|
|
* Inserts a template.
|
|
* Layout deserialization must be handled afterwards.
|
|
* NOTE: Modifies the template in-place, be sure you cloned it beforehand!
|
|
*/
|
|
insertTemplate(template: SerializedComfyBoxTemplate, pos: Vector2): Record<NodeID, LGraphNode> {
|
|
const minPos = [0, 0]
|
|
|
|
const templateNodeIDToNewNode: Record<NodeID, LGraphNode> = {}
|
|
|
|
var nodes = [];
|
|
for (var i = 0; i < template.nodes.length; ++i) {
|
|
var node_data = template.nodes[i];
|
|
var node = LiteGraph.createNode(node_data.type);
|
|
|
|
let mapping: GraphIDMapping = null;
|
|
if (node_data.type === "graph/subgraph") {
|
|
mapping = reassignGraphIDs((node_data as any).subgraph as SerializedLGraph);
|
|
}
|
|
|
|
if (node) {
|
|
const prevNodeId = node_data.id;
|
|
node_data.id = uuidv4();
|
|
templateNodeIDToNewNode[prevNodeId] = node
|
|
|
|
// Strip tags from top-level nodes
|
|
if (Array.isArray(node_data.properties.tags)) {
|
|
node_data.properties.tags = []
|
|
}
|
|
|
|
node.configure(node_data);
|
|
|
|
if (mapping) {
|
|
for (const subnode of (node as Subgraph).subgraph.iterateNodesInOrderRecursive()) {
|
|
const oldNodeID = mapping.nodeIDs[subnode.id];
|
|
templateNodeIDToNewNode[oldNodeID] = subnode;
|
|
}
|
|
}
|
|
|
|
node.pos[0] += pos[0] - minPos[0]; //+= 5;
|
|
node.pos[1] += pos[1] - minPos[1]; //+= 5;
|
|
|
|
this.add(node, { doProcessChange: false, addedBy: "template" as any });
|
|
|
|
nodes.push(node);
|
|
}
|
|
}
|
|
|
|
//create links
|
|
for (var i = 0; i < template.links.length; ++i) {
|
|
var link_info = template.links[i];
|
|
var origin_node = templateNodeIDToNewNode[link_info[0]];
|
|
var target_node = templateNodeIDToNewNode[link_info[2]];
|
|
if (origin_node && target_node)
|
|
origin_node.connect(link_info[1], target_node, link_info[3]);
|
|
else
|
|
console.error("[ComfyGraphCanvas] nodes missing on template insertion!", link_info);
|
|
}
|
|
|
|
this.afterChange();
|
|
|
|
return templateNodeIDToNewNode;
|
|
}
|
|
}
|