import { randomInt, randomUUID } from "node:crypto"; import { setTimeout as sleep } from "node:timers/promises"; import type { AdapterDescriptor, GeneratedMedia, GpuStatus, ImageGenerationCapabilities, ImageGenerationProvider, ImageGenerationRequest, MediaFile, MediaProgressListener, VideoGenerationCapabilities, VideoGenerationProvider, VideoGenerationRequest, } from "@rakazo/adapter-kit"; import type { AttachmentKind } from "@rakazo/contracts"; import { ARTIFACT_MAX_BYTES_BY_KIND, attachmentKindOf, MEDIA_VIDEO_MAX_SECONDS, } from "@rakazo/contracts"; import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; import type { ComfyWorkflow, WorkflowTemplate, WorkflowValues } from "@rakazo/core/media/workflows"; import { DEFAULT_WORKFLOWS, fillWorkflow, parseWorkflowTemplate, } from "@rakazo/core/media/workflows"; import { assertAllowedOpenAiCompatibleRequestUrl, assertHttpsForKeyedOpenAiCompatibleUrl, } from "./openai-compatible-url.js"; import type { OpenWebSocket } from "./pi-openai-compatible-provider.js"; import { createOpenAiCompatibleFetch, createOpenAiCompatibleWebSocket, } from "./pi-openai-compatible-provider.js"; import { readBodyCapped } from "./web-ssrf.js"; export const COMFYUI_PROVIDER_ID = "comfyui"; /** Frame rate of generated video; frame counts follow LTX-Video's 8n+1 rule. */ export const COMFYUI_VIDEO_FPS = 24; const MAX_IMAGE_SIDE = 2048; const REQUEST_TIMEOUT_MS = 30_000; const TRANSFER_TIMEOUT_MS = 120_000; const CANCEL_TIMEOUT_MS = 5_000; const IMAGE_PROMPT_TIMEOUT_MS = 10 * 60_000; const VIDEO_PROMPT_TIMEOUT_MS = 30 * 60_000; const POLL_INTERVAL_MS = 1_000; /** Live progress stops short of a finished prompt, which only counts once its outputs are in. */ const RUNNING_PROMPT_MAX_PROGRESS = 0.99; const MAX_JSON_BYTES = 8 * 1024 * 1024; const MAX_ERROR_CHARS = 400; const MIB = 1024 * 1024; const IMAGE_STEPS = 25; const VIDEO_STEPS = 30; const DEFAULT_STRENGTH = 0.6; const DEFAULT_NEGATIVE_PROMPTS = { image: "blurry, low quality, distorted, watermark, text", video: "worst quality, inconsistent motion, blurry, jittery, distorted", }; const JSON_HEADERS = { Accept: "application/json", "Content-Type": "application/json" }; export interface ComfyUiProviderConfig { baseUrl: string; apiKey?: string; /** Checkpoint for `{{model}}`; the first installed image checkpoint when omitted. */ imageModel?: string; /** API-format workflow JSON that replaces the matching default image workflow. */ imageWorkflow?: string; /** API-format workflow JSON that replaces the matching default video workflow. */ videoWorkflow?: string; fetch?: typeof globalThis.fetch; /** * Opens the progress websocket; tests pass the emulator's. A replaced `fetch` without one gets * no live progress, so the socket never takes a different route than the requests. */ webSocket?: OpenWebSocket; /** How long to wait between checks for a finished prompt; tests shorten it. */ pollIntervalMs?: number; } export interface ComfyUiProvider extends ImageGenerationProvider, VideoGenerationProvider { describe(): AdapterDescriptor; gpuStatus(signal?: AbortSignal): Promise; } type MediaKind = "image" | "video"; type ReadBody = (response: Response, signal: AbortSignal) => Promise; interface OutputFile { filename: string; subfolder: string; type: string; mimeType: string; kind: AttachmentKind; } interface GenerationJob { kind: MediaKind; template: WorkflowTemplate; prompt: string; negativePrompt?: string; width: number; height: number; seed?: number; count: number; source?: MediaFile; strength?: number; frames?: number; signal?: AbortSignal; onProgress?: MediaProgressListener; } /** * Image and video generation on a self-hosted ComfyUI server through its HTTP API: submit an * API-format workflow, poll its history, download the outputs. While a prompt runs, step * progress comes from ComfyUI's websocket when it is reachable. Requests and the websocket go * through the same address checks as OpenAI-compatible model servers. */ export function createComfyUiProvider(config: ComfyUiProviderConfig): ComfyUiProvider { return new ComfyUiClient(config); } class ComfyUiClient implements ComfyUiProvider { private readonly baseUrl: URL; private readonly apiKey: string | undefined; private readonly imageModel: string | undefined; private readonly imageWorkflow: WorkflowTemplate | undefined; private readonly videoWorkflow: WorkflowTemplate | undefined; private readonly fetch: typeof globalThis.fetch; private readonly openSocket: ReturnType | undefined; private readonly pollIntervalMs: number; constructor(config: ComfyUiProviderConfig) { this.baseUrl = comfyUiBaseUrl(config.baseUrl); this.apiKey = config.apiKey?.trim() || undefined; assertHttpsForKeyedOpenAiCompatibleUrl(this.baseUrl, this.apiKey); this.imageModel = config.imageModel?.trim() || undefined; this.imageWorkflow = config.imageWorkflow?.trim() ? parseWorkflowTemplate(config.imageWorkflow) : undefined; this.videoWorkflow = config.videoWorkflow?.trim() ? parseWorkflowTemplate(config.videoWorkflow) : undefined; this.fetch = createOpenAiCompatibleFetch(config.fetch); this.openSocket = config.fetch && !config.webSocket ? undefined : createOpenAiCompatibleWebSocket(config.webSocket); this.pollIntervalMs = config.pollIntervalMs ?? POLL_INTERVAL_MS; } describe() { return { id: COMFYUI_PROVIDER_ID, contractVersion: "1", adapterVersion: "0.1.0", capabilities: { textToImage: true, imageToImage: true, inpaint: false, maxWidth: MAX_IMAGE_SIDE, maxHeight: MAX_IMAGE_SIDE, textToVideo: true, imageToVideo: true, maxDurationSeconds: MEDIA_VIDEO_MAX_SECONDS, }, }; } async gpuStatus(signal?: AbortSignal): Promise { return parseGpuStatus(await this.getJson("system_stats", signal)); } async generateImage(request: ImageGenerationRequest): Promise { if (request.mask) throw new Error("ComfyUI inpainting is not supported"); return this.generate({ ...request, kind: "image", template: pickTemplate( this.imageWorkflow, Boolean(request.source), DEFAULT_WORKFLOWS.textToImage, DEFAULT_WORKFLOWS.imageToImage, ), width: roundSide(request.width, 8), height: roundSide(request.height, 8), count: request.count ?? 1, }); } async generateVideo(request: VideoGenerationRequest): Promise { const seconds = Math.min(Math.max(request.durationSeconds, 1), MEDIA_VIDEO_MAX_SECONDS); return this.generate({ ...request, kind: "video", template: pickTemplate( this.videoWorkflow, Boolean(request.source), DEFAULT_WORKFLOWS.textToVideo, DEFAULT_WORKFLOWS.imageToVideo, ), width: roundSide(request.width, 32), height: roundSide(request.height, 32), frames: 8 * Math.round((seconds * COMFYUI_VIDEO_FPS) / 8) + 1, count: 1, }); } private async generate(job: GenerationJob): Promise { const { template, signal } = job; const values: WorkflowValues = { prompt: job.prompt, negative_prompt: job.negativePrompt ?? DEFAULT_NEGATIVE_PROMPTS[job.kind], width: job.width, height: job.height, steps: job.kind === "image" ? IMAGE_STEPS : VIDEO_STEPS, strength: job.strength ?? DEFAULT_STRENGTH, frames: job.frames ?? 1, fps: COMFYUI_VIDEO_FPS, }; if (template.placeholders.has("model")) values.model = await this.imageCheckpoint(signal); if (template.placeholders.has("image")) { if (!job.source) throw new Error("ComfyUI workflow needs an input image"); values.image = await this.upload(job.source, signal); } const sized = template.placeholders.has("width") && template.placeholders.has("height"); const timed = template.placeholders.has("frames") && template.placeholders.has("fps"); const firstSeed = job.seed ?? randomInt(2 ** 48 - 1); const results: GeneratedMedia[] = []; // One client id per generation, so concurrent jobs each get their own websocket messages. const clientId = randomUUID(); const progress = new JobProgress(job.count, job.onProgress); const stopSocket = new AbortController(); if (job.onProgress) { const socketSignal = signal ? AbortSignal.any([signal, stopSocket.signal]) : stopSocket.signal; this.followProgress(clientId, progress, socketSignal); } try { for (let index = 0; index < job.count; index += 1) { const graph = fillWorkflow(template, { ...values, seed: firstSeed + index }); const outputs = await this.runPrompt(graph, job.kind, clientId, signal, (promptId) => progress.start(promptId, index), ); const files = outputFiles(outputs, job.kind); if (files.length === 0) throw new Error(`ComfyUI produced no ${job.kind}`); for (const file of files) { results.push({ bytes: await this.download(file, signal), mimeType: file.mimeType, name: file.filename, ...(sized && file.kind === job.kind ? { width: job.width, height: job.height } : {}), ...(timed && file.kind === "video" ? { durationMs: Math.round(((values.frames ?? 1) / COMFYUI_VIDEO_FPS) * 1000) } : {}), }); } progress.finish(index); } return results; } finally { stopSocket.abort(); } } /** * Best effort: listens on ComfyUI's websocket for step progress. The history poll stays the * source of truth, so a socket that fails, closes or stays silent changes nothing else. */ private followProgress(clientId: string, progress: JobProgress, signal: AbortSignal): void { if (!this.openSocket) return; const url = new URL("ws", this.baseUrl); url.searchParams.set("clientId", clientId); try { const socket = this.openSocket(url.href, { headers: this.authHeaders(), signal }); socket.addEventListener("message", (event) => progress.receive(event.data)); } catch { // Without the socket, progress still advances as each prompt finishes. } } /** Submits a prompt and waits for its outputs; a cancelled or timed-out prompt is dequeued. */ private async runPrompt( graph: ComfyWorkflow, kind: MediaKind, clientId: string, signal: AbortSignal | undefined, onQueued: (promptId: string) => void, ): Promise> { const deadline = AbortSignal.timeout( kind === "image" ? IMAGE_PROMPT_TIMEOUT_MS : VIDEO_PROMPT_TIMEOUT_MS, ); const runSignal = signal ? AbortSignal.any([signal, deadline]) : deadline; const promptId = await this.submit(graph, clientId, runSignal); onQueued(promptId); try { for (;;) { const outputs = await this.finishedOutputs(promptId, runSignal); if (outputs) return outputs; await sleep(this.pollIntervalMs, undefined, { signal: runSignal }); } } catch (error) { if (!runSignal.aborted) throw error; await this.cancelPrompt(promptId); if (signal?.aborted) throw signal.reason ?? error; throw new Error("ComfyUI timed out"); } } private async submit( graph: ComfyWorkflow, clientId: string, signal: AbortSignal, ): Promise { const payload = { prompt: graph, client_id: 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; }); } private async finishedOutputs( promptId: string, signal: AbortSignal, ): Promise | undefined> { const history = await this.getJson(`history/${encodeURIComponent(promptId)}`, signal); const entry = isRecord(history) ? history[promptId] : undefined; if (!isRecord(entry)) return undefined; const status = isRecord(entry.status) ? entry.status : {}; if (status.status_str === "error") throw new Error(describeExecutionError(status.messages)); if (status.completed === false) return undefined; return isRecord(entry.outputs) ? entry.outputs : {}; } /** Best effort: drops the prompt from the queue, or interrupts it when it is ours and running. */ private async cancelPrompt(promptId: string): Promise { try { await this.postJson( "queue", { delete: [promptId] }, undefined, 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 = { 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. } } private async imageCheckpoint(signal: AbortSignal | undefined): Promise { if (this.imageModel) return this.imageModel; const info = await this.getJson("object_info/CheckpointLoaderSimple", signal); const videoCheckpoints = new Set( [DEFAULT_WORKFLOWS.textToVideo, DEFAULT_WORKFLOWS.imageToVideo, this.videoWorkflow].flatMap( (template) => (template ? checkpointNames(template.graph) : []), ), ); const checkpoint = parseCheckpointList(info).find((name) => !videoCheckpoints.has(name)); if (!checkpoint) throw new Error("ComfyUI has no image checkpoint installed"); return checkpoint; } private async upload(file: MediaFile, signal: AbortSignal | undefined): Promise { const form = new FormData(); const extension = attachmentExtensionForMimeType(file.mimeType) || ".png"; form.append( "image", new Blob([Buffer.from(file.bytes)], { type: file.mimeType }), `evoo-${randomUUID()}${extension}`, ); form.append("type", "input"); form.append("overwrite", "true"); // Serialize here: the address-checked fetch runs on the undici package, whose FormData // class differs from the global one. const encoded = new Response(form); const headers = { Accept: "application/json", "Content-Type": encoded.headers.get("content-type") ?? "", }; const body = new Uint8Array(await encoded.arrayBuffer()); return this.send("upload/image", { method: "POST", headers, body }, signal, { timeoutMs: TRANSFER_TIMEOUT_MS, read: async (response, readSignal) => { const result = await readOkJson(response, readSignal); const name = isRecord(result) && typeof result.name === "string" ? result.name : ""; if (!name) throw new Error("ComfyUI did not accept the image"); const subfolder = isRecord(result) && typeof result.subfolder === "string" ? result.subfolder : ""; return subfolder ? `${subfolder}/${name}` : name; }, }); } private download(file: OutputFile, signal: AbortSignal | undefined): Promise { const query = new URLSearchParams({ filename: file.filename, subfolder: file.subfolder, type: file.type, }); return this.send(`view?${query}`, {}, signal, { timeoutMs: TRANSFER_TIMEOUT_MS, read: async (response, readSignal) => { await assertOk(response); return readBodyCapped(response, ARTIFACT_MAX_BYTES_BY_KIND[file.kind], readSignal); }, }); } 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. */ private async send( path: string, init: { method?: string; headers?: Record; body?: string | Uint8Array; }, signal: AbortSignal | undefined, options: { timeoutMs?: number; read: ReadBody }, ): Promise { const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS); const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; const headers = { ...init.headers, ...this.authHeaders() }; try { const response = await this.fetch(new URL(path, this.baseUrl).href, { method: init.method ?? "GET", headers, body: init.body, redirect: "error", signal: requestSignal, }); return await options.read(response, requestSignal); } catch (error) { if (signal?.aborted) throw error; if (timeout.aborted) throw new Error("ComfyUI timed out"); // undici reports refused or reset connections as a bare "fetch failed". if (error instanceof TypeError && error.message === "fetch failed") { throw new Error("ComfyUI is unreachable"); } throw error; } } private authHeaders(): Record { return this.apiKey ? { Authorization: `Bearer ${this.apiKey}` } : {}; } } /** * Maps prompt progress to the job's overall fraction across its prompts: never backwards, and * short of each prompt's share until its outputs are downloaded. Only ComfyUI's `progress` * messages for the running prompt count; the history poll decides completion and errors. */ class JobProgress { private promptId: string | undefined; private index = 0; private reported = 0; constructor( private readonly count: number, private readonly listener: MediaProgressListener | undefined, ) {} start(promptId: string, index: number): void { this.promptId = promptId; this.index = index; } finish(index: number): void { this.promptId = undefined; this.report((index + 1) / this.count); } /** A websocket message; anything malformed or about another prompt is ignored. */ receive(data: unknown): void { if (typeof data !== "string" || !this.promptId) return; let message: unknown; try { message = JSON.parse(data); } catch { return; } if (!isRecord(message) || message.type !== "progress" || !isRecord(message.data)) return; const { value, max, prompt_id: promptId } = message.data; if (promptId !== this.promptId) return; const done = finiteNumber(value); const total = finiteNumber(max); if (done === undefined || total === undefined || total <= 0) return; const fraction = Math.min(Math.max(done / total, 0), RUNNING_PROMPT_MAX_PROGRESS); this.report((this.index + fraction) / this.count); } private report(fraction: number): void { if (fraction <= this.reported) return; this.reported = fraction; try { this.listener?.(fraction); } catch { // A failing progress listener must not fail the generation. } } } function comfyUiBaseUrl(raw: string): URL { const url = assertAllowedOpenAiCompatibleRequestUrl(raw.trim()); url.search = ""; url.hash = ""; url.pathname = `${url.pathname.replace(/\/+$/, "")}/`; return url; } function isRecord(value: unknown): value is Record { return Boolean(value) && typeof value === "object" && !Array.isArray(value); } function truncate(text: string): string { return text.length > MAX_ERROR_CHARS ? `${text.slice(0, MAX_ERROR_CHARS - 1)}…` : text; } function roundSide(side: number, multiple: number): number { const rounded = Math.round(side / multiple) * multiple; return Math.min(Math.max(rounded, multiple * Math.ceil(64 / multiple)), MAX_IMAGE_SIDE); } /** A custom workflow replaces the default that matches whether it takes an input image. */ function pickTemplate( custom: WorkflowTemplate | undefined, withImage: boolean, textDefault: WorkflowTemplate, imageDefault: WorkflowTemplate, ): WorkflowTemplate { if (custom && custom.placeholders.has("image") === withImage) return custom; return withImage ? imageDefault : textDefault; } function checkpointNames(graph: ComfyWorkflow): string[] { return Object.values(graph).flatMap((node) => typeof node.inputs.ckpt_name === "string" ? [node.inputs.ckpt_name] : [], ); } async function discardBody(response: Response): Promise { await response.body?.cancel().catch(() => undefined); } async function assertOk(response: Response): Promise { if (response.ok) return; await discardBody(response); throw new Error(`ComfyUI returned ${response.status}`); } async function readJson(response: Response, signal: AbortSignal): Promise { const bytes = await readBodyCapped(response, MAX_JSON_BYTES, signal); try { return JSON.parse(new TextDecoder().decode(bytes)); } catch { throw new Error("ComfyUI returned invalid JSON"); } } async function readOkJson(response: Response, signal: AbortSignal): Promise { await assertOk(response); return readJson(response, signal); } /** * ComfyUI can evict its own models, so memory its PyTorch process holds counts as available; * only memory other processes use (another model server on the same GPU) does not. */ export function parseGpuStatus(body: unknown): GpuStatus | undefined { const devices = isRecord(body) && Array.isArray(body.devices) ? body.devices : []; for (const device of devices) { if (!isRecord(device) || device.type === "cpu") continue; const total = finiteNumber(device.vram_total); const free = finiteNumber(device.vram_free); if (!total || free === undefined) continue; const reserved = finiteNumber(device.torch_vram_total) ?? 0; const reservedFree = finiteNumber(device.torch_vram_free) ?? 0; const available = Math.min(total, Math.max(0, free - reservedFree + reserved)); return { vramFreeMb: Math.floor(available / MIB), vramTotalMb: Math.floor(total / MIB) }; } return undefined; } function finiteNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) ? value : undefined; } /** Older ComfyUI lists combo values as `[[...names], {...}]`, newer as `["COMBO", {options}]`. */ export function parseCheckpointList(body: unknown): string[] { const loader = isRecord(body) ? body.CheckpointLoaderSimple : undefined; const input = isRecord(loader) ? loader.input : undefined; const required = isRecord(input) ? input.required : undefined; const spec = isRecord(required) ? required.ckpt_name : undefined; if (!Array.isArray(spec)) return []; const names = Array.isArray(spec[0]) ? spec[0] : isRecord(spec[1]) && Array.isArray(spec[1].options) ? spec[1].options : []; return names.filter((name): name is string => typeof name === "string" && name.length > 0); } /** Saved files win over previews; video runs return clips (or animations), not frames. */ function outputFiles(outputs: Record, kind: MediaKind): OutputFile[] { const files: OutputFile[] = []; for (const node of Object.values(outputs)) { if (!isRecord(node)) continue; for (const key of ["images", "gifs", "videos"]) { const entries = node[key]; if (!Array.isArray(entries)) continue; for (const entry of entries) { if (!isRecord(entry) || typeof entry.filename !== "string" || !entry.filename) continue; const mimeType = inferAttachmentMimeType(entry.filename); if (!mimeType) continue; files.push({ filename: entry.filename, subfolder: typeof entry.subfolder === "string" ? entry.subfolder : "", type: typeof entry.type === "string" ? entry.type : "output", mimeType, kind: attachmentKindOf(mimeType), }); } } } const videos = files.filter((file) => file.kind === "video"); const wanted = kind === "image" ? files.filter((file) => file.kind === "image") : videos.length > 0 ? videos : files.filter((file) => file.mimeType === "image/gif" || file.mimeType === "image/webp"); const saved = wanted.filter((file) => file.type === "output"); return saved.length > 0 ? saved : wanted; } export function describeWorkflowRejection(body: unknown): string { const problems: string[] = []; const nodeErrors = isRecord(body) && isRecord(body.node_errors) ? body.node_errors : {}; for (const node of Object.values(nodeErrors)) { if (!isRecord(node) || !Array.isArray(node.errors)) continue; for (const error of node.errors) { if (!isRecord(error)) continue; const parts = [node.class_type, error.message, error.details]; problems.push(parts.filter((part) => typeof part === "string" && part).join(": ")); } } if (problems.length === 0 && isRecord(body) && isRecord(body.error)) { if (typeof body.error.message === "string") problems.push(body.error.message); } return truncate( problems.length > 0 ? `ComfyUI rejected the workflow: ${problems.join("; ")}` : "ComfyUI rejected the workflow", ); } export function describeExecutionError(messages: unknown): string { const events = Array.isArray(messages) ? messages : []; for (const event of events) { if (!Array.isArray(event)) continue; const [name, data] = event; if (name === "execution_interrupted") return "ComfyUI interrupted the run"; if (name !== "execution_error" || !isRecord(data)) continue; const node = typeof data.node_type === "string" ? data.node_type : "a node"; const message = typeof data.exception_message === "string" ? data.exception_message.trim().split("\n")[0] : ""; return truncate(`ComfyUI failed in ${node}${message ? `: ${message}` : ""}`); } return "ComfyUI failed"; }