From 733d71aaac9958ce864a961694218523093986fd Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Tue, 1 Apr 2025 20:48:32 -0400 Subject: [PATCH] [Refactor] Split node constructor logic (#3307) --- src/services/litegraphService.ts | 122 +++++++++++++++++++------------ 1 file changed, 74 insertions(+), 48 deletions(-) diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts index fb078344d0..1124c93634 100644 --- a/src/services/litegraphService.ts +++ b/src/services/litegraphService.ts @@ -14,7 +14,11 @@ import { useNodeImage, useNodeVideo } from '@/composables/node/useNodeImage' import { st, t } from '@/i18n' import type { NodeId } from '@/schemas/comfyWorkflowSchema' import { transformInputSpecV2ToV1 } from '@/schemas/nodeDef/migration' -import type { ComfyNodeDef as ComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' +import type { + ComfyNodeDef as ComfyNodeDefV2, + InputSpec, + OutputSpec +} from '@/schemas/nodeDef/nodeDefSchemaV2' import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import { ComfyApp, app } from '@/scripts/app' import { $el } from '@/scripts/ui' @@ -39,13 +43,36 @@ export const useLitegraphService = () => { async function registerNodeDef(nodeId: string, nodeDefV1: ComfyNodeDefV1) { const node = class ComfyNode extends LGraphNode { - static comfyClass?: string - static title?: string - static nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2 - static category?: string + static comfyClass: string + static title: string + static category: string + static nodeData: ComfyNodeDefV1 & ComfyNodeDefV2 + + /** + * @internal The initial minimum size of the node. + */ + #initialMinSize = { width: 1, height: 1 } + /** + * @internal The key for the node definition in the i18n file. + */ + get #nodeKey(): string { + return `nodeDefs.${normalizeI18nKey(ComfyNode.nodeData.name)}` + } constructor(title: string) { super(title) + this.#setupStrokeStyles() + this.#addInputs(ComfyNode.nodeData.inputs) + this.#addOutputs(ComfyNode.nodeData.outputs) + this.#setInitialSize() + this.serialize_widgets = true + extensionService.invokeExtensionsAsync('nodeCreated', this) + } + + /** + * @internal Setup stroke styles for the node under various conditions. + */ + #setupStrokeStyles() { this.strokeStyles['running'] = function (this: LGraphNode) { if (this.id == app.runningNodeId) { return { color: '#0f0' } @@ -66,12 +93,15 @@ export const useLitegraphService = () => { return { color: '#f0f', lineWidth: 2 } } } + } - const nodeMinSize = { width: 1, height: 1 } - // Process inputs using V2 schema - for (const [inputName, inputSpec] of Object.entries(nodeDef.inputs)) { + /** + * @internal Add inputs to the node. + */ + #addInputs(inputs: Record) { + for (const [inputName, inputSpec] of Object.entries(inputs)) { const inputType = inputSpec.type - const nameKey = `nodeDefs.${normalizeI18nKey(nodeDef.name)}.inputs.${normalizeI18nKey(inputName)}.name` + const nameKey = `${this.#nodeKey}.inputs.${normalizeI18nKey(inputName)}.name` const widgetConstructor = widgetStore.widgets[inputType] if (widgetConstructor) { @@ -87,29 +117,25 @@ export const useLitegraphService = () => { ) ?? {} if (widget) { - const fallback = widget.label ?? inputName - widget.label = st(nameKey, fallback) - + widget.label = st(nameKey, widget.label ?? inputName) widget.options ??= {} - if (inputSpec.isOptional) { - widget.options.inputIsOptional = true - } - if (inputSpec.forceInput) { - widget.options.forceInput = true - } - if (inputSpec.defaultInput) { - widget.options.defaultInput = true - } - if (inputSpec.advanced) { - widget.advanced = true - } - if (inputSpec.hidden) { - widget.hidden = true - } + Object.assign(widget.options, { + inputIsOptional: inputSpec.isOptional, + forceInput: inputSpec.forceInput, + defaultInput: inputSpec.defaultInput, + advanced: inputSpec.advanced, + hidden: inputSpec.hidden + }) } - nodeMinSize.width = Math.max(nodeMinSize.width, minWidth) - nodeMinSize.height = Math.max(nodeMinSize.height, minHeight) + this.#initialMinSize.width = Math.max( + this.#initialMinSize.width, + minWidth + ) + this.#initialMinSize.height = Math.max( + this.#initialMinSize.height, + minHeight + ) } else { // Node connection inputs const shapeOptions = inputSpec.isOptional @@ -122,17 +148,17 @@ export const useLitegraphService = () => { }) } } + } - // 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(nodeDef.name)}.outputs.${output.index}.name` - const typeKey = `dataTypes.${normalizeI18nKey(outputType)}` + /** + * @internal Add outputs to the node. + */ + #addOutputs(outputs: OutputSpec[]) { + for (const output of outputs) { + const { name, type, is_list } = output + const shapeOptions = is_list ? { shape: LiteGraph.GRID_SHAPE } : {} + const nameKey = `${this.#nodeKey}.outputs.${output.index}.name` + const typeKey = `dataTypes.${normalizeI18nKey(type)}` const outputOptions = { ...shapeOptions, // If the output name is different from the output type, use the output name. @@ -140,20 +166,20 @@ export const useLitegraphService = () => { // - type ("INT"); name ("Positive") => translate name // - type ("FLOAT"); name ("FLOAT") => translate type localized_name: - outputType !== outputName - ? st(nameKey, outputName) - : st(typeKey, outputName) + type !== name ? st(nameKey, name) : st(typeKey, name) } - this.addOutput(outputName, outputType, outputOptions) + this.addOutput(name, type, outputOptions) } + } + /** + * @internal Set the initial size of the node. + */ + #setInitialSize() { const s = this.computeSize() - s[0] = Math.max(nodeMinSize.width, s[0] * 1.5) - s[1] = Math.max(nodeMinSize.height, s[1]) + s[0] = Math.max(this.#initialMinSize.width, s[0] * 1.5) + s[1] = Math.max(this.#initialMinSize.height, s[1]) this.setSize(s) - this.serialize_widgets = true - - extensionService.invokeExtensionsAsync('nodeCreated', this) } configure(data: any) {