353 lines
15 KiB
TypeScript
353 lines
15 KiB
TypeScript
import sharp from "sharp";
|
|
import { describe, expect, it } from "vitest";
|
|
import type { ComfyUiProviderConfig } from "./comfyui-provider.js";
|
|
import {
|
|
createComfyUiProvider,
|
|
describeExecutionError,
|
|
parseGpuStatus,
|
|
} from "./comfyui-provider.js";
|
|
import { ComfyUiEmulator } from "./testing/comfyui-emulator.js";
|
|
|
|
const BASE_URL = "http://127.0.0.1:8188";
|
|
const IMAGE = { prompt: "a lighthouse", width: 512, height: 512 };
|
|
|
|
function setup(config: Partial<ComfyUiProviderConfig> = {}, emulator = new ComfyUiEmulator()) {
|
|
const provider = createComfyUiProvider({
|
|
baseUrl: BASE_URL,
|
|
fetch: emulator.fetch,
|
|
webSocket: emulator.openWebSocket,
|
|
pollIntervalMs: 1,
|
|
...config,
|
|
});
|
|
return { provider, emulator };
|
|
}
|
|
|
|
/** Progress rounded so float sums compare exactly. */
|
|
function rounded(progress: number[]): number[] {
|
|
return progress.map((fraction) => Number(fraction.toFixed(4)));
|
|
}
|
|
|
|
function submitted(emulator: ComfyUiEmulator, index = 0) {
|
|
const graph = emulator.prompts[index]?.graph ?? {};
|
|
const byClass = (classType: string) =>
|
|
Object.values(graph).find((node) => node.class_type === classType)?.inputs ?? {};
|
|
return { graph, byClass };
|
|
}
|
|
|
|
describe("createComfyUiProvider", () => {
|
|
it("refuses metadata hosts, public hosts and keys over public http", () => {
|
|
expect(() => createComfyUiProvider({ baseUrl: "http://169.254.169.254" })).toThrow(
|
|
"blocked metadata or link-local host",
|
|
);
|
|
expect(() => createComfyUiProvider({ baseUrl: "http://comfy.example.com" })).toThrow(
|
|
"Public model endpoints are blocked",
|
|
);
|
|
expect(() => createComfyUiProvider({ baseUrl: "ftp://127.0.0.1" })).toThrow("http or https");
|
|
});
|
|
|
|
it("validates custom workflows up front", () => {
|
|
expect(() => createComfyUiProvider({ baseUrl: BASE_URL, imageWorkflow: "{" })).toThrow(
|
|
"Workflow is not valid JSON",
|
|
);
|
|
expect(() =>
|
|
createComfyUiProvider({ baseUrl: BASE_URL, videoWorkflow: '{"nodes": [], "links": []}' }),
|
|
).toThrow("API format");
|
|
});
|
|
|
|
it("keeps a reverse-proxy path prefix and sends the API key", async () => {
|
|
const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" });
|
|
const { provider } = setup({ baseUrl: `${BASE_URL}/comfy/`, apiKey: "comfy-key" }, emulator);
|
|
await provider.generateImage(IMAGE);
|
|
expect(emulator.requests.map((request) => request.path)).toContain("prompt");
|
|
expect(emulator.requests.every((request) => request.authorization === "Bearer comfy-key")).toBe(
|
|
true,
|
|
);
|
|
const { provider: keyless } = setup({ baseUrl: `${BASE_URL}/comfy` }, emulator);
|
|
await expect(keyless.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 401");
|
|
});
|
|
|
|
it("fills the default text-to-image workflow", async () => {
|
|
const { provider, emulator } = setup();
|
|
await provider.generateImage({ ...IMAGE, width: 1001, negativePrompt: "fog", seed: 42 });
|
|
const { byClass } = submitted(emulator);
|
|
expect(byClass("CheckpointLoaderSimple").ckpt_name).toBe("sd_xl_base_1.0.safetensors");
|
|
expect(byClass("EmptyLatentImage")).toMatchObject({ width: 1000, height: 512 });
|
|
expect(byClass("KSampler")).toMatchObject({ seed: 42, denoise: 1 });
|
|
const texts = emulator.prompts[0]!.graph;
|
|
expect(Object.values(texts).map((node) => node.inputs.text)).toEqual(
|
|
expect.arrayContaining(["a lighthouse", "fog"]),
|
|
);
|
|
await provider.generateImage(IMAGE);
|
|
const defaultNegativeTexts = Object.values(emulator.prompts[1]!.graph).map((node) =>
|
|
String(node.inputs.text ?? ""),
|
|
);
|
|
expect(defaultNegativeTexts.some((text) => text.includes("amateur clipart"))).toBe(true);
|
|
expect(defaultNegativeTexts.some((text) => text.includes("artificial symmetry"))).toBe(false);
|
|
});
|
|
|
|
it("uses the configured model, else the first checkpoint the video workflows do not use", async () => {
|
|
const emulator = new ComfyUiEmulator();
|
|
emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors", "v1-5-pruned.safetensors"];
|
|
emulator.checkpointListFormat = "legacy";
|
|
await setup({}, emulator).provider.generateImage(IMAGE);
|
|
expect(submitted(emulator, 0).byClass("CheckpointLoaderSimple").ckpt_name).toBe(
|
|
"v1-5-pruned.safetensors",
|
|
);
|
|
emulator.checkpoints.push("dreamshaper.safetensors");
|
|
await setup({ imageModel: "dreamshaper.safetensors" }, emulator).provider.generateImage(IMAGE);
|
|
expect(submitted(emulator, 1).byClass("CheckpointLoaderSimple").ckpt_name).toBe(
|
|
"dreamshaper.safetensors",
|
|
);
|
|
emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors"];
|
|
await expect(setup({}, emulator).provider.generateImage(IMAGE)).rejects.toThrow(
|
|
"ComfyUI has no image checkpoint installed",
|
|
);
|
|
});
|
|
|
|
it("uploads the source image for image-to-image", async () => {
|
|
const { provider, emulator } = setup();
|
|
const bytes = new Uint8Array(Buffer.from((await sourcePng()).bytes));
|
|
await provider.generateImage({
|
|
...IMAGE,
|
|
source: { bytes, mimeType: "image/png" },
|
|
strength: 0.3,
|
|
});
|
|
const { byClass } = submitted(emulator);
|
|
const uploaded = String(byClass("LoadImage").image);
|
|
expect(uploaded).toMatch(/^evoo-.+\.png$/);
|
|
expect(Buffer.from(emulator.uploads.get(uploaded)!).equals(Buffer.from(bytes))).toBe(true);
|
|
expect(byClass("KSampler").denoise).toBe(0.3);
|
|
});
|
|
|
|
it("lets a custom workflow replace only the default it matches", async () => {
|
|
const custom = {
|
|
"1": {
|
|
class_type: "CustomTextToImage",
|
|
inputs: { text: "style: {{prompt}}", seed: "{{seed}}" },
|
|
},
|
|
"2": { class_type: "SaveImage", inputs: { images: ["1", 0], filename_prefix: "x" } },
|
|
};
|
|
const { provider, emulator } = setup({ imageWorkflow: JSON.stringify(custom) });
|
|
await provider.generateImage(IMAGE);
|
|
expect(submitted(emulator, 0).byClass("CustomTextToImage").text).toBe("style: a lighthouse");
|
|
await provider.generateImage({ ...IMAGE, source: await sourcePng() });
|
|
expect(submitted(emulator, 1).byClass("LoadImage").image).toBeDefined();
|
|
});
|
|
|
|
it("fills the LTX-Video workflow with an 8n+1 frame count and reports the clip", async () => {
|
|
const { provider, emulator } = setup();
|
|
const [clip] = await provider.generateVideo({
|
|
prompt: "a paper boat",
|
|
width: 470,
|
|
height: 836,
|
|
durationSeconds: 5,
|
|
});
|
|
const { byClass } = submitted(emulator);
|
|
expect(byClass("EmptyLTXVLatentVideo")).toMatchObject({ width: 480, height: 832, length: 121 });
|
|
expect(byClass("CreateVideo").fps).toBe(24);
|
|
expect(clip).toMatchObject({
|
|
mimeType: "video/mp4",
|
|
width: 480,
|
|
height: 832,
|
|
durationMs: 5042,
|
|
});
|
|
});
|
|
|
|
it("reads video from a custom workflow's gifs output and skips frame previews", async () => {
|
|
const custom = {
|
|
"1": { class_type: "TextToFrames", inputs: { text: "{{prompt}}" } },
|
|
"2": { class_type: "PreviewImage", inputs: { images: ["1", 0] } },
|
|
"3": { class_type: "VHS_VideoCombine", inputs: { images: ["1", 0], frame_rate: 16 } },
|
|
};
|
|
const { provider } = setup({ videoWorkflow: JSON.stringify(custom) });
|
|
const clips = await provider.generateVideo({ ...IMAGE, durationSeconds: 2 });
|
|
expect(clips.map((clip) => [clip.mimeType, clip.durationMs])).toEqual([
|
|
["video/mp4", undefined],
|
|
]);
|
|
});
|
|
|
|
it("explains rejected workflows and failed runs", async () => {
|
|
const emulator = new ComfyUiEmulator();
|
|
const { provider } = setup({ imageModel: "missing.safetensors" }, emulator);
|
|
await expect(provider.generateImage(IMAGE)).rejects.toThrow(
|
|
"ComfyUI rejected the workflow: CheckpointLoaderSimple: Value not in list: ckpt_name: 'missing.safetensors' not in list",
|
|
);
|
|
const working = setup({}, emulator).provider;
|
|
emulator.failNextExecution = "CUDA out of memory. Tried to allocate 2.00 GiB\nTraceback...";
|
|
await expect(working.generateImage(IMAGE)).rejects.toThrow(
|
|
/^ComfyUI failed in KSampler: CUDA out of memory\. Tried to allocate 2\.00 GiB$/,
|
|
);
|
|
emulator.failNextStatus = 500;
|
|
await expect(working.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 500");
|
|
});
|
|
|
|
it("reports an unreachable server", async () => {
|
|
const provider = createComfyUiProvider({
|
|
baseUrl: BASE_URL,
|
|
fetch: async () => {
|
|
throw new TypeError("fetch failed");
|
|
},
|
|
});
|
|
await expect(provider.generateImage(IMAGE)).rejects.toThrow("ComfyUI is unreachable");
|
|
await expect(provider.gpuStatus()).rejects.toThrow("ComfyUI is unreachable");
|
|
});
|
|
|
|
it("interrupts its own running prompt when the caller cancels", async () => {
|
|
const emulator = new ComfyUiEmulator();
|
|
emulator.pendingPolls = 1_000_000;
|
|
const { provider } = setup({ pollIntervalMs: 5 }, emulator);
|
|
const controller = new AbortController();
|
|
const run = provider.generateImage({ ...IMAGE, signal: controller.signal });
|
|
while (!emulator.requests.some((request) => request.path.startsWith("history/"))) {
|
|
await new Promise((resolve) => setTimeout(resolve, 1));
|
|
}
|
|
controller.abort(new Error("cancelled by user"));
|
|
await expect(run).rejects.toThrow("cancelled by user");
|
|
expect(emulator.deleted).toEqual(["prompt-1"]);
|
|
expect(emulator.interrupted).toEqual(["prompt-1"]);
|
|
});
|
|
|
|
it("reports live step progress in order, short of done until the outputs are in", async () => {
|
|
const { provider, emulator } = setup();
|
|
emulator.pendingPolls = 4;
|
|
const progress: number[] = [];
|
|
await provider.generateImage({ ...IMAGE, onProgress: (fraction) => progress.push(fraction) });
|
|
expect(progress).toEqual([0.25, 0.5, 0.75, 0.99, 1]);
|
|
const [socket] = emulator.sockets;
|
|
expect(socket?.url.href).toBe(
|
|
`ws://127.0.0.1:8188/ws?clientId=${emulator.prompts[0]?.clientId}`,
|
|
);
|
|
expect(socket?.closedByClient).toBe(true);
|
|
});
|
|
|
|
it("maps progress across the prompts of a job", async () => {
|
|
const { provider, emulator } = setup();
|
|
emulator.pendingPolls = 2;
|
|
const progress: number[] = [];
|
|
await provider.generateImage({
|
|
...IMAGE,
|
|
count: 2,
|
|
onProgress: (fraction) => progress.push(fraction),
|
|
});
|
|
expect(rounded(progress)).toEqual([0.25, 0.495, 0.5, 0.75, 0.995, 1]);
|
|
expect(emulator.sockets).toHaveLength(1);
|
|
});
|
|
|
|
it("sends the key and keeps the path prefix on the websocket", async () => {
|
|
const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" });
|
|
const { provider } = setup(
|
|
{ baseUrl: "https://127.0.0.1:8188/comfy", apiKey: "comfy-key" },
|
|
emulator,
|
|
);
|
|
const progress: number[] = [];
|
|
await provider.generateImage({ ...IMAGE, onProgress: (fraction) => progress.push(fraction) });
|
|
const socket = emulator.sockets[0]!;
|
|
expect(`${socket.url.origin}${socket.url.pathname}`).toBe("wss://127.0.0.1:8188/comfy/ws");
|
|
expect(socket.authorization).toBe("Bearer comfy-key");
|
|
expect(progress).toEqual([0.99, 1]);
|
|
});
|
|
|
|
it("generates as before when the websocket is unavailable", async () => {
|
|
const emulator = new ComfyUiEmulator();
|
|
emulator.webSocketAvailable = false;
|
|
const progress: number[] = [];
|
|
const onProgress = (fraction: number) => progress.push(fraction);
|
|
await setup({}, emulator).provider.generateImage({ ...IMAGE, onProgress });
|
|
const throwing = setup({
|
|
webSocket: () => {
|
|
throw new Error("no websocket");
|
|
},
|
|
});
|
|
await throwing.provider.generateImage({ ...IMAGE, onProgress });
|
|
// A substituted fetch without a matching socket opens none, rather than a real connection.
|
|
const fetchOnly = new ComfyUiEmulator();
|
|
await setup({ webSocket: undefined }, fetchOnly).provider.generateImage({
|
|
...IMAGE,
|
|
onProgress,
|
|
});
|
|
expect(progress).toEqual([1, 1, 1]);
|
|
expect(emulator.sockets.map((socket) => socket.readyState)).toEqual([3]);
|
|
expect(fetchOnly.sockets).toEqual([]);
|
|
});
|
|
|
|
it("ignores stray messages, a socket that drops mid-run and a failing listener", async () => {
|
|
const { provider, emulator } = setup();
|
|
emulator.pendingPolls = 3;
|
|
const progress: number[] = [];
|
|
await provider.generateImage({
|
|
...IMAGE,
|
|
onProgress: (fraction) => {
|
|
progress.push(fraction);
|
|
const socket = emulator.sockets[0]!;
|
|
socket.push("not json");
|
|
socket.push({ type: "progress", data: { value: "2", max: 3, prompt_id: "prompt-1" } });
|
|
socket.push({ type: "progress", data: { value: 2, max: 0, prompt_id: "prompt-1" } });
|
|
socket.push({ type: "progress", data: { value: 3, max: 3, prompt_id: "prompt-9" } });
|
|
socket.drop();
|
|
},
|
|
});
|
|
expect(rounded(progress)).toEqual([0.3333, 1]);
|
|
const failing = setup().provider;
|
|
await expect(
|
|
failing.generateImage({
|
|
...IMAGE,
|
|
onProgress: () => {
|
|
throw new Error("listener failed");
|
|
},
|
|
}),
|
|
).resolves.toHaveLength(1);
|
|
});
|
|
|
|
it("closes the websocket when the caller cancels or the run fails", async () => {
|
|
const emulator = new ComfyUiEmulator();
|
|
emulator.pendingPolls = 1_000_000;
|
|
const { provider } = setup({ pollIntervalMs: 5 }, emulator);
|
|
const controller = new AbortController();
|
|
const onProgress = () => undefined;
|
|
const run = provider.generateImage({ ...IMAGE, signal: controller.signal, onProgress });
|
|
while (!emulator.requests.some((request) => request.path.startsWith("history/"))) {
|
|
await new Promise((resolve) => setTimeout(resolve, 1));
|
|
}
|
|
expect(emulator.sockets[0]?.readyState).toBe(1);
|
|
controller.abort(new Error("cancelled by user"));
|
|
expect(emulator.sockets[0]?.closedByClient).toBe(true);
|
|
await expect(run).rejects.toThrow("cancelled by user");
|
|
emulator.pendingPolls = 1;
|
|
emulator.failNextExecution = "CUDA out of memory";
|
|
await expect(provider.generateImage({ ...IMAGE, onProgress })).rejects.toThrow(
|
|
"CUDA out of memory",
|
|
);
|
|
expect(emulator.sockets.map((socket) => socket.readyState)).toEqual([3, 3]);
|
|
});
|
|
|
|
it("reports GPU memory other processes leave to ComfyUI", async () => {
|
|
const emulator = new ComfyUiEmulator();
|
|
emulator.gpu = { totalMb: 12_288, freeMb: 3_000, torchReservedMb: 7_000, torchFreeMb: 1_000 };
|
|
await expect(setup({}, emulator).provider.gpuStatus()).resolves.toEqual({
|
|
vramFreeMb: 9_000,
|
|
vramTotalMb: 12_288,
|
|
});
|
|
emulator.gpu = null;
|
|
await expect(setup({}, emulator).provider.gpuStatus()).resolves.toBeUndefined();
|
|
expect(parseGpuStatus({ devices: [{ type: "cuda", vram_total: 0, vram_free: 0 }] })).toBe(
|
|
undefined,
|
|
);
|
|
});
|
|
|
|
it("names interrupted runs", () => {
|
|
expect(describeExecutionError([["execution_interrupted", {}]])).toBe(
|
|
"ComfyUI interrupted the run",
|
|
);
|
|
expect(describeExecutionError(undefined)).toBe("ComfyUI failed");
|
|
});
|
|
});
|
|
|
|
async function sourcePng() {
|
|
const png = await sharp({
|
|
create: { width: 32, height: 32, channels: 3, background: { r: 1, g: 2, b: 3 } },
|
|
})
|
|
.png()
|
|
.toBuffer();
|
|
return { bytes: new Uint8Array(png), mimeType: "image/png" };
|
|
}
|