Save panel ordering to workflow file

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>