Basic graph template saving
This commit is contained in:
space-nuko
335
src/lib/ComfyBoxTemplate.ts
Normal file
335
src/lib/ComfyBoxTemplate.ts
Normal file
@@ -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("&", "&").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("</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,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user