Files
ComfyBox/src/lib/stores/workflowState.ts
2023-06-01 18:43:45 -05:00

504 lines
16 KiB
TypeScript

import type { SerializedGraphCanvasState } from '$lib/ComfyGraphCanvas';
import { clamp, LGraphNode, type LGraphCanvas, type NodeID, type SerializedLGraph, type UUID, LGraph, LiteGraph, type SlotType, NodeMode } from '@litegraph-ts/core';
import { get, writable } from 'svelte/store';
import type { Readable, Writable } from 'svelte/store';
import { defaultWorkflowAttributes, type SerializedLayoutState, type WritableLayoutStateStore } from './layoutStates';
import ComfyGraph from '$lib/ComfyGraph';
import layoutStates from './layoutStates';
import { v4 as uuidv4 } from "uuid";
import type ComfyGraphCanvas from '$lib/ComfyGraphCanvas';
import { blankGraph } from '$lib/defaultGraph';
import type { SerializedAppState, SerializedPrompt } from '$lib/components/ComfyApp';
import type ComfyReceiveOutputNode from '$lib/nodes/actions/ComfyReceiveOutputNode';
import type { ComfyBoxPromptExtraData, PromptID } from '$lib/api';
import type { ComfyAPIPromptErrorResponse, ComfyExecutionError } from '$lib/apiErrors';
type ActiveCanvas = {
canvas: LGraphCanvas | null;
canvasHandler: () => void | null;
state: SerializedGraphCanvasState;
}
export type SerializedWorkflowState = {
graph: SerializedLGraph,
layout: SerializedLayoutState,
attrs: WorkflowAttributes
}
/*
* ID for an opened workflow.
*
* Unlike NodeID and PromptID, these are *not* saved to the workflow itself.
* They are only used for identifying an open workflow in the program. If the
* workflow is closed and reopened, a different workflow ID will be assigned to
* it.
*/
export type WorkflowInstID = UUID;
/*
* Global workflow attributes
*/
export type WorkflowAttributes = {
/*
* Title of the workflow.
*/
title: string,
/*
* Name of the "Queue Prompt" button. Set to blank to hide the button.
*/
queuePromptButtonName: string,
/*
* If true, clicking the "Queue Prompt" button will run the default
* subgraph. Set this to false if you need special behavior before running
* any subgraphs, and instead use the `onDefaultQueueAction` event of the
* Comfy.QueueEvents node.
*/
queuePromptButtonRunWorkflow: boolean,
/*
* Default subgraph to run if `queuePromptButtonRunWorkflow` is `true`. Set
* to blank to run the default subgraph (tagless).
*/
queuePromptButtonDefaultWorkflow: string,
/*
* If true, notifications will be shown when a prompt is queued and
* completed. Set to false if you need more detailed control over the
* notification type/contents, and use the `ComfyNotifyAction` node instead.
*/
showDefaultNotifications: boolean,
}
export type WorkflowValidationError = {
type: "validation"
workflowID: WorkflowInstID,
error: ComfyAPIPromptErrorResponse,
prompt: SerializedPrompt,
extraData: ComfyBoxPromptExtraData
}
export type WorkflowExecutionError = {
type: "execution"
error: ComfyExecutionError,
}
export type WorkflowError = WorkflowValidationError | WorkflowExecutionError;
export class ComfyBoxWorkflow {
/*
* Used for uniquely identifying the instance of the opened workflow in the frontend.
*/
id: WorkflowInstID;
/*
* Graph of this workflow, whose nodes are bound to the UI layout
*/
graph: ComfyGraph;
/*
* Global workflow attributes
*/
attrs: WorkflowAttributes;
/*
* True if an unsaved modification has been detected on this workflow
*/
isModified: boolean = false;
/*
* Missing node types encountered when deserializing the graph
*/
missingNodeTypes: Set<string> = new Set();
/*
* Completed queue entry ID that holds the last validation/execution error.
*/
lastError?: PromptID
get layout(): WritableLayoutStateStore | null {
return layoutStates.getLayout(this.id)
}
/*
* Graph canvases attached to the graph of this workflow
*/
canvases: Record<string, ActiveCanvas> = {};
constructor(title: string) {
this.id = uuidv4();
this.attrs = {
...defaultWorkflowAttributes,
title,
}
this.graph = new ComfyGraph(this.id);
}
notifyModified() {
this.isModified = true;
store.set(get(store));
}
setAttribute<K extends keyof WorkflowAttributes>(key: K, value: WorkflowAttributes[K]) {
this.attrs[key] = value;
this.notifyModified();
}
start(key: string, canvas: ComfyGraphCanvas) {
if (this.canvases[key] != null)
throw new Error(`This workflow is already being displayed on canvas ${key}`)
const canvasHandler = () => canvas.draw(true);
this.canvases[key] = {
canvas,
canvasHandler,
state: {
// TODO
offset: [0, 0],
scale: 1
}
}
this.graph.attachCanvas(canvas);
this.graph.eventBus.on("afterExecute", canvasHandler)
if (Object.keys(this.canvases).length === 1)
this.graph.start();
}
stop(key: string) {
const canvas = this.canvases[key]
if (canvas == null) {
console.debug("This workflow is not being displayed on canvas ${key}")
return;
}
canvas.canvas.closeAllSubgraphs();
this.graph.detachCanvas(canvas.canvas);
this.graph.eventBus.removeListener("afterExecute", canvas.canvasHandler)
delete this.canvases[key]
if (Object.keys(this.canvases).length === 0)
this.graph.stop();
}
stopAll() {
for (const key of Object.keys(this.canvases))
this.stop(key)
this.graph.stop()
}
serialize(layoutState: WritableLayoutStateStore): SerializedWorkflowState {
const graph = this.graph;
const serializedGraph = graph.serialize()
const serializedLayout = layoutState.serialize()
return {
graph: serializedGraph,
layout: serializedLayout,
attrs: this.attrs
}
}
/*
* Creates a workflow and layout.
*
* NOTE: The layout will be attached to the global store, but the workflow
* will not. If you change your mind later be sure to call
* layoutStates.remove(workflow.id)!
*/
static create(title: string = "New Workflow"): [ComfyBoxWorkflow, WritableLayoutStateStore] {
const workflow = new ComfyBoxWorkflow(title);
const layoutState = layoutStates.create(workflow);
return [workflow, layoutState]
}
deserialize(layoutState: WritableLayoutStateStore, data: SerializedWorkflowState) {
this.missingNodeTypes.clear();
for (let n of data.graph.nodes) {
// Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
if (n.type == "T2IAdapterLoader") n.type = "ControlNetLoader";
// Find missing node types
if (!(n.type in LiteGraph.registered_node_types)) {
this.missingNodeTypes.add(n.type);
}
}
// 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();
// Patch T2IAdapterLoader to ControlNetLoader since they are the same node now
for (let n of data.graph.nodes) {
if (n.type == "T2IAdapterLoader") n.type = "ControlNetLoader";
}
this.graph.configure(data.graph);
this.graph.workflowID = this.id;
for (const node of this.graph._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);
}
this.attrs = { ...defaultWorkflowAttributes, ...data.attrs };
// Now restore the layout
// Subsequent added nodes will add the UI data to layoutState
// TODO
layoutState.deserialize(data.layout, this.graph)
}
}
export type WorkflowState = {
openedWorkflows: ComfyBoxWorkflow[],
openedWorkflowsByID: Record<WorkflowInstID, ComfyBoxWorkflow>,
activeWorkflowID: WorkflowInstID | null,
activeWorkflow: ComfyBoxWorkflow | null,
}
export type WorkflowReceiveOutputTargets = {
workflow: ComfyBoxWorkflow,
targetNodes: ComfyReceiveOutputNode[]
}
type WorkflowStateOps = {
getWorkflow: (id: WorkflowInstID) => ComfyBoxWorkflow | null
getWorkflowByGraph: (graph: LGraph) => ComfyBoxWorkflow | null
getWorkflowByNode: (node: LGraphNode) => ComfyBoxWorkflow | null
getWorkflowByNodeID: (id: NodeID) => ComfyBoxWorkflow | null
getActiveWorkflow: () => ComfyBoxWorkflow | null
createNewWorkflow: (canvas: ComfyGraphCanvas, title?: string, setActive?: boolean) => ComfyBoxWorkflow,
openWorkflow: (canvas: ComfyGraphCanvas, data: SerializedAppState, setActive?: boolean) => ComfyBoxWorkflow,
addWorkflow: (canvas: ComfyGraphCanvas, data: ComfyBoxWorkflow, setActive?: boolean) => void,
closeWorkflow: (canvas: ComfyGraphCanvas, index: number) => void,
closeAllWorkflows: (canvas: ComfyGraphCanvas) => void,
setActiveWorkflow: (canvas: ComfyGraphCanvas, index: number | WorkflowInstID) => ComfyBoxWorkflow | null,
findReceiveOutputTargets: (type: SlotType | SlotType[]) => WorkflowReceiveOutputTargets[],
afterQueued: (id: WorkflowInstID, promptID: PromptID) => void
promptError: (id: WorkflowInstID, promptID: PromptID) => void
executionError: (id: WorkflowInstID, promptID: PromptID) => void
}
export type WritableWorkflowStateStore = Writable<WorkflowState> & WorkflowStateOps;
const store: Writable<WorkflowState> = writable(
{
openedWorkflows: [],
openedWorkflowsByID: {},
activeWorkflowID: null,
activeWorkflow: null
})
function getWorkflow(id: WorkflowInstID): ComfyBoxWorkflow | null {
return get(store).openedWorkflowsByID[id];
}
function getWorkflowByGraph(graph: LGraph): ComfyBoxWorkflow | null {
const workflowID = (graph.getRootGraph() as ComfyGraph)?.workflowID;
if (workflowID == null)
return null;
return getWorkflow(workflowID);
}
function getWorkflowByNode(node: LGraphNode): ComfyBoxWorkflow | null {
return getWorkflowByGraph(node.graph);
}
function getWorkflowByNodeID(id: NodeID): ComfyBoxWorkflow | null {
return Object.values(get(store).openedWorkflows).find(w => {
return w.graph.getNodeByIdRecursive(id) != null
})
}
function getActiveWorkflow(): ComfyBoxWorkflow | null {
const state = get(store);
if (state.activeWorkflowID == null)
return null;
return state.openedWorkflowsByID[state.activeWorkflowID];
}
function createNewWorkflow(canvas: ComfyGraphCanvas, title: string = "New Workflow", setActive: boolean = false): ComfyBoxWorkflow {
const workflow = new ComfyBoxWorkflow(title);
const layoutState = layoutStates.create(workflow);
layoutState.initDefaultLayout();
const state = get(store);
state.openedWorkflows.push(workflow);
state.openedWorkflowsByID[workflow.id] = workflow;
if (setActive || state.activeWorkflowID == null)
setActiveWorkflow(canvas, state.openedWorkflows.length - 1)
store.set(state)
return workflow;
}
function openWorkflow(canvas: ComfyGraphCanvas, data: SerializedAppState, setActive: boolean = true): ComfyBoxWorkflow {
const [workflow, layoutState] = ComfyBoxWorkflow.create("Workflow")
workflow.deserialize(layoutState, { graph: data.workflow, layout: data.layout, attrs: data.attrs })
addWorkflow(canvas, workflow, setActive);
return workflow;
}
function addWorkflow(canvas: ComfyGraphCanvas, workflow: ComfyBoxWorkflow, setActive: boolean = true) {
const state = get(store);
state.openedWorkflows.push(workflow);
state.openedWorkflowsByID[workflow.id] = workflow;
if (setActive || state.activeWorkflowID == null)
setActiveWorkflow(canvas, state.openedWorkflows.length - 1)
store.set(state)
return workflow;
}
function closeWorkflow(canvas: ComfyGraphCanvas, index: number) {
const state = get(store);
if (index < 0 || index >= state.openedWorkflows.length)
return;
const workflow = state.openedWorkflows[index];
workflow.stopAll();
layoutStates.remove(workflow.id)
state.openedWorkflows.splice(index, 1)
delete state.openedWorkflowsByID[workflow.id]
let newIndex = clamp(index, 0, state.openedWorkflows.length - 1)
setActiveWorkflow(canvas, newIndex);
store.set(state);
}
function closeAllWorkflows(canvas: ComfyGraphCanvas) {
const state = get(store)
while (state.openedWorkflows.length > 0)
closeWorkflow(canvas, 0)
}
function setActiveWorkflow(canvas: ComfyGraphCanvas, index: number | WorkflowInstID): ComfyBoxWorkflow | null {
const state = get(store);
if (state.openedWorkflows.length === 0) {
state.activeWorkflowID = null;
state.activeWorkflow = null
return null;
}
if (typeof index === "string") {
index = state.openedWorkflows.findIndex(w => w.id === index)
}
if (index < 0 || index >= state.openedWorkflows.length)
return state.activeWorkflow;
const workflow = state.openedWorkflows[index]
if (workflow.id === state.activeWorkflowID)
return state.activeWorkflow;
if (state.activeWorkflow != null)
state.activeWorkflow.stop("app")
state.activeWorkflowID = workflow.id;
state.activeWorkflow = workflow;
workflow.start("app", canvas);
canvas.deserialize(workflow.canvases["app"].state)
store.set(state)
return workflow;
}
function findReceiveOutputTargets(type: SlotType | SlotType[]): WorkflowReceiveOutputTargets[] {
let result = []
const state = get(store);
if (!Array.isArray(type))
type = [type]
const types = new Set(type);
for (const workflow of state.openedWorkflows) {
const targetNodes = workflow.graph
// can't use class here because of circular import
.findNodesByTypeRecursive("events/receive_output")
.filter(n => {
return types.has(n.properties.type) && n.mode === NodeMode.ALWAYS
})
if (targetNodes.length > 0)
result.push({ workflow, targetNodes });
}
return result;
}
function afterQueued(id: WorkflowInstID, promptID: PromptID) {
const workflow = getWorkflow(id);
if (workflow == null) {
console.warn("[workflowState] afterQueued: workflow not found", id, promptID)
return
}
workflow.lastError = null;
}
function promptError(id: WorkflowInstID, promptID: PromptID) {
const workflow = getWorkflow(id);
if (workflow == null) {
console.warn("[workflowState] promptError: workflow not found", id, promptID)
return
}
workflow.lastError = promptID;
}
function executionError(id: WorkflowInstID, promptID: PromptID) {
const workflow = getWorkflow(id);
if (workflow == null) {
console.warn("[workflowState] executionError: workflow not found", id, promptID)
return
}
workflow.lastError = promptID;
}
const workflowStateStore: WritableWorkflowStateStore =
{
...store,
getWorkflow,
getWorkflowByGraph,
getWorkflowByNode,
getWorkflowByNodeID,
getActiveWorkflow,
createNewWorkflow,
openWorkflow,
addWorkflow,
closeWorkflow,
closeAllWorkflows,
setActiveWorkflow,
findReceiveOutputTargets,
afterQueued,
promptError,
executionError,
}
export default workflowStateStore;