Transforms ComfyInputsSpec on nodes (#220)

* Convert input spec defs

* Fix test

* Add combo test

* import metadata

index.html

@@ -14,6 +14,7 @@
 		</style> -->
 		<script src="node_modules/reflect-metadata/Reflect.js"></script>
 		<script type="module">
+			import 'reflect-metadata';
 			window["__COMFYUI_FRONTEND_VERSION__"] = __COMFYUI_FRONTEND_VERSION__;
 			console.log("ComfyUI Front-end version:", __COMFYUI_FRONTEND_VERSION__);
 		</script>

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<T = any> {
+  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<number> {
+  @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<boolean> {
+  type: 'BOOLEAN' = 'BOOLEAN'
+
+  labelOn?: string
+  labelOff?: string
+}
+
+export class StringInputSpec extends BaseInputSpec<string> {
+  type: 'STRING' = 'STRING'
+
+  @Type(() => Boolean)
+  multiline?: boolean
+
+  @Type(() => Boolean)
+  dynamicPrompts?: boolean
+}
+
+export class ComboInputSpec extends BaseInputSpec<any> {
+  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<string, BaseInputSpec>
+
+  @Transform(({ value }) => ComfyInputsSpec.transformInputSpecRecord(value))
+  optional?: Record<string, BaseInputSpec>
+
+  hidden?: Record<string, any>
+
+  private static transformInputSpecRecord(
+    record: Record<string, any>
+  ): Record<string, BaseInputSpec> {
+    if (!record) return record
+    const result: Record<string, BaseInputSpec> = {}
+    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[] = [
   {

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<typeof zInputSpec>
-export type ComfyOutputSpec = z.infer<typeof zComfyOutputSpec>
+export type ComfyInputsSpec = z.infer<typeof zComfyInputsSpec>
+export type ComfyOutputTypesSpec = z.infer<typeof zComfyOutputTypesSpec>
 export type ComfyNodeDef = z.infer<typeof zComfyNodeDef>
 
 export function validateComfyNodeDef(data: any): ComfyNodeDef {

tests-ui/globalSetup.ts

@@ -5,6 +5,7 @@ module.exports = async function () {
     disconnect() {}
   }
 
+  require('reflect-metadata')
   const { nop } = require('./utils/nopProxy')
   global.enableWebGLCanvas = nop
 

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()
+  })
+})

tsconfig.json

@@ -8,6 +8,8 @@
     "sourceMap": true,
     "esModuleInterop": true,
     "moduleResolution": "Node",
+    "experimentalDecorators": true,
+    "emitDecoratorMetadata": true,
 
     /* Linting */
     "strict": false,