From c3ab3aa69afac96e89f23f16e672605cd133a368 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Sun, 21 May 2023 11:10:10 -0500
Subject: [PATCH] Start vanilla workflow conversion, better PNG parser based on
 catbox userscript code

---
 src/lib/api.ts                               |  12 +-
 src/lib/components/ComfyApp.ts               |  54 ++++--
 src/lib/components/ComfyQueue.svelte         |   7 +-
 src/lib/components/ComfyWorkflowsView.svelte |   4 +-
 src/lib/components/PromptDisplay.svelte      |   7 +-
 src/lib/convertVanillaWorkflow.ts            |  51 ++++++
 src/lib/pnginfo.ts                           | 183 ++++++++++++++-----
 src/lib/stores/layoutStates.ts               |  22 ++-
 tsconfig.json                                |   1 +
 9 files changed, 267 insertions(+), 74 deletions(-)
 create mode 100644 src/lib/convertVanillaWorkflow.ts

diff --git a/src/lib/api.ts b/src/lib/api.ts
index 3d0a960..3c1f413 100644
--- a/src/lib/api.ts
+++ b/src/lib/api.ts
@@ -1,4 +1,4 @@
-import type { Progress, SerializedPrompt, SerializedPromptInputsForNode, SerializedPromptInputsAll, SerializedPromptOutputs } from "./components/ComfyApp";
+import type { Progress, SerializedPrompt, SerializedPromptInputsForNode, SerializedPromptInputsAll, SerializedPromptOutputs, SerializedAppState } from "./components/ComfyApp";
 import type TypedEmitter from "typed-emitter";
 import EventEmitter from "events";
 import type { ComfyImageLocation } from "$lib/utils";
@@ -57,10 +57,14 @@ export type ComfyAPIHistoryResponse = {
     error?: string
 }
 
+export type SerializedComfyBoxPromptData = {
+    subgraphs: string[]
+}
+
 export type ComfyPromptPNGInfo = {
-    workflow: SerializedLGraph,
-    comfyBoxLayout: SerializedLayoutState,
-    comfyBoxSubgraphs: string[],
+    workflow?: SerializedLGraph, // ComfyUI format
+    comfyBoxWorkflow: SerializedAppState,
+    comfyBoxPrompt: SerializedComfyBoxPromptData,
 }
 
 export type ComfyBoxPromptExtraData = ComfyUIPromptExtraData & {
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index a713d31..173cfe4 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -1,7 +1,7 @@
 import { LiteGraph, LGraph, LGraphCanvas, LGraphNode, type LGraphNodeConstructor, type LGraphNodeExecutable, type SerializedLGraph, type SerializedLGraphGroup, type SerializedLGraphNode, type SerializedLLink, NodeMode, type Vector2, BuiltInSlotType, type INodeInputSlot, type NodeID, type NodeTypeSpec, type NodeTypeOpts, type SlotIndex, type UUID } from "@litegraph-ts/core";
 import type { LConnectionKind, INodeSlot } from "@litegraph-ts/core";
 import ComfyAPI, { type ComfyAPIStatusResponse, type ComfyBoxPromptExtraData, type ComfyPromptRequest, type ComfyNodeID, type PromptID } from "$lib/api"
-import { getPngMetadata, importA1111 } from "$lib/pnginfo";
+import { getPngMetadata, importA1111, parsePNGMetadata } from "$lib/pnginfo";
 import EventEmitter from "events";
 import type TypedEmitter from "typed-emitter";
 
@@ -44,6 +44,7 @@ import selectionState from "$lib/stores/selectionState";
 import layoutStates from "$lib/stores/layoutStates";
 import { ComfyWorkflow, type WorkflowAttributes, type WorkflowInstID } from "$lib/stores/workflowState";
 import workflowState from "$lib/stores/workflowState";
+import convertVanillaWorkflow from "$lib/convertVanillaWorkflow";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
 
@@ -72,8 +73,10 @@ export type A1111PromptAndInfo = {
  * Represents a single workflow that can be loaded into the program from JSON.
  */
 export type SerializedAppState = {
-    /** Program identifier, should always be "ComfyBox" */
-    createdBy: "ComfyBox",
+    /** For easy structural typing use */
+    comfyBoxWorkflow: true,
+    /** Program identifier, should be something like "ComfyBox" or "ComfyUI" */
+    createdBy: string,
     /** Serial version, should be incremented on breaking changes */
     version: number,
     /** Commit hash if found */
@@ -138,6 +141,14 @@ type CanvasState = {
     canvas: ComfyGraphCanvas,
 }
 
+function isComfyBoxWorkflow(data: any): data is SerializedAppState {
+    return data != null && (typeof data === "object") && data.comfyBoxWorkflow;
+}
+
+function isVanillaWorkflow(data: any): data is SerializedLGraph {
+    return data != null && (typeof data === "object") && data.last_node_id != null;
+}
+
 export default class ComfyApp {
     api: ComfyAPI;
 
@@ -237,6 +248,7 @@ export default class ComfyApp {
         const canvas = this.lCanvas.serialize();
 
         return {
+            comfyBoxWorkflow: true,
             createdBy: "ComfyBox",
             version: COMFYBOX_SERIAL_VERSION,
             commitHash: __GIT_COMMIT_HASH__,
@@ -399,7 +411,7 @@ export default class ComfyApp {
                 } catch (error) { }
             }
 
-            if (workflow && workflow.createdBy === "ComfyBox") {
+            if (workflow && typeof workflow.createdBy === "string") {
                 this.openWorkflow(workflow);
             }
             else {
@@ -537,6 +549,13 @@ export default class ComfyApp {
         return workflow;
     }
 
+    async openVanillaWorkflow(data: SerializedLGraph) {
+        const converted = convertVanillaWorkflow(data)
+        console.info("WORKFLWO", converted)
+        notify("Converted ComfyUI workflow to ComfyBox format.", { type: "info" })
+        // await this.openWorkflow(JSON.parse(pngInfo.workflow));
+    }
+
     setActiveWorkflow(id: WorkflowInstID) {
         const index = get(workflowState).openedWorkflows.findIndex(w => w.id === id)
         if (index === -1)
@@ -695,7 +714,7 @@ export default class ComfyApp {
                     }
 
                     const p = this.graphToPrompt(workflow, tag);
-                    const l = workflow.layout.serialize();
+                    const wf = this.serialize(workflow)
                     console.debug(graphToGraphVis(workflow.graph))
                     console.debug(promptToGraphVis(p))
 
@@ -704,9 +723,10 @@ export default class ComfyApp {
 
                     const extraData: ComfyBoxPromptExtraData = {
                         extra_pnginfo: {
-                            workflow: p.workflow,
-                            comfyBoxLayout: l,
-                            comfyBoxSubgraphs: [tag],
+                            comfyBoxWorkflow: wf,
+                            comfyBoxPrompt: {
+                                subgraphs: [tag]
+                            }
                         },
                         thumbnails
                     }
@@ -761,10 +781,14 @@ export default class ComfyApp {
      */
     async handleFile(file: File) {
         if (file.type === "image/png") {
-            const pngInfo = await getPngMetadata(file);
+            const buffer = await file.arrayBuffer();
+            const pngInfo = await parsePNGMetadata(buffer);
             if (pngInfo) {
-                if (pngInfo.comfyBoxConfig) {
-                    await this.openWorkflow(JSON.parse(pngInfo.comfyBoxConfig));
+                if (pngInfo.comfyBoxWorkflow) {
+                    await this.openWorkflow(JSON.parse(pngInfo.comfyBoxWorkflow));
+                } else if (pngInfo.workflow) {
+                    const workflow = JSON.parse(pngInfo.workflow);
+                    await this.openVanillaWorkflow(workflow);
                 } else if (pngInfo.parameters) {
                     const parsed = parseA1111(pngInfo.parameters)
                     if ("error" in parsed) {
@@ -787,7 +811,13 @@ export default class ComfyApp {
         } else if (file.type === "application/json" || file.name.endsWith(".json")) {
             const reader = new FileReader();
             reader.onload = async () => {
-                await this.openWorkflow(JSON.parse(reader.result as string));
+                const result = JSON.parse(reader.result as string)
+                if (isComfyBoxWorkflow(result)) {
+                    await this.openWorkflow(result);
+                }
+                else if (isVanillaWorkflow(result)) {
+                    await this.openVanillaWorkflow(result);
+                }
             };
             reader.readAsText(file);
         }
diff --git a/src/lib/components/ComfyQueue.svelte b/src/lib/components/ComfyQueue.svelte
index bdf0b2b..60315c9 100644
--- a/src/lib/components/ComfyQueue.svelte
+++ b/src/lib/components/ComfyQueue.svelte
@@ -69,7 +69,7 @@
          dateStr = formatDate(date);
      }
 
-     const subgraphs: string[] | null = entry.extraData?.extra_pnginfo?.comfyBoxSubgraphs;
+     const subgraphs: string[] | null = entry.extraData?.extra_pnginfo?.comfyBoxPrompt?.subgraphs;
 
      let message = "Prompt";
      if (entry.workflowID != null) {
@@ -208,7 +208,7 @@
 </Modal>
 
 <div class="queue">
-    <DropZone {app} />
+    <!-- <DropZone {app} /> -->
     <div class="queue-entries {mode}-mode" bind:this={queueList}>
         {#if _entries.length > 0}
             {#each _entries as entry}
@@ -305,8 +305,9 @@
 <style lang="scss">
  $pending-height: 200px;
  $bottom-bar-height: 70px;
+ $workflow-tabs-height: 2.5rem;
  $mode-buttons-height: 30px;
- $queue-height: calc(100vh - #{$pending-height} - #{$mode-buttons-height} - #{$bottom-bar-height});
+ $queue-height: calc(100vh - #{$pending-height} - #{$mode-buttons-height} - #{$bottom-bar-height} - #{$workflow-tabs-height} - 0.9rem);
 
  .prompt-modal-header {
      padding-left: 0.2rem;
diff --git a/src/lib/components/ComfyWorkflowsView.svelte b/src/lib/components/ComfyWorkflowsView.svelte
index 8dfa784..6f2de2c 100644
--- a/src/lib/components/ComfyWorkflowsView.svelte
+++ b/src/lib/components/ComfyWorkflowsView.svelte
@@ -203,7 +203,7 @@
                     {#if $workflowState.activeWorkflow != null}
                         <ComfyWorkflowView {app} workflow={$workflowState.activeWorkflow} />
                     {:else}
-                        <span>No workflow loaded</span>
+                        <span style:color="var(--body-text-color)">No workflow loaded</span>
                     {/if}
                 </Pane>
                 <Pane bind:size={graphSize}>
@@ -313,7 +313,7 @@
         </div>
     </div>
 </div>
-<input bind:this={fileInput} id="comfy-file-input" type="file" accept=".json" on:change={loadWorkflow} />
+<input bind:this={fileInput} id="comfy-file-input" type="file" accept="application/json,image/png" on:change={loadWorkflow} />
 
 {#if appSetupPromise}
     {#await appSetupPromise}
diff --git a/src/lib/components/PromptDisplay.svelte b/src/lib/components/PromptDisplay.svelte
index a11e245..7504d1a 100644
--- a/src/lib/components/PromptDisplay.svelte
+++ b/src/lib/components/PromptDisplay.svelte
@@ -18,7 +18,7 @@
 
  let galleryStyle: Styles = {
      grid_cols: [2],
-     object_fit: "cover",
+     object_fit: "contain",
      height: "var(--size-96)"
  }
 
@@ -165,6 +165,11 @@
 
          > :global(.block) {
              height: 100%;
+
+             :global(> .preview) {
+                 height: 100%;
+                 max-height: none !important;
+             }
          }
      }
 
diff --git a/src/lib/convertVanillaWorkflow.ts b/src/lib/convertVanillaWorkflow.ts
new file mode 100644
index 0000000..3af0cf9
--- /dev/null
+++ b/src/lib/convertVanillaWorkflow.ts
@@ -0,0 +1,51 @@
+import { LGraph, type SerializedLGraph } from "@litegraph-ts/core";
+import type { SerializedAppState } from "./components/ComfyApp";
+import layoutStates, { defaultWorkflowAttributes, type DragItemID, type SerializedDragEntry, type SerializedLayoutState } from "./stores/layoutStates";
+import type { WorkflowAttributes } from "./stores/workflowState";
+import type { SerializedGraphCanvasState } from "./ComfyGraphCanvas";
+
+/*
+ * The workflow type used by base ComfyUI
+ */
+export type ComfyVanillaWorkflow = SerializedLGraph;
+
+function addLayoutToVanillaWorkflow(workflow: ComfyVanillaWorkflow): SerializedLayoutState {
+    // easier to create a real layout first and then serialize it, then have to
+    // deal with manually constructing the serialized state from the ground up
+    const layoutState = layoutStates.createRaw();
+    const graph = new LGraph();
+    graph.configure(workflow)
+
+    for (const node of graph.iterateNodesInOrder()) {
+        console.warn("NODE", node)
+    }
+
+    return layoutState.serialize()
+}
+
+export default function convertVanillaWorkflow(workflow: ComfyVanillaWorkflow): SerializedAppState {
+    const attrs: WorkflowAttributes = {
+        ...defaultWorkflowAttributes,
+        title: "ComfyUI Workflow"
+    }
+
+    const canvas: SerializedGraphCanvasState = {
+        offset: [0, 0],
+        scale: 1
+    }
+
+    const layout = addLayoutToVanillaWorkflow(workflow);
+
+    const appState: SerializedAppState = {
+        comfyBoxWorkflow: true,
+        createdBy: "ComfyUI", // ???
+        version: 1,
+        workflow,
+        attrs,
+        canvas,
+        layout,
+        commitHash: null
+    }
+
+    return appState
+}
diff --git a/src/lib/pnginfo.ts b/src/lib/pnginfo.ts
index 4a02d72..e221941 100644
--- a/src/lib/pnginfo.ts
+++ b/src/lib/pnginfo.ts
@@ -1,65 +1,158 @@
 import { LiteGraph, LGraph, LGraphNode } from "@litegraph-ts/core"
 import type ComfyAPI from "$lib/api"
 
-class PNGMetadataPromise extends Promise<Record<string, string>> {
-    public cancelMethod: () => void;
-    constructor(executor: (resolve: (value?: Record<string, string>) => void, reject: (reason?: any) => void) => void) {
-        super(executor);
+class Lazy<T> {
+    thunk: () => T;
+    cache: T | null;
 
+    constructor(thunk: () => T) {
+        this.thunk = thunk;
+        this.cache = null;
     }
 
-    //cancel the operation
-    public cancel() {
-        if (this.cancelMethod) {
-            this.cancelMethod();
+    get(): T {
+        if (this.cache === null) {
+            this.cache = this.thunk();
         }
+
+        return this.cache;
     }
 }
 
-export function getPngMetadata(file: File): PNGMetadataPromise {
-    return new PNGMetadataPromise((r, _) => {
-        const reader = new FileReader();
-        reader.onload = (event: Event) => {
-            // Get the PNG data as a Uint8Array
-            const pngData = new Uint8Array((event.target as any).result);
-            const dataView = new DataView(pngData.buffer);
+interface PNGChunk {
+    type: string,
+    data: Uint8Array,
+}
 
-            // Check that the PNG signature is present
-            if (dataView.getUint32(0) !== 0x89504e47) {
-                console.error("Not a valid PNG file");
-                r();
-                return;
-            }
+enum ParseErrorKind {
+    PNGMalformedHeader,
+    PNGMalformedTextChunk,
+    JPEGNoEXIF,
+    JPEGNoUserComment,
+    JPEGMalformedUserComment,
+    UnexpectedEOF,
+    UnsupportedFileType,
+    Other,
+}
 
-            // Start searching for chunks after the PNG signature
-            let offset = 8;
-            let txt_chunks = {};
-            // Loop through the chunks in the PNG file
-            while (offset < pngData.length) {
-                // Get the length of the chunk
-                const length = dataView.getUint32(offset);
-                // Get the chunk type
-                const type = String.fromCharCode(...pngData.slice(offset + 4, offset + 8));
-                if (type === "tEXt") {
-                    // Get the keyword
-                    let keyword_end = offset + 8;
-                    while (pngData[keyword_end] !== 0) {
-                        keyword_end++;
-                    }
-                    const keyword = String.fromCharCode(...pngData.slice(offset + 8, keyword_end));
-                    // Get the text
-                    const text = String.fromCharCode(...pngData.slice(keyword_end + 1, offset + 8 + length));
-                    txt_chunks[keyword] = text;
+interface ParseError {
+    kind: ParseErrorKind,
+    error: any,
+}
+
+/*
+ * This function was taken from the hdg userscript
+ */
+function* pngChunks(bytes: Uint8Array): Generator<PNGChunk, ParseError | null> {
+    const HEADER: number[] = [137, 80, 78, 71, 13, 10, 26, 10];
+    const LENGTH_LEN = 4;
+    const TYPE_LEN = 4;
+    const CRC_LEN = 4;
+
+    let view = new DataView(bytes.buffer);
+    let decoder = new TextDecoder("utf-8", { fatal: true });
+    let pos = 0;
+
+    for (let i = 0; i < HEADER.length; i++) {
+        if (bytes[i] != HEADER[i]) {
+            return {
+                kind: ParseErrorKind.PNGMalformedHeader,
+                error: `wrong PNG header: ${bytes.slice(0, HEADER.length)}`,
+            };
+        }
+    }
+    pos += HEADER.length;
+
+    while (pos < bytes.byteLength) {
+        try {
+            let len = view.getUint32(pos, false);
+            let type = decoder.decode(bytes.subarray(pos + LENGTH_LEN, pos + LENGTH_LEN + TYPE_LEN));
+            if (type.length < 4) {
+                return {
+                    kind: ParseErrorKind.UnexpectedEOF,
+                    error: "PNG parse error: unexpected EOF when parsing chunk type",
                 }
+            }
+            let start = pos + LENGTH_LEN + TYPE_LEN;
 
-                offset += 12 + length;
+            yield {
+                type,
+                data: bytes.subarray(start, start + len),
             }
 
-            r(txt_chunks);
-        };
+            pos = start + len + CRC_LEN;
+        } catch (err) {
+            return {
+                kind: ParseErrorKind.Other,
+                error: err,
+            };
+        }
+    }
 
-        reader.readAsArrayBuffer(file);
-    });
+    return null;
+}
+
+type PNGTextChunk = {
+    keyword: string,
+    text: string
+}
+
+function parsePNGTextChunk(data: Uint8Array): PNGTextChunk | ParseError {
+    let decoder = new TextDecoder("utf-8", { fatal: true });
+
+    let sep = data.findIndex(v => v === 0);
+    if (sep < 0) {
+        return {
+            kind: ParseErrorKind.PNGMalformedTextChunk,
+            error: "PNG parse error: no null separator in tEXt chunk",
+        }
+    }
+
+    try {
+        let keyword = decoder.decode(data.subarray(0, sep));
+        let text = decoder.decode(data.subarray(sep + 1, data.byteLength));
+        return { keyword, text }
+    } catch (err) {
+        return {
+            kind: ParseErrorKind.Other,
+            error: err,
+        };
+    }
+}
+
+export async function parsePNGMetadata(buf: ArrayBuffer): Promise<Record<string, string>> {
+    let bytes = new Uint8Array(buf);
+    const metadata = {}
+    let next: IteratorResult<PNGChunk, ParseError>;
+
+    let chunks = pngChunks(bytes);
+    do {
+        next = chunks.next();
+        if (!next.done) {
+            let chunk = next.value;
+            if ("kind" in chunk) {
+                console.warn("ignored a malformed PNG text chunk");
+                console.error(chunk.error);
+                continue;
+            }
+
+            if (chunk.type !== "tEXt") {
+                continue;
+            }
+
+            let result = parsePNGTextChunk(chunk.data);
+            if ("kind" in result) {
+                console.warn("ignored a malformed PNG text chunk");
+                console.error(result.error);
+                continue;
+            }
+
+            let textChunk = result;
+            metadata[textChunk.keyword] = textChunk.text
+        }
+    } while (next.value != null && !next.done)
+
+    return metadata;
 }
 
 type NodeIndex = { node: LGraphNode, index: number }
diff --git a/src/lib/stores/layoutStates.ts b/src/lib/stores/layoutStates.ts
index 3d0af3f..58b5770 100644
--- a/src/lib/stores/layoutStates.ts
+++ b/src/lib/stores/layoutStates.ts
@@ -660,7 +660,7 @@ export interface WidgetLayout extends IDragItem {
 export type DragItemID = UUID;
 
 type LayoutStateOps = {
-    workflow: ComfyWorkflow,
+    workflow: ComfyWorkflow | null,
 
     addContainer: (parent: ContainerLayout | null, attrs: Partial<Attributes>, index?: number) => ContainerLayout,
     addWidget: (parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partial<Attributes>, index?: number) => WidgetLayout,
@@ -700,11 +700,7 @@ export type SerializedDragItem = {
 
 export type WritableLayoutStateStore = Writable<LayoutState> & LayoutStateOps;
 
-function create(workflow: ComfyWorkflow): WritableLayoutStateStore {
-    if (get(layoutStates).all[workflow.id] != null) {
-        throw new Error(`Layout state already created! ${id}`)
-    }
-
+function createRaw(workflow: ComfyWorkflow | null = null): WritableLayoutStateStore {
     const store: Writable<LayoutState> = writable({
         root: null,
         allItems: {},
@@ -1160,7 +1156,7 @@ function create(workflow: ComfyWorkflow): WritableLayoutStateStore {
 
     function notifyWorkflowModified() {
         if (!get(store).isConfiguring)
-            workflow.notifyModified();
+            workflow?.notifyModified();
     }
 
     const layoutStateStore: WritableLayoutStateStore =
@@ -1185,6 +1181,16 @@ function create(workflow: ComfyWorkflow): WritableLayoutStateStore {
         notifyWorkflowModified
     }
 
+    return layoutStateStore
+}
+
+function create(workflow: ComfyWorkflow): WritableLayoutStateStore {
+    if (get(layoutStates).all[workflow.id] != null) {
+        throw new Error(`Layout state already created! ${id}`)
+    }
+
+    const layoutStateStore = createRaw(workflow);
+
     layoutStates.update(s => {
         s.all[workflow.id] = layoutStateStore;
         return s;
@@ -1233,6 +1239,7 @@ export type LayoutStateStores = {
 
 export type LayoutStateStoresOps = {
     create: (workflow: ComfyWorkflow) => WritableLayoutStateStore,
+    createRaw: (workflow?: ComfyWorkflow | null) => WritableLayoutStateStore,
     remove: (workflowID: WorkflowInstID) => void,
     getLayout: (workflowID: WorkflowInstID) => WritableLayoutStateStore | null,
     getLayoutByGraph: (graph: LGraph) => WritableLayoutStateStore | null,
@@ -1249,6 +1256,7 @@ const store = writable({
 const layoutStates: WritableLayoutStateStores = {
     ...store,
     create,
+    createRaw,
     remove,
     getLayout,
     getLayoutByGraph,
diff --git a/tsconfig.json b/tsconfig.json
index 10daf22..21380bb 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -5,6 +5,7 @@
     "useDefineForClassFields": true,
     "module": "ESNext",
     "resolveJsonModule": true,
+    "esModuleInterop": true,
     "allowJs": true,
     "checkJs": true,
     "strict": false,