feat: add ComfyUI integration for image generation and management
- Implemented ComfyGenerateRequest type and associated functions for generating images using various architectures and modes. - Added functions for listing generation options and handling image uploads. - Created workflows for different generation modes including SDXL, Z-Image, Z-Image Turbo, and Anima. - Introduced GenerationJobStatus component to display the status of ongoing generation jobs. - Developed MaskControls for managing mask operations and displaying mask analysis. - Created palette items for tool selection, layer management, and generation settings.
This commit is contained in:
@@ -8,13 +8,15 @@ import { selectionCommands } from "@commands/selection";
|
||||
import { toolCommands } from "@commands/tool";
|
||||
import { transformCommands } from "@commands/transform";
|
||||
import { viewportCommands } from "@commands/viewport";
|
||||
import { workspaceCommands } from "@commands/workspace";
|
||||
import { editorCommands } from "@commands/editor";
|
||||
import { createInitialAppState } from "@editor/initial-state";
|
||||
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, ...generationCommands, ...transformCommands, ...historyCommands, ...commandPaletteCommands]);
|
||||
const registry = createCommandRegistry([...viewportCommands, ...selectionCommands, ...documentCommands, ...toolCommands, ...generationCommands, ...transformCommands, ...historyCommands, ...commandPaletteCommands, ...workspaceCommands, ...editorCommands]);
|
||||
const store = createAppStore(createInitialAppState(options?.documentName), registry);
|
||||
|
||||
if (options?.createDefaultArtboard !== false) {
|
||||
|
||||
@@ -1,153 +0,0 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { buildAnimaWorkflow, buildSdxlWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, handleComfyApi, 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]);
|
||||
});
|
||||
|
||||
test("builds Z-Image text-to-image with separated model loaders", () => {
|
||||
const workflow = buildZImageWorkflow(textRequest({ architecture: "z-image", model: "z_image_bf16.safetensors", steps: 30, cfg: 4 }));
|
||||
|
||||
expect(workflow["1"]).toMatchObject({ class_type: "UNETLoader", inputs: { unet_name: "z_image_bf16.safetensors", weight_dtype: "default" } });
|
||||
expect(workflow["2"]).toMatchObject({ class_type: "CLIPLoader", inputs: { clip_name: "qwen_3_4b.safetensors", type: "lumina2" } });
|
||||
expect(workflow["6"]?.class_type).toBe("EmptySD3LatentImage");
|
||||
expect(workflow["7"]).toMatchObject({ class_type: "ModelSamplingAuraFlow", inputs: { shift: 3 } });
|
||||
expect(workflow["8"]?.inputs).toMatchObject({ steps: 30, cfg: 4, sampler_name: "res_multistep", scheduler: "simple", model: ["7", 0] });
|
||||
expect(workflow["10"]?.class_type).toBe("SaveImage");
|
||||
});
|
||||
|
||||
test("builds Z-Image Turbo with zeroed negative conditioning", () => {
|
||||
const workflow = buildZImageTurboWorkflow(textRequest({ architecture: "z-image-turbo", model: "z_image_turbo_bf16.safetensors", negativePrompt: "ignored" }));
|
||||
|
||||
expect(workflow["1"]?.inputs.unet_name).toBe("z_image_turbo_bf16.safetensors");
|
||||
expect(workflow["5"]).toMatchObject({ class_type: "ConditioningZeroOut", inputs: { conditioning: ["4", 0] } });
|
||||
expect(workflow["8"]?.inputs).toMatchObject({ steps: 8, cfg: 1, sampler_name: "res_multistep", scheduler: "simple" });
|
||||
});
|
||||
|
||||
test("builds Anima text-to-image workflow", () => {
|
||||
const workflow = buildAnimaWorkflow(textRequest({ architecture: "anima", model: "anima-base-v1.0.safetensors" }));
|
||||
|
||||
expect(workflow["1"]).toMatchObject({ class_type: "UNETLoader", inputs: { unet_name: "anima-base-v1.0.safetensors" } });
|
||||
expect(workflow["2"]).toMatchObject({ class_type: "CLIPLoader", inputs: { clip_name: "qwen_3_06b_base.safetensors", type: "stable_diffusion" } });
|
||||
expect(workflow["3"]).toMatchObject({ class_type: "VAELoader", inputs: { vae_name: "qwen_image_vae.safetensors" } });
|
||||
expect(workflow["6"]?.class_type).toBe("EmptyLatentImage");
|
||||
expect(workflow["8"]?.inputs).toMatchObject({ steps: 30, cfg: 4, sampler_name: "er_sde", scheduler: "simple", model: ["1", 0] });
|
||||
});
|
||||
|
||||
test("builds Anima with model-specific text encoder and VAE", () => {
|
||||
const workflow = buildAnimaWorkflow(textRequest({
|
||||
architecture: "anima",
|
||||
model: "miaomiaoHarem_anima13.safetensors",
|
||||
textEncoder: "miaomiaoHarem_anima13_txt.safetensors",
|
||||
vae: "qwen_image_vae.safetensors",
|
||||
}));
|
||||
|
||||
expect(workflow["1"]?.inputs.unet_name).toBe("miaomiaoHarem_anima13.safetensors");
|
||||
expect(workflow["2"]?.inputs.clip_name).toBe("miaomiaoHarem_anima13_txt.safetensors");
|
||||
expect(workflow["3"]?.inputs.vae_name).toBe("qwen_image_vae.safetensors");
|
||||
});
|
||||
|
||||
test("lists branded non-Z diffusion models under Anima", async () => {
|
||||
const originalFetch = globalThis.fetch;
|
||||
const mockFetch: typeof fetch = Object.assign(async () => new Response(JSON.stringify({
|
||||
CheckpointLoaderSimple: { input: { required: { ckpt_name: [["sd_xl_base_1.0.safetensors"]] } } },
|
||||
KSampler: { input: { required: { sampler_name: [["euler"]], scheduler: [["normal"]] } } },
|
||||
UNETLoader: {
|
||||
input: {
|
||||
required: {
|
||||
unet_name: [[
|
||||
"anima-base-v1.0.safetensors",
|
||||
"miaomiaoHarem_anima13.safetensors",
|
||||
"novaAnimeAM_v30.safetensors",
|
||||
"z_image_bf16.safetensors",
|
||||
"z_image_turbo_bf16.safetensors",
|
||||
]],
|
||||
},
|
||||
},
|
||||
},
|
||||
CLIPLoader: { input: { required: { clip_name: [["qwen_3_06b_base.safetensors"]] } } },
|
||||
VAELoader: { input: { required: { vae_name: [["qwen_image_vae.safetensors"]] } } },
|
||||
}), { headers: { "content-type": "application/json" } }), { preconnect: originalFetch.preconnect });
|
||||
globalThis.fetch = mockFetch;
|
||||
|
||||
try {
|
||||
const response = await handleComfyApi(new Request("http://image-studio.test/api/comfy/models"));
|
||||
const body = await response.json() as { architectures: { value: string; models: string[] }[] };
|
||||
const anima = body.architectures.find((architecture) => architecture.value === "anima");
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(anima?.models).toContain("novaAnimeAM_v30.safetensors");
|
||||
expect(anima?.models).toContain("miaomiaoHarem_anima13.safetensors");
|
||||
expect(anima?.models).not.toContain("z_image_bf16.safetensors");
|
||||
expect(anima?.models).not.toContain("z_image_turbo_bf16.safetensors");
|
||||
} finally {
|
||||
globalThis.fetch = originalFetch;
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
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,
|
||||
};
|
||||
}
|
||||
|
||||
function textRequest(overrides: Partial<Parameters<typeof buildSdxlWorkflow>[0]> = {}) {
|
||||
return {
|
||||
architecture: "sdxl" as const,
|
||||
mode: "text-to-image" as const,
|
||||
model: "model.safetensors",
|
||||
prompt: "a studio portrait",
|
||||
negativePrompt: "low quality",
|
||||
width: 1024,
|
||||
height: 1024,
|
||||
seed: 123,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
475
app/comfy.ts
475
app/comfy.ts
@@ -1,475 +0,0 @@
|
||||
type GenerateArchitecture = "sdxl" | "z-image" | "z-image-turbo" | "anima";
|
||||
type GenerateMode = "text-to-image" | "image-to-image" | "inpaint" | "outpaint";
|
||||
type Workflow = Record<string, { class_type: string; inputs: Record<string, unknown> }>;
|
||||
|
||||
type ComfyGenerateRequest = {
|
||||
architecture?: GenerateArchitecture;
|
||||
mode: GenerateMode;
|
||||
model: string;
|
||||
textEncoder?: string;
|
||||
vae?: 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;
|
||||
};
|
||||
|
||||
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[]] } } };
|
||||
};
|
||||
|
||||
const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188";
|
||||
const comfyHistoryTimeoutMs = parsePositiveInteger(process.env.COMFYUI_HISTORY_TIMEOUT_MS, 20 * 60 * 1000);
|
||||
const comfyHistoryPollIntervalMs = parsePositiveInteger(process.env.COMFYUI_HISTORY_POLL_INTERVAL_MS, 1000);
|
||||
const defaultModels: Record<GenerateArchitecture, string> = {
|
||||
sdxl: "sd_xl_base_1.0.safetensors",
|
||||
"z-image": "z_image_bf16.safetensors",
|
||||
"z-image-turbo": "z_image_turbo_bf16.safetensors",
|
||||
anima: "anima-base-v1.0.safetensors",
|
||||
};
|
||||
|
||||
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 ComfyObjectInfo;
|
||||
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] ?? [];
|
||||
const vaes = info.VAELoader?.input?.required?.vae_name?.[0] ?? [];
|
||||
|
||||
return {
|
||||
models: checkpointModels,
|
||||
samplers: info.KSampler?.input?.required?.sampler_name?.[0] ?? [],
|
||||
schedulers: info.KSampler?.input?.required?.scheduler?.[0] ?? [],
|
||||
diffusionModels,
|
||||
textEncoders,
|
||||
vaes,
|
||||
architectures: [
|
||||
{
|
||||
value: "sdxl",
|
||||
label: "SDXL",
|
||||
defaultModel: checkpointModels[0] ?? defaultModels.sdxl,
|
||||
models: checkpointModels,
|
||||
supportedModes: ["text-to-image", "image-to-image", "inpaint", "outpaint"],
|
||||
},
|
||||
{
|
||||
value: "z-image",
|
||||
label: "Z-Image",
|
||||
defaultModel: defaultModels["z-image"],
|
||||
models: modelsForArchitecture(diffusionModels, defaultModels["z-image"], [/z[_-]?image(?!.*turbo)/i]),
|
||||
supportedModes: ["text-to-image"],
|
||||
},
|
||||
{
|
||||
value: "z-image-turbo",
|
||||
label: "Z-Image Turbo",
|
||||
defaultModel: defaultModels["z-image-turbo"],
|
||||
models: modelsForArchitecture(diffusionModels, defaultModels["z-image-turbo"], [/z[_-]?image.*turbo/i]),
|
||||
supportedModes: ["text-to-image"],
|
||||
},
|
||||
{
|
||||
value: "anima",
|
||||
label: "Anima",
|
||||
defaultModel: defaultModels.anima,
|
||||
models: modelsForArchitecture(diffusionModels, defaultModels.anima, [], { exclude: [/z[_-]?image/i] }),
|
||||
supportedModes: ["text-to-image"],
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
async function listCheckpointModels() {
|
||||
return (await listGenerationOptions()).models;
|
||||
}
|
||||
|
||||
async function generate(request: ComfyGenerateRequest) {
|
||||
if (!request.prompt?.trim()) throw new Error("Prompt is required");
|
||||
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);
|
||||
|
||||
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 = buildComfyWorkflow({ ...request, architecture, 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) {
|
||||
const startedAt = Date.now();
|
||||
let attempts = 0;
|
||||
while (Date.now() - startedAt < comfyHistoryTimeoutMs) {
|
||||
const response = await fetch(`${comfyBaseUrl}/history/${promptId}`);
|
||||
attempts += 1;
|
||||
if (response.ok) {
|
||||
const history = await response.json() as Record<string, unknown>;
|
||||
if (history[promptId]) return history[promptId];
|
||||
}
|
||||
await Bun.sleep(comfyHistoryPollIntervalMs);
|
||||
}
|
||||
throw new Error(`Timed out waiting for ComfyUI prompt ${promptId} after ${Math.round(comfyHistoryTimeoutMs / 1000)} seconds and ${attempts} checks`);
|
||||
}
|
||||
|
||||
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 buildComfyWorkflow(request: ComfyGenerateRequest): Workflow {
|
||||
switch (normalizeArchitecture(request.architecture)) {
|
||||
case "z-image":
|
||||
return buildZImageWorkflow(request);
|
||||
case "z-image-turbo":
|
||||
return buildZImageTurboWorkflow(request);
|
||||
case "anima":
|
||||
return buildAnimaWorkflow(request);
|
||||
case "sdxl":
|
||||
default:
|
||||
return buildSdxlWorkflow(request);
|
||||
}
|
||||
}
|
||||
|
||||
export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow {
|
||||
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 samplerInputs: Record<string, unknown> = { seed, steps, cfg, sampler_name: sampler, scheduler, denoise, model: ["1", 0], positive: ["2", 0], negative: ["3", 0], latent_image: ["5", 0] };
|
||||
const workflow: Workflow = {
|
||||
"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: samplerInputs },
|
||||
"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] } };
|
||||
samplerInputs.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;
|
||||
}
|
||||
|
||||
export function buildZImageWorkflow(request: ComfyGenerateRequest): Workflow {
|
||||
return buildSeparatedTextToImageWorkflow(request, {
|
||||
architecture: "z-image",
|
||||
filenamePrefix: "image-studio-z-image",
|
||||
model: defaultModels["z-image"],
|
||||
textEncoder: "qwen_3_4b.safetensors",
|
||||
vae: "ae.safetensors",
|
||||
clipType: "lumina2",
|
||||
latentNode: "EmptySD3LatentImage",
|
||||
modelSamplingAuraFlow: true,
|
||||
negativeMode: "prompt",
|
||||
steps: 30,
|
||||
cfg: 4,
|
||||
sampler: "res_multistep",
|
||||
scheduler: "simple",
|
||||
});
|
||||
}
|
||||
|
||||
export function buildZImageTurboWorkflow(request: ComfyGenerateRequest): Workflow {
|
||||
return buildSeparatedTextToImageWorkflow(request, {
|
||||
architecture: "z-image-turbo",
|
||||
filenamePrefix: "image-studio-z-image-turbo",
|
||||
model: defaultModels["z-image-turbo"],
|
||||
textEncoder: "qwen_3_4b.safetensors",
|
||||
vae: "ae.safetensors",
|
||||
clipType: "lumina2",
|
||||
latentNode: "EmptySD3LatentImage",
|
||||
modelSamplingAuraFlow: true,
|
||||
negativeMode: "zero",
|
||||
steps: 8,
|
||||
cfg: 1,
|
||||
sampler: "res_multistep",
|
||||
scheduler: "simple",
|
||||
});
|
||||
}
|
||||
|
||||
export function buildAnimaWorkflow(request: ComfyGenerateRequest): Workflow {
|
||||
return buildSeparatedTextToImageWorkflow(request, {
|
||||
architecture: "anima",
|
||||
filenamePrefix: "image-studio-anima",
|
||||
model: defaultModels.anima,
|
||||
textEncoder: "qwen_3_06b_base.safetensors",
|
||||
vae: "qwen_image_vae.safetensors",
|
||||
clipType: "stable_diffusion",
|
||||
latentNode: "EmptyLatentImage",
|
||||
modelSamplingAuraFlow: false,
|
||||
negativeMode: "prompt",
|
||||
steps: 30,
|
||||
cfg: 4,
|
||||
sampler: "er_sde",
|
||||
scheduler: "simple",
|
||||
});
|
||||
}
|
||||
|
||||
function buildSeparatedTextToImageWorkflow(request: ComfyGenerateRequest, config: {
|
||||
architecture: GenerateArchitecture;
|
||||
filenamePrefix: string;
|
||||
model: string;
|
||||
textEncoder: string;
|
||||
vae: string;
|
||||
clipType: string;
|
||||
latentNode: "EmptyLatentImage" | "EmptySD3LatentImage";
|
||||
modelSamplingAuraFlow: boolean;
|
||||
negativeMode: "prompt" | "zero";
|
||||
steps: number;
|
||||
cfg: number;
|
||||
sampler: string;
|
||||
scheduler: string;
|
||||
}): Workflow {
|
||||
if (request.mode !== "text-to-image") throw new Error(`${architectureLabel(config.architecture)} currently supports text-to-image only`);
|
||||
const options = resolveSamplerOptions(request, config);
|
||||
const modelOutput: [string, number] = config.modelSamplingAuraFlow ? ["7", 0] : ["1", 0];
|
||||
const workflow: Workflow = {
|
||||
"1": { class_type: "UNETLoader", inputs: { unet_name: request.model && request.model !== "auto" ? request.model : config.model, weight_dtype: "default" } },
|
||||
"2": { class_type: "CLIPLoader", inputs: { clip_name: resolveSupportModelName(request.textEncoder, config.textEncoder), type: config.clipType, device: "default" } },
|
||||
"3": { class_type: "VAELoader", inputs: { vae_name: resolveSupportModelName(request.vae, config.vae) } },
|
||||
"4": { class_type: "CLIPTextEncode", inputs: { text: request.prompt, clip: ["2", 0] } },
|
||||
"6": { class_type: config.latentNode, inputs: { width: options.width, height: options.height, batch_size: 1 } },
|
||||
"8": { class_type: "KSampler", inputs: { seed: options.seed, steps: options.steps, cfg: options.cfg, sampler_name: options.sampler, scheduler: options.scheduler, denoise: 1, model: modelOutput, positive: ["4", 0], negative: ["5", 0], latent_image: ["6", 0] } },
|
||||
"9": { class_type: "VAEDecode", inputs: { samples: ["8", 0], vae: ["3", 0] } },
|
||||
"10": { class_type: "SaveImage", inputs: { filename_prefix: config.filenamePrefix, images: ["9", 0] } },
|
||||
};
|
||||
|
||||
if (config.modelSamplingAuraFlow) workflow["7"] = { class_type: "ModelSamplingAuraFlow", inputs: { model: ["1", 0], shift: 3 } };
|
||||
workflow["5"] = config.negativeMode === "zero"
|
||||
? { class_type: "ConditioningZeroOut", inputs: { conditioning: ["4", 0] } }
|
||||
: { class_type: "CLIPTextEncode", inputs: { text: request.negativePrompt ?? "", clip: ["2", 0] } };
|
||||
|
||||
return workflow;
|
||||
}
|
||||
|
||||
function resolveSamplerOptions(request: ComfyGenerateRequest, defaults: { steps: number; cfg: number; sampler: string; scheduler: string }) {
|
||||
if (!request.width || !request.height) throw new Error("Generation width and height are required");
|
||||
return {
|
||||
width: Math.max(64, Math.round(request.width)),
|
||||
height: Math.max(64, Math.round(request.height)),
|
||||
seed: request.seed === undefined || request.seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(request.seed),
|
||||
steps: Math.max(1, Math.round(request.steps ?? defaults.steps)),
|
||||
cfg: Math.max(0, request.cfg ?? defaults.cfg),
|
||||
sampler: request.sampler ?? defaults.sampler,
|
||||
scheduler: request.scheduler ?? defaults.scheduler,
|
||||
};
|
||||
}
|
||||
|
||||
function resolveSupportModelName(value: string | undefined, fallback: string) {
|
||||
return value && value !== "auto" ? value : fallback;
|
||||
}
|
||||
|
||||
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";
|
||||
}
|
||||
|
||||
async function defaultModelForArchitecture(architecture: GenerateArchitecture) {
|
||||
if (architecture !== "sdxl") return defaultModels[architecture];
|
||||
const models = await listCheckpointModels();
|
||||
return models[0] ?? defaultModels.sdxl;
|
||||
}
|
||||
|
||||
function normalizeArchitecture(architecture: ComfyGenerateRequest["architecture"]): GenerateArchitecture {
|
||||
if (architecture === "z-image" || architecture === "z-image-turbo" || architecture === "anima") return architecture;
|
||||
return "sdxl";
|
||||
}
|
||||
|
||||
function architectureLabel(architecture: GenerateArchitecture): string {
|
||||
switch (architecture) {
|
||||
case "z-image":
|
||||
return "Z-Image";
|
||||
case "z-image-turbo":
|
||||
return "Z-Image Turbo";
|
||||
case "anima":
|
||||
return "Anima";
|
||||
case "sdxl":
|
||||
default:
|
||||
return "SDXL";
|
||||
}
|
||||
}
|
||||
|
||||
function modelsForArchitecture(models: string[], defaultModel: string, matchers: RegExp[], options: { exclude?: RegExp[] } = {}) {
|
||||
return unique([defaultModel, ...models.filter((model) => {
|
||||
if (model === defaultModel) return true;
|
||||
if (options.exclude?.some((matcher) => matcher.test(model))) return false;
|
||||
if (matchers.length === 0) return true;
|
||||
return matchers.some((matcher) => matcher.test(model));
|
||||
})]);
|
||||
}
|
||||
|
||||
function unique<T>(values: T[]) {
|
||||
return Array.from(new Set(values));
|
||||
}
|
||||
|
||||
function parsePositiveInteger(value: string | undefined, fallback: number) {
|
||||
if (!value) return fallback;
|
||||
const parsed = Number(value);
|
||||
return Number.isFinite(parsed) && parsed > 0 ? Math.round(parsed) : fallback;
|
||||
}
|
||||
|
||||
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