From ee6788a35e8e95d0c0c2c417c31c6439a5667a82 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sun, 30 Jun 2024 10:14:16 -0400 Subject: [PATCH] Add object_info schema (#67) --- package.json | 2 +- src/extensions/core/uploadAudio.ts | 3 +- src/extensions/core/uploadImage.ts | 3 +- src/scripts/api.ts | 4 +- src/scripts/app.ts | 5 +- src/scripts/widgets.ts | 41 ++++------- src/types/apiTypes.ts | 107 ++++++++++++++++++++++++++++- 7 files changed, 128 insertions(+), 37 deletions(-) diff --git a/package.json b/package.json index 96b5ddd2ce..8c7b685c14 100644 --- a/package.json +++ b/package.json @@ -5,7 +5,7 @@ "type": "module", "scripts": { "dev": "vite", - "build": "npm run typecheck && vite build && npm run zipdist", + "build": "npm run typecheck && vite build", "zipdist": "node scripts/zipdist.js", "typecheck": "tsc --noEmit", "test": "npm run build && jest", diff --git a/src/extensions/core/uploadAudio.ts b/src/extensions/core/uploadAudio.ts index 6cb7d6421c..cec7d74cbd 100644 --- a/src/extensions/core/uploadAudio.ts +++ b/src/extensions/core/uploadAudio.ts @@ -2,6 +2,7 @@ import { app } from "../../scripts/app"; import { api } from "../../scripts/api"; import type { IWidget } from "/types/litegraph"; import type { DOMWidget } from "/scripts/domWidget"; +import { ComfyNodeDef } from "/types/apiTypes"; type FolderType = "input" | "output" | "temp"; @@ -120,7 +121,7 @@ app.registerExtension({ app.registerExtension({ name: "Comfy.UploadAudio", - async beforeRegisterNodeDef(nodeType, nodeData) { + async beforeRegisterNodeDef(nodeType, nodeData: ComfyNodeDef) { if (nodeData?.input?.required?.audio?.[1]?.audio_upload === true) { nodeData.input.required.upload = ["AUDIOUPLOAD"]; } diff --git a/src/extensions/core/uploadImage.ts b/src/extensions/core/uploadImage.ts index 7e6186c64a..7e28026abc 100644 --- a/src/extensions/core/uploadImage.ts +++ b/src/extensions/core/uploadImage.ts @@ -1,10 +1,11 @@ import { app } from "../../scripts/app"; +import { ComfyNodeDef } from "/types/apiTypes"; // Adds an upload button to the nodes app.registerExtension({ name: "Comfy.UploadImage", - async beforeRegisterNodeDef(nodeType, nodeData, app) { + async beforeRegisterNodeDef(nodeType, nodeData: ComfyNodeDef, app) { if (nodeData?.input?.required?.image?.[1]?.image_upload === true) { nodeData.input.required.upload = ["IMAGEUPLOAD"]; } diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 6a0aa32913..581268e24a 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -1,4 +1,4 @@ -import { HistoryTaskItem, PendingTaskItem, RunningTaskItem } from "/types/apiTypes"; +import { HistoryTaskItem, PendingTaskItem, RunningTaskItem, ComfyNodeDef } from "/types/apiTypes"; interface QueuePromptRequestBody { @@ -215,7 +215,7 @@ class ComfyApi extends EventTarget { * Loads node object definitions for the graph * @returns The node definitions */ - async getNodeDefs() { + async getNodeDefs(): Promise> { const resp = await this.fetchApi("/object_info", { cache: "no-store" }); return await resp.json(); } diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 62b9ab733c..a2cad8da25 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -11,6 +11,7 @@ import { applyTextReplacements, addStylesheet } from "./utils"; import type { ComfyExtension } from "/types/comfy"; import type { LGraph, LGraphCanvas, LGraphNode } from "/types/litegraph"; import { type ComfyWorkflow, parseComfyWorkflow } from "../types/comfyWorkflow"; +import { ComfyNodeDef } from "/types/apiTypes"; export const ANIM_PREVIEW_WIDGET = "$$comfy_animation_preview" @@ -1728,7 +1729,7 @@ export class ComfyApp { } } - async registerNodeDef(nodeId, nodeData) { + async registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) { const self = this; const node = Object.assign( function ComfyNode() { @@ -1805,7 +1806,7 @@ export class ComfyApp { node.category = nodeData.category; } - async registerNodesFromDefs(defs) { + async registerNodesFromDefs(defs: Record) { await this.#invokeExtensionsAsync("addCustomNodeDefs", defs); // Generate list of known widgets diff --git a/src/scripts/widgets.ts b/src/scripts/widgets.ts index 1ee75180cc..cb8210dfa5 100644 --- a/src/scripts/widgets.ts +++ b/src/scripts/widgets.ts @@ -2,27 +2,10 @@ import { api } from "./api" import "./domWidget"; import type { ComfyApp } from "./app"; import type { IWidget, LGraphNode } from "/types/litegraph"; - -interface InputDataOptions { - display?: string; - default?: any; - label_on?: boolean; - label_off?: boolean; - multiline?: boolean; - // TODO: infer type - dynamicPrompts?: any; - // TODO: infer type - control_after_generate?: any; - // Name of widget. - widget?: string -} - -export type InputData = [ - string, InputDataOptions -]; +import { ComfyNodeDef } from "/types/apiTypes"; export type ComfyWidgetConstructor = ( - node: LGraphNode, inputName: string, inputData: InputData, app?: ComfyApp, widgetName?: string) => + node: LGraphNode, inputName: string, inputData: ComfyNodeDef, app?: ComfyApp, widgetName?: string) => {widget: IWidget, minWidth?: number; minHeight?: number }; @@ -39,7 +22,7 @@ export function updateControlWidgetLabel(widget) { const IS_CONTROL_WIDGET = Symbol(); const HAS_EXECUTED = Symbol(); -function getNumberDefaults(inputData, defaultStep, precision, enable_rounding) { +function getNumberDefaults(inputData: ComfyNodeDef, defaultStep, precision, enable_rounding) { let defaultVal = inputData[1]["default"]; let { min, max, step, round} = inputData[1]; @@ -61,7 +44,7 @@ function getNumberDefaults(inputData, defaultStep, precision, enable_rounding) { return { val: defaultVal, config: { min, max, step: 10.0 * step, round, precision } }; } -export function addValueControlWidget(node, targetWidget, defaultValue = "randomize", values, widgetName, inputData) { +export function addValueControlWidget(node, targetWidget, defaultValue = "randomize", values, widgetName, inputData: ComfyNodeDef) { let name = inputData[1]?.control_after_generate; if(typeof name !== "string") { name = widgetName; @@ -73,7 +56,7 @@ export function addValueControlWidget(node, targetWidget, defaultValue = "random return widgets[0]; } -export function addValueControlWidgets(node, targetWidget, defaultValue = "randomize", options, inputData) { +export function addValueControlWidgets(node, targetWidget, defaultValue = "randomize", options, inputData: ComfyNodeDef) { if (!defaultValue) defaultValue = "randomize"; if (!options) options = {}; @@ -232,7 +215,7 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando return widgets; }; -function seedWidget(node, inputName, inputData, app, widgetName) { +function seedWidget(node, inputName, inputData: ComfyNodeDef, app, widgetName) { const seed = createIntWidget(node, inputName, inputData, app, true); const seedControl = addValueControlWidget(node, seed.widget, "randomize", undefined, widgetName, inputData); @@ -240,7 +223,7 @@ function seedWidget(node, inputName, inputData, app, widgetName) { return seed; } -function createIntWidget(node, inputName, inputData, app, isSeedInput: boolean = false) { +function createIntWidget(node, inputName, inputData: ComfyNodeDef, app, isSeedInput: boolean = false) { const control = inputData[1]?.control_after_generate; if (!isSeedInput && control) { return seedWidget(node, inputName, inputData, app, typeof control === "string" ? control : undefined); @@ -329,7 +312,7 @@ export function initWidgets(app) { export const ComfyWidgets: Record = { "INT:seed": seedWidget, "INT:noise_seed": seedWidget, - FLOAT(node, inputName, inputData, app) { + FLOAT(node, inputName, inputData: ComfyNodeDef, app) { let widgetType: "number" | "slider" = isSlider(inputData[1]["display"], app); let precision = app.ui.settings.getSettingValue("Comfy.FloatRoundingPrecision"); let disable_rounding = app.ui.settings.getSettingValue("Comfy.DisableFloatRounding") @@ -346,7 +329,7 @@ export const ComfyWidgets: Record = { } }, config) }; }, - INT(node, inputName, inputData, app) { + INT(node, inputName, inputData: ComfyNodeDef, app) { return createIntWidget(node, inputName, inputData, app); }, BOOLEAN(node, inputName, inputData) { @@ -370,7 +353,7 @@ export const ComfyWidgets: Record = { ) }; }, - STRING(node, inputName, inputData, app) { + STRING(node, inputName, inputData: ComfyNodeDef, app) { const defaultVal = inputData[1].default || ""; const multiline = !!inputData[1].multiline; @@ -386,7 +369,7 @@ export const ComfyWidgets: Record = { return res; }, - COMBO(node, inputName, inputData) { + COMBO(node, inputName, inputData: ComfyNodeDef) { const type = inputData[0]; let defaultValue = type[0]; if (inputData[1] && inputData[1].default) { @@ -400,7 +383,7 @@ export const ComfyWidgets: Record = { } return res; }, - IMAGEUPLOAD(node: LGraphNode, inputName: string, inputData, app) { + IMAGEUPLOAD(node: LGraphNode, inputName: string, inputData: ComfyNodeDef, app) { // TODO make image upload handle a custom node type? // @ts-ignore const imageWidget = node.widgets.find((w) => w.name === (inputData[1]?.widget ?? "image")); diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 05ad5a5d0b..b2ac38e148 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -1,4 +1,4 @@ -import { z } from "zod"; +import { ZodType, z } from "zod"; import { zComfyWorkflow } from "./comfyWorkflow"; const zNodeId = z.number(); @@ -109,9 +109,114 @@ const zHistoryTaskItem = z.object({ const zTaskItem = z.union([zRunningTaskItem, zPendingTaskItem, zHistoryTaskItem]); +// `/queue` export type RunningTaskItem = z.infer; export type PendingTaskItem = z.infer; +// `/history` export type HistoryTaskItem = z.infer; export type TaskItem = z.infer; // TODO: validate `/history` `/queue` API endpoint responses. + +function inputSpec(spec: [ZodType, ZodType]): ZodType { + const [inputType, inputSpec] = spec; + return z.union([ + z.tuple([inputType, inputSpec]), + z.tuple([inputType]), + ]); +} + +const zIntInputSpec = inputSpec([ + z.literal("INT"), + z.object({ + min: z.number().optional(), + max: z.number().optional(), + step: z.number().optional(), + default: z.number().optional(), + forceInput: z.boolean().optional(), + }), +]); + +const zFloatInputSpec = inputSpec([ + z.literal("FLOAT"), + z.object({ + min: z.number().optional(), + max: z.number().optional(), + step: z.number().optional(), + round: z.number().optional(), + default: z.number().optional(), + forceInput: z.boolean().optional(), + }), +]); + +const zBooleanInputSpec = inputSpec([ + z.literal("BOOLEAN"), + z.object({ + label_on: z.string().optional(), + label_off: z.string().optional(), + default: z.boolean().optional(), + forceInput: z.boolean().optional(), + }) +]); + +const zStringInputSpec = inputSpec([ + z.literal("STRING"), + z.object({ + default: z.string().optional(), + multiline: z.boolean().optional(), + dynamicPrompts: z.boolean().optional(), + forceInput: z.boolean().optional(), + }), +]); + +// 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 zCustomInputSpec = inputSpec([ + z.string(), + z.object({ + default: z.any().optional(), + forceInput: z.boolean().optional(), + }), +]); + +const zInputSpec = z.union([ + zIntInputSpec, + zFloatInputSpec, + zBooleanInputSpec, + zStringInputSpec, + zComboInputSpec, + zCustomInputSpec, +]); + +const zComfyNodeDataType = z.string(); +const zComfyComboOutput = z.array(z.any()); +const zComfyOutputSpec = z.array(z.union([zComfyNodeDataType, zComfyComboOutput])); + +const zComfyNodeDef = z.object({ + input: z.object({ + required: z.record(zInputSpec).optional(), + optional: z.record(zInputSpec).optional(), + }), + output: zComfyOutputSpec, + output_is_list: z.array(z.boolean()), + output_name: z.array(z.string()), + name: z.string(), + display_name: z.string(), + description: z.string(), + category: z.string(), + output_node: z.boolean(), +}); + +// `/object_info` +export type ComfyNodeDef = z.infer; + +// TODO: validate `/object_info` API endpoint responses.