feat: implement inpainting functionality with mask handling
- Added `createMaskedPixelReplacementSource` function to handle pixel replacement using inpainting. - Introduced `buildInpaintBundle` to prepare inpainting data including mask generation and validation. - Created utility functions for mask operations such as `applyMaskedContentModeToRgba`, `expandRectWithinBounds`, and others for mask manipulation. - Developed tests for inpainting preparation and mask raster utilities to ensure functionality and correctness. - Implemented mask raster operations including inversion, feathering, blurring, and more.
This commit is contained in:
syntaxbullet
@@ -1,4 +1,5 @@
|
||||
import { documentCommands } from "@commands/document";
|
||||
import { generationCommands } from "@commands/generation";
|
||||
import { historyCommands } from "@commands/history";
|
||||
import { commandIds } from "@commands/ids";
|
||||
import { createCommandRegistry } from "@commands/registry";
|
||||
@@ -12,7 +13,7 @@ import { createAppStore } from "@editor/store";
|
||||
export type ImageStudioApp = ReturnType<typeof createImageStudioApp>;
|
||||
|
||||
export function createImageStudioApp(options?: { documentName?: string; createDefaultArtboard?: boolean }) {
|
||||
const registry = createCommandRegistry([...viewportCommands, ...selectionCommands, ...documentCommands, ...toolCommands, ...transformCommands, ...historyCommands]);
|
||||
const registry = createCommandRegistry([...viewportCommands, ...selectionCommands, ...documentCommands, ...toolCommands, ...generationCommands, ...transformCommands, ...historyCommands]);
|
||||
const store = createAppStore(createInitialAppState(options?.documentName), registry);
|
||||
|
||||
if (options?.createDefaultArtboard !== false) {
|
||||
|
||||
59
app/comfy.test.ts
Normal file
59
app/comfy.test.ts
Normal file
@@ -0,0 +1,59 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { buildSdxlWorkflow, selectGeneratedOutputImage } from "./comfy";
|
||||
|
||||
describe("Comfy adapter", () => {
|
||||
test("selects SaveImage output instead of uploaded input or mask images", () => {
|
||||
const image = selectGeneratedOutputImage({
|
||||
outputs: {
|
||||
"4": { images: [{ filename: "image-studio-input.png", type: "input" }] },
|
||||
"9": { images: [{ filename: "image-studio-mask.png", type: "input" }] },
|
||||
"8": { images: [{ filename: "image-studio-inpaint_00001_.png", subfolder: "", type: "output" }] },
|
||||
},
|
||||
});
|
||||
|
||||
expect(image).toEqual({ filename: "image-studio-inpaint_00001_.png", subfolder: "", type: "output" });
|
||||
});
|
||||
|
||||
test("falls back to generated filename prefixes when node ids differ", () => {
|
||||
const image = selectGeneratedOutputImage({
|
||||
outputs: {
|
||||
"12": { images: [{ filename: "image-studio-inpaint_00002_.png", type: "output" }] },
|
||||
"4": { images: [{ filename: "image-studio-input.png", type: "input" }] },
|
||||
},
|
||||
});
|
||||
|
||||
expect(image?.filename).toBe("image-studio-inpaint_00002_.png");
|
||||
});
|
||||
|
||||
test("builds neutral inpaint with VAEEncodeForInpaint", () => {
|
||||
const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "neutral" }));
|
||||
|
||||
expect(workflow["5"]?.class_type).toBe("VAEEncodeForInpaint");
|
||||
expect(workflow["5"]?.inputs).toMatchObject({ grow_mask_by: 6, mask: ["11", 0] });
|
||||
expect(workflow["6"]?.inputs.latent_image).toEqual(["5", 0]);
|
||||
});
|
||||
|
||||
test("builds original-content inpaint with a latent noise mask", () => {
|
||||
const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "original", growMaskBy: 12 }));
|
||||
|
||||
expect(workflow["5"]?.class_type).toBe("VAEEncode");
|
||||
expect(workflow["12"]?.class_type).toBe("GrowMask");
|
||||
expect(workflow["12"]?.inputs).toMatchObject({ mask: ["11", 0], expand: 12 });
|
||||
expect(workflow["13"]?.class_type).toBe("SetLatentNoiseMask");
|
||||
expect(workflow["13"]?.inputs).toMatchObject({ samples: ["5", 0], mask: ["12", 0] });
|
||||
expect(workflow["6"]?.inputs.latent_image).toEqual(["13", 0]);
|
||||
});
|
||||
});
|
||||
|
||||
function inpaintRequest(inpaint: { maskedContent: "neutral" | "original"; growMaskBy?: number }) {
|
||||
return {
|
||||
mode: "inpaint" as const,
|
||||
model: "model.safetensors",
|
||||
prompt: "replace garment",
|
||||
width: 128,
|
||||
height: 128,
|
||||
inputImage: "input.png",
|
||||
maskImage: "mask.png",
|
||||
inpaint,
|
||||
};
|
||||
}
|
||||
122
app/comfy.ts
122
app/comfy.ts
@@ -20,6 +20,17 @@ type ComfyGenerateRequest = {
|
||||
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;
|
||||
};
|
||||
@@ -27,10 +38,14 @@ type ComfyGenerateRequest = {
|
||||
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 });
|
||||
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() {
|
||||
@@ -61,6 +76,7 @@ async function generate(request: ComfyGenerateRequest) {
|
||||
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`, {
|
||||
@@ -69,9 +85,15 @@ async function generate(request: ComfyGenerateRequest) {
|
||||
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 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 image = firstOutputImage(history);
|
||||
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" })}`);
|
||||
@@ -82,7 +104,7 @@ async function generate(request: ComfyGenerateRequest) {
|
||||
|
||||
async function uploadDataUrl(dataUrl: string, filename: string) {
|
||||
const match = /^data:([^;]+);base64,(.+)$/.exec(dataUrl);
|
||||
if (!match) return undefined;
|
||||
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();
|
||||
@@ -103,19 +125,25 @@ async function waitForHistory(promptId: string) {
|
||||
}
|
||||
await Bun.sleep(500);
|
||||
}
|
||||
throw new Error("Timed out waiting for ComfyUI");
|
||||
throw new Error(`Timed out waiting for ComfyUI prompt ${promptId}`);
|
||||
}
|
||||
|
||||
function firstOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
|
||||
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 ?? {};
|
||||
for (const output of Object.values(outputs)) {
|
||||
const image = output.images?.[0];
|
||||
if (image) return image;
|
||||
}
|
||||
return undefined;
|
||||
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 buildSdxlWorkflow(request: ComfyGenerateRequest) {
|
||||
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));
|
||||
@@ -127,7 +155,7 @@ function buildSdxlWorkflow(request: ComfyGenerateRequest) {
|
||||
const scheduler = request.scheduler ?? "normal";
|
||||
const positive = request.prompt;
|
||||
const negative = request.negativePrompt ?? "";
|
||||
const workflow: Record<string, unknown> = {
|
||||
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] } },
|
||||
@@ -145,13 +173,21 @@ function buildSdxlWorkflow(request: ComfyGenerateRequest) {
|
||||
|
||||
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 } };
|
||||
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: 6 } };
|
||||
workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["10", 0], vae: ["1", 2], mask: ["10", 1], grow_mask_by: resolveGrowMaskBy(request) } };
|
||||
return workflow;
|
||||
}
|
||||
|
||||
@@ -159,6 +195,56 @@ function buildSdxlWorkflow(request: ComfyGenerateRequest) {
|
||||
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" } });
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user