Start vanilla workflow conversion, better PNG parser based on catbox

userscript code
2023-05-21 11:10:10 -05:00
parent affe973375
commit c3ab3aa69a
9 changed files with 267 additions and 74 deletions
--- 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 & {
--- 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);
        }
--- 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;
--- 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}
--- 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;
+             }
         }
     }

--- a/src/lib/convertVanillaWorkflow.ts
+++ 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
+}
--- 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 }
--- 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,