From 0143b6430ff7132bcc67df04ac97a19a2e14a3af Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Tue, 16 May 2023 17:02:11 -0500 Subject: [PATCH] Subgraph ergonomic improvements & node def types --- litegraph | 2 +- src/lib/ComfyNodeDef.ts | 52 +++++++++++++++++++++++++++++++ src/lib/api.ts | 3 +- src/lib/components/ComfyApp.ts | 44 +++++++++++++++++--------- src/lib/init.ts | 19 +++++++++++ src/lib/nodes/ComfyBackendNode.ts | 28 +++++++---------- 6 files changed, 115 insertions(+), 33 deletions(-) create mode 100644 src/lib/ComfyNodeDef.ts create mode 100644 src/lib/init.ts diff --git a/litegraph b/litegraph index edc0ccb..b9762fa 160000 --- a/litegraph +++ b/litegraph @@ -1 +1 @@ -Subproject commit edc0ccbb086cf44db847e48f371b2793a19b1340 +Subproject commit b9762fa6d56bf6b1f4f61361b6f9cd2e3972de1d diff --git a/src/lib/ComfyNodeDef.ts b/src/lib/ComfyNodeDef.ts new file mode 100644 index 0000000..d0fb5d0 --- /dev/null +++ b/src/lib/ComfyNodeDef.ts @@ -0,0 +1,52 @@ +import { range } from "./utils" +import ComfyWidgets from "./widgets" + +export type ComfyNodeDef = { + name: string + display_name?: string + category: string + input: ComfyNodeDefInputs + /** Output type like "LATENT" or "IMAGE" */ + output: string[] + output_name: string[] + output_is_list: boolean[] +} + +export type ComfyNodeDefInputs = { + required: Record, + optional?: Record +} +export type ComfyNodeDefInput = [ComfyNodeDefInputType, ComfyNodeDefInputOptions | null] +export type ComfyNodeDefInputType = string[] | keyof typeof ComfyWidgets | string +export type ComfyNodeDefInputOptions = { + forceInput?: boolean +} + +// TODO when comfy refactors +export type ComfyNodeDefOutput = { + type: string, + name: string, + is_list?: boolean +} + +export function isBackendNodeDefInputType(inputName: string, type: ComfyNodeDefInputType): type is string { + return !Array.isArray(type) && !(type in ComfyWidgets) && !(`${type}:${inputName}` in ComfyWidgets); +} + +export function iterateNodeDefInputs(def: ComfyNodeDef): Iterable<[string, ComfyNodeDefInput]> { + var inputs = def.input.required + if (def.input.optional != null) { + inputs = Object.assign({}, def.input.required, def.input.optional) + } + return Object.entries(inputs); +} + +export function iterateNodeDefOutputs(def: ComfyNodeDef): Iterable { + return range(def.output.length).map(i => { + return { + type: def.output[i], + name: def.output_name[i] || def.output[i], + is_list: def.output_is_list[i], + } + }) +} diff --git a/src/lib/api.ts b/src/lib/api.ts index aaf256a..a9c8581 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -4,6 +4,7 @@ import EventEmitter from "events"; import type { ComfyExecutionResult, ComfyImageLocation } from "./nodes/ComfyWidgetNodes"; import type { SerializedLGraph, UUID } from "@litegraph-ts/core"; import type { SerializedLayoutState } from "./stores/layoutState"; +import type { ComfyNodeDef } from "./ComfyNodeDef"; export type ComfyPromptRequest = { client_id?: string, @@ -226,7 +227,7 @@ export default class ComfyAPI { * Loads node object definitions for the graph * @returns The node definitions */ - async getNodeDefs(): Promise { + async getNodeDefs(): Promise> { return fetch(this.getBackendUrl() + "/object_info", { cache: "no-store" }) .then(resp => resp.json()) } diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts index fa094b3..c18a955 100644 --- a/src/lib/components/ComfyApp.ts +++ b/src/lib/components/ComfyApp.ts @@ -1,4 +1,4 @@ -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 } from "@litegraph-ts/core"; +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 } 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"; @@ -34,13 +34,10 @@ import configState from "$lib/stores/configState"; import { blankGraph } from "$lib/defaultGraph"; import type { ComfyExecutionResult } from "$lib/nodes/ComfyWidgetNodes"; import ComfyPromptSerializer from "./ComfyPromptSerializer"; +import { iterateNodeDefInputs, type ComfyNodeDef, isBackendNodeDefInputType, iterateNodeDefOutputs } from "$lib/ComfyNodeDef"; export const COMFYBOX_SERIAL_VERSION = 1; -LiteGraph.catch_exceptions = false; -LiteGraph.CANVAS_GRID_SIZE = 32; -LiteGraph.default_subgraph_lgraph_factory = () => new ComfyGraph(); - if (typeof window !== "undefined") { // Load default visibility nodes.ComfyReroute.setDefaultTextVisibility(!!localStorage["Comfy.ComfyReroute.DefaultVisibility"]); @@ -130,8 +127,6 @@ export default class ComfyApp { this.lCanvas.allow_dragnodes = uiUnlocked; this.lCanvas.allow_interaction = uiUnlocked; - (window as any).LiteGraph = LiteGraph; - // await this.#invokeExtensionsAsync("init"); await this.registerNodes(); @@ -205,15 +200,11 @@ export default class ComfyApp { static widget_type_overrides: Record = {} private async registerNodes() { - const app = this; - // Load node definitions from the backend const defs = await this.api.getNodeDefs(); // Register a node for each definition - for (const nodeId in defs) { - const nodeData = defs[nodeId]; - + for (const [nodeId, nodeDef] of Object.entries(defs)) { const typeOverride = ComfyApp.node_type_overrides[nodeId] if (typeOverride) console.debug("Attaching custom type to received node:", nodeId, typeOverride) @@ -221,19 +212,42 @@ export default class ComfyApp { const ctor = class extends baseClass { constructor(title?: string) { - super(title, nodeId, nodeData); + super(title, nodeId, nodeDef); } } const node: LGraphNodeConstructor = { class: ctor, - title: nodeData.display_name || nodeData.name, + title: nodeDef.display_name || nodeDef.name, type: nodeId, desc: `ComfyNode: ${nodeId}` } LiteGraph.registerNodeType(node); - node.category = nodeData.category; + node.category = nodeDef.category; + + ComfyApp.registerDefaultSlotHandlers(nodeId, nodeDef) + } + } + + static registerDefaultSlotHandlers(nodeId: string, nodeDef: ComfyNodeDef) { + const nodeTypeSpec: NodeTypeOpts = { + node: nodeId, + title: nodeDef.display_name || nodeDef.name, + properties: null, + inputs: null, + outputs: null + } + + for (const [inputName, [inputType, _inputOpts]] of iterateNodeDefInputs(nodeDef)) { + if (isBackendNodeDefInputType(inputName, inputType)) { + LiteGraph.slot_types_default_out[inputType] ||= [] + LiteGraph.slot_types_default_out[inputType].push(nodeTypeSpec) + } + } + for (const output of iterateNodeDefOutputs(nodeDef)) { + LiteGraph.slot_types_default_in[output.type] ||= [] + LiteGraph.slot_types_default_in[output.type].push(nodeTypeSpec) } } diff --git a/src/lib/init.ts b/src/lib/init.ts new file mode 100644 index 0000000..913dfe5 --- /dev/null +++ b/src/lib/init.ts @@ -0,0 +1,19 @@ +import ComfyGraph from '$lib/ComfyGraph'; +import { LGraphCanvas, LiteGraph, Subgraph } from '@litegraph-ts/core'; + +export function configureLitegraph(isMobile: boolean = false) { + LiteGraph.catch_exceptions = false; + LiteGraph.use_uuids = true; + LiteGraph.CANVAS_GRID_SIZE = 32; + + if (isMobile) { + LiteGraph.dialog_close_on_mouse_leave = false; + LiteGraph.search_hide_on_mouse_leave = false; + LiteGraph.pointerevents_method = "pointer"; + } + + Subgraph.default_lgraph_factory = () => new ComfyGraph; + + (window as any).LiteGraph = LiteGraph; + (window as any).LGraphCanvas = LGraphCanvas; +} diff --git a/src/lib/nodes/ComfyBackendNode.ts b/src/lib/nodes/ComfyBackendNode.ts index fd17f00..1bf0f5c 100644 --- a/src/lib/nodes/ComfyBackendNode.ts +++ b/src/lib/nodes/ComfyBackendNode.ts @@ -5,6 +5,7 @@ import type { ComfyWidgetNode, ComfyExecutionResult } from "./ComfyWidgetNodes"; import { BuiltInSlotShape, BuiltInSlotType, type SerializedLGraphNode } from "@litegraph-ts/core"; import type IComfyInputSlot from "$lib/IComfyInputSlot"; import type { ComfyInputConfig } from "$lib/IComfyInputSlot"; +import { iterateNodeDefOutputs, type ComfyNodeDef, iterateNodeDefInputs } from "$lib/ComfyNodeDef"; /* * Base class for any node with configuration sent by the backend. @@ -37,22 +38,19 @@ export class ComfyBackendNode extends ComfyGraphNode { // comfy class -> input name -> input config private static defaultInputConfigs: Record> = {} - private setup(nodeData: any) { - var inputs = nodeData["input"]["required"]; - if (nodeData["input"]["optional"] != undefined) { - inputs = Object.assign({}, nodeData["input"]["required"], nodeData["input"]["optional"]) - } - + private setup(nodeDef: ComfyNodeDef) { ComfyBackendNode.defaultInputConfigs[this.type] = {} - for (const inputName in inputs) { + for (const [inputName, inputData] of iterateNodeDefInputs(nodeDef)) { const config: Partial = {}; - const inputData = inputs[inputName]; - const type = inputData[0]; + const [type, opts] = inputData; - if (inputData[1]?.forceInput) { - this.addInput(inputName, type); + if (opts?.forceInput) { + if (Array.isArray(type)) { + throw new Error(`Can't have forceInput set to true for an enum type! ${type}`) + } + this.addInput(inputName, type as string); } else { if (Array.isArray(type)) { // Enums @@ -73,11 +71,9 @@ export class ComfyBackendNode extends ComfyGraphNode { ComfyBackendNode.defaultInputConfigs[this.type][inputName] = (config as IComfyInputSlot).config } - for (const o in nodeData["output"]) { - const output = nodeData["output"][o]; - const outputName = nodeData["output_name"][o] || output; - const outputShape = nodeData["output_is_list"][o] ? BuiltInSlotShape.GRID_SHAPE : BuiltInSlotShape.CIRCLE_SHAPE; - this.addOutput(outputName, output, { shape: outputShape }); + for (const output of iterateNodeDefOutputs(nodeDef)) { + const outputShape = output.is_list ? BuiltInSlotShape.GRID_SHAPE : BuiltInSlotShape.CIRCLE_SHAPE; + this.addOutput(output.name, output.type, { shape: outputShape }); } this.serialize_widgets = false;