diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts index c2647b0ca..6614517e2 100644 --- a/browser_tests/ComfyPage.ts +++ b/browser_tests/ComfyPage.ts @@ -8,6 +8,11 @@ interface Position { y: number; } +interface Size { + width: number; + height: number; +} + class ComfyNodeSearchBox { public readonly input: Locator; public readonly dropdown: Locator; @@ -214,6 +219,65 @@ export class ComfyPage { await this.page.keyboard.up("Control"); await this.nextFrame(); } + + async closeMenu() { + await this.page.click("button.comfy-close-menu-btn"); + await this.nextFrame(); + } + + async resizeNode(nodePos: Position, nodeSize: Size, ratioX: number, ratioY: number, revertAfter: boolean = false) { + const bottomRight = { + x: nodePos.x + nodeSize.width, + y: nodePos.y + nodeSize.height, + } + const target = { + x: nodePos.x + nodeSize.width * ratioX, + y: nodePos.y + nodeSize.height * ratioY, + } + await this.dragAndDrop(bottomRight, target); + await this.nextFrame(); + if (revertAfter) { + await this.dragAndDrop(target, bottomRight); + await this.nextFrame(); + } + } + + async resizeKsamplerNode(percentX: number, percentY: number, revertAfter: boolean = false) { + const ksamplerPos = { + x: 864, + y: 157 + } + const ksamplerSize = { + width: 315, + height: 292, + } + this.resizeNode(ksamplerPos, ksamplerSize, percentX, percentY, revertAfter); + } + + async resizeLoadCheckpointNode(percentX: number, percentY: number, revertAfter: boolean = false) { + const loadCheckpointPos = { + x: 25, + y: 440, + } + const loadCheckpointSize = { + width: 320, + height: 120, + } + this.resizeNode(loadCheckpointPos, loadCheckpointSize, percentX, percentY, revertAfter); + } + + async resizeEmptyLatentNode(percentX: number, percentY: number, revertAfter: boolean = false) { + const emptyLatentPos = { + x: 475, + y: 580, + } + const emptyLatentSize = { + width: 303, + height: 132, + } + this.resizeNode(emptyLatentPos, emptyLatentSize, percentX, percentY, revertAfter); + } + } export const comfyPageFixture = base.extend<{ comfyPage: ComfyPage }>({ diff --git a/browser_tests/textWidgetTruncate.spec.ts b/browser_tests/textWidgetTruncate.spec.ts new file mode 100644 index 000000000..db55ca5cf --- /dev/null +++ b/browser_tests/textWidgetTruncate.spec.ts @@ -0,0 +1,28 @@ +import { expect } from "@playwright/test"; +import { comfyPageFixture as test } from "./ComfyPage"; + +test.describe("Combo text widget", () => { + test("Truncates text when resized", async ({ comfyPage }) => { + await comfyPage.resizeLoadCheckpointNode(0.2, 1); + await expect(comfyPage.canvas).toHaveScreenshot( + "load-checkpoint-resized-min-width.png" + ); + await comfyPage.closeMenu(); + await comfyPage.resizeKsamplerNode(0.2, 1); + await expect(comfyPage.canvas).toHaveScreenshot( + `ksampler-resized-min-width.png` + ); + }); + + test("Doesn't truncate when space still available", async ({ comfyPage }) => { + await comfyPage.resizeEmptyLatentNode(0.8, 0.8); + await expect(comfyPage.canvas).toHaveScreenshot( + "empty-latent-resized-80-percent.png" + ); + }); + + test("Can revert to full text", async ({ comfyPage }) => { + await comfyPage.resizeLoadCheckpointNode(0.8, 1, true); + await expect(comfyPage.canvas).toHaveScreenshot("resized-to-original.png"); + }); +}); diff --git a/browser_tests/textWidgetTruncate.spec.ts-snapshots/empty-latent-resized-80-percent-chromium-linux.png b/browser_tests/textWidgetTruncate.spec.ts-snapshots/empty-latent-resized-80-percent-chromium-linux.png new file mode 100644 index 000000000..2cef5fa75 Binary files /dev/null and b/browser_tests/textWidgetTruncate.spec.ts-snapshots/empty-latent-resized-80-percent-chromium-linux.png differ diff --git a/browser_tests/textWidgetTruncate.spec.ts-snapshots/ksampler-resized-min-width-chromium-linux.png b/browser_tests/textWidgetTruncate.spec.ts-snapshots/ksampler-resized-min-width-chromium-linux.png new file mode 100644 index 000000000..a30715e33 Binary files /dev/null and b/browser_tests/textWidgetTruncate.spec.ts-snapshots/ksampler-resized-min-width-chromium-linux.png differ diff --git a/browser_tests/textWidgetTruncate.spec.ts-snapshots/load-checkpoint-resized-min-width-chromium-linux.png b/browser_tests/textWidgetTruncate.spec.ts-snapshots/load-checkpoint-resized-min-width-chromium-linux.png new file mode 100644 index 000000000..aa49d2a3f Binary files /dev/null and b/browser_tests/textWidgetTruncate.spec.ts-snapshots/load-checkpoint-resized-min-width-chromium-linux.png differ diff --git a/browser_tests/textWidgetTruncate.spec.ts-snapshots/resized-to-original-chromium-linux.png b/browser_tests/textWidgetTruncate.spec.ts-snapshots/resized-to-original-chromium-linux.png new file mode 100644 index 000000000..7d42dc96c Binary files /dev/null and b/browser_tests/textWidgetTruncate.spec.ts-snapshots/resized-to-original-chromium-linux.png differ diff --git a/index.html b/index.html index b90b22876..30d4d0a45 100644 --- a/index.html +++ b/index.html @@ -12,6 +12,7 @@ font-family: 'Roboto Mono', 'Noto Color Emoji'; } --> + + + + + diff --git a/src/components/sidebar/SideBarIcon.vue b/src/components/sidebar/SideBarIcon.vue new file mode 100644 index 000000000..075d3a31d --- /dev/null +++ b/src/components/sidebar/SideBarIcon.vue @@ -0,0 +1,50 @@ + + + + + + + diff --git a/src/components/sidebar/SideBarSettingsToggleIcon.vue b/src/components/sidebar/SideBarSettingsToggleIcon.vue new file mode 100644 index 000000000..00f346bea --- /dev/null +++ b/src/components/sidebar/SideBarSettingsToggleIcon.vue @@ -0,0 +1,12 @@ + + + diff --git a/src/components/sidebar/SideBarThemeToggleIcon.vue b/src/components/sidebar/SideBarThemeToggleIcon.vue new file mode 100644 index 000000000..4a0c2bebe --- /dev/null +++ b/src/components/sidebar/SideBarThemeToggleIcon.vue @@ -0,0 +1,20 @@ + + + diff --git a/src/components/sidebar/SideToolBar.vue b/src/components/sidebar/SideToolBar.vue new file mode 100644 index 000000000..2754e376b --- /dev/null +++ b/src/components/sidebar/SideToolBar.vue @@ -0,0 +1,90 @@ + + + + + + + diff --git a/src/components/sidebar/items/NodeDetailSideBarItem.vue b/src/components/sidebar/items/NodeDetailSideBarItem.vue new file mode 100644 index 000000000..ad19897db --- /dev/null +++ b/src/components/sidebar/items/NodeDetailSideBarItem.vue @@ -0,0 +1,20 @@ + + + diff --git a/src/components/sidebar/items/QueueSideBarItem.vue b/src/components/sidebar/items/QueueSideBarItem.vue new file mode 100644 index 000000000..156020d19 --- /dev/null +++ b/src/components/sidebar/items/QueueSideBarItem.vue @@ -0,0 +1,176 @@ + + + + + + + diff --git a/src/main.ts b/src/main.ts index aee15636f..3d0c5ba49 100644 --- a/src/main.ts +++ b/src/main.ts @@ -1,12 +1,24 @@ import { createApp } from "vue"; import PrimeVue from "primevue/config"; import Aura from "@primevue/themes/aura"; +import { definePreset } from "@primevue/themes"; +import ConfirmationService from "primevue/confirmationservice"; +import ToastService from "primevue/toastservice"; import "primeicons/primeicons.css"; import App from "./App.vue"; import { app as comfyApp } from "@/scripts/app"; +import { createPinia } from "pinia"; + +const ComfyUIPreset = definePreset(Aura, { + semantic: { + // @ts-ignore + primary: Aura.primitive.blue, + }, +}); const app = createApp(App); +const pinia = createPinia(); comfyApp.setup().then(() => { window["app"] = comfyApp; @@ -15,7 +27,7 @@ comfyApp.setup().then(() => { app .use(PrimeVue, { theme: { - preset: Aura, + preset: ComfyUIPreset, options: { prefix: "p", cssLayer: false, @@ -25,5 +37,8 @@ comfyApp.setup().then(() => { }, }, }) + .use(ConfirmationService) + .use(ToastService) + .use(pinia) .mount("#vue-app"); }); diff --git a/src/scripts/api.ts b/src/scripts/api.ts index 3f8bc3d14..4be4b0627 100644 --- a/src/scripts/api.ts +++ b/src/scripts/api.ts @@ -184,6 +184,11 @@ class ComfyApi extends EventTarget { new CustomEvent("execution_start", { detail: msg.data }) ); break; + case "execution_success": + this.dispatchEvent( + new CustomEvent("execution_success", { detail: msg.data }) + ); + break; case "execution_error": this.dispatchEvent( new CustomEvent("execution_error", { detail: msg.data }) @@ -315,10 +320,14 @@ class ComfyApi extends EventTarget { return { // Running action uses a different endpoint for cancelling Running: data.queue_running.map((prompt) => ({ + taskType: "Running", prompt, remove: { name: "Cancel", cb: () => api.interrupt() }, })), - Pending: data.queue_pending.map((prompt) => ({ prompt })), + Pending: data.queue_pending.map((prompt) => ({ + taskType: "Pending", + prompt, + })), }; } catch (error) { console.error(error); @@ -335,7 +344,11 @@ class ComfyApi extends EventTarget { ): Promise<{ History: HistoryTaskItem[] }> { try { const res = await this.fetchApi(`/history?max_items=${max_items}`); - return { History: Object.values(await res.json()) }; + return { + History: Object.values(await res.json()).map( + (item: HistoryTaskItem) => ({ ...item, taskType: "History" }) + ), + }; } catch (error) { console.error(error); return { History: [] }; diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 91c52e017..76d243351 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -1845,13 +1845,18 @@ export class ComfyApp { await this.#setUser(); // Create and mount the LiteGraph in the DOM + const canvasContainer = document.createElement("div"); + canvasContainer.id = "graph-canvas-container"; + const mainCanvas = document.createElement("canvas"); mainCanvas.style.touchAction = "none"; const canvasEl = (this.canvasEl = Object.assign(mainCanvas, { id: "graph-canvas", })); canvasEl.tabIndex = 1; - document.body.prepend(canvasEl); + canvasContainer.prepend(canvasEl); + document.body.prepend(canvasContainer); + this.resizeCanvas(); await Promise.all([ diff --git a/src/scripts/ui/menu/index.ts b/src/scripts/ui/menu/index.ts index 7e2135f78..f593de0ad 100644 --- a/src/scripts/ui/menu/index.ts +++ b/src/scripts/ui/menu/index.ts @@ -5,10 +5,8 @@ import { downloadBlob } from "../../utils"; import { ComfyButton } from "../components/button"; import { ComfyButtonGroup } from "../components/buttonGroup"; import { ComfySplitButton } from "../components/splitButton"; -import { ComfyViewHistoryButton } from "./viewHistory"; import { ComfyQueueButton } from "./queueButton"; import { ComfyWorkflowsMenu } from "./workflows"; -import { ComfyViewQueueButton } from "./viewQueue"; import { getInteruptButton } from "./interruptButton"; import "./menu.css"; import type { ComfySettingsDialog } from "../settings"; @@ -128,19 +126,10 @@ export class ComfyAppMenu { }, }) ); - this.settingsGroup = new ComfyButtonGroup( - new ComfyButton({ - icon: "cog", - content: "Settings", - tooltip: "Open settings", - action: () => { - app.ui.settings.show(); - }, - }) - ); + // Keep the settings group as there are custom scripts attaching extra + // elements to it. + this.settingsGroup = new ComfyButtonGroup(); this.viewGroup = new ComfyButtonGroup( - new ComfyViewHistoryButton(app).element, - new ComfyViewQueueButton(app).element, getInteruptButton("nlg-hide").element ); this.mobileMenuButton = new ComfyButton({ @@ -193,6 +182,7 @@ export class ComfyAppMenu { app.ui.menuContainer.style.removeProperty("display"); this.element.style.display = "none"; app.ui.restoreMenuPosition(); + document.dispatchEvent(new Event("comfy:setting:beta-menu-disabled")); } window.dispatchEvent(new Event("resize")); }, diff --git a/src/scripts/ui/menu/viewHistory.ts b/src/scripts/ui/menu/viewHistory.ts deleted file mode 100644 index e76194f8f..000000000 --- a/src/scripts/ui/menu/viewHistory.ts +++ /dev/null @@ -1,26 +0,0 @@ -import type { ComfyApp } from "@/scripts/app"; -import { ComfyButton } from "../components/button"; -import { ComfyViewList, ComfyViewListButton } from "./viewList"; - -export class ComfyViewHistoryButton extends ComfyViewListButton { - constructor(app: ComfyApp) { - super(app, { - button: new ComfyButton({ - content: "View History", - icon: "history", - tooltip: "View history", - classList: "comfyui-button comfyui-history-button", - }), - list: ComfyViewHistoryList, - mode: "History", - }); - } -} - -export class ComfyViewHistoryList extends ComfyViewList { - async loadItems() { - const items = await super.loadItems(); - items["History"].reverse(); - return items; - } -} diff --git a/src/scripts/ui/menu/viewList.ts b/src/scripts/ui/menu/viewList.ts deleted file mode 100644 index cf2385ec0..000000000 --- a/src/scripts/ui/menu/viewList.ts +++ /dev/null @@ -1,231 +0,0 @@ -import { ComfyButton } from "../components/button"; -import { $el } from "../../ui"; -import { api } from "../../api"; -import { ComfyPopup } from "../components/popup"; -import type { ComfyApp } from "@/scripts/app"; - -type ViewListMode = "Queue" | "History"; - -export class ComfyViewListButton { - popup: ComfyPopup; - app: ComfyApp; - button: ComfyButton; - element: HTMLDivElement; - list: ComfyViewList; - - get open() { - return this.popup.open; - } - - set open(open) { - this.popup.open = open; - } - - constructor( - app: ComfyApp, - { - button, - list, - mode, - }: { button: ComfyButton; list: typeof ComfyViewList; mode: ViewListMode } - ) { - this.app = app; - this.button = button; - this.element = $el( - "div.comfyui-button-wrapper", - this.button.element - ) as HTMLDivElement; - this.popup = new ComfyPopup({ - target: this.element, - container: this.element, - horizontal: "right", - }); - this.list = new (list ?? ComfyViewList)(app, mode, this.popup); - this.popup.children = [this.list.element]; - this.popup.addEventListener("open", () => { - this.list.update(); - }); - this.popup.addEventListener("close", () => { - this.list.close(); - }); - this.button.withPopup(this.popup); - - api.addEventListener("status", () => { - if (this.popup.open) { - this.popup.update(); - } - }); - } -} - -export class ComfyViewList { - app: ComfyApp; - mode: ViewListMode; - popup: ComfyPopup; - type: string; - items: HTMLElement; - clear: ComfyButton; - refresh: ComfyButton; - element: HTMLElement; - constructor(app: ComfyApp, mode: ViewListMode, popup: ComfyPopup) { - this.app = app; - this.mode = mode; - this.popup = popup; - this.type = mode.toLowerCase(); - - this.items = $el(`div.comfyui-${this.type}-items.comfyui-view-list-items`); - this.clear = new ComfyButton({ - icon: "cancel", - content: "Clear", - action: async () => { - this.showSpinner(false); - await api.clearItems(this.type); - await this.update(); - }, - }); - - this.refresh = new ComfyButton({ - icon: "refresh", - content: "Refresh", - action: async () => { - await this.update(false); - }, - }); - - this.element = $el( - `div.comfyui-${this.type}-popup.comfyui-view-list-popup`, - [ - $el("h3", mode), - $el("header", [this.clear.element, this.refresh.element]), - this.items, - ] - ); - - api.addEventListener("status", () => { - if (this.popup.open) { - this.update(); - } - }); - } - - async close() { - this.items.replaceChildren(); - } - - async update(resize = true) { - this.showSpinner(resize); - const res = await this.loadItems(); - let any = false; - - const names = Object.keys(res); - const sections = names - .map((section) => { - const items = res[section]; - if (items?.length) { - any = true; - } else { - return; - } - - const rows = []; - if (names.length > 1) { - rows.push($el("h5", section)); - } - rows.push(...items.flatMap((item) => this.createRow(item, section))); - return $el("section", rows); - }) - .filter(Boolean); - - if (any) { - this.items.replaceChildren(...sections); - } else { - this.items.replaceChildren($el("h5", "None")); - } - - this.popup.update(); - this.clear.enabled = this.refresh.enabled = true; - this.element.style.removeProperty("height"); - } - - showSpinner(resize = true) { - // if (!this.spinner) { - // this.spinner = createSpinner(); - // } - // if (!resize) { - // this.element.style.height = this.element.clientHeight + "px"; - // } - // this.clear.enabled = this.refresh.enabled = false; - // this.items.replaceChildren( - // $el( - // "div", - // { - // style: { - // fontSize: "18px", - // }, - // }, - // this.spinner - // ) - // ); - // this.popup.update(); - } - - async loadItems() { - return await api.getItems(this.type); - } - - getRow(item, section) { - return { - text: item.prompt[0] + "", - actions: [ - { - text: "Load", - action: async () => { - try { - await this.app.loadGraphData( - item.prompt[3].extra_pnginfo.workflow - ); - if (item.outputs) { - this.app.nodeOutputs = item.outputs; - } - } catch (error) { - alert("Error loading workflow: " + error.message); - console.error(error); - } - }, - }, - { - text: "Delete", - action: async () => { - try { - await api.deleteItem(this.type, item.prompt[1]); - this.update(); - } catch (error) {} - }, - }, - ], - }; - } - - createRow = (item, section) => { - const row = this.getRow(item, section); - return [ - $el("span", row.text), - ...row.actions.map( - (a) => - new ComfyButton({ - content: a.text, - action: async (e, btn) => { - btn.enabled = false; - try { - await a.action(); - } catch (error) { - throw error; - } finally { - btn.enabled = true; - } - }, - }).element - ), - ]; - }; -} diff --git a/src/scripts/ui/menu/viewQueue.ts b/src/scripts/ui/menu/viewQueue.ts deleted file mode 100644 index bc18ca306..000000000 --- a/src/scripts/ui/menu/viewQueue.ts +++ /dev/null @@ -1,56 +0,0 @@ -import { ComfyButton } from "../components/button"; -import { ComfyViewList, ComfyViewListButton } from "./viewList"; -import { api } from "../../api"; -import type { ComfyApp } from "@/scripts/app"; - -export class ComfyViewQueueButton extends ComfyViewListButton { - constructor(app: ComfyApp) { - super(app, { - button: new ComfyButton({ - content: "View Queue", - icon: "format-list-numbered", - tooltip: "View queue", - classList: "comfyui-button comfyui-queue-button", - }), - list: ComfyViewQueueList, - mode: "Queue", - }); - } -} - -export class ComfyViewQueueList extends ComfyViewList { - getRow = (item, section) => { - if (section !== "Running") { - return super.getRow(item, section); - } - return { - text: item.prompt[0] + "", - actions: [ - { - text: "Load", - action: async () => { - try { - await this.app.loadGraphData( - item.prompt[3].extra_pnginfo.workflow - ); - if (item.outputs) { - this.app.nodeOutputs = item.outputs; - } - } catch (error) { - alert("Error loading workflow: " + error.message); - console.error(error); - } - }, - }, - { - text: "Cancel", - action: async () => { - try { - await api.interrupt(); - } catch (error) {} - }, - }, - ], - }; - }; -} diff --git a/src/stores/queueStore.ts b/src/stores/queueStore.ts new file mode 100644 index 000000000..82e6d7996 --- /dev/null +++ b/src/stores/queueStore.ts @@ -0,0 +1,206 @@ +import { api } from "@/scripts/api"; +import { app } from "@/scripts/app"; +import { + validateTaskItem, + TaskItem, + TaskType, + TaskPrompt, + TaskStatus, + TaskOutput, +} from "@/types/apiTypes"; +import { plainToClass } from "class-transformer"; +import _ from "lodash"; +import { defineStore } from "pinia"; +import { toRaw } from "vue"; + +// Task type used in the API. +export type APITaskType = "queue" | "history"; + +export enum TaskItemDisplayStatus { + Running = "Running", + Pending = "Pending", + Completed = "Completed", + Failed = "Failed", + Cancelled = "Cancelled", +} + +export class TaskItemImpl { + taskType: TaskType; + prompt: TaskPrompt; + status?: TaskStatus; + outputs?: TaskOutput; + + get apiTaskType(): APITaskType { + switch (this.taskType) { + case "Running": + case "Pending": + return "queue"; + case "History": + return "history"; + } + } + + get queueIndex() { + return this.prompt[0]; + } + + get promptId() { + return this.prompt[1]; + } + + get promptInputs() { + return this.prompt[2]; + } + + get extraData() { + return this.prompt[3]; + } + + get outputsToExecute() { + return this.prompt[4]; + } + + get extraPngInfo() { + return this.extraData.extra_pnginfo; + } + + get clientId() { + return this.extraData.client_id; + } + + get workflow() { + return this.extraPngInfo.workflow; + } + + get messages() { + return this.status?.messages || []; + } + + get interrupted() { + return _.some( + this.messages, + (message) => message[0] === "execution_interrupted" + ); + } + + get isHistory() { + return this.taskType === "History"; + } + + get isRunning() { + return this.taskType === "Running"; + } + + get displayStatus(): TaskItemDisplayStatus { + switch (this.taskType) { + case "Running": + return TaskItemDisplayStatus.Running; + case "Pending": + return TaskItemDisplayStatus.Pending; + case "History": + switch (this.status!.status_str) { + case "success": + return TaskItemDisplayStatus.Completed; + case "error": + return this.interrupted + ? TaskItemDisplayStatus.Cancelled + : TaskItemDisplayStatus.Failed; + } + } + } + + get executionStartTimestamp() { + const message = this.messages.find( + (message) => message[0] === "execution_start" + ); + return message ? message[1].timestamp : undefined; + } + + get executionEndTimestamp() { + const messages = this.messages.filter((message) => + [ + "execution_success", + "execution_interrupted", + "execution_error", + ].includes(message[0]) + ); + if (!messages.length) { + return undefined; + } + return _.max(messages.map((message) => message[1].timestamp)); + } + + get executionTime() { + if (!this.executionStartTimestamp || !this.executionEndTimestamp) { + return undefined; + } + return this.executionEndTimestamp - this.executionStartTimestamp; + } + + get executionTimeInSeconds() { + return this.executionTime !== undefined + ? this.executionTime / 1000 + : undefined; + } + + public async loadWorkflow() { + await app.loadGraphData(toRaw(this.workflow)); + if (this.outputs) { + app.nodeOutputs = toRaw(this.outputs); + } + } +} + +interface State { + runningTasks: TaskItemImpl[]; + pendingTasks: TaskItemImpl[]; + historyTasks: TaskItemImpl[]; +} + +export const useQueueStore = defineStore("queue", { + state: (): State => ({ + runningTasks: [], + pendingTasks: [], + historyTasks: [], + }), + getters: { + tasks(state) { + return [ + ...state.pendingTasks, + ...state.runningTasks, + ...state.historyTasks, + ]; + }, + }, + actions: { + // Fetch the queue data from the API + async update() { + const [queue, history] = await Promise.all([ + api.getQueue(), + api.getHistory(), + ]); + + const toClassAll = (tasks: TaskItem[]): TaskItemImpl[] => + tasks + .map((task) => validateTaskItem(task)) + .filter((result) => result.success) + .map((result) => plainToClass(TaskItemImpl, result.data)) + // Desc order to show the latest tasks first + .sort((a, b) => b.queueIndex - a.queueIndex); + + this.runningTasks = toClassAll(queue.Running); + this.pendingTasks = toClassAll(queue.Pending); + this.historyTasks = toClassAll(history.History); + }, + async clear() { + await Promise.all( + ["queue", "history"].map((type) => api.clearItems(type)) + ); + await this.update(); + }, + async delete(task: TaskItemImpl) { + await api.deleteItem(task.apiTaskType, task.promptId); + await this.update(); + }, + }, +}); diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 98c3bec37..918a02bd9 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -2,17 +2,17 @@ import { ZodType, z } from "zod"; import { zComfyWorkflow } from "./comfyWorkflow"; import { fromZodError } from "zod-validation-error"; -const zNodeId = z.number(); +const zNodeId = z.union([z.number(), z.string()]); const zNodeType = z.string(); const zQueueIndex = z.number(); const zPromptId = z.string(); -const zPromptItem = z.object({ +const zPromptInputItem = z.object({ inputs: z.record(z.string(), z.any()), class_type: zNodeType, }); -const zPrompt = z.array(zPromptItem); +const zPromptInputs = z.record(zPromptInputItem); const zExtraPngInfo = z .object({ @@ -26,26 +26,32 @@ const zExtraData = z.object({ }); const zOutputsToExecute = z.array(zNodeId); +const zMessageDetailBase = z.object({ + prompt_id: zPromptId, + timestamp: z.number(), +}); + const zExecutionStartMessage = z.tuple([ z.literal("execution_start"), - z.object({ - prompt_id: zPromptId, - }), + zMessageDetailBase, +]); + +const zExecutionSuccessMessage = z.tuple([ + z.literal("execution_success"), + zMessageDetailBase, ]); const zExecutionCachedMessage = z.tuple([ z.literal("execution_cached"), - z.object({ - prompt_id: zPromptId, + zMessageDetailBase.extend({ nodes: z.array(zNodeId), }), ]); const zExecutionInterruptedMessage = z.tuple([ z.literal("execution_interrupted"), - z.object({ + zMessageDetailBase.extend({ // InterruptProcessingException - prompt_id: zPromptId, node_id: zNodeId, node_type: zNodeType, executed: z.array(zNodeId), @@ -54,8 +60,7 @@ const zExecutionInterruptedMessage = z.tuple([ const zExecutionErrorMessage = z.tuple([ z.literal("execution_error"), - z.object({ - prompt_id: zPromptId, + zMessageDetailBase.extend({ node_id: zNodeId, node_type: zNodeType, executed: z.array(zNodeId), @@ -70,6 +75,7 @@ const zExecutionErrorMessage = z.tuple([ const zStatusMessage = z.union([ zExecutionStartMessage, + zExecutionSuccessMessage, zExecutionCachedMessage, zExecutionInterruptedMessage, zExecutionErrorMessage, @@ -87,13 +93,15 @@ const zOutput = z.any(); const zTaskPrompt = z.tuple([ zQueueIndex, zPromptId, - zPrompt, + zPromptInputs, zExtraData, zOutputsToExecute, ]); const zRunningTaskItem = z.object({ + taskType: z.literal("Running"), prompt: zTaskPrompt, + // @Deprecated remove: z.object({ name: z.literal("Cancel"), cb: z.function(), @@ -101,13 +109,17 @@ const zRunningTaskItem = z.object({ }); const zPendingTaskItem = z.object({ + taskType: z.literal("Pending"), prompt: zTaskPrompt, }); +const zTaskOutput = z.record(zNodeId, zOutput); + const zHistoryTaskItem = z.object({ + taskType: z.literal("History"), prompt: zTaskPrompt, status: zStatus.optional(), - outputs: z.record(zNodeId, zOutput), + outputs: zTaskOutput, }); const zTaskItem = z.union([ @@ -116,6 +128,17 @@ const zTaskItem = z.union([ zHistoryTaskItem, ]); +const zTaskType = z.union([ + z.literal("Running"), + z.literal("Pending"), + z.literal("History"), +]); + +export type TaskType = z.infer; +export type TaskPrompt = z.infer; +export type TaskStatus = z.infer; +export type TaskOutput = z.infer; + // `/queue` export type RunningTaskItem = z.infer; export type PendingTaskItem = z.infer; @@ -123,7 +146,17 @@ export type PendingTaskItem = z.infer; export type HistoryTaskItem = z.infer; export type TaskItem = z.infer; -// TODO: validate `/history` `/queue` API endpoint responses. +export function validateTaskItem(taskItem: unknown) { + const result = zTaskItem.safeParse(taskItem); + if (!result.success) { + const zodError = fromZodError(result.error); + // TODO accept a callback to report error. + console.warn( + `Invalid TaskItem: ${JSON.stringify(taskItem)}\n${zodError.message}` + ); + } + return result; +} function inputSpec( spec: [ZodType, ZodType],