image-studio/app/comfy.ts

type GenerateMode = "text-to-image" | "image-to-image" | "inpaint" | "outpaint";

type ComfyGenerateRequest = {
  mode: GenerateMode;
  model: string;
  prompt: string;
  negativePrompt?: string;
  strength?: number;
  steps?: number;
  cfg?: number;
  seed?: number;
  sampler?: string;
  scheduler?: string;
  width?: number;
  height?: number;
  outpaint?: {
    left?: number;
    top?: number;
    right?: number;
    bottom?: number;
    feathering?: number;
  };
  inpaint?: {
    growMaskBy?: number;
    maskBlur?: number;
    maskFeather?: number;
    maskExpand?: number;
    cropPadding?: number;
    maskPolarity?: "hidden" | "revealed";
    maskedContent?: "neutral" | "original" | "originalColor" | "edges";
    crop?: unknown;
    placement?: unknown;
  };
  inputImage?: string;
  maskImage?: string;
};

const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188";

export async function handleComfyApi(request: Request) {
  try {
    const url = new URL(request.url);
    if (url.pathname === "/api/comfy/models" && request.method === "GET") return json(await listGenerationOptions());
    if (url.pathname === "/api/comfy/generate" && request.method === "POST") return json(await generate(await request.json() as ComfyGenerateRequest));
    return new Response("Not found", { status: 404 });
  } catch (error) {
    return new Response(error instanceof Error ? error.message : "ComfyUI request failed", { status: 500 });
  }
}

async function listGenerationOptions() {
  const response = await fetch(`${comfyBaseUrl}/object_info`);
  if (!response.ok) throw new Error(`ComfyUI option lookup failed: ${response.status}`);
  const info = await response.json() as {
    CheckpointLoaderSimple?: { input?: { required?: { ckpt_name?: [string[]] } } };
    KSampler?: { input?: { required?: { sampler_name?: [string[]]; scheduler?: [string[]] } } };
  };
  return {
    models: info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? [],
    samplers: info.KSampler?.input?.required?.sampler_name?.[0] ?? [],
    schedulers: info.KSampler?.input?.required?.scheduler?.[0] ?? [],
  };
}

async function listCheckpointModels() {
  return (await listGenerationOptions()).models;
}

async function generate(request: ComfyGenerateRequest) {
  if (!request.prompt?.trim()) throw new Error("Prompt is required");
  if (!request.model || request.model === "auto") {
    const models = await listCheckpointModels();
    request.model = models[0] ?? "sd_xl_base_1.0.safetensors";
  }

  const clientId = crypto.randomUUID();
  const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`) : undefined;
  const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`) : undefined;
  if (request.mode === "inpaint" && (!uploaded || !mask)) throw new Error("Inpaint requires normalized input and mask images");
  const prompt = buildSdxlWorkflow({ ...request, inputImage: uploaded, maskImage: mask });

  const queued = await fetch(`${comfyBaseUrl}/prompt`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ client_id: clientId, prompt }),
  });
  if (!queued.ok) throw new Error(`ComfyUI prompt failed: ${queued.status} ${await queued.text()}`);
  const queuedBody = await queued.json() as { prompt_id?: string; node_errors?: unknown };
  const nodeError = nodeErrorsMessage(queuedBody.node_errors);
  if (nodeError) throw new Error(`ComfyUI rejected the workflow: ${nodeError}`);
  if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a prompt id");
  const prompt_id = queuedBody.prompt_id;
  const history = await waitForHistory(prompt_id);
  const historyError = historyErrorMessage(history);
  if (historyError) throw new Error(`ComfyUI generation failed: ${historyError}`);
  const image = selectGeneratedOutputImage(history);
  if (!image) throw new Error("ComfyUI did not return an image");

  const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`);
  if (!imageResponse.ok) throw new Error(`ComfyUI image fetch failed: ${imageResponse.status}`);
  const bytes = Buffer.from(await imageResponse.arrayBuffer());
  return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png" };
}

async function uploadDataUrl(dataUrl: string, filename: string) {
  const match = /^data:([^;]+);base64,(.+)$/.exec(dataUrl);
  if (!match) throw new Error("Expected a base64 data URL image");
  const mimeType = match[1] ?? "image/png";
  const base64 = match[2] ?? "";
  const form = new FormData();
  form.append("image", new File([new Uint8Array(Buffer.from(base64, "base64"))], filename, { type: mimeType }));
  form.append("overwrite", "true");
  const response = await fetch(`${comfyBaseUrl}/upload/image`, { method: "POST", body: form });
  if (!response.ok) throw new Error(`ComfyUI upload failed: ${response.status}`);
  const uploaded = await response.json() as { name: string };
  return uploaded.name;
}

async function waitForHistory(promptId: string) {
  for (let attempt = 0; attempt < 240; attempt++) {
    const response = await fetch(`${comfyBaseUrl}/history/${promptId}`);
    if (response.ok) {
      const history = await response.json() as Record<string, unknown>;
      if (history[promptId]) return history[promptId];
    }
    await Bun.sleep(500);
  }
  throw new Error(`Timed out waiting for ComfyUI prompt ${promptId}`);
}

export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
  const outputs = (history as { outputs?: Record<string, { images?: { filename: string; subfolder?: string; type?: string }[] }> }).outputs ?? {};
  const saveImageOutput = outputs["8"]?.images?.find(isGeneratedImage);
  if (saveImageOutput) return saveImageOutput;

  const prefixedOutput = Object.values(outputs).flatMap((output) => output.images ?? []).find((image) => image.filename.startsWith("image-studio-") && image.type !== "input");
  if (prefixedOutput) return prefixedOutput;

  return Object.values(outputs).flatMap((output) => output.images ?? []).find(isGeneratedImage);
}

function isGeneratedImage(image: { filename: string; subfolder?: string; type?: string }) {
  return image.type === undefined || image.type === "output";
}

export function buildSdxlWorkflow(request: ComfyGenerateRequest) {
  if (!request.width || !request.height) throw new Error("Generation width and height are required");
  const width = Math.max(64, Math.round(request.width));
  const height = Math.max(64, Math.round(request.height));
  const denoise = Math.max(0, Math.min(1, (request.strength ?? 75) / 100));
  const seed = request.seed === undefined || request.seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(request.seed);
  const steps = Math.max(1, Math.round(request.steps ?? 30));
  const cfg = Math.max(0, request.cfg ?? 7);
  const sampler = request.sampler ?? "euler";
  const scheduler = request.scheduler ?? "normal";
  const positive = request.prompt;
  const negative = request.negativePrompt ?? "";
  const workflow: Record<string, { class_type: string; inputs: Record<string, unknown> }> = {
    "1": { class_type: "CheckpointLoaderSimple", inputs: { ckpt_name: request.model } },
    "2": { class_type: "CLIPTextEncode", inputs: { text: positive, clip: ["1", 1] } },
    "3": { class_type: "CLIPTextEncode", inputs: { text: negative, clip: ["1", 1] } },
    "6": { class_type: "KSampler", inputs: { seed, steps, cfg, sampler_name: sampler, scheduler, denoise, model: ["1", 0], positive: ["2", 0], negative: ["3", 0], latent_image: ["5", 0] } },
    "7": { class_type: "VAEDecode", inputs: { samples: ["6", 0], vae: ["1", 2] } },
    "8": { class_type: "SaveImage", inputs: { filename_prefix: `image-studio-${request.mode}`, images: ["7", 0] } },
  };

  if (request.mode === "text-to-image" || !request.inputImage) {
    workflow["5"] = { class_type: "EmptyLatentImage", inputs: { width, height, batch_size: 1 } };
    return workflow;
  }

  workflow["4"] = { class_type: "LoadImage", inputs: { image: request.inputImage } };

  if (request.mode === "inpaint" && request.maskImage) {
    workflow["9"] = { class_type: "LoadImage", inputs: { image: request.maskImage } };
    workflow["11"] = { class_type: "ImageToMask", inputs: { image: ["9", 0], channel: "red" } };
    if (usesOriginalLatentContent(request)) {
      workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
      workflow["12"] = { class_type: "GrowMask", inputs: { mask: ["11", 0], expand: resolveGrowMaskBy(request), tapered_corners: true } };
      workflow["13"] = { class_type: "SetLatentNoiseMask", inputs: { samples: ["5", 0], mask: ["12", 0] } };
      workflow["6"].inputs.latent_image = ["13", 0];
    } else {
      workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["4", 0], vae: ["1", 2], mask: ["11", 0], grow_mask_by: resolveGrowMaskBy(request) } };
    }
    return workflow;
  }

  if (request.mode === "outpaint") {
    workflow["10"] = { class_type: "ImagePadForOutpaint", inputs: { image: ["4", 0], left: Math.round(request.outpaint?.left ?? 0), top: Math.round(request.outpaint?.top ?? 0), right: Math.round(request.outpaint?.right ?? 0), bottom: Math.round(request.outpaint?.bottom ?? 0), feathering: Math.round(request.outpaint?.feathering ?? 0) } };
    workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["10", 0], vae: ["1", 2], mask: ["10", 1], grow_mask_by: resolveGrowMaskBy(request) } };
    return workflow;
  }

  workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
  return workflow;
}

function resolveGrowMaskBy(request: ComfyGenerateRequest): number {
  const value = request.inpaint?.growMaskBy ?? 6;
  if (!Number.isFinite(value)) return 6;
  return Math.round(Math.max(0, Math.min(256, value)));
}

function usesOriginalLatentContent(request: ComfyGenerateRequest): boolean {
  return request.inpaint?.maskedContent === "original" || request.inpaint?.maskedContent === "originalColor" || request.inpaint?.maskedContent === "edges";
}

function nodeErrorsMessage(nodeErrors: unknown): string | undefined {
  if (!nodeErrors) return undefined;
  if (Array.isArray(nodeErrors) && nodeErrors.length === 0) return undefined;
  if (typeof nodeErrors === "object" && Object.keys(nodeErrors).length === 0) return undefined;

  if (typeof nodeErrors === "string") return nodeErrors;
  try {
    return JSON.stringify(nodeErrors);
  } catch {
    return "Unknown node validation error";
  }
}

function historyErrorMessage(history: unknown): string | undefined {
  const status = (history as { status?: { status_str?: string; completed?: boolean; messages?: unknown[] } }).status;
  if (!status) return undefined;
  if (status.status_str && status.status_str !== "success") return statusMessage(status);
  if (status.completed === false) return statusMessage(status);
  return undefined;
}

function statusMessage(status: { status_str?: string; messages?: unknown[] }) {
  const message = status.messages?.map(formatHistoryMessage).filter(Boolean).join("; ");
  return message || status.status_str || "Unknown execution error";
}

function formatHistoryMessage(message: unknown): string | undefined {
  if (!Array.isArray(message)) return undefined;
  const eventName = typeof message[0] === "string" ? message[0] : undefined;
  const payload = message[1];
  if (payload && typeof payload === "object") {
    const detail = payload as { exception_message?: string; node_type?: string; node_id?: string | number };
    if (detail.exception_message) {
      const node = detail.node_type ? ` in ${detail.node_type}${detail.node_id !== undefined ? ` ${detail.node_id}` : ""}` : "";
      return `${eventName ?? "error"}${node}: ${detail.exception_message}`;
    }
  }
  return eventName;
}

function json(value: unknown) {
  return new Response(JSON.stringify(value), { headers: { "content-type": "application/json" } });
}