Restore/load workflow button in prompt details modal

src/lib/components/ComfyApp.ts

@@ -28,7 +28,7 @@ import modalState from "$lib/stores/modalState";
 import queueState from "$lib/stores/queueState";
 import selectionState from "$lib/stores/selectionState";
 import uiState from "$lib/stores/uiState";
-import workflowState, { ComfyBoxWorkflow, type WorkflowAttributes, type WorkflowInstID } from "$lib/stores/workflowState";
+import workflowState, { ComfyBoxWorkflow, OpenWorkflowMode, type WorkflowAttributes, type WorkflowInstID } from "$lib/stores/workflowState";
 import { playSound, readFileToText, type SerializedPromptOutput } from "$lib/utils";
 import { basename, capitalize, download, graphToGraphVis, jsonToJsObject, promptToGraphVis, range } from "$lib/utils";
 import { tick } from "svelte";
@@ -49,7 +49,7 @@ if (typeof window !== "undefined") {
 }
 
 export type OpenWorkflowOptions = {
-    setActive?: boolean,
+    mode: OpenWorkflowMode,
     refreshCombos?: boolean | Record<string, ComfyNodeDef>,
     warnMissingNodeTypes?: boolean,
 }
@@ -234,7 +234,7 @@ export default class ComfyApp {
         if (!restored) {
             const options: OpenWorkflowOptions = {
                 refreshCombos: defs,
-                setActive: false
+                mode: OpenWorkflowMode.Append
             }
             await this.initDefaultWorkflow("defaultWorkflow", options);
             await this.initDefaultWorkflow("upscaleByModel", options);
@@ -408,7 +408,7 @@ export default class ComfyApp {
             return false;
 
         await Promise.all(workflows.map(w => {
-            return this.openWorkflow(w, { refreshCombos: defs, warnMissingNodeTypes: false, setActive: false }).catch(error => {
+            return this.openWorkflow(w, { refreshCombos: defs, warnMissingNodeTypes: false, mode: OpenWorkflowMode.Append }).catch(error => {
                 console.error("Failed restoring previous workflow", error)
                 notify(`Failed restoring previous workflow: ${error}`, { type: "error" })
             })
@@ -778,7 +778,7 @@ export default class ComfyApp {
     }
 
     async openWorkflow(data: SerializedAppState, options: OpenWorkflowOptions = {
-        setActive: true,
+        mode: OpenWorkflowMode.AppendAndSetActive,
         refreshCombos: true,
         warnMissingNodeTypes: true
     }
@@ -793,7 +793,7 @@ export default class ComfyApp {
 
         let workflow: ComfyBoxWorkflow;
         try {
-            workflow = workflowState.openWorkflow(this.lCanvas, data, options.setActive);
+            workflow = workflowState.openWorkflow(this.lCanvas, data, options.mode);
         }
         catch (error) {
             modalState.pushModal({
@@ -827,7 +827,11 @@ export default class ComfyApp {
         return workflow;
     }
 
-    async openVanillaWorkflow(data: SerializedLGraph, filename: string) {
+    async openVanillaWorkflow(data: SerializedLGraph, filename: string, options: OpenWorkflowOptions = {
+        mode: OpenWorkflowMode.AppendAndSetActive,
+        refreshCombos: true,
+        warnMissingNodeTypes: true
+    }) {
         const title = basename(filename)
 
         const attrs: WorkflowAttributes = {
@@ -844,7 +848,7 @@ export default class ComfyApp {
 
         const addWorkflow = () => {
             notify("Converted ComfyUI workflow to ComfyBox format.", { type: "info" })
-            workflowState.addWorkflow(this.lCanvas, comfyBoxWorkflow)
+            workflowState.addWorkflow(this.lCanvas, comfyBoxWorkflow, options.mode)
             this.lCanvas.deserialize(canvas);
         }
 
@@ -886,7 +890,7 @@ export default class ComfyApp {
     }
 
     createNewWorkflow() {
-        workflowState.createNewWorkflow(this.lCanvas, undefined, true);
+        workflowState.createNewWorkflow(this.lCanvas, undefined, OpenWorkflowMode.AppendAndSetActive);
         selectionState.clear();
     }
 
@@ -1146,13 +1150,13 @@ export default class ComfyApp {
     /**
      * Loads workflow data from the specified file
      */
-    async handleFile(file: File) {
+    async handleFile(file: File, openOptions?: OpenWorkflowOptions) {
         if (file.type === "image/png") {
             const buffer = await file.arrayBuffer();
             const pngInfo = await parsePNGMetadata(buffer);
             if (pngInfo) {
                 if (pngInfo.comfyBoxWorkflow) {
-                    await this.openWorkflow(JSON.parse(pngInfo.comfyBoxWorkflow));
+                    await this.openWorkflow(JSON.parse(pngInfo.comfyBoxWorkflow), openOptions);
                 } else if (pngInfo.workflow) {
                     const workflow = JSON.parse(pngInfo.workflow);
                     await this.openVanillaWorkflow(workflow, file.name);
@@ -1187,7 +1191,7 @@ export default class ComfyApp {
             reader.onload = async () => {
                 const result = JSON.parse(reader.result as string)
                 if (isComfyBoxWorkflow(result)) {
-                    await this.openWorkflow(result);
+                    await this.openWorkflow(result, openOptions);
                 }
                 else if (isVanillaWorkflow(result)) {
                     await this.openVanillaWorkflow(result, file.name);

src/lib/components/ComfyQueue.svelte

@@ -12,11 +12,12 @@
  import type ComfyApp from "./ComfyApp";
  import { getContext, tick } from "svelte";
  import Modal from "./Modal.svelte";
- import { type WorkflowError } from "$lib/stores/workflowState";
  import ComfyQueueListDisplay from "./ComfyQueueListDisplay.svelte";
  import ComfyQueueGridDisplay from "./ComfyQueueGridDisplay.svelte";
  import { WORKFLOWS_VIEW } from "./ComfyBoxWorkflowsView.svelte";
  import uiQueueState, { type QueueUIEntry } from "$lib/stores/uiQueueState";
+	import type { SerializedPromptInputsAll } from "./ComfyApp";
+	import { OpenWorkflowMode } from "$lib/stores/workflowState";
 
  export let app: ComfyApp;
 
@@ -109,18 +110,21 @@
 
  let showModal = false;
  let expandAll = false;
- let selectedPrompt = null;
+ let selectedQueueEntry: QueueUIEntry = null;
+ let selectedPrompt: SerializedPromptInputsAll = null;
  let selectedImages: ComfyImageLocation[] = [];
  function showPrompt(entry: QueueUIEntry) {
      if (entry.error != null) {
          showModal = false;
          expandAll = false;
+         selectedQueueEntry = null;
          selectedPrompt = null;
          selectedImages = [];
 
          showError(entry.entry.promptID);
      }
      else {
+         selectedQueueEntry = entry;
          selectedPrompt = entry.entry.prompt;
          selectedImages = entry.images;
          showModal = true;
@@ -129,6 +133,7 @@
  }
 
  function closeModal() {
+     selectedQueueEntry = null;
      selectedPrompt = null
      selectedImages = []
      showModal = false;
@@ -136,6 +141,24 @@
      console.warn("CLOSEMODAL")
  }
 
+ async function restoreWorkflow(closeDialog: () => void, mode: OpenWorkflowMode) {
+     if (selectedQueueEntry == null) {
+         console.error("No active prompt!");
+         return;
+     }
+
+     let workflow = selectedQueueEntry.entry.extraData.extra_pnginfo.comfyBoxWorkflow;
+     if (workflow == null) {
+         console.error("No workflow found in PNG info!");
+         return;
+     }
+
+     closeDialog();
+     closeModal();
+
+     await app.openWorkflow(structuredClone(workflow), { mode: mode, refreshCombos: true, warnMissingNodeTypes: true });
+ }
+
  let queued = false
  $: queued = Boolean($queueState.runningNodeID || $queueState.progress);
 
@@ -160,6 +183,12 @@
         <Button variant="secondary" on:click={() => (expandAll = !expandAll)}>
             Expand All
         </Button>
+        <Button variant="primary" on:click={() => restoreWorkflow(closeDialog, OpenWorkflowMode.ReplaceActive)}>
+            Restore
+        </Button>
+        <Button variant="primary" on:click={() => restoreWorkflow(closeDialog, OpenWorkflowMode.AppendAndSetActive)}>
+            Load
+        </Button>
     </div>
 </Modal>
 

src/lib/stores/workflowState.ts

@@ -7,7 +7,6 @@ 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';
@@ -25,6 +24,12 @@ export type SerializedWorkflowState = {
     attrs: WorkflowAttributes
 }
 
+export enum OpenWorkflowMode {
+    Append,
+    AppendAndSetActive,
+    ReplaceActive,
+}
+
 /*
  * ID for an opened workflow.
  *
@@ -265,6 +270,7 @@ export type WorkflowState = {
     openedWorkflows: ComfyBoxWorkflow[],
     openedWorkflowsByID: Record<WorkflowInstID, ComfyBoxWorkflow>,
     activeWorkflowID: WorkflowInstID | null,
+    activeWorkflowIndex: number | null,
     activeWorkflow: ComfyBoxWorkflow | null,
 }
 
@@ -279,9 +285,9 @@ type WorkflowStateOps = {
     getWorkflowByNode: (node: LGraphNode) => ComfyBoxWorkflow | null
     getWorkflowByNodeID: (id: NodeID) => ComfyBoxWorkflow | null
     getActiveWorkflow: () => ComfyBoxWorkflow | null
-    createNewWorkflow: (canvas: ComfyGraphCanvas, title?: string, setActive?: boolean) => ComfyBoxWorkflow,
+    createNewWorkflow: (canvas: ComfyGraphCanvas, title?: string, mode?: OpenWorkflowMode) => ComfyBoxWorkflow,
-    openWorkflow: (canvas: ComfyGraphCanvas, data: SerializedAppState, setActive?: boolean) => ComfyBoxWorkflow,
+    openWorkflow: (canvas: ComfyGraphCanvas, data: SerializedAppState, mode?: OpenWorkflowMode) => ComfyBoxWorkflow,
-    addWorkflow: (canvas: ComfyGraphCanvas, data: ComfyBoxWorkflow, setActive?: boolean) => void,
+    addWorkflow: (canvas: ComfyGraphCanvas, data: ComfyBoxWorkflow, mode?: OpenWorkflowMode) => void,
     closeWorkflow: (canvas: ComfyGraphCanvas, index: number) => void,
     closeAllWorkflows: (canvas: ComfyGraphCanvas) => void,
     setActiveWorkflow: (canvas: ComfyGraphCanvas, index: number | WorkflowInstID) => ComfyBoxWorkflow | null,
@@ -297,6 +303,7 @@ const store: Writable<WorkflowState> = writable(
         openedWorkflows: [],
         openedWorkflowsByID: {},
         activeWorkflowID: null,
+        activeWorkflowIndex: null,
         activeWorkflow: null
     })
 
@@ -328,39 +335,51 @@ function getActiveWorkflow(): ComfyBoxWorkflow | null {
     return state.openedWorkflowsByID[state.activeWorkflowID];
 }
 
-function createNewWorkflow(canvas: ComfyGraphCanvas, title: string = "New Workflow", setActive: boolean = false): ComfyBoxWorkflow {
+function createNewWorkflow(canvas: ComfyGraphCanvas, title: string = "New Workflow", mode: OpenWorkflowMode = OpenWorkflowMode.AppendAndSetActive): ComfyBoxWorkflow {
     const workflow = new ComfyBoxWorkflow(title);
     const layoutState = layoutStates.create(workflow);
     layoutState.initDefaultLayout();
 
-    const state = get(store);
+    return addWorkflow(canvas, workflow, mode);
-    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 {
+function openWorkflow(canvas: ComfyGraphCanvas, data: SerializedAppState, mode: OpenWorkflowMode = OpenWorkflowMode.AppendAndSetActive): ComfyBoxWorkflow {
     const [workflow, layoutState] = ComfyBoxWorkflow.create("Workflow")
     workflow.deserialize(layoutState, { graph: data.workflow, layout: data.layout, attrs: data.attrs })
 
-    addWorkflow(canvas, workflow, setActive);
+    addWorkflow(canvas, workflow, mode);
 
     return workflow;
 }
 
-function addWorkflow(canvas: ComfyGraphCanvas, workflow: ComfyBoxWorkflow, setActive: boolean = true) {
+function addWorkflow(canvas: ComfyGraphCanvas, workflow: ComfyBoxWorkflow, mode: OpenWorkflowMode = OpenWorkflowMode.AppendAndSetActive) {
     const state = get(store);
+
+    let resultIndex: number = state.openedWorkflows.length;
+    if (mode == OpenWorkflowMode.ReplaceActive && state.activeWorkflowIndex != null) {
+        resultIndex = state.activeWorkflowIndex;
+        closeWorkflow(canvas, resultIndex);
+        state.openedWorkflows.splice(resultIndex, 0, workflow);
+    }
+    else {
         state.openedWorkflows.push(workflow);
+    }
+
     state.openedWorkflowsByID[workflow.id] = workflow;
 
+    let setActive: boolean;
+    switch (mode) {
+        case OpenWorkflowMode.Append:
+            setActive = false;
+            break;
+        case OpenWorkflowMode.AppendAndSetActive:
+        case OpenWorkflowMode.ReplaceActive:
+            setActive = true;
+            break;
+    }
+
     if (setActive || state.activeWorkflowID == null)
-        setActiveWorkflow(canvas, state.openedWorkflows.length - 1)
+        setActiveWorkflow(canvas, resultIndex)
 
     store.set(state)
 
@@ -397,6 +416,7 @@ function setActiveWorkflow(canvas: ComfyGraphCanvas, index: number | WorkflowIns
 
     if (state.openedWorkflows.length === 0) {
         state.activeWorkflowID = null;
+        state.activeWorkflowIndex = null;
         state.activeWorkflow = null
         return null;
     }
@@ -416,6 +436,7 @@ function setActiveWorkflow(canvas: ComfyGraphCanvas, index: number | WorkflowIns
         state.activeWorkflow.stop("app")
 
     state.activeWorkflowID = workflow.id;
+    state.activeWorkflowIndex = index;
     state.activeWorkflow = workflow;
 
     workflow.start("app", canvas);