From fe5964ceb66ab09e3ffd7e6b4dc3c561f17b8cce Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Tue, 4 Mar 2025 11:14:14 -0500 Subject: [PATCH] [Refactor] Use V2 node def in ComfyApp (#2854) --- src/schemas/nodeDef/nodeDefSchemaV2.ts | 30 +++++++ src/scripts/app.ts | 110 ++++++++++++++----------- src/stores/nodeDefStore.ts | 18 ++-- 3 files changed, 101 insertions(+), 57 deletions(-) diff --git a/src/schemas/nodeDef/nodeDefSchemaV2.ts b/src/schemas/nodeDef/nodeDefSchemaV2.ts index 8e3d2e746..e1aa24b08 100644 --- a/src/schemas/nodeDef/nodeDefSchemaV2.ts +++ b/src/schemas/nodeDef/nodeDefSchemaV2.ts @@ -91,3 +91,33 @@ export type CustomInputSpec = z.infer export type InputSpec = z.infer export type OutputSpec = z.infer export type ComfyNodeDef = z.infer + +export const isIntInputSpec = ( + inputSpec: InputSpec +): inputSpec is IntInputSpec => { + return inputSpec.type === 'INT' +} + +export const isFloatInputSpec = ( + inputSpec: InputSpec +): inputSpec is FloatInputSpec => { + return inputSpec.type === 'FLOAT' +} + +export const isBooleanInputSpec = ( + inputSpec: InputSpec +): inputSpec is BooleanInputSpec => { + return inputSpec.type === 'BOOLEAN' +} + +export const isStringInputSpec = ( + inputSpec: InputSpec +): inputSpec is StringInputSpec => { + return inputSpec.type === 'STRING' +} + +export const isComboInputSpec = ( + inputSpec: InputSpec +): inputSpec is ComboInputSpec => { + return inputSpec.type === 'COMBO' +} diff --git a/src/scripts/app.ts b/src/scripts/app.ts index e4aece79b..114c60c0b 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -19,7 +19,11 @@ import { type NodeId, validateComfyWorkflow } from '@/schemas/comfyWorkflowSchema' -import type { ComfyNodeDef } from '@/schemas/nodeDefSchema' +import { + type ComfyNodeDef as ComfyNodeDefV2, + isComboInputSpec +} from '@/schemas/nodeDef/nodeDefSchemaV2' +import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import { getFromWebmFile } from '@/scripts/metadata/ebml' import { useDialogService } from '@/services/dialogService' import { useExtensionService } from '@/services/extensionService' @@ -30,7 +34,11 @@ import { useExecutionStore } from '@/stores/executionStore' import { useExtensionStore } from '@/stores/extensionStore' import { KeyComboImpl, useKeybindingStore } from '@/stores/keybindingStore' import { useModelStore } from '@/stores/modelStore' -import { SYSTEM_NODE_DEFS, useNodeDefStore } from '@/stores/nodeDefStore' +import { + ComfyNodeDefImpl, + SYSTEM_NODE_DEFS, + useNodeDefStore +} from '@/stores/nodeDefStore' import { useSettingStore } from '@/stores/settingStore' import { useToastStore } from '@/stores/toastStore' import { useWidgetStore } from '@/stores/widgetStore' @@ -841,10 +849,12 @@ export class ComfyApp { this.canvas?.draw(true, true) } - private updateVueAppNodeDefs(defs: Record) { + private updateVueAppNodeDefs( + defs: Record + ) { // Frontend only nodes registered by custom nodes. // Example: https://github.com/rgthree/rgthree-comfy/blob/dd534e5384be8cf0c0fa35865afe2126ba75ac55/src_web/comfyui/fast_groups_bypasser.ts#L10 - const rawDefs = Object.fromEntries( + const rawDefs: Record = Object.fromEntries( Object.entries(LiteGraph.registered_node_types).map(([name, node]) => [ name, { @@ -868,7 +878,7 @@ export class ComfyApp { } const nodeDefStore = useNodeDefStore() - const nodeDefArray: ComfyNodeDef[] = Object.values(allNodeDefs) + const nodeDefArray: ComfyNodeDefV1[] = Object.values(allNodeDefs) useExtensionService().invokeExtensions( 'beforeRegisterVueAppNodeDefs', nodeDefArray, @@ -877,33 +887,29 @@ export class ComfyApp { nodeDefStore.updateNodeDefs(nodeDefArray) } - #translateNodeDefs(defs: Record) { - return Object.fromEntries( - Object.entries(defs).map(([name, def]) => [ - name, - { - ...def, - display_name: st( - `nodeDefs.${name}.display_name`, - def.display_name ?? def.name - ), - description: def.description - ? st(`nodeDefs.${name}.description`, def.description) - : undefined, - category: def.category - .split('/') - .map((category) => st(`nodeCategories.${category}`, category)) - .join('/') - } - ]) - ) - } + async #getNodeDefs(): Promise< + Record + > { + const translateNodeDef = (def: ComfyNodeDefV1): ComfyNodeDefV1 => ({ + ...def, + display_name: st( + `nodeDefs.${def.name}.display_name`, + def.display_name ?? def.name + ), + description: def.description + ? st(`nodeDefs.${def.name}.description`, def.description) + : undefined, + category: def.category + .split('/') + .map((category: string) => st(`nodeCategories.${category}`, category)) + .join('/') + }) - async #getNodeDefs() { - return this.#translateNodeDefs( + return _.mapValues( await api.getNodeDefs({ validate: useSettingStore().get('Comfy.Validation.NodeDefs') - }) + }), + (def) => new ComfyNodeDefImpl(translateNodeDef(def)) ) } @@ -938,11 +944,16 @@ export class ComfyApp { } } - async registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) { - return await useLitegraphService().registerNodeDef(nodeId, nodeData) + async registerNodeDef( + nodeId: string, + nodeDef: ComfyNodeDefV1 & ComfyNodeDefV2 + ) { + return await useLitegraphService().registerNodeDef(nodeId, nodeDef) } - async registerNodesFromDefs(defs: Record) { + async registerNodesFromDefs( + defs: Record + ) { await useExtensionService().invokeExtensionsAsync('addCustomNodeDefs', defs) // Register a node for each definition @@ -1590,28 +1601,31 @@ export class ComfyApp { useToastStore().add(requestToastMessage) } - const defs = await this.#getNodeDefs() - for (const nodeId in defs) { - this.registerNodeDef(nodeId, defs[nodeId]) + const defs: Record = + await this.#getNodeDefs() + + for (const [nodeId, nodeDef] of Object.entries(defs)) { + this.registerNodeDef(nodeId, nodeDef) } - for (let nodeNum in this.graph.nodes) { - const node = this.graph.nodes[nodeNum] + + for (const node of this.graph.nodes) { const def = defs[node.type] // Allow primitive nodes to handle refresh node.refreshComboInNode?.(defs) - if (!def?.input) continue + if (!def) continue - for (const widgetNum in node.widgets) { - const widget = node.widgets[widgetNum] - if (widget.type === 'combo') { - if (def['input'].required?.[widget.name] !== undefined) { - // @ts-expect-error InputSpec is not typed correctly - widget.options.values = def['input'].required[widget.name][0] - } else if (def['input'].optional?.[widget.name] !== undefined) { - // @ts-expect-error InputSpec is not typed correctly - widget.options.values = def['input'].optional[widget.name][0] - } + // Update combo options in combo widgets + for (const widget of node.widgets) { + const inputSpec = def.inputs[widget.name] + if ( + inputSpec && + isComboInputSpec(inputSpec) && + widget.type === 'combo' + ) { + widget.options.values = inputSpec.options.map((o) => + typeof o === 'string' ? o : o.toString() + ) } } } diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index 4d4ce1c54..474573afb 100644 --- a/src/stores/nodeDefStore.ts +++ b/src/stores/nodeDefStore.ts @@ -228,17 +228,17 @@ export const useNodeDefStore = defineStore('nodeDef', () => { function updateNodeDefs(nodeDefs: ComfyNodeDefV1[]) { const newNodeDefsByName: Record = {} const newNodeDefsByDisplayName: Record = {} + for (const nodeDef of nodeDefs) { - try { - const nodeDefImpl = new ComfyNodeDefImpl(nodeDef) - newNodeDefsByName[nodeDef.name] = nodeDefImpl - newNodeDefsByDisplayName[nodeDef.display_name] = nodeDefImpl - } catch (e) { - // Avoid breaking the app for invalid nodeDefs - // NodeDef validation is now optional for performance reasons - console.error('Error adding nodeDef:', e) - } + const nodeDefImpl = + nodeDef instanceof ComfyNodeDefImpl + ? nodeDef + : new ComfyNodeDefImpl(nodeDef) + + newNodeDefsByName[nodeDef.name] = nodeDefImpl + newNodeDefsByDisplayName[nodeDef.display_name] = nodeDefImpl } + nodeDefsByName.value = newNodeDefsByName nodeDefsByDisplayName.value = newNodeDefsByDisplayName }