Files
ComfyBox/src/lib/components/ComfyApp.ts
2023-05-07 11:54:54 -05:00

705 lines
25 KiB
TypeScript

import { LiteGraph, LGraph, LGraphCanvas, LGraphNode, type LGraphNodeConstructor, type LGraphNodeExecutable, type SerializedLGraph, type SerializedLGraphGroup, type SerializedLGraphNode, type SerializedLLink, NodeMode, type Vector2, BuiltInSlotType } from "@litegraph-ts/core";
import type { LConnectionKind, INodeSlot } from "@litegraph-ts/core";
import ComfyAPI from "$lib/api"
import defaultGraph from "$lib/defaultGraph"
import { getPngMetadata, importA1111 } from "$lib/pnginfo";
import EventEmitter from "events";
import type TypedEmitter from "typed-emitter";
// Import nodes
import "@litegraph-ts/nodes-basic"
import "@litegraph-ts/nodes-events"
import "@litegraph-ts/nodes-math"
import "@litegraph-ts/nodes-strings"
import "$lib/nodes/index"
import * as nodes from "$lib/nodes/index"
import ComfyGraphCanvas, { type SerializedGraphCanvasState } from "$lib/ComfyGraphCanvas";
import type ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
import * as widgets from "$lib/widgets/index"
import queueState from "$lib/stores/queueState";
import type { SvelteComponentDev } from "svelte/internal";
import type IComfyInputSlot from "$lib/IComfyInputSlot";
import type { SerializedLayoutState } from "$lib/stores/layoutState";
import layoutState from "$lib/stores/layoutState";
import { toast } from '@zerodevx/svelte-toast'
import ComfyGraph from "$lib/ComfyGraph";
import { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode";
import { get } from "svelte/store";
import uiState from "$lib/stores/uiState";
import { promptToGraphVis, workflowToGraphVis } from "$lib/utils";
import notify from "$lib/notify";
export const COMFYBOX_SERIAL_VERSION = 1;
LiteGraph.catch_exceptions = false;
LiteGraph.CANVAS_GRID_SIZE = 32;
if (typeof window !== "undefined") {
// Load default visibility
nodes.ComfyReroute.setDefaultTextVisibility(!!localStorage["Comfy.ComfyReroute.DefaultVisibility"]);
}
type QueueItem = { num: number, batchCount: number }
export type SerializedAppState = {
createdBy: "ComfyBox",
version: number,
workflow: SerializedLGraph,
layout: SerializedLayoutState,
canvas: SerializedGraphCanvasState
}
/** [link origin, link index] | value */
export type SerializedPromptInput = [string, number] | any
export type SerializedPromptInputs = {
inputs: Record<string, SerializedPromptInput>,
class_type: string
}
export type SerializedPromptOutput = Record<string, SerializedPromptInputs>
export type SerializedPrompt = {
workflow: SerializedLGraph,
output: SerializedPromptOutput
}
export type Progress = {
value: number,
max: number
}
export default class ComfyApp {
api: ComfyAPI;
rootEl: HTMLDivElement | null = null;
canvasEl: HTMLCanvasElement | null = null;
canvasCtx: CanvasRenderingContext2D | null = null;
lGraph: ComfyGraph | null = null;
lCanvas: ComfyGraphCanvas | null = null;
dropZone: HTMLElement | null = null;
nodeOutputs: Record<string, any> = {};
dragOverNode: LGraphNode | null = null;
shiftDown: boolean = false;
selectedGroupMoving: boolean = false;
private queueItems: QueueItem[] = [];
private processingQueue: boolean = false;
private alreadySetup = false;
constructor() {
this.api = new ComfyAPI();
}
async setup(): Promise<void> {
if (this.alreadySetup) {
console.error("Already setup")
return;
}
this.rootEl = document.getElementById("main") as HTMLDivElement;
this.canvasEl = document.getElementById("graph-canvas") as HTMLCanvasElement;
this.lGraph = new ComfyGraph();
this.lCanvas = new ComfyGraphCanvas(this, this.canvasEl);
this.canvasCtx = this.canvasEl.getContext("2d");
LiteGraph.release_link_on_empty_shows_menu = true;
LiteGraph.alt_drag_do_clone_nodes = true;
(window as any).LiteGraph = LiteGraph;
// await this.#invokeExtensionsAsync("init");
await this.registerNodes();
// Load previous workflow
let restored = false;
try {
const json = localStorage.getItem("workflow");
if (json) {
const state = JSON.parse(json) as SerializedAppState;
await this.deserialize(state)
restored = true;
}
} catch (err) {
console.error("Error loading previous workflow", err);
}
// We failed to restore a workflow so load the default
if (!restored) {
await this.initDefaultGraph();
}
// Save current workflow automatically
// setInterval(this.saveStateToLocalStorage.bind(this), 1000);
this.addApiUpdateHandlers();
this.addDropHandler();
this.addPasteHandler();
this.addKeyboardHandler();
this.setupColorScheme()
// await this.#invokeExtensionsAsync("setup");
// Ensure the canvas fills the window
this.resizeCanvas();
window.addEventListener("resize", this.resizeCanvas.bind(this));
this.alreadySetup = true;
return Promise.resolve();
}
resizeCanvas() {
this.canvasEl.width = this.canvasEl.parentElement.offsetWidth;
this.canvasEl.height = this.canvasEl.parentElement.offsetHeight;
this.canvasEl.style.width = ""
this.canvasEl.style.height = ""
this.lCanvas.draw(true, true);
}
saveStateToLocalStorage() {
const savedWorkflow = this.serialize();
const json = JSON.stringify(savedWorkflow);
localStorage.setItem("workflow", json)
}
static node_type_overrides: Record<string, typeof ComfyBackendNode> = {}
static widget_type_overrides: Record<string, typeof SvelteComponentDev> = {}
private async registerNodes() {
const app = this;
// Load node definitions from the backend
const defs = await this.api.getNodeDefs();
// Register a node for each definition
for (const nodeId in defs) {
const nodeData = defs[nodeId];
const typeOverride = ComfyApp.node_type_overrides[nodeId]
if (typeOverride)
console.debug("Attaching custom type to received node:", nodeId, typeOverride)
const baseClass: typeof ComfyBackendNode = typeOverride || ComfyBackendNode;
const ctor = class extends baseClass {
constructor(title?: string) {
super(title, nodeId, nodeData);
}
}
const node: LGraphNodeConstructor = {
class: ctor,
title: nodeData.name,
type: nodeId,
desc: `ComfyNode: ${nodeId}`
}
LiteGraph.registerNodeType(node);
node.category = nodeData.category;
}
}
private showDropZone() {
if (this.dropZone)
this.dropZone.style.display = "block";
}
private hideDropZone() {
if (this.dropZone)
this.dropZone.style.display = "none";
}
private allowDrag(event: DragEvent) {
if (event.dataTransfer.items?.length > 0) {
event.dataTransfer.dropEffect = 'copy';
this.showDropZone();
event.preventDefault();
}
}
private async handleDrop(event: DragEvent) {
event.preventDefault();
event.stopPropagation();
this.hideDropZone();
if (event.dataTransfer.files.length > 0) {
await this.handleFile(event.dataTransfer.files[0]);
}
}
private addDropHandler() {
// this.dropZone = document.getElementById("dropzone");
// if (this.dropZone) {
// window.addEventListener('dragenter', this.allowDrag.bind(this));
// this.dropZone.addEventListener('dragover', this.allowDrag.bind(this));
// this.dropZone.addEventListener('dragleave', this.hideDropZone.bind(this));
// this.dropZone.addEventListener('drop', this.handleDrop.bind(this));
// }
// else {
// console.warn("No dropzone detected (probably on mobile).")
// }
}
/**
* Adds a handler on paste that extracts and loads workflows from pasted JSON data
*/
private addPasteHandler() {
// document.addEventListener("paste", (e) => {
// let data = (e.clipboardData || (window as any).clipboardData).getData("text/plain");
// let workflow;
// try {
// data = data.slice(data.indexOf("{"));
// workflow = JSON.parse(data);
// } catch (err) {
// try {
// data = data.slice(data.indexOf("workflow\n"));
// data = data.slice(data.indexOf("{"));
// workflow = JSON.parse(data);
// } catch (error) { }
// }
// if (workflow && workflow.version && workflow.nodes && workflow.extra) {
// this.loadGraphData(workflow);
// }
// });
}
/**
* Handles updates from the API socket
*/
private addApiUpdateHandlers() {
this.api.addEventListener("status", ({ detail: ComfyAPIStatus }: CustomEvent) => {
// this.ui.setStatus(detail);
});
this.api.addEventListener("reconnecting", () => {
// this.ui.dialog.show("Reconnecting...");
});
this.api.addEventListener("reconnected", () => {
// this.ui.dialog.close();
});
this.api.addEventListener("progress", ({ detail }: CustomEvent) => {
queueState.progressUpdated(detail);
this.lGraph.setDirtyCanvas(true, false);
});
this.api.addEventListener("executing", ({ detail }: CustomEvent) => {
queueState.executingUpdated(detail.node);
this.lGraph.setDirtyCanvas(true, false);
});
this.api.addEventListener("executed", ({ detail }: CustomEvent) => {
this.nodeOutputs[detail.node] = detail.output;
const node = this.lGraph.getNodeById(detail.node) as ComfyGraphNode;
if (node?.onExecuted) {
node.onExecuted(detail.output);
}
});
this.api.init();
}
private addKeyboardHandler() {
window.addEventListener("keydown", (e) => {
this.shiftDown = e.shiftKey;
// Queue prompt using ctrl or command + enter
if ((e.ctrlKey || e.metaKey) && (e.key === "Enter" || e.keyCode === 13 || e.keyCode === 10)) {
this.queuePrompt(e.shiftKey ? -1 : 0);
}
});
window.addEventListener("keyup", (e) => {
this.shiftDown = e.shiftKey;
});
}
private setupColorScheme() {
const setColor = (type: any, color: string) => {
this.lCanvas.link_type_colors[type] = color
this.lCanvas.default_connection_color_byType[type] = color
}
// Distinguish frontend/backend connections
const BACKEND_TYPES = ["CLIP", "CLIP_VISION", "CLIP_VISION_OUTPUT", "CONDITIONING", "CONTROL_NET", "LATENT", "MASK", "MODEL", "STYLE_MODEL", "VAE"]
for (const type of BACKEND_TYPES) {
setColor(type, "orange")
}
setColor("IMAGE", "rebeccapurple")
setColor(BuiltInSlotType.EVENT, "lightseagreen")
setColor(BuiltInSlotType.ACTION, "lightseagreen")
}
serialize(): SerializedAppState {
const graph = this.lGraph;
const serializedGraph = graph.serialize()
const serializedLayout = layoutState.serialize()
const serializedCanvas = this.lCanvas.serialize();
return {
createdBy: "ComfyBox",
version: COMFYBOX_SERIAL_VERSION,
workflow: serializedGraph,
layout: serializedLayout,
canvas: serializedCanvas
}
}
async deserialize(data: SerializedAppState) {
if (data.version !== COMFYBOX_SERIAL_VERSION) {
throw `Invalid ComfyBox saved data format: ${data.version}`
}
// Ensure loadGraphData does not trigger any state changes in layoutState
// (isConfiguring is set to true here)
// lGraph.configure will add new nodes, triggering onNodeAdded, but we
// want to restore the layoutState ourselves
layoutState.onStartConfigure();
this.loadGraphData(data.workflow)
// Now restore the layout
// Subsequent added nodes will add the UI data to layoutState
layoutState.deserialize(data.layout, this.lGraph)
// Restore canvas offset/zoom
this.lCanvas.deserialize(data.canvas)
await this.refreshComboInNodes();
this.lGraph.start();
this.lGraph.eventBus.on("afterExecute", () => this.lCanvas.draw(true))
}
async initDefaultGraph() {
const state = structuredClone(defaultGraph)
await this.deserialize(state)
}
/**
* Populates the graph with the specified workflow data
* @param {*} graphData A serialized graph object
*/
loadGraphData(graphData: SerializedLGraph) {
this.clean();
// Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
for (let n of graphData.nodes) {
if (n.type == "T2IAdapterLoader") n.type = "ControlNetLoader";
}
this.lGraph.configure(graphData);
for (const node of this.lGraph._nodes) {
const size = node.computeSize();
size[0] = Math.max(node.size[0], size[0]);
size[1] = Math.max(node.size[1], size[1]);
node.size = size;
// this.#invokeExtensions("loadedGraphNode", node);
}
}
clear() {
this.clean();
const blankGraph: SerializedLGraph = {
last_node_id: 0,
last_link_id: 0,
nodes: [],
links: [],
groups: [],
config: {},
extra: {},
version: 0
}
layoutState.onStartConfigure();
this.lGraph.configure(blankGraph)
layoutState.initDefaultLayout();
}
/**
* Converts the current graph workflow for sending to the API
* @returns The workflow and node links
*/
async graphToPrompt(tag: string | null = null): Promise<SerializedPrompt> {
// Run frontend-only logic
this.lGraph.runStep(1)
const workflow = this.lGraph.serialize();
const output = {};
// Process nodes in order of execution
for (const node_ of this.lGraph.computeExecutionOrder<ComfyGraphNode>(false, null)) {
const n = workflow.nodes.find((n) => n.id === node_.id);
if (!node_.isBackendNode) {
// console.debug("Not serializing node: ", node_.type)
continue;
}
const node = node_ as ComfyBackendNode;
if (tag && node.tags.indexOf(tag) === -1) {
console.debug("Skipping tagged node", tag, node.tags)
continue;
}
if (node.mode === NodeMode.NEVER) {
// Don't serialize muted nodes
continue;
}
const inputs = {};
// Store all link values
if (node.inputs) {
for (let i = 0; i < node.inputs.length; i++) {
const inp = node.inputs[i];
const inputLink = node.getInputLink(i)
const inputNode = node.getInputNode(i)
if (inputNode && tag && "tags" in inputNode && (inputNode.tags as string[]).indexOf(tag) === -1) {
continue;
}
if (!inputLink || !inputNode) {
if ("config" in inp) {
const defaultValue = (inp as IComfyInputSlot).config?.defaultValue
if (defaultValue !== null && defaultValue !== undefined)
inputs[inp.name] = defaultValue
}
continue;
}
let serialize = true;
if ("config" in inp)
serialize = (inp as IComfyInputSlot).serialize
let isBackendNode = node.isBackendNode;
let isInputBackendNode = false;
if ("isBackendNode" in inputNode)
isInputBackendNode = (inputNode as ComfyGraphNode).isBackendNode;
// The reasoning behind this check:
// We only want to serialize inputs to nodes with backend equivalents.
// And in ComfyBox, the nodes in litegraph *never* have widgets, instead they're all inputs.
// All values are passed by separate frontend-only nodes,
// either UI-bound or something like ConstantInteger.
// So we know that any value passed into a backend node *must* come from
// a frontend node.
// The rest (links between backend nodes) will be serialized after this bit runs.
if (serialize && isBackendNode && !isInputBackendNode) {
inputs[inp.name] = inputLink.data
}
}
}
// Store all links between nodes
for (let i = 0; i < node.inputs.length; i++) {
let parent: ComfyGraphNode = node.getInputNode(i) as ComfyGraphNode;
if (parent) {
const seen = {}
let link = node.getInputLink(i);
const isValidParent = (parent: ComfyGraphNode) => {
if (!parent || parent.isBackendNode)
return false;
if ("tags" in parent && (parent.tags as string[]).indexOf(tag) === -1)
return false;
return true;
}
while (isValidParent(parent)) {
link = parent.getInputLink(link.origin_slot);
if (link && !seen[link.id]) {
seen[link.id] = true
const inputNode = parent.getInputNode(link.origin_slot) as ComfyGraphNode;
if (inputNode && "tags" in inputNode && tag && (inputNode.tags as string[]).indexOf(tag) === -1) {
console.debug("Skipping tagged parent node", tag, node.tags)
parent = null;
}
else {
parent = inputNode;
}
} else {
parent = null;
}
}
if (link && parent && parent.isBackendNode) {
if ("tags" in parent && tag && (parent.tags as string[]).indexOf(tag) === -1)
continue;
const input = node.inputs[i]
// TODO can null be a legitimate value in some cases?
// Nodes like CLIPLoader will never have a value in the frontend, hence "null".
if (!(input.name in inputs))
inputs[input.name] = [String(link.origin_id), link.origin_slot];
}
}
}
output[String(node.id)] = {
inputs,
class_type: node.comfyClass,
};
}
// Remove inputs connected to removed nodes
for (const o in output) {
for (const i in output[o].inputs) {
if (Array.isArray(output[o].inputs[i])
&& output[o].inputs[i].length === 2
&& !output[output[o].inputs[i][0]]) {
console.debug("Prune removed node link", o, i, output[o].inputs[i])
delete output[o].inputs[i];
}
}
}
// console.debug({ workflow, output })
// console.debug(promptToGraphVis({ workflow, output }))
return { workflow, output };
}
async queuePrompt(num: number, batchCount: number = 1, tag: string | null = null) {
this.queueItems.push({ num, batchCount });
// Only have one action process the items so each one gets a unique seed correctly
if (this.processingQueue) {
return;
}
if (tag === "")
tag = null;
this.processingQueue = true;
try {
while (this.queueItems.length) {
({ num, batchCount } = this.queueItems.pop());
console.debug(`Queue get! ${num} ${batchCount} ${tag}`);
for (let i = 0; i < batchCount; i++) {
for (const node of this.lGraph._nodes_in_order) {
if ("beforeQueued" in node) {
(node as ComfyGraphNode).beforeQueued(tag);
}
}
const p = await this.graphToPrompt(tag);
try {
await this.api.queuePrompt(num, p);
} catch (error) {
// this.ui.dialog.show(error.response || error.toString());
const mes = error.response || error.toString()
notify(`Error queuing prompt:\n${mes}`, null, "error")
console.error(promptToGraphVis(p))
console.error("Error queuing prompt", mes, num, p)
break;
}
for (const n of p.workflow.nodes) {
const node = this.lGraph.getNodeById(n.id);
if ("afterQueued" in node) {
(node as ComfyGraphNode).afterQueued(p, tag);
}
}
this.lCanvas.draw(true, true);
// await this.ui.queue.update();
}
}
} finally {
console.log("Queue finished!");
this.processingQueue = false;
}
}
/**
* Loads workflow data from the specified file
*/
async handleFile(file: File) {
if (file.type === "image/png") {
const pngInfo = await getPngMetadata(file);
if (pngInfo) {
if (pngInfo.comfyBoxConfig) {
this.deserialize(JSON.parse(pngInfo.comfyBoxConfig));
} else if (pngInfo.parameters) {
throw "TODO import A111 import!"
// importA1111(this.lGraph, pngInfo.parameters, this.api);
}
else {
console.error("No metadata found in image file.", pngInfo)
notify("No metadata found in image file.")
}
}
} else if (file.type === "application/json" || file.name.endsWith(".json")) {
const reader = new FileReader();
reader.onload = () => {
this.deserialize(JSON.parse(reader.result as string));
};
reader.readAsText(file);
}
}
// registerExtension(extension) {
// if (!extension.name) {
// throw new Error("Extensions must have a 'name' property.");
// }
// if (this.extensions.find((ext) => ext.name === extension.name)) {
// throw new Error(`Extension named '${extension.name}' already registered.`);
// }
// this.extensions.push(extension);
// }
/**
* Refresh combo list on whole nodes
*/
async refreshComboInNodes() {
const defs = await this.api.getNodeDefs();
for (let nodeNum in this.lGraph._nodes) {
const node = this.lGraph._nodes[nodeNum];
const def = defs[node.type];
for (let index = 0; index < node.inputs.length; index++) {
const input = node.inputs[index];
if ("config" in input) {
const comfyInput = input as IComfyInputSlot;
if (comfyInput.defaultWidgetNode == nodes.ComfyComboNode && def["input"]["required"][comfyInput.name] !== undefined) {
comfyInput.config.values = def["input"]["required"][comfyInput.name][0];
const inputNode = node.getInputNode(index)
if (inputNode && "doAutoConfig" in inputNode) {
const comfyInputNode = inputNode as nodes.ComfyWidgetNode;
comfyInputNode.doAutoConfig(comfyInput)
if (!comfyInput.config.values.includes(get(comfyInputNode.value))) {
comfyInputNode.setValue(comfyInput.config.defaultValue || comfyInput.config.values[0])
}
}
}
}
}
}
}
/**
* Clean current state
*/
clean() {
this.nodeOutputs = {};
}
}