diff --git a/src/lib/components/ComfyApp.svelte b/src/lib/components/ComfyApp.svelte index c2d0955..0463f3c 100644 --- a/src/lib/components/ComfyApp.svelte +++ b/src/lib/components/ComfyApp.svelte @@ -4,10 +4,10 @@ import { Pane, Splitpanes } from 'svelte-splitpanes'; import { Button } from "@gradio/button"; import ComfyUIPane from "./ComfyUIPane.svelte"; - import ComfyApp from "./ComfyApp"; + import ComfyApp, { type SerializedAppState } from "./ComfyApp"; import widgetState from "$lib/stores/widgetState"; - import { LGraphNode } from "@litegraph-ts/core"; + import { LGraph, LGraphNode } from "@litegraph-ts/core"; let app: ComfyApp = undefined; let uiPane: ComfyUIPane = undefined; @@ -51,6 +51,22 @@ let graphResizeTimer: typeof Timer = -1; + function doAutosave(graph: LGraph): void { + const serializedGraph = graph.serialize() + const serializedPaneOrder = uiPane.serialize() + + const savedWorkflow = { + graph: serializedGraph, + panes: serializedPaneOrder + } + + localStorage.setItem("workflow", JSON.stringify(savedWorkflow)) + } + + function doRestore(workflow: SerializedAppState) { + uiPane.restore(workflow.panes); + } + onMount(async () => { app = new ComfyApp(); @@ -58,11 +74,14 @@ app.eventBus.on("nodeRemoved", widgetState.nodeRemoved); app.eventBus.on("configured", widgetState.configureFinished); app.eventBus.on("cleared", widgetState.clear); + app.eventBus.on("autosave", doAutosave); + app.eventBus.on("restored", doRestore); await app.setup(); refreshView(); (window as any).app = app; + (window as any).appPane = uiPane; let wrappers = containerElem.querySelectorAll(".pane-wrapper") for (const wrapper of wrappers) { diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts index ea7be82..a21dbb6 100644 --- a/src/lib/components/ComfyApp.ts +++ b/src/lib/components/ComfyApp.ts @@ -22,12 +22,25 @@ if (typeof window !== "undefined") { type QueueItem = { num: number, batchCount: number } +export type SerializedPanes = { + panels: { nodeId: number }[][] +} + +export type SerializedAppState = { + panes: SerializedPanes, + workflow: SerializedLGraph +} + type ComfyAppEvents = { configured: (graph: LGraph) => void nodeAdded: (node: LGraphNode) => void nodeRemoved: (node: LGraphNode) => void nodeConnectionChanged: (kind: LConnectionKind, node: LGraphNode, slot: INodeSlot, targetNode: LGraphNode, targetSlot: INodeSlot) => void cleared: () => void + beforeChange: (graph: LGraph, param: any) => void + afterChange: (graph: LGraph, param: any) => void + autosave: (graph: LGraph) => void + restored: (workflow: SerializedAppState) => void } interface ComfyGraphNodeExecutable extends LGraphNodeExecutable { @@ -85,8 +98,9 @@ export default class ComfyApp { try { const json = localStorage.getItem("workflow"); if (json) { - const workflow = JSON.parse(json); - this.loadGraphData(workflow); + const workflow = JSON.parse(json) as SerializedAppState; + this.loadGraphData(workflow["workflow"]); + this.eventBus.emit("restored", workflow); restored = true; } } catch (err) { @@ -99,7 +113,7 @@ export default class ComfyApp { } // Save current workflow automatically - setInterval(() => localStorage.setItem("workflow", JSON.stringify(this.lGraph.serialize())), 1000); + setInterval(this.requestAutosave.bind(this), 1000); this.addApiUpdateHandlers(); this.addDropHandler(); @@ -158,6 +172,10 @@ export default class ComfyApp { this.eventBus.emit("cleared"); } + private requestAutosave() { + this.eventBus.emit("autosave", this.lGraph); + } + private addGraphLifecycleHooks() { this.lGraph.onConfigure = this.graphOnConfigure.bind(this); this.lGraph.onBeforeChange = this.graphOnBeforeChange.bind(this); @@ -372,11 +390,11 @@ export default class ComfyApp { * Populates the graph with the specified workflow data * @param {*} graphData A serialized graph object */ - loadGraphData(graphData: any = null) { + loadGraphData(graphData?: SerializedLGraph) { this.clean(); if (!graphData) { - graphData = structuredClone(defaultGraph); + graphData = structuredClone(defaultGraph) as SerializedLGraph; } // Patch T2IAdapterLoader to ControlNetLoader since they are the same node now diff --git a/src/lib/components/ComfyUIPane.svelte b/src/lib/components/ComfyUIPane.svelte index 617930f..500a699 100644 --- a/src/lib/components/ComfyUIPane.svelte +++ b/src/lib/components/ComfyUIPane.svelte @@ -2,24 +2,30 @@ import { get } from "svelte/store"; import { LGraphNode, LGraph } from "@litegraph-ts/core"; import type { IWidget } from "@litegraph-ts/core"; - import ComfyApp from "./ComfyApp"; + import ComfyApp from "./ComfyApp"; + import type { SerializedPanes } from "./ComfyApp" import ComfyPane from "./ComfyPane.svelte"; - import widgetState from "$lib/stores/widgetState"; + import widgetState, { type WidgetUIState } from "$lib/stores/widgetState"; + + type DragItem = { + id: number, + node: LGraphNode + } export let app: ComfyApp; let dragConfigured: boolean = false; - export let dragItemss: any[][] = [] + export let dragItems: DragItem[][] = [] export let totalId = 0; - function addUIForNewNode(node: LGraphNode) { + function findLeastPopulatedPaneIndex(): number { let minWidgetCount = 2 ** 64; let minIndex = 0; let state = get(widgetState); - for (let i = 0; i < dragItemss.length; i++) { + for (let i = 0; i < dragItems.length; i++) { let widgetCount = 0; - for (let j = 0; j < dragItemss[i].length; j++) { - const nodeID = dragItemss[i][j].node.id; + for (let j = 0; j < dragItems[i].length; j++) { + const nodeID = dragItems[i][j].node.id; widgetCount += state[nodeID].length; } if (widgetCount < minWidgetCount) { @@ -27,7 +33,13 @@ minIndex = i; } } - dragItemss[minIndex].push({ id: totalId++, node: node }); + return minIndex + } + + function addUIForNewNode(node: LGraphNode, paneIndex?: number) { + if (!paneIndex) + paneIndex = findLeastPopulatedPaneIndex(); + dragItems[paneIndex].push({ id: totalId++, node: node }); } $: if(app && !dragConfigured) { @@ -38,15 +50,54 @@ /* * Serialize UI panel order so it can be restored when workflow is loaded */ - function getUIState(): any { + export function serialize(): any { + let panels = [] + for (let i = 0; i < dragItems.length; i++) { + panels[i] = []; + for (let j = 0; j < dragItems[i].length; j++) { + panels[i].push({ nodeId: dragItems[i][j].node.id }); + } + } + return { + panels + } + } + export function restore(panels: SerializedPanes) { + let nodeIdToDragItem: Record = {}; + for (let i = 0; i < dragItems.length; i++) { + for (const dragItem of dragItems[i]) { + nodeIdToDragItem[dragItem.node.id] = dragItem + } + } + + for (let i = 0; i < panels.panels.length; i++) { + dragItems[i].length = 0; + for (const panel of panels.panels[i]) { + const dragItem = nodeIdToDragItem[panel.nodeId]; + if (dragItem) { + delete nodeIdToDragItem[panel.nodeId]; + dragItems[i].push(dragItem) + } + } + } + + // Put everything left over into other columns + if (Object.keys(nodeIdToDragItem).length > 0) { + console.warn("Extra panels without ordering found", nodeIdToDragItem) + for (const nodeId in nodeIdToDragItem) { + const dragItem = nodeIdToDragItem[nodeId]; + const paneIndex = findLeastPopulatedPaneIndex(); + dragItems[paneIndex].push(dragItem); + } + } }
- - - + + +