From 2d179ad6321070f0950593823db68eba1cce30f0 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Tue, 4 Mar 2025 12:07:13 -0500 Subject: [PATCH] [Refactor] Use node def v2 in registerNodeDef (#2856) --- src/extensions/core/groupNode.ts | 1 - src/schemas/nodeDef/nodeDefSchemaV2.ts | 16 ++++ src/scripts/app.ts | 23 +++--- src/services/litegraphService.ts | 103 ++++++++++++------------- src/types/litegraph-augmentation.d.ts | 5 +- 5 files changed, 84 insertions(+), 64 deletions(-) diff --git a/src/extensions/core/groupNode.ts b/src/extensions/core/groupNode.ts index 7d1ba3730..3a165c4e9 100644 --- a/src/extensions/core/groupNode.ts +++ b/src/extensions/core/groupNode.ts @@ -1162,7 +1162,6 @@ export class GroupNodeHandler { def?.input?.optional?.[old.inputName] if (!input) continue - // @ts-expect-error InputSpec is not typed correctly widget.options.values = input[0] if ( diff --git a/src/schemas/nodeDef/nodeDefSchemaV2.ts b/src/schemas/nodeDef/nodeDefSchemaV2.ts index e1aa24b08..095ca757c 100644 --- a/src/schemas/nodeDef/nodeDefSchemaV2.ts +++ b/src/schemas/nodeDef/nodeDefSchemaV2.ts @@ -121,3 +121,19 @@ export const isComboInputSpec = ( ): inputSpec is ComboInputSpec => { return inputSpec.type === 'COMBO' } + +/** + * Check if a node definition is a valid ComfyUI node definition. + * + * Note: This is just a simple check against the V1 schema. + * + * @param nodeDef - The node definition to check. + * @returns True if the node definition is valid, false otherwise. + */ +export const isComfyNodeDef = (nodeDef: unknown): nodeDef is ComfyNodeDef => { + return ( + !!nodeDef && + typeof nodeDef === 'object' && + ['inputs', 'outputs'].every((key) => key in nodeDef) + ) +} diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 114c60c0b..80fb6b8fa 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -19,9 +19,11 @@ import { type NodeId, validateComfyWorkflow } from '@/schemas/comfyWorkflowSchema' +import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration' import { type ComfyNodeDef as ComfyNodeDefV2, - isComboInputSpec + isComboInputSpec, + isComfyNodeDef as isComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import { getFromWebmFile } from '@/scripts/metadata/ebml' @@ -944,16 +946,19 @@ export class ComfyApp { } } - async registerNodeDef( - nodeId: string, - nodeDef: ComfyNodeDefV1 & ComfyNodeDefV2 - ) { - return await useLitegraphService().registerNodeDef(nodeId, nodeDef) + async registerNodeDef(nodeId: string, nodeDef: ComfyNodeDefV1) { + return await useLitegraphService().registerNodeDef( + nodeId, + isComfyNodeDefV2(nodeDef) + ? nodeDef + : { + ...(nodeDef as ComfyNodeDefV1), + ...transformNodeDefV1ToV2(nodeDef) + } + ) } - async registerNodesFromDefs( - defs: Record - ) { + async registerNodesFromDefs(defs: Record) { await useExtensionService().invokeExtensionsAsync('addCustomNodeDefs', defs) // Register a node for each definition diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts index a613d8d74..cfa5affaf 100644 --- a/src/services/litegraphService.ts +++ b/src/services/litegraphService.ts @@ -13,11 +13,8 @@ import { IBaseWidget, IWidget } from '@comfyorg/litegraph/dist/types/widgets' import { useNodeImage, useNodeVideo } from '@/composables/node/useNodeImage' import { st } from '@/i18n' import type { NodeId } from '@/schemas/comfyWorkflowSchema' -import { - type ComfyNodeDef, - InputSpec, - getInputSpecType -} from '@/schemas/nodeDefSchema' +import type { ComfyNodeDef as ComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' +import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import { ANIM_PREVIEW_WIDGET, ComfyApp, app } from '@/scripts/app' import { $el } from '@/scripts/ui' import { calculateImageGrid, createImageHost } from '@/scripts/ui/imagePreview' @@ -38,54 +35,55 @@ export const useLitegraphService = () => { const toastStore = useToastStore() const canvasStore = useCanvasStore() - async function registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) { + async function registerNodeDef( + nodeId: string, + nodeDef: ComfyNodeDefV2 & ComfyNodeDefV1 + ) { const node = class ComfyNode extends LGraphNode { - static comfyClass? = nodeData.name - static title? = nodeData.display_name || nodeData.name - static nodeData? = nodeData + static comfyClass?: string = nodeDef.name + static title?: string = nodeDef.display_name || nodeDef.name + static nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2 = nodeDef static category?: string constructor(title?: string) { super(title) - const requiredInputs = nodeData.input.required - let inputs = nodeData['input']['required'] - if (nodeData['input']['optional'] != undefined) { - inputs = Object.assign( - {}, - nodeData['input']['required'], - nodeData['input']['optional'] - ) - } const config: { minWidth: number minHeight: number widget?: IBaseWidget } = { minWidth: 1, minHeight: 1 } - for (const inputName in inputs) { - const inputData = [ - inputs[inputName][0], - inputs[inputName][1] ?? {} - ] as InputSpec - const inputType = getInputSpecType(inputData) - const inputOptions = inputData[1] - const nameKey = `nodeDefs.${normalizeI18nKey(nodeData.name)}.inputs.${normalizeI18nKey(inputName)}.name` - - const inputIsRequired = requiredInputs && inputName in requiredInputs + // Process inputs using V2 schema + for (const [inputName, inputSpec] of Object.entries(nodeDef.inputs)) { + const inputType = inputSpec.type + const nameKey = `nodeDefs.${normalizeI18nKey(nodeDef.name)}.inputs.${normalizeI18nKey(inputName)}.name` let widgetCreated = true - const widgetType = app.getWidgetType(inputData, inputName) + const widgetType = app.getWidgetType( + [inputType, inputSpec], + inputName + ) if (widgetType) { if (widgetType === 'COMBO') { Object.assign( config, - app.widgets.COMBO(this, inputName, inputData, app) || {} + app.widgets.COMBO( + this, + inputName, + [inputType, inputSpec], + app + ) || {} ) } else { Object.assign( config, - app.widgets[widgetType](this, inputName, inputData, app) || {} + app.widgets[widgetType]( + this, + inputName, + [inputType, inputSpec], + app + ) || {} ) } if (config.widget) { @@ -94,9 +92,9 @@ export const useLitegraphService = () => { } } else { // Node connection inputs - const shapeOptions = inputIsRequired - ? {} - : { shape: RenderShape.HollowCircle } + const shapeOptions = inputSpec.isOptional + ? { shape: RenderShape.HollowCircle } + : {} this.addInput(inputName, inputType, { ...shapeOptions, @@ -107,34 +105,34 @@ export const useLitegraphService = () => { if (widgetCreated && config?.widget) { config.widget.options ??= {} - if (!inputIsRequired) { + if (inputSpec.isOptional) { config.widget.options.inputIsOptional = true } - if (inputOptions.forceInput) { + if (inputSpec.forceInput) { config.widget.options.forceInput = true } - if (inputOptions.defaultInput) { + if (inputSpec.defaultInput) { config.widget.options.defaultInput = true } - if (inputOptions.advanced) { + if (inputSpec.advanced) { config.widget.advanced = true } - if (inputOptions.hidden) { + if (inputSpec.hidden) { config.widget.hidden = true } } } - for (const o in nodeData['output']) { - let output = nodeData['output'][o] - if (output instanceof Array) output = 'COMBO' - const outputName = nodeData['output_name'][o] || output - const outputIsList = nodeData['output_is_list'][o] + // Process outputs using V2 schema + for (const output of nodeDef.outputs) { + const outputName = output.name + const outputType = output.type + const outputIsList = output.is_list const shapeOptions = outputIsList ? { shape: LiteGraph.GRID_SHAPE } : {} - const nameKey = `nodeDefs.${normalizeI18nKey(nodeData.name)}.outputs.${o}.name` - const typeKey = `dataTypes.${normalizeI18nKey(output)}` + const nameKey = `nodeDefs.${normalizeI18nKey(nodeDef.name)}.outputs.${output.index}.name` + const typeKey = `dataTypes.${normalizeI18nKey(outputType)}` const outputOptions = { ...shapeOptions, // If the output name is different from the output type, use the output name. @@ -142,11 +140,11 @@ export const useLitegraphService = () => { // - type ("INT"); name ("Positive") => translate name // - type ("FLOAT"); name ("FLOAT") => translate type localized_name: - output !== outputName + outputType !== outputName ? st(nameKey, outputName) : st(typeKey, outputName) } - this.addOutput(outputName, output, outputOptions) + this.addOutput(outputName, outputType, outputOptions) } const s = this.computeSize() @@ -186,7 +184,7 @@ export const useLitegraphService = () => { super.configure(data) } } - node.prototype.comfyClass = nodeData.name + node.prototype.comfyClass = nodeDef.name addNodeContextMenuHandler(node) addDrawBackgroundHandler(node) @@ -195,11 +193,12 @@ export const useLitegraphService = () => { await extensionService.invokeExtensionsAsync( 'beforeRegisterNodeDef', node, - nodeData + nodeDef // Receives V1 NodeDef ) + LiteGraph.registerNodeType(nodeId, node) // Note: Do not move this to the class definition, it will be overwritten - node.category = nodeData.category + node.category = nodeDef.category } /** @@ -722,7 +721,7 @@ export const useLitegraphService = () => { } function addNodeOnGraph( - nodeDef: ComfyNodeDef, + nodeDef: ComfyNodeDefV1 | ComfyNodeDefV2, options: Record = {} ): LGraphNode { options.pos ??= getCanvasCenter() diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts index daca9f64d..70172770d 100644 --- a/src/types/litegraph-augmentation.d.ts +++ b/src/types/litegraph-augmentation.d.ts @@ -1,7 +1,8 @@ import '@comfyorg/litegraph' import type { LLink, Size } from '@comfyorg/litegraph' -import type { ComfyNodeDef } from '@/schemas/nodeDefSchema' +import type { ComfyNodeDef as ComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' +import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import type { DOMWidget, DOMWidgetOptions } from '@/scripts/domWidget' import type { NodeId } from '../schemas/comfyWorkflowSchema' @@ -63,7 +64,7 @@ declare module '@comfyorg/litegraph' { type?: string comfyClass: string title: string - nodeData?: ComfyNodeDef + nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2 category?: string new (): T }