feat: add ComfyUI integration for image generation
- Implemented ComfyUI API for generating images with various modes (text-to-image, image-to-image, inpaint, outpaint). - Created GenerateSheet and associated controls for user input on generation settings. - Added subtle scrollbar styles for improved UI experience. - Enhanced canvas input handling to ignore key events when focused on editable elements. - Optimized canvas resizing logic to prevent unnecessary dispatches. - Introduced error handling for generation failures and loading models. - Added functionality to upload images and masks for inpainting.
This commit is contained in:
syntaxbullet
164
app/comfy.ts
Normal file
164
app/comfy.ts
Normal file
@@ -0,0 +1,164 @@
|
||||
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;
|
||||
};
|
||||
inputImage?: string;
|
||||
maskImage?: string;
|
||||
};
|
||||
|
||||
const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188";
|
||||
|
||||
export async function handleComfyApi(request: Request) {
|
||||
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 });
|
||||
}
|
||||
|
||||
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;
|
||||
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 { prompt_id } = await queued.json() as { prompt_id: string };
|
||||
const history = await waitForHistory(prompt_id);
|
||||
const image = firstOutputImage(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) return undefined;
|
||||
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");
|
||||
}
|
||||
|
||||
function firstOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
|
||||
const outputs = (history as { outputs?: Record<string, { images?: { filename: string; subfolder?: string; type?: string }[] }> }).outputs ?? {};
|
||||
for (const output of Object.values(outputs)) {
|
||||
const image = output.images?.[0];
|
||||
if (image) return image;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
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, 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["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["4", 0], vae: ["1", 2], mask: ["9", 1], grow_mask_by: 6 } };
|
||||
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: 6 } };
|
||||
return workflow;
|
||||
}
|
||||
|
||||
workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
|
||||
return workflow;
|
||||
}
|
||||
|
||||
function json(value: unknown) {
|
||||
return new Response(JSON.stringify(value), { headers: { "content-type": "application/json" } });
|
||||
}
|
||||
Reference in New Issue
Block a user