422 lines
16 KiB
TypeScript
422 lines
16 KiB
TypeScript
import { z } from "zod";
|
|
import { selectPreset, presets, defaultPromptPrefix } from "./presets.ts";
|
|
import { test } from "node:test";
|
|
import assert from "node:assert/strict";
|
|
import { spawn } from "node:child_process";
|
|
import { createHash } from "node:crypto";
|
|
import { mkdtemp, mkdir, writeFile, readFile, readdir, rm } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { basename, join, resolve } from "node:path";
|
|
import { setTimeout as delay } from "node:timers/promises";
|
|
|
|
const payloadSchema = z.object({
|
|
seed: z.number(),
|
|
steps: z.number(),
|
|
cfg_scale: z.number(),
|
|
sampler_name: z.string(),
|
|
width: z.number(),
|
|
height: z.number(),
|
|
scheduler: z.string(),
|
|
negative_prompt: z.string(),
|
|
prompt: z.string(),
|
|
});
|
|
const metadataSchema = z.object({
|
|
prompt: z.string(),
|
|
inputPrompt: z.string(),
|
|
settings: z.object({
|
|
seed: z.number(),
|
|
cache: z.object({ mode: z.literal("ucache"), options: z.string() }).nullable(),
|
|
preset: z.string(),
|
|
decoder: z.string(),
|
|
decoderSha256: z.string().nullable(),
|
|
promptPrefix: z.string(),
|
|
}),
|
|
});
|
|
|
|
const generations = resolve("public/assets/generations");
|
|
const metadataDirectory = resolve(".cache/generations");
|
|
const testName = (tools: string) => `test-image-${basename(tools)}`;
|
|
const metadataFolder = (filename: string) =>
|
|
join(metadataDirectory, createHash("sha256").update(filename).digest("hex"));
|
|
const metadataPath = (output: string) =>
|
|
join(metadataFolder(basename(output)), `${basename(output)}.json`);
|
|
const script = resolve("scripts/image-generation/generate.ts");
|
|
const png =
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4f+A0AAVMAotDMHBLAAAAAElFTkSuQmCC";
|
|
async function fixture(mode = "success") {
|
|
const tools = await mkdtemp(join(tmpdir(), "image-runner-"));
|
|
await mkdir(join(tools, "stable-diffusion.cpp/build/bin"), { recursive: true });
|
|
await mkdir(join(tools, "models"));
|
|
const model = `test-model-${tools}`;
|
|
await writeFile(join(tools, "models/dreamshaper_8_lcm.safetensors"), model);
|
|
await writeFile(join(tools, "models/cyberrealistic_v5_fp16.safetensors"), model);
|
|
await writeFile(join(tools, "models/juggernaut_xi_lightning.safetensors"), model);
|
|
await writeFile(join(tools, "models/taesd.safetensors"), "test-decoder");
|
|
await writeFile(join(tools, "models/realisticVision_v51_hyper.safetensors"), model);
|
|
const binary = join(tools, "stable-diffusion.cpp/build/bin/sd-server");
|
|
await writeFile(
|
|
binary,
|
|
`#!/usr/bin/env node
|
|
const http=require('node:http'),fs=require('node:fs');
|
|
const port=Number(process.argv[process.argv.indexOf('--listen-port')+1]);
|
|
fs.writeFileSync(${JSON.stringify(join(tools, "args.json"))},JSON.stringify(process.argv));
|
|
const server=http.createServer(async(req,res)=>{
|
|
if(req.method==='GET'){res.setHeader('Content-Type','application/json');res.end(JSON.stringify({sd_model_checkpoint:'dreamshaper_8_lcm'}));return;}
|
|
let body='';for await(const part of req){body+=part.toString();}
|
|
fs.writeFileSync(${JSON.stringify(join(tools, "payload.json"))},body);
|
|
fs.writeFileSync(${JSON.stringify(join(tools, "requested"))},'yes');
|
|
if(${JSON.stringify(mode)}==='slow')await new Promise(r=>setTimeout(r,60000));
|
|
res.statusCode=${mode === "error" ? 503 : 200};res.setHeader('Content-Type','application/json');res.end(JSON.stringify({images:[${JSON.stringify(png)}]}));
|
|
});
|
|
server.listen(port,'127.0.0.1',()=>{fs.writeFileSync(${JSON.stringify(join(tools, "pid"))},String(process.pid));fs.writeFileSync(${JSON.stringify(join(tools, "port"))},String(server.address().port));});
|
|
`,
|
|
{ mode: 0o755 },
|
|
);
|
|
const cleanup = async () => {
|
|
for (const entry of await readdir(generations).catch(() => [])) {
|
|
if (entry.startsWith(`${testName(tools)}.`))
|
|
await rm(join(generations, entry), { force: true });
|
|
}
|
|
for (const entry of await readdir(metadataDirectory).catch(() => [])) {
|
|
const folder = join(metadataDirectory, entry);
|
|
const files = await readdir(folder);
|
|
if (files.some((file) => file.startsWith(`${testName(tools)}.`)))
|
|
await rm(folder, { recursive: true });
|
|
}
|
|
await rm(tools, { recursive: true, force: true });
|
|
};
|
|
return { tools, binary, cleanup };
|
|
}
|
|
function run(tools: string, extra: string[] = []) {
|
|
const child = spawn(process.execPath, [
|
|
script,
|
|
"--tools",
|
|
tools,
|
|
"--prompt",
|
|
"pink ink",
|
|
"--name",
|
|
testName(tools),
|
|
...extra,
|
|
]);
|
|
let stdout = "",
|
|
stderr = "";
|
|
child.stdout.on("data", (d: Buffer) => {
|
|
stdout += d.toString();
|
|
});
|
|
child.stderr.on("data", (d: Buffer) => {
|
|
stderr += d.toString();
|
|
});
|
|
const done = new Promise<{ code: number | null; stdout: string; stderr: string }>(
|
|
(done, reject) => {
|
|
child.once("error", reject);
|
|
child.once("exit", (code) => done({ code, stdout, stderr }));
|
|
},
|
|
);
|
|
return { child, done };
|
|
}
|
|
async function waitFile(path: string) {
|
|
for (let i = 0; i < 100; i++) {
|
|
try {
|
|
return await readFile(path, "utf8");
|
|
} catch {
|
|
await delay(50);
|
|
}
|
|
}
|
|
throw new Error(`Timed out waiting for ${path}`);
|
|
}
|
|
function gone(pid: number) {
|
|
assert.throws(() => process.kill(pid, 0), { code: "ESRCH" });
|
|
}
|
|
|
|
test("generates flat public WebP names with optional prompt hashes, rejects overwrite, stops owned server", async () => {
|
|
const f = await fixture();
|
|
try {
|
|
const result = await run(f.tools).done;
|
|
assert.equal(result.code, 0, result.stderr);
|
|
const output = result.stdout.split("\n").find((x) => x.endsWith(".webp"));
|
|
assert.ok(output);
|
|
assert.equal(output, join(generations, `${testName(f.tools)}.dreamshaper-8-lcm.webp`));
|
|
assert.equal((await readFile(output)).subarray(8, 12).toString(), "WEBP");
|
|
const initialMetadata = metadataSchema.parse(
|
|
JSON.parse(await readFile(metadataPath(output), "utf8")),
|
|
);
|
|
const initialPayload = payloadSchema.parse(
|
|
JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8")),
|
|
);
|
|
assert.ok(
|
|
Number.isInteger(initialPayload.seed) &&
|
|
initialPayload.seed >= 0 &&
|
|
initialPayload.seed <= 2147483647,
|
|
);
|
|
assert.equal(initialMetadata.settings.seed, initialPayload.seed);
|
|
gone(Number(await readFile(join(f.tools, "pid"), "utf8")));
|
|
const second = await run(f.tools, ["--prompt", "blue ink", "--with-hash"]).done;
|
|
const secondOutput = second.stdout.split("\n").find((x) => x.endsWith(".webp"));
|
|
assert.ok(secondOutput);
|
|
const secondMetadata = metadataSchema.parse(
|
|
JSON.parse(await readFile(metadataPath(secondOutput), "utf8")),
|
|
);
|
|
const effectiveSecondPrompt = `${secondMetadata.settings.promptPrefix} blue ink`;
|
|
assert.equal(secondMetadata.prompt, effectiveSecondPrompt);
|
|
assert.equal(secondMetadata.inputPrompt, "blue ink");
|
|
assert.equal(
|
|
secondOutput,
|
|
join(
|
|
generations,
|
|
`${testName(f.tools)}.dreamshaper-8-lcm.${createHash("sha256").update(effectiveSecondPrompt).digest("hex").slice(0, 8)}.webp`,
|
|
),
|
|
);
|
|
assert.equal(second.code, 0, second.stderr);
|
|
assert.equal(
|
|
secondOutput.split("/").slice(0, -1).join("/"),
|
|
output.split("/").slice(0, -1).join("/"),
|
|
);
|
|
const duplicate = await run(f.tools).done;
|
|
assert.equal(duplicate.code, 1);
|
|
assert.match(duplicate.stderr, /already exists/);
|
|
await writeFile(output, "old image");
|
|
const forced = await run(f.tools, ["--force", "--prompt", "replacement ink", "--seed", "73"])
|
|
.done;
|
|
assert.equal(forced.code, 0, forced.stderr);
|
|
assert.equal(
|
|
forced.stdout.split("\n").find((x) => x.endsWith(".webp")),
|
|
output,
|
|
);
|
|
assert.equal((await readFile(output)).subarray(8, 12).toString(), "WEBP");
|
|
const forcedMetadata = metadataSchema.parse(
|
|
JSON.parse(await readFile(metadataPath(output), "utf8")),
|
|
);
|
|
assert.equal(forcedMetadata.prompt, `${forcedMetadata.settings.promptPrefix} replacement ink`);
|
|
assert.equal(forcedMetadata.inputPrompt, "replacement ink");
|
|
assert.equal(
|
|
metadataSchema.parse(JSON.parse(await readFile(metadataPath(output), "utf8"))).settings.seed,
|
|
73,
|
|
);
|
|
assert.equal(
|
|
payloadSchema.parse(JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8"))).seed,
|
|
73,
|
|
);
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("forced generation failure preserves the existing image", async () => {
|
|
const f = await fixture("error");
|
|
const output = join(generations, `${testName(f.tools)}.dreamshaper-8-lcm.webp`);
|
|
try {
|
|
await mkdir(generations, { recursive: true });
|
|
await writeFile(output, "previous image");
|
|
const result = await run(f.tools, ["--force"]).done;
|
|
assert.equal(result.code, 1);
|
|
assert.match(result.stderr, /HTTP 503/);
|
|
assert.equal(await readFile(output, "utf8"), "previous image");
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("generation failure stops owned server", async () => {
|
|
const f = await fixture("error");
|
|
try {
|
|
const result = await run(f.tools).done;
|
|
assert.equal(result.code, 1);
|
|
assert.match(result.stderr, /HTTP 503/);
|
|
gone(Number(await readFile(join(f.tools, "pid"), "utf8")));
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("Ctrl+C during generation stops owned server and removes reservation", async () => {
|
|
const f = await fixture("slow");
|
|
try {
|
|
const job = run(f.tools);
|
|
await waitFile(join(f.tools, "requested"));
|
|
job.child.kill("SIGINT");
|
|
const result = await job.done;
|
|
assert.equal(result.code, 130);
|
|
gone(Number(await readFile(join(f.tools, "pid"), "utf8")));
|
|
assert.equal(
|
|
(await readdir(metadataFolder(`${testName(f.tools)}.dreamshaper-8-lcm.webp`))).some((entry) =>
|
|
entry.endsWith(".lock"),
|
|
),
|
|
false,
|
|
);
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("rejects filename traversal before starting server", async () => {
|
|
const f = await fixture();
|
|
try {
|
|
const result = await run(f.tools, ["--name", "../escape"]).done;
|
|
assert.equal(result.code, 1);
|
|
assert.match(result.stderr, /filename/);
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("reuses an external server without shutting it down", async () => {
|
|
const f = await fixture();
|
|
const server = spawn(f.binary, ["--listen-port", "0"]);
|
|
try {
|
|
const port = await waitFile(join(f.tools, "port"));
|
|
const result = await run(f.tools, ["--server", `http://127.0.0.1:${port}`]).done;
|
|
assert.equal(result.code, 0, result.stderr);
|
|
assert.match(result.stdout, /Reusing/);
|
|
assert.ok(server.pid);
|
|
process.kill(server.pid, 0);
|
|
} finally {
|
|
server.kill("SIGTERM");
|
|
await new Promise<void>((done) => server.once("exit", () => done()));
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("DreamShaper is default; each preset carries its own model, sampling and cache configuration", async () => {
|
|
assert.equal(selectPreset().id, "dreamshaper-8-lcm");
|
|
assert.equal(selectPreset("CyberRealistic").config.cfgScale, 7);
|
|
assert.equal(selectPreset("dreamshaper-8-lcm").config.cache, null);
|
|
assert.equal(selectPreset("cyberrealistic-no-cache").config.cache, null);
|
|
assert.throws(() => selectPreset("missing"), /Available/);
|
|
const f = await fixture();
|
|
try {
|
|
const result = await run(f.tools, ["--preset", "CyberRealistic"]).done;
|
|
assert.equal(result.code, 0, result.stderr);
|
|
const args = z
|
|
.array(z.string())
|
|
.parse(JSON.parse(await readFile(join(f.tools, "args.json"), "utf8")));
|
|
assert.equal(
|
|
args[args.indexOf("-m") + 1],
|
|
join(f.tools, "models/cyberrealistic_v5_fp16.safetensors"),
|
|
);
|
|
assert.equal(args[args.indexOf("--cache-mode") + 1], "ucache");
|
|
assert.equal(args[args.indexOf("--taesd") + 1], join(f.tools, "models/taesd.safetensors"));
|
|
const payload = payloadSchema.parse(
|
|
JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8")),
|
|
);
|
|
assert.equal(payload.steps, 20);
|
|
assert.equal(payload.cfg_scale, 7);
|
|
assert.equal(payload.sampler_name, "dpm++2m");
|
|
const output = result.stdout.split("\n").find((x) => x.endsWith(".webp"));
|
|
assert.ok(output);
|
|
assert.equal(output, join(generations, `${testName(f.tools)}.cyberrealistic.webp`));
|
|
const settings = metadataSchema.parse(
|
|
JSON.parse(await readFile(metadataPath(output), "utf8")),
|
|
).settings;
|
|
assert.ok(settings.cache);
|
|
assert.equal(settings.cache.mode, "ucache");
|
|
assert.equal(settings.preset, "cyberrealistic");
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("Realistic Vision uses four-step settings and hashes negative-prompt overrides", async () => {
|
|
const f = await fixture();
|
|
try {
|
|
const first = await run(f.tools, ["--preset", "realistic-vision-hyper"]).done;
|
|
assert.equal(first.code, 0, first.stderr);
|
|
const payload = payloadSchema.parse(
|
|
JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8")),
|
|
);
|
|
assert.equal(payload.steps, 4);
|
|
assert.equal(payload.cfg_scale, 1.8);
|
|
assert.equal(payload.sampler_name, "dpm++2m");
|
|
assert.equal(payload.scheduler, "discrete");
|
|
assert.equal(payload.negative_prompt, "ugly, deformed, bad anatomy, blurry, cartoon");
|
|
const args = z
|
|
.array(z.string())
|
|
.parse(JSON.parse(await readFile(join(f.tools, "args.json"), "utf8")));
|
|
assert.equal(
|
|
args[args.indexOf("-m") + 1],
|
|
join(f.tools, "models/realisticVision_v51_hyper.safetensors"),
|
|
);
|
|
assert.equal(args.includes("--cache-mode"), false);
|
|
const second = await run(f.tools, [
|
|
"--preset",
|
|
"realistic-vision-hyper",
|
|
"--negative-prompt",
|
|
"",
|
|
"--name",
|
|
`${testName(f.tools)}.negative-override`,
|
|
]).done;
|
|
assert.equal(second.code, 0, second.stderr);
|
|
assert.equal(
|
|
payloadSchema.parse(JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8")))
|
|
.negative_prompt,
|
|
"",
|
|
);
|
|
assert.notEqual(
|
|
first.stdout.split("\n").find((x) => x.endsWith(".webp")),
|
|
second.stdout.split("\n").find((x) => x.endsWith(".webp")),
|
|
);
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("Juggernaut uses SDXL dimensions and embedded VAE; removed presets are rejected", async () => {
|
|
assert.throws(() => selectPreset("sd-turbo"), /Unknown preset/);
|
|
const f = await fixture();
|
|
try {
|
|
const result = await run(f.tools, ["--preset", "juggernaut-xi-lightning"]).done;
|
|
assert.equal(result.code, 0, result.stderr);
|
|
const args = z
|
|
.array(z.string())
|
|
.parse(JSON.parse(await readFile(join(f.tools, "args.json"), "utf8")));
|
|
assert.equal(args.includes("--taesd"), false);
|
|
const payload = payloadSchema.parse(
|
|
JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8")),
|
|
);
|
|
assert.equal(payload.width, 1024);
|
|
assert.equal(payload.height, 1024);
|
|
assert.equal(payload.steps, 6);
|
|
assert.equal(payload.sampler_name, "dpm++2m_sde");
|
|
const output = result.stdout.split("\n").find((x) => x.endsWith(".webp"));
|
|
assert.ok(output);
|
|
const settings = metadataSchema.parse(
|
|
JSON.parse(await readFile(metadataPath(output), "utf8")),
|
|
).settings;
|
|
assert.equal(settings.decoder, "embedded-vae");
|
|
assert.equal(settings.decoderSha256, null);
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|
|
|
|
test("shared prompt prefix applies to all presets; CLI overrides or disables it and hashes the effective prompt", async () => {
|
|
for (const preset of Object.values(presets))
|
|
assert.equal(preset.promptPrefix, defaultPromptPrefix);
|
|
const f = await fixture();
|
|
try {
|
|
for (const prefix of ["watercolor illustration,", ""]) {
|
|
const result = await run(f.tools, ["--prompt-prefix", prefix, "--with-hash"]).done;
|
|
assert.equal(result.code, 0, result.stderr);
|
|
const prompt = prefix ? `${prefix} pink ink` : "pink ink";
|
|
const output = result.stdout.split("\n").find((x) => x.endsWith(".webp"));
|
|
assert.ok(output);
|
|
const expectedHash = createHash("sha256").update(prompt).digest("hex").slice(0, 8);
|
|
assert.equal(
|
|
output,
|
|
join(generations, `${testName(f.tools)}.dreamshaper-8-lcm.${expectedHash}.webp`),
|
|
);
|
|
assert.equal(
|
|
payloadSchema.parse(JSON.parse(await readFile(join(f.tools, "payload.json"), "utf8")))
|
|
.prompt,
|
|
prompt,
|
|
);
|
|
const metadata = metadataSchema.parse(
|
|
JSON.parse(await readFile(metadataPath(output), "utf8")),
|
|
);
|
|
assert.equal(metadata.prompt, prompt);
|
|
assert.equal(metadata.inputPrompt, "pink ink");
|
|
assert.equal(metadata.settings.promptPrefix, prefix);
|
|
}
|
|
} finally {
|
|
await f.cleanup();
|
|
}
|
|
});
|