evOo-local/packages/adapters/src/comfyui-provider.ts

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 { createOpenAiCompatibleFetch } 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;
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;
  /** How long to wait between checks for a finished prompt; tests shorten it. */
  pollIntervalMs?: number;
}

export interface ComfyUiProvider extends ImageGenerationProvider, VideoGenerationProvider {
  describe(): AdapterDescriptor<ImageGenerationCapabilities & VideoGenerationCapabilities>;
  gpuStatus(signal?: AbortSignal): Promise<GpuStatus | undefined>;
}

type MediaKind = "image" | "video";
type ReadBody<T> = (response: Response, signal: AbortSignal) => Promise<T>;

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. Requests go through the same
 * address-checked fetch 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 pollIntervalMs: number;
  private readonly clientId = randomUUID();

  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.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<GpuStatus | undefined> {
    return parseGpuStatus(await this.getJson("system_stats", signal));
  }

  async generateImage(request: ImageGenerationRequest): Promise<GeneratedMedia[]> {
    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<GeneratedMedia[]> {
    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<GeneratedMedia[]> {
    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[] = [];
    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, signal);
      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) }
            : {}),
        });
      }
      job.onProgress?.((index + 1) / job.count);
    }
    return results;
  }

  /** Submits a prompt and waits for its outputs; a cancelled or timed-out prompt is dequeued. */
  private async runPrompt(
    graph: ComfyWorkflow,
    kind: MediaKind,
    signal: AbortSignal | undefined,
  ): Promise<Record<string, unknown>> {
    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, runSignal);
    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, signal: AbortSignal): Promise<string> {
    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;
    });
  }

  private async finishedOutputs(
    promptId: string,
    signal: AbortSignal,
  ): Promise<Record<string, unknown> | 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<void> {
    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<string> {
    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<string> {
    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<Uint8Array> {
    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<T>(
    path: string,
    payload: unknown,
    signal: AbortSignal | undefined,
    read: ReadBody<T>,
    timeoutMs?: number,
  ): Promise<T> {
    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<T>(
    path: string,
    init: {
      method?: string;
      headers?: Record<string, string>;
      body?: string | Uint8Array<ArrayBuffer>;
    },
    signal: AbortSignal | undefined,
    options: { timeoutMs?: number; read: ReadBody<T> },
  ): Promise<T> {
    const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS);
    const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout;
    const headers: Record<string, string> = { ...init.headers };
    if (this.apiKey) headers.Authorization = `Bearer ${this.apiKey}`;
    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;
    }
  }
}

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<string, unknown> {
  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<void> {
  await response.body?.cancel().catch(() => undefined);
}

async function assertOk(response: Response): Promise<void> {
  if (response.ok) return;
  await discardBody(response);
  throw new Error(`ComfyUI returned ${response.status}`);
}

async function readJson(response: Response, signal: AbortSignal): Promise<unknown> {
  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<unknown> {
  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<string, unknown>, 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";
}