diff --git a/litegraph b/litegraph index a57af62..a1bf4cb 160000 --- a/litegraph +++ b/litegraph @@ -1 +1 @@ -Subproject commit a57af62d4648bc85b264329ff193c61d9ba84d6d +Subproject commit a1bf4cb511c46831b371d79d1495b76052beecce diff --git a/package.json b/package.json index 5dbcbd0..7c8340c 100644 --- a/package.json +++ b/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", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2d7d1f8..6f5ac91 100644 --- a/pnpm-lock.yaml +++ b/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'} diff --git a/src/lib/ComfyBoxTemplate.ts b/src/lib/ComfyBoxTemplate.ts new file mode 100644 index 0000000..5210b00 --- /dev/null +++ b/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 = 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("&", "&").replaceAll("<", "<").replaceAll(">", ">"); +} +function unescapeXml(safe) { + return safe.replaceAll("&", "&").replaceAll("<", "<").replaceAll(">", ">"); +} + +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("", `${escapeXml(json)}`); + + 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, + } + } +} diff --git a/src/lib/ComfyGraphCanvas.ts b/src/lib/ComfyGraphCanvas.ts index ad71005..f70ecc0 100644 --- a/src/lib/ComfyGraphCanvas.ts +++ b/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 } diff --git a/src/lib/convertVanillaWorkflow.ts b/src/lib/convertVanillaWorkflow.ts index 4b4ca7a..fe0553d 100644 --- a/src/lib/convertVanillaWorkflow.ts +++ b/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", diff --git a/src/lib/nodes/ComfyBackendNode.ts b/src/lib/nodes/ComfyBackendNode.ts index 466a642..1678326 100644 --- a/src/lib/nodes/ComfyBackendNode.ts +++ b/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" }); } } diff --git a/src/lib/nodes/ComfyGraphNode.ts b/src/lib/nodes/ComfyGraphNode.ts index 6ac2d45..f75a55b 100644 --- a/src/lib/nodes/ComfyGraphNode.ts +++ b/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); } diff --git a/src/lib/stores/layoutStates.ts b/src/lib/stores/layoutStates.ts index 5831094..7fcb2a8 100644 --- a/src/lib/stores/layoutStates.ts +++ b/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, 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 { + 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 = {} + + 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 & LayoutStateStoresOps; @@ -1308,7 +1383,9 @@ const layoutStates: WritableLayoutStateStores = { remove, getLayout, getLayoutByGraph, - getLayoutByNode + getLayoutByNode, + getLayoutByDragItemID, + getDragItemByNode } export default layoutStates;