diff --git a/packages/adapters/src/comfyui-provider.ts b/packages/adapters/src/comfyui-provider.ts index ec35de7..e2410a1 100644 --- a/packages/adapters/src/comfyui-provider.ts +++ b/packages/adapters/src/comfyui-provider.ts @@ -75,6 +75,7 @@ export interface ComfyUiProvider extends ImageGenerationProvider, VideoGeneratio } type MediaKind = "image" | "video"; +type ReadBody = (response: Response, signal: AbortSignal) => Promise; interface OutputFile { filename: string; @@ -259,20 +260,18 @@ class ComfyUiClient implements ComfyUiProvider { } private async submit(graph: ComfyWorkflow, signal: AbortSignal): Promise { - const body = JSON.stringify({ prompt: graph, client_id: this.clientId }); - return this.send("prompt", { method: "POST", headers: JSON_HEADERS, body }, signal, { - read: async (response, readSignal) => { - if (response.status === 400) { - const rejection = await readJson(response, readSignal).catch(() => undefined); - throw new Error(describeWorkflowRejection(rejection)); - } - const result = await readOkJson(response, readSignal); - const promptId = isRecord(result) ? result.prompt_id : undefined; - if (typeof promptId !== "string" || !promptId) { - throw new Error("ComfyUI did not accept the workflow"); - } - return promptId; - }, + const payload = { prompt: graph, client_id: this.clientId }; + return this.postJson("prompt", payload, signal, async (response, readSignal) => { + if (response.status === 400) { + const rejection = await readJson(response, readSignal).catch(() => undefined); + throw new Error(describeWorkflowRejection(rejection)); + } + const result = await readOkJson(response, readSignal); + const promptId = isRecord(result) ? result.prompt_id : undefined; + if (typeof promptId !== "string" || !promptId) { + throw new Error("ComfyUI did not accept the workflow"); + } + return promptId; }); } @@ -291,29 +290,20 @@ class ComfyUiClient implements ComfyUiProvider { /** Best effort: drops the prompt from the queue, or interrupts it when it is ours and running. */ private async cancelPrompt(promptId: string): Promise { - const options = { timeoutMs: CANCEL_TIMEOUT_MS, read: discardBody }; try { - const body = JSON.stringify({ delete: [promptId] }); - await this.send("queue", { method: "POST", headers: JSON_HEADERS, body }, undefined, options); - const queue = await this.send( + await this.postJson( "queue", - { headers: { Accept: "application/json" } }, + { delete: [promptId] }, undefined, - { - timeoutMs: CANCEL_TIMEOUT_MS, - read: readOkJson, - }, + discardBody, + CANCEL_TIMEOUT_MS, ); + const queue = await this.getJson("queue", undefined, CANCEL_TIMEOUT_MS); const running = isRecord(queue) && Array.isArray(queue.queue_running) ? queue.queue_running : []; if (running.some((entry) => Array.isArray(entry) && entry[1] === promptId)) { - const interrupt = JSON.stringify({ prompt_id: promptId }); - await this.send( - "interrupt", - { method: "POST", headers: JSON_HEADERS, body: interrupt }, - undefined, - options, - ); + const interrupt = { prompt_id: promptId }; + await this.postJson("interrupt", interrupt, undefined, discardBody, CANCEL_TIMEOUT_MS); } } catch { // The prompt then finishes on its own; nothing waits for it. @@ -379,10 +369,20 @@ class ComfyUiClient implements ComfyUiProvider { }); } - private getJson(path: string, signal: AbortSignal | undefined): Promise { - return this.send(path, { headers: { Accept: "application/json" } }, signal, { - read: readOkJson, - }); + private getJson(path: string, signal: AbortSignal | undefined, timeoutMs?: number) { + const init = { headers: { Accept: "application/json" } }; + return this.send(path, init, signal, { timeoutMs, read: readOkJson }); + } + + private postJson( + path: string, + payload: unknown, + signal: AbortSignal | undefined, + read: ReadBody, + timeoutMs?: number, + ): Promise { + const init = { method: "POST", headers: JSON_HEADERS, body: JSON.stringify(payload) }; + return this.send(path, init, signal, { timeoutMs, read }); } /** One request, with the timeout covering the response body too. */ @@ -394,10 +394,7 @@ class ComfyUiClient implements ComfyUiProvider { body?: string | Uint8Array; }, signal: AbortSignal | undefined, - options: { - timeoutMs?: number; - read: (response: Response, signal: AbortSignal) => Promise; - }, + options: { timeoutMs?: number; read: ReadBody }, ): Promise { const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS); const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout;