diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 5ca51c9be..4246f5b2f 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -1,8 +1,7 @@ import { ZodType, z } from 'zod' -import { zComfyWorkflow } from './comfyWorkflow' +import { zComfyWorkflow, zNodeId } from './comfyWorkflow' import { fromZodError } from 'zod-validation-error' -const zNodeId = z.union([z.number(), z.string()]) const zNodeType = z.string() const zQueueIndex = z.number() const zPromptId = z.string() diff --git a/src/types/comfyWorkflow.ts b/src/types/comfyWorkflow.ts index 519a1c69d..6e3b6453f 100644 --- a/src/types/comfyWorkflow.ts +++ b/src/types/comfyWorkflow.ts @@ -1,6 +1,20 @@ import { z } from 'zod' import { fromZodError } from 'zod-validation-error' +// GroupNode is hacking node id to be a string, so we need to allow that. +// innerNode.id = `${this.node.id}:${i}` +// Remove it after GroupNode is redesigned. +export const zNodeId = z.union([z.number().int(), z.string()]) +export const zSlotIndex = z.union([ + z.number().int(), + z + .string() + .transform((val) => parseInt(val)) + .refine((val) => !isNaN(val), { + message: 'Invalid number' + }) +]) + // Definition of an AI model file used in the workflow. const zModelFile = z.object({ name: z.string(), @@ -12,10 +26,10 @@ const zModelFile = z.object({ const zComfyLink = z.tuple([ z.number(), // Link id - z.number(), // Node id of source node - z.number(), // Output slot# of source node - z.number(), // Node id of destination node - z.number(), // Input slot# of destination node + zNodeId, // Node id of source node + zSlotIndex, // Output slot# of source node + zNodeId, // Node id of destination node + zSlotIndex, // Input slot# of destination node z.string() // Data type ]) @@ -24,7 +38,7 @@ const zNodeOutput = z name: z.string(), type: z.string(), links: z.array(z.number()).nullable(), - slot_index: z.number().optional() + slot_index: zSlotIndex.optional() }) .passthrough() @@ -33,7 +47,7 @@ const zNodeInput = z name: z.string(), type: z.string(), link: z.number().nullable(), - slot_index: z.number().optional() + slot_index: zSlotIndex.optional() }) .passthrough() @@ -62,7 +76,7 @@ const zWidgetValues = z.union([z.array(z.any()), z.record(z.any())]) const zComfyNode = z .object({ - id: z.number(), + id: zNodeId, type: z.string(), pos: zVector2, size: zVector2, @@ -123,7 +137,7 @@ const zExtra = z export const zComfyWorkflow = z .object({ - last_node_id: z.number(), + last_node_id: zNodeId, last_link_id: z.number(), nodes: z.array(zComfyNode), links: z.array(zComfyLink), diff --git a/tests-ui/tests/comfyWorkflow.test.ts b/tests-ui/tests/comfyWorkflow.test.ts index 597ff1374..bcb729d33 100644 --- a/tests-ui/tests/comfyWorkflow.test.ts +++ b/tests-ui/tests/comfyWorkflow.test.ts @@ -88,4 +88,20 @@ describe('parseComfyWorkflow', () => { const validatedWorkflow = await validateComfyWorkflow(workflow) expect(validatedWorkflow.nodes[0].widgets_values).toEqual({ foo: 'bar' }) }) + + it('workflow.links', async () => { + const workflow = JSON.parse(JSON.stringify(defaultGraph)) + + workflow.links = [ + [ + 1, // Link id + '100:1', // Node id of source node + '12', // Output slot# of source node + '100:2', // Node id of destination node + 15, // Input slot# of destination node + 'INT' // Data type + ] + ] + expect(await validateComfyWorkflow(workflow)).not.toBeNull() + }) })