From c6d2767af116538bebd1d757574dbf0bdba536a8 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Thu, 25 Jul 2024 13:50:55 -0400 Subject: [PATCH] Transforms ComfyInputsSpec on nodes (#220) * Convert input spec defs * Fix test * Add combo test * import metadata --- index.html | 1 + src/stores/nodeDefStore.ts | 119 +++++++++++++++++++++++++++ src/types/apiTypes.ts | 24 +++--- tests-ui/globalSetup.ts | 1 + tests-ui/tests/nodeDef.test.ts | 142 +++++++++++++++++++++++++++++++++ tsconfig.json | 2 + 6 files changed, 278 insertions(+), 11 deletions(-) create mode 100644 tests-ui/tests/nodeDef.test.ts diff --git a/index.html b/index.html index 30d4d0a45..072de61ff 100644 --- a/index.html +++ b/index.html @@ -14,6 +14,7 @@ --> diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index 2caf85830..e0c615159 100644 --- a/src/stores/nodeDefStore.ts +++ b/src/stores/nodeDefStore.ts @@ -1,6 +1,125 @@ import { NodeSearchService } from '@/services/nodeSearchService' import { ComfyNodeDef } from '@/types/apiTypes' import { defineStore } from 'pinia' +import { Type, Transform, plainToClass } from 'class-transformer' + +export class BaseInputSpec { + type: string + + default?: T + + @Type(() => Boolean) + forceInput?: boolean + + static isInputSpec(obj: any): boolean { + return ( + Array.isArray(obj) && + obj.length >= 1 && + (typeof obj[0] === 'string' || Array.isArray(obj[0])) + ) + } +} + +export class NumericInputSpec extends BaseInputSpec { + @Type(() => Number) + min?: number + + @Type(() => Number) + max?: number + + @Type(() => Number) + step?: number +} + +export class IntInputSpec extends NumericInputSpec { + type: 'INT' = 'INT' +} + +export class FloatInputSpec extends NumericInputSpec { + type: 'FLOAT' = 'FLOAT' + + @Type(() => Number) + round?: number +} + +export class BooleanInputSpec extends BaseInputSpec { + type: 'BOOLEAN' = 'BOOLEAN' + + labelOn?: string + labelOff?: string +} + +export class StringInputSpec extends BaseInputSpec { + type: 'STRING' = 'STRING' + + @Type(() => Boolean) + multiline?: boolean + + @Type(() => Boolean) + dynamicPrompts?: boolean +} + +export class ComboInputSpec extends BaseInputSpec { + type: string = 'COMBO' + + @Transform(({ value }) => value[0]) + comboOptions: any[] + + @Type(() => Boolean) + controlAfterGenerate?: boolean + + @Type(() => Boolean) + imageUpload?: boolean +} + +export class CustomInputSpec extends BaseInputSpec {} + +export class ComfyInputsSpec { + @Transform(({ value }) => ComfyInputsSpec.transformInputSpecRecord(value)) + required?: Record + + @Transform(({ value }) => ComfyInputsSpec.transformInputSpecRecord(value)) + optional?: Record + + hidden?: Record + + private static transformInputSpecRecord( + record: Record + ): Record { + if (!record) return record + const result: Record = {} + for (const [key, value] of Object.entries(record)) { + result[key] = ComfyInputsSpec.transformSingleInputSpec(value) + } + return result + } + + private static transformSingleInputSpec(value: any): BaseInputSpec { + if (!BaseInputSpec.isInputSpec(value)) return value + + const [typeRaw, spec] = value + const type = Array.isArray(typeRaw) ? 'COMBO' : value[0] + + switch (type) { + case 'INT': + return plainToClass(IntInputSpec, { type, ...spec }) + case 'FLOAT': + return plainToClass(FloatInputSpec, { type, ...spec }) + case 'BOOLEAN': + return plainToClass(BooleanInputSpec, { type, ...spec }) + case 'STRING': + return plainToClass(StringInputSpec, { type, ...spec }) + case 'COMBO': + return plainToClass(ComboInputSpec, { + type, + ...spec, + comboOptions: typeRaw + }) + default: + return plainToClass(CustomInputSpec, { type, ...spec }) + } + } +} export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [ { diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 7004873eb..00c4638b2 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -249,21 +249,23 @@ const zInputSpec = z.union([ zCustomInputSpec ]) +const zComfyInputsSpec = z.object({ + required: z.record(zInputSpec).optional(), + optional: z.record(zInputSpec).optional(), + // Frontend repo is not using it, but some custom nodes are using the + // hidden field to pass various values. + hidden: z.record(z.any()).optional() +}) + const zComfyNodeDataType = z.string() const zComfyComboOutput = z.array(z.any()) -const zComfyOutputSpec = z.array( +const zComfyOutputTypesSpec = z.array( z.union([zComfyNodeDataType, zComfyComboOutput]) ) const zComfyNodeDef = z.object({ - input: z.object({ - required: z.record(zInputSpec).optional(), - optional: z.record(zInputSpec).optional(), - // Frontend repo is not using it, but some custom nodes are using the - // hidden field to pass various values. - hidden: z.record(z.any()).optional() - }), - output: zComfyOutputSpec, + input: zComfyInputsSpec, + output: zComfyOutputTypesSpec, output_is_list: z.array(z.boolean()), output_name: z.array(z.string()), name: z.string(), @@ -275,8 +277,8 @@ const zComfyNodeDef = z.object({ }) // `/object_info` -export type ComfyInputSpec = z.infer -export type ComfyOutputSpec = z.infer +export type ComfyInputsSpec = z.infer +export type ComfyOutputTypesSpec = z.infer export type ComfyNodeDef = z.infer export function validateComfyNodeDef(data: any): ComfyNodeDef { diff --git a/tests-ui/globalSetup.ts b/tests-ui/globalSetup.ts index 87332fef9..e3cf4c131 100644 --- a/tests-ui/globalSetup.ts +++ b/tests-ui/globalSetup.ts @@ -5,6 +5,7 @@ module.exports = async function () { disconnect() {} } + require('reflect-metadata') const { nop } = require('./utils/nopProxy') global.enableWebGLCanvas = nop diff --git a/tests-ui/tests/nodeDef.test.ts b/tests-ui/tests/nodeDef.test.ts new file mode 100644 index 000000000..2f0d933e4 --- /dev/null +++ b/tests-ui/tests/nodeDef.test.ts @@ -0,0 +1,142 @@ +import { plainToClass } from 'class-transformer' +import { + ComfyInputsSpec, + IntInputSpec, + StringInputSpec, + BooleanInputSpec, + FloatInputSpec, + CustomInputSpec, + ComboInputSpec +} from '@/stores/nodeDefStore' // Adjust the import path as needed + +describe('ComfyInputsSpec', () => { + it('should transform a plain object to ComfyInputsSpec instance', () => { + const plainObject = { + required: { + intInput: ['INT', { min: 0, max: 100, default: 50 }], + stringInput: ['STRING', { default: 'Hello', multiline: true }] + }, + optional: { + booleanInput: [ + 'BOOLEAN', + { default: true, labelOn: 'Yes', labelOff: 'No' } + ], + floatInput: ['FLOAT', { min: 0, max: 1, step: 0.1 }] + }, + hidden: { + someHiddenValue: 42 + } + } + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result).toBeInstanceOf(ComfyInputsSpec) + expect(result.required).toBeDefined() + expect(result.optional).toBeDefined() + expect(result.hidden).toBeDefined() + }) + + it('should correctly transform required input specs', () => { + const plainObject = { + required: { + intInput: ['INT', { min: 0, max: 100, default: 50 }], + stringInput: ['STRING', { default: 'Hello', multiline: true }] + } + } + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result.required?.intInput).toBeInstanceOf(IntInputSpec) + expect(result.required?.stringInput).toBeInstanceOf(StringInputSpec) + + const intInput = result.required?.intInput as IntInputSpec + const stringInput = result.required?.stringInput as StringInputSpec + + expect(intInput.min).toBe(0) + expect(intInput.max).toBe(100) + expect(intInput.default).toBe(50) + expect(stringInput.default).toBe('Hello') + expect(stringInput.multiline).toBe(true) + }) + + it('should correctly transform optional input specs', () => { + const plainObject = { + optional: { + booleanInput: [ + 'BOOLEAN', + { default: true, labelOn: 'Yes', labelOff: 'No' } + ], + floatInput: ['FLOAT', { min: 0, max: 1, step: 0.1 }] + } + } + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result.optional?.booleanInput).toBeInstanceOf(BooleanInputSpec) + expect(result.optional?.floatInput).toBeInstanceOf(FloatInputSpec) + + const booleanInput = result.optional?.booleanInput as BooleanInputSpec + const floatInput = result.optional?.floatInput as FloatInputSpec + + expect(booleanInput.default).toBe(true) + expect(booleanInput.labelOn).toBe('Yes') + expect(booleanInput.labelOff).toBe('No') + expect(floatInput.min).toBe(0) + expect(floatInput.max).toBe(1) + expect(floatInput.step).toBe(0.1) + }) + + it('should handle custom input specs', () => { + const plainObject = { + optional: { + comboInput: [[1, 2, 3], { default: 2 }] + } + } + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result.optional?.comboInput).toBeInstanceOf(ComboInputSpec) + expect(result.optional?.comboInput.type).toBe('COMBO') + expect(result.optional?.comboInput.default).toBe(2) + }) + + it('should handle custom input specs', () => { + const plainObject = { + optional: { + customInput: ['CUSTOM_TYPE', { default: 'custom value' }] + } + } + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result.optional?.customInput).toBeInstanceOf(CustomInputSpec) + expect(result.optional?.customInput.type).toBe('CUSTOM_TYPE') + expect(result.optional?.customInput.default).toBe('custom value') + }) + + it('should not transform hidden fields', () => { + const plainObject = { + hidden: { + someHiddenValue: 42, + anotherHiddenValue: { nested: 'object' } + } + } + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result.hidden).toEqual(plainObject.hidden) + expect(result.hidden?.someHiddenValue).toBe(42) + expect(result.hidden?.anotherHiddenValue).toEqual({ nested: 'object' }) + }) + + it('should handle empty or undefined fields', () => { + const plainObject = {} + + const result = plainToClass(ComfyInputsSpec, plainObject) + + expect(result).toBeInstanceOf(ComfyInputsSpec) + expect(result.required).toBeUndefined() + expect(result.optional).toBeUndefined() + expect(result.hidden).toBeUndefined() + }) +}) diff --git a/tsconfig.json b/tsconfig.json index ef159ef27..e3d9b16a2 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -8,6 +8,8 @@ "sourceMap": true, "esModuleInterop": true, "moduleResolution": "Node", + "experimentalDecorators": true, + "emitDecoratorMetadata": true, /* Linting */ "strict": false,