feat: add inpaint region functionality and related tools

- Enhanced cursor behavior for new tools: semantic select, mask lasso, and mask rectangle.
- Updated mask edit state to include mask asset ID and kind.
- Implemented inpaint region commands for adding, applying, and removing inpaint regions.
- Introduced new operations for lasso and semantic selection tools.
- Created UI components for candidate review and inpaint region management.
- Added tests for inpaint region commands to ensure functionality.
- Updated various components to support new inpaint features and improve user experience.
This commit is contained in:
syntaxbullet
2026-07-11 16:41:22 +02:00
parent f4e13b80e7
commit ff762b8f17
78 changed files with 1632 additions and 301 deletions

View File

@@ -18,6 +18,9 @@ export type ComfyGenerateRequest = {
scheduler?: string;
width?: number;
height?: number;
batchSize?: number;
refinePass?: boolean;
refineStrength?: number;
outpaint?: {
left?: number;
top?: number;
@@ -35,17 +38,30 @@ export type ComfyGenerateRequest = {
maskedContent?: "neutral" | "original" | "originalColor" | "edges";
crop?: unknown;
placement?: unknown;
structureControl?: "none" | "canny" | "depth" | "pose";
controlStrength?: number;
controlModel?: string;
};
inputImage?: string;
maskImage?: string;
};
export type ComfyProgress = { progress: number; detail: string };
export type ComfySegmentRequest = { inputImage: string; x: number; y: number; model?: string };
type ComfyObjectInfo = {
CheckpointLoaderSimple?: { input?: { required?: { ckpt_name?: [string[]] } } };
KSampler?: { input?: { required?: { sampler_name?: [string[]]; scheduler?: [string[]] } } };
UNETLoader?: { input?: { required?: { unet_name?: [string[]] } } };
CLIPLoader?: { input?: { required?: { clip_name?: [string[]] } } };
VAELoader?: { input?: { required?: { vae_name?: [string[]] } } };
ControlNetLoader?: { input?: { required?: { control_net_name?: [string[]] } } };
Canny?: unknown;
"MiDaS-DepthMapPreprocessor"?: unknown;
OpenposePreprocessor?: unknown;
ControlNetApplyAdvanced?: unknown;
SAM3_Detect?: unknown;
MaskToImage?: unknown;
};
const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188";
@@ -59,9 +75,7 @@ const defaultModels: Record<GenerateArchitecture, string> = {
};
export 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 ComfyObjectInfo;
const info = await fetchObjectInfo();
const checkpointModels = info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? [];
const diffusionModels = info.UNETLoader?.input?.required?.unet_name?.[0] ?? [];
const textEncoders = info.CLIPLoader?.input?.required?.clip_name?.[0] ?? [];
@@ -69,11 +83,20 @@ export async function listGenerationOptions() {
return {
models: checkpointModels,
inpaintModels: checkpointModels.filter((model) => /inpaint|fill/i.test(model)),
samplers: info.KSampler?.input?.required?.sampler_name?.[0] ?? [],
schedulers: info.KSampler?.input?.required?.scheduler?.[0] ?? [],
diffusionModels,
textEncoders,
vaes,
controlModels: info.ControlNetLoader?.input?.required?.control_net_name?.[0] ?? [],
structureControls: [
...(info.Canny && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["canny"] : []),
...(info["MiDaS-DepthMapPreprocessor"] && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["depth"] : []),
...(info.OpenposePreprocessor && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["pose"] : []),
],
semanticSelection: Boolean(info.SAM3_Detect && info.MaskToImage && diffusionModels.some((model) => /sam.?3/i.test(model))),
sam3Models: diffusionModels.filter((model) => /sam.?3/i.test(model)),
architectures: [
{
value: "sdxl",
@@ -107,21 +130,60 @@ export async function listGenerationOptions() {
};
}
export async function segment(request: ComfySegmentRequest, signal?: AbortSignal) {
if (!request.inputImage) throw new Error("Semantic selection requires an image");
if (!Number.isFinite(request.x) || !Number.isFinite(request.y)) throw new Error("Semantic selection requires a valid point");
const info = await fetchObjectInfo();
const models = info.UNETLoader?.input?.required?.unet_name?.[0] ?? [];
const model = request.model && request.model !== "auto" ? request.model : models.find((candidate) => /sam.?3/i.test(candidate));
if (!info.SAM3_Detect || !info.MaskToImage || !model || !models.includes(model)) throw new Error("SAM3 semantic selection is not installed in ComfyUI. Install a SAM3 model and enable the native SAM3 nodes.");
const uploaded = await uploadDataUrl(request.inputImage, `image-studio-segment-${crypto.randomUUID()}.png`, signal);
const prompt = buildSemanticSelectionWorkflow({ ...request, model, inputImage: uploaded });
const queued = await fetch(`${comfyBaseUrl}/prompt`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ client_id: crypto.randomUUID(), prompt }), signal });
if (!queued.ok) throw new Error(`ComfyUI semantic selection 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 semantic selection: ${nodeError}`);
if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a semantic selection prompt id");
const history = await waitForHistory(queuedBody.prompt_id, signal);
const historyError = historyErrorMessage(history);
if (historyError) throw new Error(`ComfyUI semantic selection failed: ${historyError}`);
const image = selectGeneratedOutputImage(history);
if (!image) throw new Error("ComfyUI did not return a semantic selection mask");
const response = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal });
if (!response.ok) throw new Error(`ComfyUI mask fetch failed: ${response.status}`);
const bytes = Buffer.from(await response.arrayBuffer());
return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png" };
}
export function buildSemanticSelectionWorkflow(request: ComfySegmentRequest & { model: string }): Workflow {
return {
"1": { class_type: "UNETLoader", inputs: { unet_name: request.model, weight_dtype: "default" } },
"2": { class_type: "LoadImage", inputs: { image: request.inputImage } },
"3": { class_type: "SAM3_Detect", inputs: { model: ["1", 0], image: ["2", 0], positive_coords: JSON.stringify([{ x: Math.round(request.x), y: Math.round(request.y) }]), threshold: 0.5, refine_iterations: 2, individual_masks: false } },
"4": { class_type: "MaskToImage", inputs: { mask: ["3", 0] } },
"8": { class_type: "SaveImage", inputs: { filename_prefix: "image-studio-segment", images: ["4", 0] } },
};
}
async function listCheckpointModels() {
return (await listGenerationOptions()).models;
}
export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal) {
export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal, onProgress?: (event: ComfyProgress) => void) {
if (!request.prompt?.trim()) throw new Error("Prompt is required");
onProgress?.({ progress: 0.03, detail: "Preparing workflow" });
const architecture = normalizeArchitecture(request.architecture);
if (request.mode !== "text-to-image" && architecture !== "sdxl") throw new Error(`${architectureLabel(architecture)} currently supports text-to-image only`);
if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture);
if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture, request.mode);
const clientId = crypto.randomUUID();
const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`, signal) : undefined;
const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`, signal) : undefined;
onProgress?.({ progress: 0.16, detail: "Inputs uploaded" });
if (request.mode === "inpaint" && (!uploaded || !mask)) throw new Error("Inpaint requires normalized input and mask images");
const prompt = buildComfyWorkflow({ ...request, architecture, inputImage: uploaded, maskImage: mask });
const resolvedRequest = await resolveStructureControl({ ...request, architecture, inputImage: uploaded, maskImage: mask });
const prompt = buildComfyWorkflow(resolvedRequest);
const queued = await fetch(`${comfyBaseUrl}/prompt`, {
method: "POST",
@@ -134,23 +196,28 @@ export async function generate(request: ComfyGenerateRequest, signal?: AbortSign
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");
onProgress?.({ progress: 0.24, detail: "Queued in ComfyUI" });
const prompt_id = queuedBody.prompt_id;
let history: unknown;
try {
history = await waitForHistory(prompt_id, signal);
history = await waitForHistory(prompt_id, signal, onProgress);
} catch (error) {
if (signal?.aborted) await cancelComfyPrompt(prompt_id);
throw error;
}
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" })}`, { signal });
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" };
const images = selectGeneratedOutputImages(history);
if (images.length === 0) throw new Error("ComfyUI did not return an image");
onProgress?.({ progress: 0.9, detail: "Downloading results" });
const results = await Promise.all(images.map(async (image, index) => {
const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal });
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", seed: Math.round((resolvedRequest.seed ?? 0) + index) };
}));
onProgress?.({ progress: 1, detail: "Results ready" });
return { results };
}
async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSignal) {
@@ -167,7 +234,7 @@ async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSi
return uploaded.name;
}
async function waitForHistory(promptId: string, signal?: AbortSignal) {
async function waitForHistory(promptId: string, signal?: AbortSignal, onProgress?: (event: ComfyProgress) => void) {
const startedAt = Date.now();
let attempts = 0;
while (Date.now() - startedAt < comfyHistoryTimeoutMs) {
@@ -177,6 +244,8 @@ async function waitForHistory(promptId: string, signal?: AbortSignal) {
}
const response = await fetch(`${comfyBaseUrl}/history/${promptId}`, { signal });
attempts += 1;
const elapsedRatio = Math.min(1, (Date.now() - startedAt) / comfyHistoryTimeoutMs);
onProgress?.({ progress: 0.25 + elapsedRatio * 0.6, detail: attempts <= 1 ? "Generating" : `Generating · check ${attempts}` });
if (response.ok) {
const history = await response.json() as Record<string, unknown>;
if (history[promptId]) return history[promptId];
@@ -209,14 +278,18 @@ function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
}
export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
return selectGeneratedOutputImages(history)[0];
}
export function selectGeneratedOutputImages(history: unknown): Array<{ filename: string; subfolder?: string; type?: string }> {
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 saveImageOutput = outputs["8"]?.images?.filter(isGeneratedImage) ?? [];
if (saveImageOutput.length > 0) 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;
const prefixedOutput = Object.values(outputs).flatMap((output) => output.images ?? []).filter((image) => image.filename.startsWith("image-studio-") && image.type !== "input");
if (prefixedOutput.length > 0) return prefixedOutput;
return Object.values(outputs).flatMap((output) => output.images ?? []).find(isGeneratedImage);
return Object.values(outputs).flatMap((output) => output.images ?? []).filter(isGeneratedImage);
}
function isGeneratedImage(image: { filename: string; subfolder?: string; type?: string }) {
@@ -260,8 +333,8 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow {
};
if (request.mode === "text-to-image" || !request.inputImage) {
workflow["5"] = { class_type: "EmptyLatentImage", inputs: { width, height, batch_size: 1 } };
return workflow;
workflow["5"] = { class_type: "EmptyLatentImage", inputs: { width, height, batch_size: resolveBatchSize(request) } };
return finalizeSdxlWorkflow(workflow, samplerInputs, request);
}
workflow["4"] = { class_type: "LoadImage", inputs: { image: request.inputImage } };
@@ -277,19 +350,96 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow {
} else {
workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["4", 0], vae: ["1", 2], mask: ["11", 0], grow_mask_by: resolveGrowMaskBy(request) } };
}
return workflow;
return finalizeSdxlWorkflow(workflow, samplerInputs, request);
}
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;
return finalizeSdxlWorkflow(workflow, samplerInputs, request);
}
workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } };
return finalizeSdxlWorkflow(workflow, samplerInputs, request);
}
function finalizeSdxlWorkflow(workflow: Workflow, samplerInputs: Record<string, unknown>, request: ComfyGenerateRequest): Workflow {
const batchSize = resolveBatchSize(request);
if (batchSize > 1 && workflow["5"]?.class_type !== "EmptyLatentImage") {
const latent = samplerInputs.latent_image;
workflow["19"] = { class_type: "RepeatLatentBatch", inputs: { samples: latent, amount: batchSize } };
samplerInputs.latent_image = ["19", 0];
}
const control = request.inpaint?.structureControl ?? "none";
if (request.mode !== "inpaint" || control === "none" || !request.inpaint?.controlModel || request.inpaint.controlModel === "auto") return addRefinementPass(workflow, samplerInputs, request);
workflow["20"] = { class_type: "ControlNetLoader", inputs: { control_net_name: request.inpaint.controlModel } };
if (control === "canny") {
workflow["21"] = { class_type: "Canny", inputs: { image: ["4", 0], low_threshold: 100, high_threshold: 200 } };
} else if (control === "depth") {
workflow["21"] = { class_type: "MiDaS-DepthMapPreprocessor", inputs: { image: ["4", 0], a: 6.283, bg_threshold: 0.1, resolution: Math.max(request.width ?? 512, request.height ?? 512) } };
} else {
workflow["21"] = { class_type: "OpenposePreprocessor", inputs: { image: ["4", 0], detect_hand: "enable", detect_body: "enable", detect_face: "enable", resolution: Math.max(request.width ?? 512, request.height ?? 512) } };
}
workflow["22"] = {
class_type: "ControlNetApplyAdvanced",
inputs: {
positive: ["2", 0],
negative: ["3", 0],
control_net: ["20", 0],
image: ["21", 0],
strength: Math.max(0, Math.min(1, request.inpaint.controlStrength ?? 0.55)),
start_percent: 0,
end_percent: 0.85,
vae: ["1", 2],
},
};
samplerInputs.positive = ["22", 0];
samplerInputs.negative = ["22", 1];
return addRefinementPass(workflow, samplerInputs, request);
}
function addRefinementPass(workflow: Workflow, samplerInputs: Record<string, unknown>, request: ComfyGenerateRequest): Workflow {
if (!request.refinePass) return workflow;
workflow["30"] = {
class_type: "KSampler",
inputs: {
...samplerInputs,
seed: Math.round((request.seed ?? 0) + 1),
steps: Math.max(6, Math.round((request.steps ?? 30) / 3)),
denoise: Math.max(0, Math.min(1, (request.refineStrength ?? 20) / 100)),
latent_image: ["6", 0],
},
};
if (workflow["7"]) workflow["7"].inputs.samples = ["30", 0];
return workflow;
}
function resolveBatchSize(request: ComfyGenerateRequest) {
return Math.round(Math.max(1, Math.min(8, request.batchSize ?? 1)));
}
async function fetchObjectInfo(): Promise<ComfyObjectInfo> {
const response = await fetch(`${comfyBaseUrl}/object_info`);
if (!response.ok) throw new Error(`ComfyUI option lookup failed: ${response.status}`);
return response.json() as Promise<ComfyObjectInfo>;
}
async function resolveStructureControl(request: ComfyGenerateRequest): Promise<ComfyGenerateRequest> {
const control = request.inpaint?.structureControl ?? "none";
if (request.mode !== "inpaint" || control === "none") return request;
const info = await fetchObjectInfo();
const requiredPreprocessor = control === "canny" ? info.Canny : control === "depth" ? info["MiDaS-DepthMapPreprocessor"] : info.OpenposePreprocessor;
if (!requiredPreprocessor || !info.ControlNetLoader || !info.ControlNetApplyAdvanced) {
throw new Error(`${control === "canny" ? "Canny" : control === "depth" ? "Depth" : "Pose"} structural control is not installed in ComfyUI.`);
}
const models = info.ControlNetLoader.input?.required?.control_net_name?.[0] ?? [];
const requested = request.inpaint?.controlModel;
const model = requested && requested !== "auto" ? requested : models.find((candidate) => candidate.toLowerCase().includes(control));
if (!model || !models.includes(model)) throw new Error(`Install or select a ${control} ControlNet model before using structural control.`);
return { ...request, inpaint: { ...request.inpaint, controlModel: model } };
}
export function buildZImageWorkflow(request: ComfyGenerateRequest): Workflow {
return buildSeparatedTextToImageWorkflow(request, {
architecture: "z-image",
@@ -408,9 +558,10 @@ function usesOriginalLatentContent(request: ComfyGenerateRequest): boolean {
return request.inpaint?.maskedContent === "original" || request.inpaint?.maskedContent === "originalColor" || request.inpaint?.maskedContent === "edges";
}
async function defaultModelForArchitecture(architecture: GenerateArchitecture) {
async function defaultModelForArchitecture(architecture: GenerateArchitecture, mode: GenerateMode) {
if (architecture !== "sdxl") return defaultModels[architecture];
const models = await listCheckpointModels();
if (mode === "inpaint") return models.find((model) => /inpaint|fill/i.test(model)) ?? models[0] ?? defaultModels.sdxl;
return models[0] ?? defaultModels.sdxl;
}