Save panel ordering to workflow file

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<HTMLDivNode>(".pane-wrapper")
      for (const wrapper of wrappers) {

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

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<number, DragItem> = {};
+     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);
+         }
+     }
  }
 </script>
 
 <div id="comfy-ui-panes" >
-    <ComfyPane bind:dragItems={dragItemss[0]} />
-    <ComfyPane bind:dragItems={dragItemss[1]} />
-    <ComfyPane bind:dragItems={dragItemss[2]} />
+    <ComfyPane bind:dragItems={dragItems[0]} />
+    <ComfyPane bind:dragItems={dragItems[1]} />
+    <ComfyPane bind:dragItems={dragItems[2]} />
 </div>
 
 <style>