Validate node def from /object_info endpoint (#159)

* Validate node def

* nit

* nit

* More tests
This commit is contained in:
Chenlei Hu
2024-07-18 12:20:47 -04:00
committed by GitHub
parent 2568746071
commit 9961be1bc7
4 changed files with 132 additions and 14 deletions

View File

@@ -1,5 +1,6 @@
import { ZodType, z } from "zod";
import { zComfyWorkflow } from "./comfyWorkflow";
import { fromZodError } from "zod-validation-error";
const zNodeId = z.number();
const zNodeType = z.string();
@@ -124,9 +125,21 @@ export type TaskItem = z.infer<typeof zTaskItem>;
// TODO: validate `/history` `/queue` API endpoint responses.
function inputSpec(spec: [ZodType, ZodType]): ZodType {
function inputSpec(
spec: [ZodType, ZodType],
allowUpcast: boolean = true
): ZodType {
const [inputType, inputSpec] = spec;
return z.union([z.tuple([inputType, inputSpec]), z.tuple([inputType])]);
// e.g. "INT" => ["INT", {}]
const upcastTypes: ZodType[] = allowUpcast
? [inputType.transform((type) => [type, {}])]
: [];
return z.union([
z.tuple([inputType, inputSpec]),
z.tuple([inputType]).transform(([type]) => [type, {}]),
...upcastTypes,
]);
}
const zIntInputSpec = inputSpec([
@@ -173,15 +186,18 @@ const zStringInputSpec = inputSpec([
]);
// Dropdown Selection.
const zComboInputSpec = inputSpec([
z.array(z.any()),
z.object({
default: z.any().optional(),
control_after_generate: z.boolean().optional(),
image_upload: z.boolean().optional(),
forceInput: z.boolean().optional(),
}),
]);
const zComboInputSpec = inputSpec(
[
z.array(z.any()),
z.object({
default: z.any().optional(),
control_after_generate: z.boolean().optional(),
image_upload: z.boolean().optional(),
forceInput: z.boolean().optional(),
}),
],
/* allowUpcast=*/ false
);
const zCustomInputSpec = inputSpec([
z.string(),
@@ -210,6 +226,9 @@ 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,
output_is_list: z.array(z.boolean()),
@@ -227,4 +246,15 @@ export type ComfyInputSpec = z.infer<typeof zInputSpec>;
export type ComfyOutputSpec = z.infer<typeof zComfyOutputSpec>;
export type ComfyNodeDef = z.infer<typeof zComfyNodeDef>;
// TODO: validate `/object_info` API endpoint responses.
export function validateComfyNodeDef(data: any): ComfyNodeDef {
const result = zComfyNodeDef.safeParse(data);
if (!result.success) {
const zodError = fromZodError(result.error);
const error = new Error(
`Invalid ComfyNodeDef: ${JSON.stringify(data)}\n${zodError.message}`
);
error.cause = zodError;
throw error;
}
return result.data;
}