[Schema] ComfyNodeDefV2 schema (#2847)

src/components/graph/NodeTooltip.vue

@@ -76,7 +76,7 @@ const onIdle = () => {
     const inputName = node.inputs[inputSlot].name
     const translatedTooltip = st(
       `nodeDefs.${normalizeI18nKey(node.type)}.inputs.${normalizeI18nKey(inputName)}.tooltip`,
-      nodeDef.inputs.getInput(inputName)?.tooltip
+      nodeDef.inputs[inputName]?.tooltip
     )
     return showTooltip(translatedTooltip)
   }
@@ -90,7 +90,7 @@ const onIdle = () => {
   if (outputSlot !== -1) {
     const translatedTooltip = st(
       `nodeDefs.${normalizeI18nKey(node.type)}.outputs.${outputSlot}.tooltip`,
-      nodeDef.outputs.all?.[outputSlot]?.tooltip
+      nodeDef.outputs[outputSlot]?.tooltip
     )
     return showTooltip(translatedTooltip)
   }
@@ -100,7 +100,7 @@ const onIdle = () => {
   if (widget && !widget.element) {
     const translatedTooltip = st(
       `nodeDefs.${normalizeI18nKey(node.type)}.inputs.${normalizeI18nKey(widget.name)}.tooltip`,
-      nodeDef.inputs.getInput(widget.name)?.tooltip
+      nodeDef.inputs[widget.name]?.tooltip
     )
     // Widget tooltip can be set dynamically, current translation collection does not support this.
     return showTooltip(widget.tooltip ?? translatedTooltip)

src/components/node/NodePreview.vue

@@ -83,16 +83,13 @@ https://github.com/Nuked88/ComfyUI-N-Sidebar/blob/7ae7da4a9761009fb6629bc04c6830
 import _ from 'lodash'
 import { computed } from 'vue'
 
-import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+import type { ComfyNodeDef as ComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import { useWidgetStore } from '@/stores/widgetStore'
 import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
 
-const props = defineProps({
-  nodeDef: {
-    type: ComfyNodeDefImpl,
-    required: true
-  }
-})
+const props = defineProps<{
+  nodeDef: ComfyNodeDefV2
+}>()
 
 const colorPaletteStore = useColorPaletteStore()
 const litegraphColors = computed(
@@ -102,8 +99,8 @@ const litegraphColors = computed(
 const widgetStore = useWidgetStore()
 
 const nodeDef = props.nodeDef
-const allInputDefs = nodeDef.inputs.all
-const allOutputDefs = nodeDef.outputs.all
+const allInputDefs = Object.values(nodeDef.inputs)
+const allOutputDefs = nodeDef.outputs
 const slotInputDefs = allInputDefs.filter(
   (input) => !widgetStore.inputIsWidget(input)
 )

src/extensions/core/widgetInputs.ts

@@ -624,7 +624,7 @@ app.registerExtension({
     app.canvas.getWidgetLinkType = function (widget, node) {
       const nodeDefStore = useNodeDefStore()
       const nodeDef = nodeDefStore.nodeDefsByName[node.type]
-      const input = nodeDef.inputs.getInput(widget.name)
+      const input = nodeDef.inputs[widget.name]
       return input?.type
     }
 

src/schemas/nodeDef/migration.ts

@@ -0,0 +1,121 @@
+import {
+  ComfyNodeDef as ComfyNodeDefV2,
+  InputSpec as InputSpecV2,
+  OutputSpec as OutputSpecV2
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import {
+  ComfyNodeDef as ComfyNodeDefV1,
+  InputSpec as InputSpecV1,
+  getComboSpecComboOptions,
+  isComboInputSpec,
+  isComboInputSpecV1
+} from '@/schemas/nodeDefSchema'
+
+/**
+ * Transforms a V1 node definition to V2 format
+ * @param nodeDefV1 The V1 node definition to transform
+ * @returns The transformed V2 node definition
+ */
+export function transformNodeDefV1ToV2(
+  nodeDefV1: ComfyNodeDefV1
+): ComfyNodeDefV2 {
+  // Transform inputs
+  const inputs: Record<string, InputSpecV2> = {}
+
+  // Process required inputs
+  if (nodeDefV1.input?.required) {
+    Object.entries(nodeDefV1.input.required).forEach(([name, inputSpecV1]) => {
+      inputs[name] = transformInputSpec(inputSpecV1, name, false)
+    })
+  }
+
+  // Process optional inputs
+  if (nodeDefV1.input?.optional) {
+    Object.entries(nodeDefV1.input.optional).forEach(([name, inputSpecV1]) => {
+      inputs[name] = transformInputSpec(inputSpecV1, name, true)
+    })
+  }
+
+  // Transform outputs
+  const outputs: OutputSpecV2[] = []
+
+  if (nodeDefV1.output) {
+    nodeDefV1.output.forEach((outputType, index) => {
+      const outputSpec: OutputSpecV2 = {
+        index,
+        name: nodeDefV1.output_name?.[index] || `output_${index}`,
+        type: Array.isArray(outputType) ? 'COMBO' : outputType,
+        is_list: nodeDefV1.output_is_list?.[index] || false,
+        tooltip: nodeDefV1.output_tooltips?.[index]
+      }
+
+      // Add options for combo outputs
+      if (Array.isArray(outputType)) {
+        outputSpec.options = outputType
+      }
+
+      outputs.push(outputSpec)
+    })
+  }
+
+  // Create the V2 node definition
+  return {
+    inputs,
+    outputs,
+    hidden: nodeDefV1.input?.hidden,
+    name: nodeDefV1.name,
+    display_name: nodeDefV1.display_name,
+    description: nodeDefV1.description,
+    category: nodeDefV1.category,
+    output_node: nodeDefV1.output_node,
+    python_module: nodeDefV1.python_module,
+    deprecated: nodeDefV1.deprecated,
+    experimental: nodeDefV1.experimental
+  }
+}
+
+/**
+ * Transforms a V1 input specification to V2 format
+ * @param inputSpecV1 The V1 input specification to transform
+ * @param name The name of the input
+ * @param isOptional Whether the input is optional
+ * @returns The transformed V2 input specification
+ */
+function transformInputSpec(
+  inputSpecV1: InputSpecV1,
+  name: string,
+  isOptional: boolean
+): InputSpecV2 {
+  // Extract options from the input spec
+  const options = inputSpecV1[1] || {}
+
+  // Base properties for all input types
+  const baseProps = {
+    name,
+    isOptional,
+    ...options
+  }
+
+  // Handle different input types
+  if (isComboInputSpec(inputSpecV1)) {
+    return {
+      type: 'COMBO',
+      ...baseProps,
+      options: isComboInputSpecV1(inputSpecV1)
+        ? inputSpecV1[0]
+        : getComboSpecComboOptions(inputSpecV1)
+    }
+  } else if (typeof inputSpecV1[0] === 'string') {
+    // Handle standard types (INT, FLOAT, BOOLEAN, STRING) and custom types
+    return {
+      type: inputSpecV1[0],
+      ...baseProps
+    }
+  }
+
+  // Fallback for any unhandled cases
+  return {
+    type: 'UNKNOWN',
+    ...baseProps
+  }
+}

src/schemas/nodeDef/nodeDefSchemaV2.ts

@@ -0,0 +1,93 @@
+import { z } from 'zod'
+
+import {
+  zBaseInputOptions,
+  zBooleanInputOptions,
+  zComboInputOptions,
+  zFloatInputOptions,
+  zIntInputOptions,
+  zStringInputOptions
+} from '@/schemas/nodeDefSchema'
+
+const zIntInputSpec = zIntInputOptions.extend({
+  type: z.literal('INT'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zFloatInputSpec = zFloatInputOptions.extend({
+  type: z.literal('FLOAT'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zBooleanInputSpec = zBooleanInputOptions.extend({
+  type: z.literal('BOOLEAN'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zStringInputSpec = zStringInputOptions.extend({
+  type: z.literal('STRING'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zComboInputSpec = zComboInputOptions.extend({
+  type: z.literal('COMBO'),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zCustomInputSpec = zBaseInputOptions.extend({
+  type: z.string(),
+  name: z.string(),
+  isOptional: z.boolean().optional()
+})
+
+const zInputSpec = z.union([
+  zIntInputSpec,
+  zFloatInputSpec,
+  zBooleanInputSpec,
+  zStringInputSpec,
+  zComboInputSpec,
+  zCustomInputSpec
+])
+
+// Output specs
+const zOutputSpec = z.object({
+  index: z.number(),
+  name: z.string(),
+  type: z.string(),
+  is_list: z.boolean(),
+  options: z.array(z.any()).optional(),
+  tooltip: z.string().optional()
+})
+
+// Main node definition schema
+const zNodeDef = z.object({
+  inputs: z.record(zInputSpec),
+  outputs: z.array(zOutputSpec),
+  hidden: z.record(z.any()).optional(),
+
+  name: z.string(),
+  display_name: z.string(),
+  description: z.string(),
+  category: z.string(),
+  output_node: z.boolean(),
+  python_module: z.string(),
+  deprecated: z.boolean().optional(),
+  experimental: z.boolean().optional()
+})
+
+// Export types
+export type IntInputSpec = z.infer<typeof zIntInputSpec>
+export type FloatInputSpec = z.infer<typeof zFloatInputSpec>
+export type BooleanInputSpec = z.infer<typeof zBooleanInputSpec>
+export type StringInputSpec = z.infer<typeof zStringInputSpec>
+export type ComboInputSpec = z.infer<typeof zComboInputSpec>
+export type CustomInputSpec = z.infer<typeof zCustomInputSpec>
+
+export type InputSpec = z.infer<typeof zInputSpec>
+export type OutputSpec = z.infer<typeof zOutputSpec>
+export type ComfyNodeDef = z.infer<typeof zNodeDef>

src/schemas/nodeDefSchema.ts

@@ -13,7 +13,7 @@ const zRemoteWidgetConfig = z.object({
   max_retries: z.number().gte(0).optional()
 })
 
-const zBaseInputOptions = z
+export const zBaseInputOptions = z
   .object({
     default: z.any().optional(),
     /** @deprecated Group node uses this field. Remove when group node feature is removed. */
@@ -28,7 +28,7 @@ const zBaseInputOptions = z
   })
   .passthrough()
 
-const zNumericInputOptions = zBaseInputOptions.extend({
+export const zNumericInputOptions = zBaseInputOptions.extend({
   min: z.number().optional(),
   max: z.number().optional(),
   step: z.number().optional(),
@@ -37,7 +37,7 @@ const zNumericInputOptions = zBaseInputOptions.extend({
   display: z.enum(['slider', 'number', 'knob']).optional()
 })
 
-const zIntInputOptions = zNumericInputOptions.extend({
+export const zIntInputOptions = zNumericInputOptions.extend({
   /**
    * If true, a linked widget will be added to the node to select the mode
    * of `control_after_generate`.
@@ -45,17 +45,17 @@ const zIntInputOptions = zNumericInputOptions.extend({
   control_after_generate: z.boolean().optional()
 })
 
-const zFloatInputOptions = zNumericInputOptions.extend({
+export const zFloatInputOptions = zNumericInputOptions.extend({
   round: z.union([z.number(), z.literal(false)]).optional()
 })
 
-const zBooleanInputOptions = zBaseInputOptions.extend({
+export const zBooleanInputOptions = zBaseInputOptions.extend({
   label_on: z.string().optional(),
   label_off: z.string().optional(),
   default: z.boolean().optional()
 })
 
-const zStringInputOptions = zBaseInputOptions.extend({
+export const zStringInputOptions = zBaseInputOptions.extend({
   default: z.string().optional(),
   multiline: z.boolean().optional(),
   dynamicPrompts: z.boolean().optional(),
@@ -65,7 +65,7 @@ const zStringInputOptions = zBaseInputOptions.extend({
   placeholder: z.string().optional()
 })
 
-const zComboInputOptions = zBaseInputOptions.extend({
+export const zComboInputOptions = zBaseInputOptions.extend({
   control_after_generate: z.boolean().optional(),
   image_upload: z.boolean().optional(),
   image_folder: z.enum(['input', 'output', 'temp']).optional(),

src/services/nodeSearchService.ts

@@ -210,7 +210,7 @@ export class NodeSearchService {
       /* name */ 'Input Type',
       /* invokeSequence */ 'i',
       /* longInvokeSequence */ 'input',
-      (node) => node.inputs.all.map((input) => input.type),
+      (node) => Object.values(node.inputs).map((input) => input.type),
       data,
       filterSearchOptions
     )
@@ -220,7 +220,7 @@ export class NodeSearchService {
       /* name */ 'Output Type',
       /* invokeSequence */ 'o',
       /* longInvokeSequence */ 'output',
-      (node) => node.outputs.all.map((output) => output.type),
+      (node) => node.outputs.map((output) => output.type),
       data,
       filterSearchOptions
     )

src/stores/nodeDefStore.ts

@@ -4,11 +4,16 @@ import { defineStore } from 'pinia'
 import type { TreeNode } from 'primevue/treenode'
 import { computed, ref } from 'vue'
 
-import {
-  type ComfyInputsSpec as ComfyInputsSpecSchema,
-  type ComfyNodeDef,
-  type ComfyOutputTypesSpec as ComfyOutputTypesSpecSchema,
-  type InputSpec
+import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration'
+import type {
+  ComfyNodeDef as ComfyNodeDefV2,
+  InputSpec as InputSpecV2,
+  OutputSpec as OutputSpecV2
+} from '@/schemas/nodeDef/nodeDefSchemaV2'
+import type {
+  ComfyInputsSpec as ComfyInputSpecV1,
+  ComfyNodeDef as ComfyNodeDefV1,
+  ComfyOutputTypesSpec as ComfyOutputSpecV1
 } from '@/schemas/nodeDefSchema'
 import {
   NodeSearchService,
@@ -21,137 +26,8 @@ import {
 } from '@/types/nodeSource'
 import { buildTree } from '@/utils/treeUtil'
 
-export interface BaseInputSpec<T = any> {
-  name: string
-  type: string
-  tooltip?: string
-  default?: T
-
-  forceInput?: boolean
-}
-
-export interface NumericInputSpec extends BaseInputSpec<number> {
-  min?: number
-  max?: number
-  step?: number
-}
-
-export interface IntInputSpec extends NumericInputSpec {
-  type: 'INT'
-}
-
-export interface FloatInputSpec extends NumericInputSpec {
-  type: 'FLOAT'
-  round?: number
-}
-
-export interface BooleanInputSpec extends BaseInputSpec<boolean> {
-  type: 'BOOLEAN'
-  labelOn?: string
-  labelOff?: string
-}
-
-export interface StringInputSpec extends BaseInputSpec<string> {
-  type: 'STRING'
-  multiline?: boolean
-  dynamicPrompts?: boolean
-}
-
-export interface ComboInputSpec extends BaseInputSpec<any> {
-  type: 'COMBO'
-  comboOptions: any[]
-  controlAfterGenerate?: boolean
-  imageUpload?: boolean
-}
-
-export class ComfyInputsSpec {
-  required: Record<string, BaseInputSpec>
-  optional: Record<string, BaseInputSpec>
-  hidden?: Record<string, any>
-
-  constructor(obj: ComfyInputsSpecSchema) {
-    this.required = ComfyInputsSpec.transformInputSpecRecord(obj.required ?? {})
-    this.optional = ComfyInputsSpec.transformInputSpecRecord(obj.optional ?? {})
-    this.hidden = obj.hidden
-  }
-
-  private static transformInputSpecRecord(
-    record: Record<string, InputSpec>
-  ): Record<string, BaseInputSpec> {
-    const result: Record<string, BaseInputSpec> = {}
-    for (const [key, value] of Object.entries(record)) {
-      result[key] = ComfyInputsSpec.transformSingleInputSpec(key, value)
-    }
-    return result
-  }
-
-  private static isInputSpec(obj: any): boolean {
-    return (
-      Array.isArray(obj) &&
-      obj.length >= 1 &&
-      (typeof obj[0] === 'string' || Array.isArray(obj[0]))
-    )
-  }
-
-  private static transformSingleInputSpec(
-    name: string,
-    value: any
-  ): BaseInputSpec {
-    if (!ComfyInputsSpec.isInputSpec(value)) return value
-
-    const [typeRaw, _spec] = value
-    const spec = _spec ?? {}
-    const type = Array.isArray(typeRaw) ? 'COMBO' : value[0]
-
-    switch (type) {
-      case 'COMBO':
-        return {
-          name,
-          type,
-          ...spec,
-          comboOptions: typeRaw,
-          default: spec.default ?? typeRaw[0]
-        } as ComboInputSpec
-      case 'INT':
-      case 'FLOAT':
-      case 'BOOLEAN':
-      case 'STRING':
-      default:
-        return { name, type, ...spec } as BaseInputSpec
-    }
-  }
-
-  get all() {
-    return [...Object.values(this.required), ...Object.values(this.optional)]
-  }
-
-  getInput(name: string): BaseInputSpec | undefined {
-    return this.required[name] ?? this.optional[name]
-  }
-}
-
-export class ComfyOutputSpec {
-  constructor(
-    public index: number,
-    // Name is not unique for output params
-    public name: string,
-    public type: string,
-    public is_list: boolean,
-    public comboOptions?: any[],
-    public tooltip?: string
-  ) {}
-}
-
-export class ComfyOutputsSpec {
-  constructor(public outputs: ComfyOutputSpec[]) {}
-
-  get all() {
-    return this.outputs
-  }
-}
-
-export class ComfyNodeDefImpl implements ComfyNodeDef {
-  // ComfyNodeDef fields
+export class ComfyNodeDefImpl implements ComfyNodeDefV1, ComfyNodeDefV2 {
+  // ComfyNodeDef fields (V1)
   readonly name: string
   readonly display_name: string
   /**
@@ -167,11 +43,11 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
   /**
    * @deprecated Use `inputs` instead
    */
-  readonly input: ComfyInputsSpecSchema
+  readonly input: ComfyInputSpecV1
   /**
    * @deprecated Use `outputs` instead
    */
-  readonly output: ComfyOutputTypesSpecSchema
+  readonly output: ComfyOutputSpecV1
   /**
    * @deprecated Use `outputs[n].is_list` instead
    */
@@ -185,12 +61,16 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
    */
   readonly output_tooltips?: string[]
 
+  // V2 fields
+  readonly inputs: Record<string, InputSpecV2>
+  readonly outputs: OutputSpecV2[]
+  readonly hidden?: Record<string, any>
+
   // ComfyNodeDefImpl fields
-  readonly inputs: ComfyInputsSpec
-  readonly outputs: ComfyOutputsSpec
   readonly nodeSource: NodeSource
 
-  constructor(obj: ComfyNodeDef) {
+  constructor(obj: ComfyNodeDefV1) {
+    // Initialize V1 fields
     this.name = obj.name
     this.display_name = obj.display_name
     this.category = obj.category
@@ -206,28 +86,16 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
     this.output_name = obj.output_name
     this.output_tooltips = obj.output_tooltips
 
-    this.inputs = new ComfyInputsSpec(obj.input ?? {})
-    this.outputs = ComfyNodeDefImpl.transformOutputSpec(obj)
+    // Initialize V2 fields
+    const defV2 = transformNodeDefV1ToV2(obj)
+    this.inputs = defV2.inputs
+    this.outputs = defV2.outputs
+    this.hidden = defV2.hidden
+
+    // Initialize node source
     this.nodeSource = getNodeSource(obj.python_module)
   }
 
-  private static transformOutputSpec(obj: any): ComfyOutputsSpec {
-    const { output, output_is_list, output_name, output_tooltips } = obj
-    const result = (output ?? []).map((type: string | any[], index: number) => {
-      const typeString = Array.isArray(type) ? 'COMBO' : type
-
-      return new ComfyOutputSpec(
-        index,
-        output_name?.[index],
-        typeString,
-        output_is_list?.[index],
-        Array.isArray(type) ? type : undefined,
-        output_tooltips?.[index]
-      )
-    })
-    return new ComfyOutputsSpec(result)
-  }
-
   get nodePath(): string {
     return (this.category ? this.category + '/' : '') + this.name
   }
@@ -253,7 +121,7 @@ export class ComfyNodeDefImpl implements ComfyNodeDef {
   }
 }
 
-export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
+export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDefV1> = {
   PrimitiveNode: {
     name: 'PrimitiveNode',
     display_name: 'Primitive',
@@ -323,7 +191,7 @@ export function createDummyFolderNodeDef(folderPath: string): ComfyNodeDefImpl {
     output_name: [],
     output_is_list: [],
     output_node: false
-  } as ComfyNodeDef)
+  } as ComfyNodeDefV1)
 }
 
 export const useNodeDefStore = defineStore('nodeDef', () => {
@@ -336,10 +204,10 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   const nodeDataTypes = computed(() => {
     const types = new Set<string>()
     for (const nodeDef of nodeDefs.value) {
-      for (const input of nodeDef.inputs.all) {
+      for (const input of Object.values(nodeDef.inputs)) {
         types.add(input.type)
       }
-      for (const output of nodeDef.outputs.all) {
+      for (const output of nodeDef.outputs) {
         types.add(output.type)
       }
     }
@@ -357,7 +225,7 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   )
   const nodeTree = computed(() => buildNodeDefTree(visibleNodeDefs.value))
 
-  function updateNodeDefs(nodeDefs: ComfyNodeDef[]) {
+  function updateNodeDefs(nodeDefs: ComfyNodeDefV1[]) {
     const newNodeDefsByName: Record<string, ComfyNodeDefImpl> = {}
     const newNodeDefsByDisplayName: Record<string, ComfyNodeDefImpl> = {}
     for (const nodeDef of nodeDefs) {
@@ -374,7 +242,7 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
     nodeDefsByName.value = newNodeDefsByName
     nodeDefsByDisplayName.value = newNodeDefsByDisplayName
   }
-  function addNodeDef(nodeDef: ComfyNodeDef) {
+  function addNodeDef(nodeDef: ComfyNodeDefV1) {
     const nodeDefImpl = new ComfyNodeDefImpl(nodeDef)
     nodeDefsByName.value[nodeDef.name] = nodeDefImpl
     nodeDefsByDisplayName.value[nodeDef.display_name] = nodeDefImpl

src/stores/widgetStore.ts

@@ -1,6 +1,7 @@
 import { defineStore } from 'pinia'
 import { computed, ref } from 'vue'
 
+import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2'
 import {
   type ComboInputSpecV2,
   type InputSpec,
@@ -8,8 +9,6 @@ import {
 } from '@/schemas/nodeDefSchema'
 import { ComfyWidgetConstructor, ComfyWidgets } from '@/scripts/widgets'
 
-import type { BaseInputSpec } from './nodeDefStore'
-
 export const useWidgetStore = defineStore('widget', () => {
   const coreWidgets = ComfyWidgets
   const customWidgets = ref<Record<string, ComfyWidgetConstructor>>({})
@@ -33,7 +32,7 @@ export const useWidgetStore = defineStore('widget', () => {
     }
   }
 
-  function inputIsWidget(spec: BaseInputSpec) {
+  function inputIsWidget(spec: InputSpecV2) {
     return getWidgetType(spec.type, spec.name) !== null
   }