From 680268bb295537f3a39a7e0882a884a9debb3213 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 4 Mar 2025 17:40:26 -0500
Subject: [PATCH] Generate json schema for node def (#2862)

---
 scripts/generate-json-schema.ts        | 27 +++++++++++++++++++++++++-
 src/schemas/nodeDef/nodeDefSchemaV2.ts |  4 ++--
 src/schemas/nodeDefSchema.ts           |  2 +-
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/scripts/generate-json-schema.ts b/scripts/generate-json-schema.ts
index 406ba538b..76aa41b2e 100644
--- a/scripts/generate-json-schema.ts
+++ b/scripts/generate-json-schema.ts
@@ -2,7 +2,12 @@ import fs from 'fs'
 import path from 'path'
 import { zodToJsonSchema } from 'zod-to-json-schema'
 
-import { zComfyWorkflow, zComfyWorkflow1 } from '../src/types/comfyWorkflow'
+import {
+  zComfyWorkflow,
+  zComfyWorkflow1
+} from '../src/schemas/comfyWorkflowSchema'
+import { zComfyNodeDef as zComfyNodeDefV2 } from '../src/schemas/nodeDef/nodeDefSchemaV2'
+import { zComfyNodeDef as zComfyNodeDefV1 } from '../src/schemas/nodeDefSchema'
 
 // Convert both workflow schemas to JSON Schema
 const workflow04Schema = zodToJsonSchema(zComfyWorkflow, {
@@ -15,6 +20,16 @@ const workflow1Schema = zodToJsonSchema(zComfyWorkflow1, {
   $refStrategy: 'none'
 })
 
+const nodeDefV1Schema = zodToJsonSchema(zComfyNodeDefV1, {
+  name: 'ComfyNodeDefV1',
+  $refStrategy: 'none'
+})
+
+const nodeDefV2Schema = zodToJsonSchema(zComfyNodeDefV2, {
+  name: 'ComfyNodeDefV2',
+  $refStrategy: 'none'
+})
+
 // Create output directory if it doesn't exist
 const outputDir = './schemas'
 if (!fs.existsSync(outputDir)) {
@@ -32,4 +47,14 @@ fs.writeFileSync(
   JSON.stringify(workflow1Schema, null, 2)
 )
 
+fs.writeFileSync(
+  path.join(outputDir, 'node-def-v1.json'),
+  JSON.stringify(nodeDefV1Schema, null, 2)
+)
+
+fs.writeFileSync(
+  path.join(outputDir, 'node-def-v2.json'),
+  JSON.stringify(nodeDefV2Schema, null, 2)
+)
+
 console.log('JSON Schemas generated successfully!')
diff --git a/src/schemas/nodeDef/nodeDefSchemaV2.ts b/src/schemas/nodeDef/nodeDefSchemaV2.ts
index 095ca757c..18a9f392e 100644
--- a/src/schemas/nodeDef/nodeDefSchemaV2.ts
+++ b/src/schemas/nodeDef/nodeDefSchemaV2.ts
@@ -65,7 +65,7 @@ const zOutputSpec = z.object({
 })
 
 // Main node definition schema
-const zNodeDef = z.object({
+export const zComfyNodeDef = z.object({
   inputs: z.record(zInputSpec),
   outputs: z.array(zOutputSpec),
   hidden: z.record(z.any()).optional(),
@@ -90,7 +90,7 @@ 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>
+export type ComfyNodeDef = z.infer<typeof zComfyNodeDef>
 
 export const isIntInputSpec = (
   inputSpec: InputSpec
diff --git a/src/schemas/nodeDefSchema.ts b/src/schemas/nodeDefSchema.ts
index 8a0f5087b..8d1a4467d 100644
--- a/src/schemas/nodeDefSchema.ts
+++ b/src/schemas/nodeDefSchema.ts
@@ -203,7 +203,7 @@ const zComfyOutputTypesSpec = z.array(
   z.union([zComfyNodeDataType, zComfyComboOutput])
 )
 
-const zComfyNodeDef = z.object({
+export const zComfyNodeDef = z.object({
   input: zComfyInputsSpec.optional(),
   output: zComfyOutputTypesSpec.optional(),
   output_is_list: z.array(z.boolean()).optional(),