diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 6b5e4b972..6a0aa3291 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -1,3 +1,6 @@ +import { HistoryTaskItem, PendingTaskItem, RunningTaskItem } from "/types/apiTypes"; + + interface QueuePromptRequestBody { client_id: string; // Mapping from node id to node info + input values @@ -268,7 +271,7 @@ class ComfyApi extends EventTarget { * Gets the current state of the queue * @returns The currently running and queued items */ - async getQueue() { + async getQueue(): Promise<{ Running: RunningTaskItem[], Pending: PendingTaskItem[] }>{ try { const res = await this.fetchApi("/queue"); const data = await res.json(); @@ -290,7 +293,7 @@ class ComfyApi extends EventTarget { * Gets the prompt execution history * @returns Prompt history including node outputs */ - async getHistory(max_items = 200) { + async getHistory(max_items: number = 200): Promise<{History: HistoryTaskItem[]}> { try { const res = await this.fetchApi(`/history?max_items=${max_items}`); return { History: Object.values(await res.json()) }; diff --git a/src/scripts/ui.ts b/src/scripts/ui.ts index d332fe9de..238a0bc45 100644 --- a/src/scripts/ui.ts +++ b/src/scripts/ui.ts @@ -3,6 +3,7 @@ import { ComfyDialog as _ComfyDialog } from "./ui/dialog"; import { toggleSwitch } from "./ui/toggleSwitch"; import { ComfySettingsDialog } from "./ui/settings"; import { ComfyApp, app } from "./app"; +import { TaskItem } from "/types/apiTypes"; export const ComfyDialog = _ComfyDialog; @@ -221,9 +222,9 @@ class ComfyList { textContent: section, }), $el("div.comfy-list-items", [ - ...(this.#reverse ? items[section].reverse() : items[section]).map((item) => { + ...(this.#reverse ? items[section].reverse() : items[section]).map((item: TaskItem) => { // Allow items to specify a custom remove action (e.g. for interrupt current prompt) - const removeAction = item.remove || { + const removeAction = "remove" in item ? item.remove : { name: "Delete", cb: () => api.deleteItem(this.#type, item.prompt[1]), }; @@ -232,7 +233,7 @@ class ComfyList { textContent: "Load", onclick: async () => { await app.loadGraphData(item.prompt[3].extra_pnginfo.workflow, true, false); - if (item.outputs) { + if ("outputs" in item) { app.nodeOutputs = item.outputs; } }, diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts new file mode 100644 index 000000000..05ad5a5d0 --- /dev/null +++ b/src/types/apiTypes.ts @@ -0,0 +1,117 @@ +import { z } from "zod"; +import { zComfyWorkflow } from "./comfyWorkflow"; + +const zNodeId = z.number(); +const zNodeType = z.string(); +const zQueueIndex = z.number(); +const zPromptId = z.string(); + +const zPromptItem = z.object({ + inputs: z.record(z.string(), z.any()), + class_type: zNodeType, +}); + +const zPrompt = z.array(zPromptItem); + +const zExtraPngInfo = z.object({ + workflow: zComfyWorkflow, +}).passthrough(); + +const zExtraData = z.object({ + extra_pnginfo: zExtraPngInfo, + client_id: z.string(), +}); +const zOutputsToExecute = z.array(zNodeId); + +const zExecutionStartMessage = z.tuple([ + z.literal("execution_start"), + z.object({ + prompt_id: zPromptId, + }), +]); + +const zExecutionCachedMessage = z.tuple([ + z.literal("execution_cached"), + z.object({ + prompt_id: zPromptId, + nodes: z.array(zNodeId), + }), +]); + +const zExecutionInterruptedMessage = z.tuple([ + z.literal("execution_interrupted"), + z.object({ + // InterruptProcessingException + prompt_id: zPromptId, + node_id: zNodeId, + node_type: zNodeType, + executed: z.array(zNodeId), + }), +]); + +const zExecutionErrorMessage = z.tuple([ + z.literal("execution_error"), + z.object({ + prompt_id: zPromptId, + node_id: zNodeId, + node_type: zNodeType, + executed: z.array(zNodeId), + + exception_message: z.string(), + exception_type: z.string(), + traceback: z.string(), + current_inputs: z.any(), + current_outputs: z.any(), + }), +]); + +const zStatusMessage = z.union([ + zExecutionStartMessage, + zExecutionCachedMessage, + zExecutionInterruptedMessage, + zExecutionErrorMessage, +]); + +const zStatus = z.object({ + status_str: z.enum(["success", "error"]), + completed: z.boolean(), + messages: z.array(zStatusMessage), +}); + +// TODO: this is a placeholder +const zOutput = z.any(); + +const zTaskPrompt = z.tuple([ + zQueueIndex, + zPromptId, + zPrompt, + zExtraData, + zOutputsToExecute, +]); + +const zRunningTaskItem = z.object({ + prompt: zTaskPrompt, + remove: z.object({ + name: z.literal("Cancel"), + cb: z.function(), + }), +}); + +const zPendingTaskItem = z.object({ + prompt: zTaskPrompt, +}); + +const zHistoryTaskItem = z.object({ + prompt: zTaskPrompt, + status: zStatus.optional(), + outputs: z.record(zNodeId, zOutput), +}); + +const zTaskItem = z.union([zRunningTaskItem, zPendingTaskItem, zHistoryTaskItem]); + +export type RunningTaskItem = z.infer; +export type PendingTaskItem = z.infer; +export type HistoryTaskItem = z.infer; +export type TaskItem = z.infer; + +// TODO: validate `/history` `/queue` API endpoint responses. diff --git a/src/types/comfyWorkflow.ts b/src/types/comfyWorkflow.ts index ef75f6046..b4c51f82f 100644 --- a/src/types/comfyWorkflow.ts +++ b/src/types/comfyWorkflow.ts @@ -90,7 +90,7 @@ const zExtra = z.object({ info: zInfo.optional(), }).passthrough(); -const zComfyWorkflow = z.object({ +export const zComfyWorkflow = z.object({ last_node_id: z.number(), last_link_id: z.number(), nodes: z.array(zComfyNode),