Basic graph template saving

package.json

@@ -66,6 +66,7 @@
     "@sveltejs/vite-plugin-svelte": "^2.1.1",
     "@tsconfig/svelte": "^4.0.1",
     "@zerodevx/svelte-json-view": "^1.0.5",
+    "canvas-to-svg": "^1.0.3",
     "events": "^3.3.0",
     "framework7": "^8.0.3",
     "framework7-svelte": "^8.0.3",

pnpm-lock.yaml

@@ -76,6 +76,9 @@ importers:
       '@zerodevx/svelte-json-view':
         specifier: ^1.0.5
         version: 1.0.5(svelte@3.58.0)
+      canvas-to-svg:
+        specifier: ^1.0.3
+        version: 1.0.3
       events:
         specifier: ^3.3.0
         version: 3.3.0
@@ -3779,6 +3782,10 @@ packages:
   /caniuse-lite@1.0.30001488:
     resolution: {integrity: sha512-NORIQuuL4xGpIy6iCCQGN4iFjlBXtfKWIenlUuyZJumLRIindLb7wXM+GO8erEhb7vXfcnf4BAg2PrSDN5TNLQ==}
 
+  /canvas-to-svg@1.0.3:
+    resolution: {integrity: sha512-qWQKPW0lRkOCClRmKtTzC5dNLtdOCzEOPYWASvW9ggkq4cRYL0o8RnC4fPc4tNE5WGm5YzJTGFztFG0JuCGXcQ==}
+    dev: false
+
   /case@1.6.3:
     resolution: {integrity: sha512-mzDSXIPaFwVDvZAHqZ9VlbyF4yyXRuX6IvB06WvPYkqJVO24kX1PPhv9bfpKNFZyxYFmmgo03HUiD8iklmJYRQ==}
     engines: {node: '>= 0.8.0'}

src/lib/ComfyBoxTemplate.ts

@@ -0,0 +1,335 @@
+import { Subgraph, type LGraphNode, type LLink, type SerializedLGraphNode, type SerializedLLink, LGraph } from "@litegraph-ts/core"
+import layoutStates, { isComfyWidgetNode, type ContainerLayout, type SerializedDragEntry, type WidgetLayout, type DragItemID, type WritableLayoutStateStore, type DragItemEntry, type SerializedLayoutState } from "./stores/layoutStates"
+import type { ComfyWidgetNode } from "./nodes/widgets"
+import type ComfyGraphCanvas from "./ComfyGraphCanvas"
+import C2S from "canvas-to-svg";
+import { download } from "./utils";
+
+/*
+ * InComfyBox a template contains a subgraph and the set of components it
+ * represents in the UI.
+ */
+export type ComfyBoxTemplate = {
+    nodes: LGraphNode[],
+    links: LLink[],
+    container?: DragItemEntry
+}
+
+/*
+ * InComfyBox a template contains a subgraph and the set of components it
+ * represents in the UI.
+ */
+export type SerializedComfyBoxTemplate = {
+    /*
+     * Serialized nodes
+     */
+    nodes: SerializedLGraphNode[]
+
+    /*
+     * Serialized inner links
+     */
+    links: SerializedLLink[],
+
+    /*
+     * Serialized container type drag item
+     */
+    layout?: SerializedLayoutState
+}
+
+export type SerializedComfyBoxTemplateData = {
+    comfyBoxTemplate: SerializedComfyBoxTemplate
+}
+
+export type ComfyBoxTemplateError = {
+    error: string
+}
+
+export type ComfyBoxTemplateResult = ComfyBoxTemplate | ComfyBoxTemplateError;
+
+function findIdealParentContainerForNodes(layout: WritableLayoutStateStore, widgets: WidgetLayout[]): DragItemEntry | null {
+    const widgetIds = new Set(widgets.map(w => w.id));
+
+    const containsExactlyTheWidgets = (container: ContainerLayout): boolean => {
+        const found: Set<DragItemID> = new Set();
+        for (const { dragItem } of layout.iterateBreadthFirst(container.id)) {
+            if (dragItem.type === "widget") {
+                if (!widgetIds.has(dragItem.id))
+                    return false;
+
+                found.add(dragItem.id);
+            }
+        }
+
+        return found.size === widgetIds.size;
+    }
+
+    for (const entry of layout.iterateBreadthFirst()) {
+        if (entry.dragItem.type === "container" && entry.children.length > 0) {
+            const container = entry.dragItem as ContainerLayout;
+            if (containsExactlyTheWidgets(container)) {
+                return entry;
+            }
+        }
+    }
+
+    return null;
+}
+
+function getWidgetNodes(nodes: LGraphNode[]): ComfyWidgetNode[] {
+    let result = []
+
+    for (const node of nodes) {
+        if (isComfyWidgetNode(node)) {
+            result.push(node)
+        }
+        else if (node.is(Subgraph)) {
+            result = result.concat(Array.from(node.subgraph.iterateNodesInOrderRecursive())
+                .filter(isComfyWidgetNode))
+        }
+    }
+
+    return result;
+}
+
+function getInnerLinks(nodes: LGraphNode[]): LLink[] {
+    const nodeIds = new Set(nodes.map(n => n.id));
+    const result = []
+
+    for (const node of nodes) {
+        for (const link of node.iterateAllLinks()) {
+            if (nodeIds.has(link.origin_id) && nodeIds.has(link.target_id))
+                result.push(link)
+        }
+    }
+
+    return result;
+}
+
+function escapeXml(unsafe) {
+    return unsafe.replaceAll("&", "&amp;").replaceAll("<", "&lt;").replaceAll(">", "&gt;");
+}
+function unescapeXml(safe) {
+    return safe.replaceAll("&amp;", "&").replaceAll("&lt;", "<").replaceAll("&gt;", ">");
+}
+
+const TEMPLATE_SVG_PADDING: number = 50;
+
+function renderSvg(canvas: ComfyGraphCanvas, graph: LGraph, extraData: any, padding: number): string {
+    // Calculate the min max bounds for the nodes on the graph
+    const bounds = graph._nodes.reduce(
+        (p, n) => {
+            if (n.pos[0] < p[0]) p[0] = n.pos[0];
+            if (n.pos[1] < p[1]) p[1] = n.pos[1];
+            const r = n.pos[0] + n.size[0];
+            const b = n.pos[1] + n.size[1];
+            if (r > p[2]) p[2] = r;
+            if (b > p[3]) p[3] = b;
+            return p;
+        },
+        [Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY,
+        Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY]
+    );
+
+    bounds[0] -= padding;
+    bounds[1] -= padding;
+    bounds[2] += padding;
+    bounds[3] += padding;
+
+    // Store current canvas values to reset after drawing
+    const ctx = canvas.ctx;
+    const scale = canvas.ds.scale;
+    const width = canvas.canvas.width;
+    const height = canvas.canvas.height;
+    const offset = canvas.ds.offset;
+    const show_info = canvas.show_info;
+    const background_image = canvas.background_image;
+    const render_canvas_border = canvas.render_canvas_border;
+    const render_subgraph_panels = canvas.render_subgraph_panels
+    const render_subgraph_stack_header = canvas.render_subgraph_stack_header
+
+    canvas.openSubgraph(graph)
+    canvas.show_info = false;
+    canvas.background_image = null;
+    canvas.render_canvas_border = false;
+    canvas.render_subgraph_panels = false;
+    canvas.render_subgraph_stack_header = false;
+
+    const svgCtx = new C2S(bounds[2] - bounds[0], bounds[3] - bounds[1]);
+
+    svgCtx.canvas.getBoundingClientRect = function() {
+        return { width: svgCtx.width, height: svgCtx.height };
+    };
+
+    // Override the c2s handling of images to draw images as canvases
+    // const drawImage = svgCtx.drawImage;
+    // svgCtx.drawImage = function(...args) {
+    //     const image = args[0];
+    //     // If we are an image node and not a datauri then we need to replace with a canvas
+    //     // we cant convert to data uri here as it is an async process
+    //     if (image.nodeName === "IMG" && !image.src.startsWith("data:image/")) {
+    //         const canvas = document.createElement("canvas");
+    //         canvas.width = image.width;
+    //         canvas.height = image.height;
+    //         const imgCtx = canvas.getContext("2d");
+    //         imgCtx.drawImage(image, 0, 0);
+    //         args[0] = canvas;
+    //     }
+
+    //     return drawImage.apply(this, args);
+    // };
+
+    // Implement missing required functions
+    svgCtx.getTransform = function() {
+        return ctx.getTransform();
+    };
+    svgCtx.resetTransform = function() {
+        return ctx.resetTransform();
+    };
+    svgCtx.roundRect = svgCtx.rect;
+
+    // Force the canvas to render the whole graph to the svg context
+    canvas.ds.scale = 1;
+    canvas.canvas.width = bounds[2] - bounds[0];
+    canvas.canvas.height = bounds[3] - bounds[1];
+    canvas.ds.offset = [-bounds[0], -bounds[1]];
+    canvas.ctx = svgCtx;
+
+    let saving = false;
+
+    // Trigger saving
+    saving = true;
+    canvas.draw(true, true);
+    saving = false;
+
+    // Restore original settings
+    canvas.closeSubgraph();
+    canvas.ds.scale = scale;
+    canvas.canvas.width = width;
+    canvas.canvas.height = height;
+    canvas.ds.offset = offset;
+    canvas.ctx = ctx;
+    canvas.show_info = show_info;
+    canvas.background_image = background_image;
+    canvas.render_canvas_border = render_canvas_border;
+    canvas.render_subgraph_panels = render_subgraph_panels;
+    canvas.render_subgraph_stack_header = render_subgraph_stack_header;
+
+    canvas.draw(true, true);
+
+    // Convert to SVG, embed graph and save
+    // const json = JSON.stringify(app.graph.serialize());
+    const json = JSON.stringify(extraData);
+    const svg = svgCtx.getSerializedSvg(true).replace("</svg>", `<desc>${escapeXml(json)}</desc></svg>`);
+
+    return svg
+}
+
+function pruneDetachedLinks(nodes: SerializedLGraphNode[], links: SerializedLLink[]): [SerializedLGraphNode[], SerializedLLink[]] {
+    const nodeIds = new Set(nodes.map(n => n.id));
+
+    for (const node of nodes) {
+        if (node.inputs) {
+            for (const input of node.inputs) {
+                if (input.link && (!nodeIds.has(input.link)))
+                    input.link = null;
+            }
+        }
+        if (node.outputs) {
+            for (const output of node.outputs) {
+                if (output.links) {
+                    output.links = output.links.filter(l => nodeIds.has(l))
+                }
+            }
+        }
+    }
+
+    links = links.filter(l => {
+        return nodeIds.has(l[1]) && nodeIds.has(l[3]);
+    })
+
+    return [nodes, links]
+}
+
+export function serializeTemplate(canvas: ComfyGraphCanvas, template: ComfyBoxTemplate): SerializedComfyBoxTemplate {
+    let graph: LGraph
+    if (template.nodes.length === 1 && template.nodes[0].is(Subgraph)) {
+        graph = template.nodes[0].subgraph
+    } else {
+        // TODO render graph portion
+        graph = canvas.graph;
+    }
+
+    const layoutState = layoutStates.getLayoutByDragItemID(template.container.dragItem.id);
+    if (layoutState == null)
+        throw "Couldn't find layout for template being serialized!"
+
+    let nodes = template.nodes.map(n => n.serialize());
+    let links = template.links.map(l => l.serialize());
+    const layout = layoutState.serializeAtRoot(template.container.dragItem.id);
+
+    [nodes, links] = pruneDetachedLinks(nodes, links);
+
+    let comfyBoxTemplate: SerializedComfyBoxTemplate = {
+        nodes: nodes,
+        links: links,
+        layout: layout
+    }
+
+    let templateData: SerializedComfyBoxTemplateData = {
+        comfyBoxTemplate
+    }
+
+    const svg = renderSvg(canvas, graph, templateData, TEMPLATE_SVG_PADDING)
+    download("workflow.svg", svg, "image/svg+xml");
+
+    return comfyBoxTemplate
+}
+
+
+export function createTemplate(nodes: LGraphNode[]): ComfyBoxTemplateResult {
+    if (nodes.length === 0) {
+        return {
+            error: "No nodes selected."
+        }
+    }
+    // Find all UI nodes in this subgraph recursively
+    const widgetNodes = getWidgetNodes(nodes)
+    const links = getInnerLinks(nodes)
+
+    const layout = layoutStates.getLayoutByNode(nodes[0])
+    if (layout == null) {
+        return {
+            error: "Subgraph not contained in a layout!"
+        }
+    }
+
+    if (widgetNodes.length > 0) {
+        // Find the highest-level container that contains all these nodes and
+        // contains no other widgets
+        const widgets = widgetNodes.map(node => node.dragItem)
+        if (!widgets.every(Boolean)) {
+            return {
+                error: "At least one widget node was missing an entry in the UI!"
+            }
+        }
+        const container = findIdealParentContainerForNodes(layout, widgets);
+        if (container == null) {
+            return {
+                error: "Couldn't find a suitable container in the UI for these nodes. Ensure all the widget nodes in the subgraph are kept inside a single container in the UI."
+            }
+        }
+
+        return {
+            nodes: nodes,
+            links: links,
+            container: container
+        }
+    }
+    else {
+        // No UI to serialize.
+        return {
+            nodes: nodes,
+            links: links,
+        }
+    }
+}

src/lib/ComfyGraphCanvas.ts

@@ -6,6 +6,8 @@ import { ComfyReroute } from "./nodes";
 import layoutStates from "./stores/layoutStates";
 import queueState from "./stores/queueState";
 import selectionState from "./stores/selectionState";
+import { createTemplate, type ComfyBoxTemplate, serializeTemplate } from "./ComfyBoxTemplate";
+import notify from "./notify";
 
 export type SerializedGraphCanvasState = {
     offset: Vector2,
@@ -452,6 +454,24 @@ export default class ComfyGraphCanvas extends LGraphCanvas {
         return options
     }
 
+    saveAsTemplate(_value: IContextMenuItem, _options, mouseEvent, prevMenu, node?: LGraphNode) {
+        if (!this.selected_nodes || Object.values(this.selected_nodes).length === 0)
+            return;
+
+        const result = createTemplate(Object.values(this.selected_nodes));
+
+        if ("error" in result) {
+            notify(`Couldn't create template: ${result.error}`, { type: "error", timeout: 5000 });
+            return;
+        }
+
+        const template = result as ComfyBoxTemplate;
+
+        console.warn("TEMPLATEFOUND", template)
+
+        const serialized = serializeTemplate(this, template);
+    }
+
     override getNodeMenuOptions(node: LGraphNode): ContextMenuItem[] {
         const options = super.getNodeMenuOptions(node);
 
@@ -464,6 +484,15 @@ export default class ComfyGraphCanvas extends LGraphCanvas {
             },
         )
 
+        options.push(
+            {
+                content: "Save as Template",
+                has_submenu: false,
+                disabled: false,
+                callback: this.saveAsTemplate.bind(this)
+            },
+        )
+
         return options
     }
 

src/lib/convertVanillaWorkflow.ts

@@ -391,8 +391,7 @@ export default function convertVanillaWorkflow(vanillaWorkflow: ComfyVanillaWork
         // Lazily create group in case there are no inputs
         let group: ContainerLayout | null = null;
 
-        // TODO needs to be generalized!
-        let isOutputNode = ["PreviewImage", "SaveImage"].indexOf(node.type) !== -1
+        let isOutputNode = def.nodeDef.output_node
 
         for (const [inputName, [inputType, inputOpts]] of iterateNodeDefInputs(def.nodeDef)) {
             // Detect if this input was a widget converted to an input
@@ -475,7 +474,7 @@ export default function convertVanillaWorkflow(vanillaWorkflow: ComfyVanillaWork
         }
 
         // Add OUTPUT event slot to output nodes
-        // TODO needs to be generalized!
+        // For now assume that all output nodes will send back images
         if (isOutputNode) {
             const newOutput: INodeOutputSlot = {
                 name: "OUTPUT",

src/lib/nodes/ComfyBackendNode.ts

@@ -32,10 +32,7 @@ export class ComfyBackendNode extends ComfyGraphNode {
 
         this.setup(nodeDef)
 
-        // ComfyUI has no obvious way to identify if a node will return outputs back to the frontend based on its properties.
-        // It just returns a hash like { "ui": { "images": results } } internally.
-        // So this will need to be hardcoded for now.
-        if (["PreviewImage", "SaveImage"].indexOf(comfyClass) !== -1) {
+        if (nodeDef.output_node) {
             this.addOutput("OUTPUT", BuiltInSlotType.EVENT, { color_off: "rebeccapurple", color_on: "rebeccapurple" });
         }
     }

src/lib/nodes/ComfyGraphNode.ts

@@ -8,7 +8,7 @@ import type IComfyInputSlot from "$lib/IComfyInputSlot";
 import uiState from "$lib/stores/uiState";
 import { get } from "svelte/store";
 import configState from "$lib/stores/configState";
-import type { WritableLayoutStateStore } from "$lib/stores/layoutStates";
+import type { WidgetLayout, WritableLayoutStateStore } from "$lib/stores/layoutStates";
 import layoutStates from "$lib/stores/layoutStates";
 import workflowStateStore, { ComfyWorkflow } from "$lib/stores/workflowState";
 
@@ -107,6 +107,10 @@ export default class ComfyGraphNode extends LGraphNode {
         return layoutStates.getLayoutByNode(this);
     }
 
+    get dragItem(): WidgetLayout | null {
+        return layoutStates.getDragItemByNode(this);
+    }
+
     get workflow(): ComfyWorkflow | null {
         return workflowStateStore.getWorkflowByNode(this);
     }

src/lib/stores/layoutStates.ts

@@ -10,11 +10,11 @@ import type ComfyGraph from '$lib/ComfyGraph';
 import type { ComfyWorkflow, WorkflowAttributes, WorkflowInstID } from './workflowState';
 import workflowState from './workflowState';
 
-function isComfyWidgetNode(node: LGraphNode): node is ComfyWidgetNode {
+export function isComfyWidgetNode(node: LGraphNode): node is ComfyWidgetNode {
     return "svelteComponentType" in node
 }
 
-type DragItemEntry = {
+export type DragItemEntry = {
     /*
      * Drag item.
      */
@@ -719,7 +719,9 @@ type LayoutStateOps = {
     ungroup: (container: ContainerLayout) => void,
     findLayoutEntryForNode: (nodeId: ComfyNodeID) => DragItemEntry | null,
     findLayoutForNode: (nodeId: ComfyNodeID) => IDragItem | null,
+    iterateBreadthFirst: (id?: DragItemID | null) => Iterable<DragItemEntry>,
     serialize: () => SerializedLayoutState,
+    serializeAtRoot: (rootID: DragItemID) => SerializedLayoutState,
     deserialize: (data: SerializedLayoutState, graph: LGraph) => void,
     initDefaultLayout: () => DefaultLayout,
     onStartConfigure: () => void
@@ -1082,6 +1084,25 @@ function createRaw(workflow: ComfyWorkflow | null = null): WritableLayoutStateSt
         return found.dragItem as WidgetLayout
     }
 
+    function* iterateBreadthFirst(id?: DragItemID | null): Iterable<DragItemEntry> {
+        const state = get(store);
+
+        id ||= state.root?.id;
+        if (id == null)
+            return;
+
+        const queue = [state.allItems[id]];
+        while (queue.length > 0) {
+            const node = queue.shift();
+            yield node;
+            if (node.children) {
+                for (const child of node.children) {
+                    queue.push(state.allItems[child.id]);
+                }
+            }
+        }
+    }
+
     function initDefaultLayout(): DefaultLayout {
         store.set({
             root: null,
@@ -1106,6 +1127,40 @@ function createRaw(workflow: ComfyWorkflow | null = null): WritableLayoutStateSt
         return { root, left, right }
     }
 
+    function serializeAtRoot(rootID: DragItemID): SerializedLayoutState {
+        const state = get(store);
+
+        if (!state.allItems[rootID])
+            throw "Root not contained in layout!"
+
+        const allItems: Record<DragItemID, SerializedDragEntry> = {}
+
+        const queue = [state.allItems[rootID]]
+        while (queue.length > 0) {
+            const entry = queue.shift();
+            allItems[entry.dragItem.id] = {
+                dragItem: {
+                    type: entry.dragItem.type,
+                    id: entry.dragItem.id,
+                    nodeId: (entry.dragItem as any).node?.id,
+                    attrs: entry.dragItem.attrs
+                },
+                children: entry.children.map((di) => di.id),
+                parent: entry.parent?.id
+            }
+            if (entry.children) {
+                for (const child of entry.children) {
+                    queue.push(state.allItems[child.id]);
+                }
+            }
+        }
+
+        return {
+            root: rootID,
+            allItems
+        }
+    }
+
     function serialize(): SerializedLayoutState {
         const state = get(store)
 
@@ -1221,10 +1276,12 @@ function createRaw(workflow: ComfyWorkflow | null = null): WritableLayoutStateSt
         groupItems,
         findLayoutEntryForNode,
         findLayoutForNode,
+        iterateBreadthFirst,
         ungroup,
         initDefaultLayout,
         onStartConfigure,
         serialize,
+        serializeAtRoot,
         deserialize,
         notifyWorkflowModified
     }
@@ -1273,6 +1330,22 @@ function getLayoutByNode(node: LGraphNode): WritableLayoutStateStore | null {
     return getLayoutByGraph(rootGraph);
 }
 
+function getLayoutByDragItemID(dragItemID: DragItemID): WritableLayoutStateStore | null {
+    return Object.values(get(layoutStates).all).find(l => get(l).allItems[dragItemID] != null)
+}
+
+function getDragItemByNode(node: LGraphNode): WidgetLayout | null {
+    const layout = getLayoutByNode(node);
+    if (layout == null)
+        return null;
+
+    const entry = get(layout).allItemsByNode[node.id]
+    if (entry && entry.dragItem.type === "widget")
+        return entry.dragItem as WidgetLayout;
+
+    return null;
+}
+
 export type LayoutStateStores = {
     /*
      * Layouts associated with opened workflows
@@ -1292,6 +1365,8 @@ export type LayoutStateStoresOps = {
     getLayout: (workflowID: WorkflowInstID) => WritableLayoutStateStore | null,
     getLayoutByGraph: (graph: LGraph) => WritableLayoutStateStore | null,
     getLayoutByNode: (node: LGraphNode) => WritableLayoutStateStore | null,
+    getLayoutByDragItemID: (dragItemID: DragItemID) => WritableLayoutStateStore | null,
+    getDragItemByNode: (node: LGraphNode) => WidgetLayout | null,
 }
 
 export type WritableLayoutStateStores = Writable<LayoutStateStores> & LayoutStateStoresOps;
@@ -1308,7 +1383,9 @@ const layoutStates: WritableLayoutStateStores = {
     remove,
     getLayout,
     getLayoutByGraph,
-    getLayoutByNode
+    getLayoutByNode,
+    getLayoutByDragItemID,
+    getDragItemByNode
 }
 
 export default layoutStates;