407 lines
16 KiB
TypeScript
407 lines
16 KiB
TypeScript
import { z } from "zod";
|
||
import { presets, defaultPreset, selectPreset } from "./presets.ts";
|
||
import { spawn, type ChildProcess } from "node:child_process";
|
||
import { createHash, randomInt } from "node:crypto";
|
||
import { createReadStream } from "node:fs";
|
||
import { access, mkdir, readFile, rename, rm, stat, writeFile } from "node:fs/promises";
|
||
import { createServer } from "node:net";
|
||
import { homedir } from "node:os";
|
||
import { basename, join, resolve } from "node:path";
|
||
import { fileURLToPath } from "node:url";
|
||
import { parseArgs } from "node:util";
|
||
import { setTimeout as delay } from "node:timers/promises";
|
||
|
||
const root = fileURLToPath(new URL("../../", import.meta.url));
|
||
const controller = new AbortController();
|
||
const children = new Set<ChildProcess>();
|
||
const hash = (value: string | Buffer) => createHash("sha256").update(value).digest("hex");
|
||
|
||
const isRunning = (child: ChildProcess) => child.exitCode === null && child.signalCode === null;
|
||
|
||
async function stop(child: ChildProcess) {
|
||
if (!isRunning(child)) return;
|
||
child.kill("SIGTERM");
|
||
await Promise.race([
|
||
new Promise<void>((done) => child.once("exit", () => done())),
|
||
delay(3000, undefined, { ref: false }),
|
||
]);
|
||
if (isRunning(child)) {
|
||
child.kill("SIGKILL");
|
||
await new Promise<void>((done) => child.once("exit", () => done()));
|
||
}
|
||
}
|
||
for (const signal of ["SIGINT", "SIGTERM", "SIGHUP"] as const) {
|
||
process.once(signal, () => {
|
||
process.exitCode = signal === "SIGINT" ? 130 : 143;
|
||
controller.abort(new Error(`Interrupted by ${signal}`));
|
||
for (const child of children) child.kill("SIGTERM");
|
||
});
|
||
}
|
||
|
||
function launch(command: string, args: string[], onOutput: (chunk: string) => void) {
|
||
controller.signal.throwIfAborted();
|
||
const child = spawn(command, args, { stdio: ["ignore", "pipe", "pipe"] });
|
||
children.add(child);
|
||
const completed = new Promise<void>((done, reject) => {
|
||
child.once("error", reject);
|
||
child.once("exit", (code, signal) => {
|
||
children.delete(child);
|
||
if (code === 0) done();
|
||
else reject(new Error(`${basename(command)} exited ${code ?? signal}`));
|
||
});
|
||
});
|
||
// Startup/generation checks consume errors later; avoid an unhandled rejection in between.
|
||
void completed.catch(() => {});
|
||
child.stdout.on("data", (chunk: Buffer) => onOutput(chunk.toString()));
|
||
child.stderr.on("data", (chunk: Buffer) => onOutput(chunk.toString()));
|
||
return { child, completed };
|
||
}
|
||
|
||
async function fingerprint(path: string) {
|
||
const info = await stat(path);
|
||
const key = hash(JSON.stringify([path, info.size, info.mtimeMs, info.ctimeMs]));
|
||
const cache = join(root, ".cache/image-generation", `${key}.sha256`);
|
||
try {
|
||
const cached = await readFile(cache, "utf8");
|
||
if (/^[a-f0-9]{64}$/.test(cached)) return cached;
|
||
} catch {
|
||
/* First use hashes the file without loading it into memory. */
|
||
}
|
||
const digest = createHash("sha256");
|
||
for await (const chunk of createReadStream(path, { signal: controller.signal })) {
|
||
if (!Buffer.isBuffer(chunk)) throw new Error("Expected binary file data.");
|
||
digest.update(chunk);
|
||
}
|
||
const result = digest.digest("hex");
|
||
await mkdir(join(root, ".cache/image-generation"), { recursive: true });
|
||
await writeFile(cache, result);
|
||
return result;
|
||
}
|
||
|
||
async function freePort() {
|
||
const server = createServer();
|
||
await new Promise<void>((done, reject) => {
|
||
server.once("error", reject);
|
||
server.listen(0, "127.0.0.1", done);
|
||
});
|
||
const address = server.address();
|
||
if (!address || typeof address === "string") throw new Error("Could not allocate a local port");
|
||
await new Promise<void>((done, reject) =>
|
||
server.close((error) => (error ? reject(error) : done())),
|
||
);
|
||
return address.port;
|
||
}
|
||
|
||
async function jsonRequest(url: string, payload?: unknown, timeout = 240_000): Promise<unknown> {
|
||
const response = await fetch(url, {
|
||
...(payload === undefined ? {} : { method: "POST", body: JSON.stringify(payload) }),
|
||
headers: { "Content-Type": "application/json" },
|
||
signal: AbortSignal.any([controller.signal, AbortSignal.timeout(timeout)]),
|
||
});
|
||
if (!response.ok)
|
||
throw new Error(
|
||
`Image server returned HTTP ${response.status}: ${(await response.text()).slice(0, 300)}`,
|
||
);
|
||
return response.json();
|
||
}
|
||
|
||
function number(
|
||
value: string | undefined,
|
||
fallback: number,
|
||
min: number,
|
||
max: number,
|
||
integer = true,
|
||
) {
|
||
const result = value === undefined ? fallback : Number(value);
|
||
if (
|
||
!Number.isFinite(result) ||
|
||
result < min ||
|
||
result > max ||
|
||
(integer && !Number.isInteger(result))
|
||
) {
|
||
throw new Error(
|
||
`Invalid number ${value}; expected ${min}–${max}${integer ? " (integer)" : ""}`,
|
||
);
|
||
}
|
||
return result;
|
||
}
|
||
|
||
async function main() {
|
||
const { values: v } = parseArgs({
|
||
options: {
|
||
prompt: { type: "string" },
|
||
"prompt-prefix": { type: "string" },
|
||
"negative-prompt": { type: "string" },
|
||
preset: { type: "string" },
|
||
"list-presets": { type: "boolean" },
|
||
name: { type: "string" },
|
||
"with-hash": { type: "boolean" },
|
||
force: { type: "boolean" },
|
||
reference: { type: "string" },
|
||
seed: { type: "string" },
|
||
steps: { type: "string" },
|
||
strength: { type: "string" },
|
||
width: { type: "string" },
|
||
height: { type: "string" },
|
||
tools: { type: "string" },
|
||
server: { type: "string" },
|
||
help: { type: "boolean" },
|
||
},
|
||
});
|
||
if (v["list-presets"]) {
|
||
for (const [id, config] of Object.entries(presets))
|
||
console.log(
|
||
`${id}${id === defaultPreset ? " (default)" : ""}: ${config.checkpoint}, ${config.steps} steps, CFG ${config.cfgScale}, ${config.sampler}/${config.scheduler}, cache ${config.cache?.mode ?? "off"}`,
|
||
);
|
||
return;
|
||
}
|
||
if (v.help) {
|
||
console.log(
|
||
'vp run generate:image --prompt "..." --name "cybelle"\nPreset: --preset dreamshaper-8-lcm (default), cyberrealistic, cyberrealistic-no-cache, realistic-vision-hyper, juggernaut-xi-lightning. Use --list-presets.\nOutput: public/assets/generations/<name>.<preset>.webp; --with-hash adds the prompt hash before .webp.\n--force overwrites an existing output after successful generation.\nSeed is random by default; --seed fixes it for reproducible generation.\nOptional: --prompt-prefix "..." (overrides shared prefix; empty disables it) --negative-prompt "..." --reference FILE --seed 73 --steps 6 --strength 0.5 --width 512 --height 896\n--tools DIR overrides ~/code/image-generation-benchmark.\n--server http://127.0.0.1:PORT reuses an existing matching-model server (left running).\nDefault: starts a private temporary server and closes it on completion, error or Ctrl+C.',
|
||
);
|
||
return;
|
||
}
|
||
if (!v.prompt?.trim() || !v.name)
|
||
throw new Error("--prompt and --name are required. Run with --help for usage.");
|
||
const name = v.name.replace(/\.webp$/i, "");
|
||
if (!/^[\p{L}\p{N}][\p{L}\p{N} _.-]*$/u.test(name) || name.length > 120 || name.includes("..")) {
|
||
throw new Error(
|
||
"--name must be a filename, not a path (letters, numbers, spaces, dots, _ and -; max 120 characters).",
|
||
);
|
||
}
|
||
const preset = selectPreset(v.preset);
|
||
const promptPrefix = (v["prompt-prefix"] ?? preset.config.promptPrefix).trim();
|
||
const prompt = promptPrefix ? `${promptPrefix} ${v.prompt}` : v.prompt;
|
||
const tools = resolve(v.tools ?? join(homedir(), "code/image-generation-benchmark"));
|
||
const binary = join(tools, "stable-diffusion.cpp/build/bin/sd-server");
|
||
const model = join(tools, preset.config.checkpoint);
|
||
const decoder = preset.config.decoder ? join(tools, preset.config.decoder) : null;
|
||
for (const path of [binary, model, ...(decoder ? [decoder] : [])]) {
|
||
try {
|
||
await access(path);
|
||
} catch {
|
||
throw new Error(
|
||
`Missing local generator file: ${path}. Use --tools to select the existing installation.`,
|
||
);
|
||
}
|
||
}
|
||
const width = number(v.width, preset.config.width, 64, 1024);
|
||
const height = number(v.height, preset.config.height, 64, 1024);
|
||
if (width % 64 || height % 64) throw new Error("Width and height must be multiples of 64.");
|
||
const reference = v.reference ? await readFile(resolve(v.reference)) : undefined;
|
||
const settings = {
|
||
model: preset.config.model,
|
||
preset: preset.id,
|
||
cache: preset.config.cache,
|
||
modelSha256: await fingerprint(model),
|
||
decoder: decoder ? "taesd" : "embedded-vae",
|
||
decoderSha256: decoder ? await fingerprint(decoder) : null,
|
||
engine: "stable-diffusion.cpp",
|
||
engineSha256: await fingerprint(binary),
|
||
precision: "f16",
|
||
seed: v.seed === undefined ? randomInt(0, 2147483648) : number(v.seed, 0, 0, 2147483647),
|
||
steps: number(v.steps, reference ? preset.config.referenceSteps : preset.config.steps, 1, 50),
|
||
width,
|
||
height,
|
||
sampler: preset.config.sampler,
|
||
scheduler: preset.config.scheduler,
|
||
cfgScale: preset.config.cfgScale,
|
||
promptPrefix,
|
||
negativePrompt: v["negative-prompt"] ?? preset.config.negativePrompt,
|
||
strength: reference ? number(v.strength, 0.5, 0.01, 1, false) : null,
|
||
referenceSha256: reference ? hash(reference) : null,
|
||
prediction: preset.config.prediction,
|
||
diffusionFlashAttention: true,
|
||
threads: 4,
|
||
webp: { quality: 90, method: 6 },
|
||
serverMode: v.server ? "external-config-not-fully-verifiable" : "owned",
|
||
};
|
||
if (v.strength && !reference) throw new Error("--strength requires --reference.");
|
||
const settingsHash = hash(JSON.stringify(settings));
|
||
const directory = join(root, "public/assets/generations");
|
||
await mkdir(directory, { recursive: true });
|
||
const promptHash = hash(prompt).slice(0, 8);
|
||
const filename = `${name}.${preset.id}${v["with-hash"] ? `.${promptHash}` : ""}.webp`;
|
||
const output = join(directory, filename);
|
||
const metadataDirectory = join(root, ".cache/generations", hash(filename));
|
||
await mkdir(metadataDirectory, { recursive: true });
|
||
const metadata = join(metadataDirectory, filename);
|
||
const lock = `${metadata}.lock`;
|
||
// Exclusive reservation prevents concurrent runs with the same output filename from overwriting.
|
||
await writeFile(lock, String(process.pid), { flag: "wx" });
|
||
let owned: ReturnType<typeof launch> | undefined;
|
||
let serverLog = "";
|
||
const png = join(metadataDirectory, `.${name}.${process.pid}.png`);
|
||
const temporary = join(metadataDirectory, `.${name}.${process.pid}.webp`);
|
||
const started = performance.now();
|
||
try {
|
||
if (!v.force) {
|
||
try {
|
||
await access(output);
|
||
throw new Error(
|
||
`Output already exists: ${output}. Use --force to overwrite, choose another --name, or use --with-hash with a different prompt.`,
|
||
);
|
||
} catch (error) {
|
||
if (!(error instanceof Error && "code" in error && error.code === "ENOENT")) throw error;
|
||
}
|
||
}
|
||
let base: string;
|
||
if (v.server) {
|
||
const url = new URL(v.server);
|
||
if (
|
||
url.protocol !== "http:" ||
|
||
!["127.0.0.1", "localhost", "[::1]"].includes(url.hostname) ||
|
||
url.username ||
|
||
url.password ||
|
||
url.pathname !== "/" ||
|
||
url.search ||
|
||
url.hash
|
||
) {
|
||
throw new Error("--server must be a loopback HTTP origin, such as http://127.0.0.1:18792.");
|
||
}
|
||
base = url.origin;
|
||
const options = z
|
||
.object({ sd_model_checkpoint: z.string().optional() })
|
||
.parse(await jsonRequest(`${base}/sdapi/v1/options`, undefined, 3000));
|
||
if (options.sd_model_checkpoint !== basename(model, ".safetensors"))
|
||
throw new Error(`Existing server is not using ${basename(model)}.`);
|
||
console.log(
|
||
`Reusing ${base}; decoder/backend settings are controlled by that server. It will remain running.`,
|
||
);
|
||
} else {
|
||
const port = await freePort();
|
||
base = `http://127.0.0.1:${port}`;
|
||
owned = launch(
|
||
binary,
|
||
[
|
||
"-m",
|
||
model,
|
||
...(decoder ? ["--taesd", decoder] : []),
|
||
"--type",
|
||
"f16",
|
||
"--prediction",
|
||
settings.prediction,
|
||
"--sampling-method",
|
||
settings.sampler,
|
||
"--cfg-scale",
|
||
String(settings.cfgScale),
|
||
...(settings.cache
|
||
? ["--cache-mode", settings.cache.mode, "--cache-option", settings.cache.options]
|
||
: []),
|
||
"--diffusion-fa",
|
||
"-t",
|
||
"4",
|
||
"--listen-ip",
|
||
"127.0.0.1",
|
||
"--listen-port",
|
||
String(port),
|
||
],
|
||
(chunk) => {
|
||
serverLog = (serverLog + chunk).slice(-100_000);
|
||
},
|
||
);
|
||
let startupFailure: unknown;
|
||
void owned.completed.then(
|
||
() => {
|
||
startupFailure = new Error("Image server exited before becoming ready.");
|
||
},
|
||
(error: unknown) => {
|
||
startupFailure = error;
|
||
},
|
||
);
|
||
console.log(`Starting local ${preset.id} + ${settings.decoder} (${base})…`);
|
||
const deadline = Date.now() + 120_000;
|
||
for (;;) {
|
||
controller.signal.throwIfAborted();
|
||
if (startupFailure) throw startupFailure;
|
||
if (owned.child.exitCode !== null || owned.child.signalCode !== null)
|
||
throw new Error("Image server exited before becoming ready.");
|
||
try {
|
||
await jsonRequest(`${base}/sdapi/v1/options`, undefined, 1000);
|
||
break;
|
||
} catch {
|
||
controller.signal.throwIfAborted();
|
||
if (Date.now() > deadline) throw new Error("Local server startup timed out.");
|
||
await delay(250, undefined, { signal: controller.signal });
|
||
}
|
||
}
|
||
}
|
||
console.log(`Generating ${width}×${height}, seed ${settings.seed}…`);
|
||
// console.log(`Prompt ${prompt}`);
|
||
const result = z.object({ images: z.array(z.unknown()).optional() }).parse(
|
||
await jsonRequest(`${base}/sdapi/v1/${reference ? "img2img" : "txt2img"}`, {
|
||
prompt,
|
||
negative_prompt: settings.negativePrompt,
|
||
width,
|
||
height,
|
||
seed: settings.seed,
|
||
steps: settings.steps,
|
||
cfg_scale: settings.cfgScale,
|
||
batch_size: 1,
|
||
sampler_name: settings.sampler,
|
||
scheduler: settings.scheduler,
|
||
...(reference
|
||
? { init_images: [reference.toString("base64")], denoising_strength: settings.strength }
|
||
: {}),
|
||
}),
|
||
);
|
||
const encoded = result.images?.[0];
|
||
if (typeof encoded !== "string") throw new Error("Image server returned no image.");
|
||
const bytes = Buffer.from(encoded, "base64");
|
||
if (!bytes.subarray(0, 8).equals(Buffer.from([137, 80, 78, 71, 13, 10, 26, 10])))
|
||
throw new Error("Image server returned invalid PNG data.");
|
||
await writeFile(png, bytes);
|
||
let conversionLog = "";
|
||
const conversion = launch(
|
||
"cwebp",
|
||
["-quiet", "-q", "90", "-m", "6", png, "-o", temporary],
|
||
(chunk) => {
|
||
conversionLog += chunk;
|
||
},
|
||
);
|
||
try {
|
||
await conversion.completed;
|
||
} catch (error) {
|
||
throw new Error(
|
||
`WebP conversion failed (install cwebp if missing): ${String(error)} ${conversionLog}`,
|
||
);
|
||
}
|
||
controller.signal.throwIfAborted();
|
||
await rename(temporary, output);
|
||
await writeFile(
|
||
`${metadata}.json`,
|
||
JSON.stringify(
|
||
{
|
||
settings,
|
||
settingsHash,
|
||
filename,
|
||
prompt,
|
||
inputPrompt: v.prompt,
|
||
promptHash,
|
||
referencePath: v.reference ? resolve(v.reference) : null,
|
||
createdAt: new Date().toISOString(),
|
||
seconds: (performance.now() - started) / 1000,
|
||
},
|
||
null,
|
||
2,
|
||
) + "\n",
|
||
);
|
||
console.log(output);
|
||
} finally {
|
||
if (owned) {
|
||
await stop(owned.child);
|
||
console.log("Local server stopped.");
|
||
}
|
||
await Promise.all([...children].map(stop));
|
||
if (serverLog) await writeFile(`${metadata}.server.log`, serverLog);
|
||
await Promise.all([png, temporary, lock].map((path) => rm(path, { force: true })));
|
||
}
|
||
}
|
||
try {
|
||
await main();
|
||
} catch (error) {
|
||
console.error(error instanceof Error ? error.message : error);
|
||
process.exitCode ||= 1;
|
||
} finally {
|
||
await Promise.all([...children].map(stop));
|
||
}
|