diff --git a/src/components/NodePreview.vue b/src/components/NodePreview.vue index 6ad7149f4..b9ab7e1ee 100644 --- a/src/components/NodePreview.vue +++ b/src/components/NodePreview.vue @@ -59,7 +59,7 @@ const nodeDefStore = useNodeDefStore() const nodeDef = props.nodeDef const allInputDefs = nodeDef.input.all -const allOutputDefs = Object.values(nodeDef.output) +const allOutputDefs = nodeDef.output.all const slotInputDefs = allInputDefs.filter( (input) => !nodeDefStore.inputIsWidget(input) ) diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index 13af2fed2..f57120a44 100644 --- a/src/stores/nodeDefStore.ts +++ b/src/stores/nodeDefStore.ts @@ -132,15 +132,21 @@ export class ComfyInputsSpec { } export class ComfyOutputSpec { - type: string - is_list: boolean - display_name: string - name?: string - comboOptions?: any[] + constructor( + public name: string, + public display_name: string, + public type: string, + public is_list: boolean, + public comboOptions?: any[] + ) {} } export class ComfyOutputsSpec { - [key: string]: ComfyOutputSpec + constructor(public outputByName: Record) {} + + get all() { + return Object.values(this.outputByName) + } } export class ComfyNodeDefImpl { @@ -158,23 +164,20 @@ export class ComfyNodeDefImpl { private static transformOutputSpec(obj: any): ComfyOutputsSpec { const { output, output_is_list, output_name } = obj - const result: ComfyOutputsSpec = {} + const result = {} output.forEach((type: string | any[], index: number) => { const typeString = Array.isArray(type) ? 'COMBO' : type const display_name = output_name[index] const name = display_name === typeString ? index.toString() : display_name - const outputSpec = { + const outputSpec = new ComfyOutputSpec( name, display_name, - type: typeString, - is_list: output_is_list[index] - } as ComfyOutputSpec - - if (Array.isArray(type)) { - outputSpec.comboOptions = type - } + typeString, + output_is_list[index], + Array.isArray(type) ? type : undefined + ) if (name in result) { throw new Error(`Duplicate output name: ${name}`) @@ -182,7 +185,7 @@ export class ComfyNodeDefImpl { result[name] = outputSpec }) - return result + return new ComfyOutputsSpec(result) } } diff --git a/tests-ui/tests/nodeDef.test.ts b/tests-ui/tests/nodeDef.test.ts index 87db782a3..e1f5c683b 100644 --- a/tests-ui/tests/nodeDef.test.ts +++ b/tests-ui/tests/nodeDef.test.ts @@ -171,7 +171,7 @@ describe('ComfyNodeDefImpl', () => { expect(result.python_module).toBe('test_module') expect(result.description).toBe('A test node') expect(result.input).toBeInstanceOf(ComfyInputsSpec) - expect(result.output).toEqual({ + expect(result.output.outputByName).toEqual({ intOutput: { name: 'intOutput', display_name: 'intOutput', @@ -196,7 +196,7 @@ describe('ComfyNodeDefImpl', () => { const result = plainToClass(ComfyNodeDefImpl, plainObject) - expect(result.output).toEqual({ + expect(result.output.outputByName).toEqual({ stringOutput: { name: 'stringOutput', display_name: 'stringOutput', @@ -234,7 +234,7 @@ describe('ComfyNodeDefImpl', () => { const result = plainToClass(ComfyNodeDefImpl, plainObject) - expect(result.output).toEqual({ + expect(result.output.outputByName).toEqual({ '0': { name: '0', display_name: 'INT', @@ -289,7 +289,7 @@ describe('ComfyNodeDefImpl', () => { const result = plainToClass(ComfyNodeDefImpl, plainObject) - expect(result.output).toEqual({}) + expect(result.output.outputByName).toEqual({}) }) it('should handle complex input specifications', () => {