204 lines
7 KiB
TypeScript
204 lines
7 KiB
TypeScript
import { mkdir, writeFile } from "node:fs/promises";
|
|
import { join } from "node:path";
|
|
import { fileURLToPath } from "node:url";
|
|
import { parseArgs } from "node:util";
|
|
import { input, search } from "@inquirer/prompts";
|
|
|
|
const root = fileURLToPath(new URL("../../", import.meta.url));
|
|
const api = "https://api.nano-gpt.com";
|
|
|
|
async function apiKey(): Promise<string> {
|
|
const direct = process.env["NANOGPT_API_KEY"]?.trim();
|
|
if (direct) return direct;
|
|
// generate.ts runs under `vp run` (env injected); plain `node script.ts`
|
|
// does not load .env, so fall back to a minimal manual parse.
|
|
const { readFile } = await import("node:fs/promises");
|
|
try {
|
|
const env = await readFile(join(root, ".env"), "utf8");
|
|
for (const line of env.split("\n")) {
|
|
const match = line.match(/^\s*NANOGPT_API_KEY\s*=\s*(.+?)\s*$/);
|
|
if (match) {
|
|
const value = match[1]!.replace(/^["']|["']$/g, "").trim();
|
|
if (value) return value;
|
|
}
|
|
}
|
|
} catch {
|
|
/* No .env file; fall through to the error below. */
|
|
}
|
|
throw new Error("Missing NANOGPT_API_KEY. Set it in .env or the environment.");
|
|
}
|
|
|
|
type ModelChoice = {
|
|
id: string;
|
|
label: string;
|
|
size: string;
|
|
estimate: number;
|
|
};
|
|
|
|
async function fetchModels(key: string): Promise<ModelChoice[]> {
|
|
const response = await fetch(`${api}/api/v1/images/models`, {
|
|
headers: { Authorization: `Bearer ${key}` },
|
|
signal: AbortSignal.timeout(30_000),
|
|
});
|
|
if (!response.ok) throw new Error(`Model list returned HTTP ${response.status}.`);
|
|
const payload = (await response.json()) as {
|
|
data?: Array<{
|
|
id: string;
|
|
name?: string;
|
|
capabilities?: { image_generation?: boolean };
|
|
pricing?: { per_image?: Record<string, number> };
|
|
}>;
|
|
};
|
|
const choices: ModelChoice[] = [];
|
|
for (const model of payload.data ?? []) {
|
|
if (!model.capabilities?.image_generation) continue;
|
|
const prices = Object.entries(model.pricing?.per_image ?? {});
|
|
if (prices.length === 0) continue;
|
|
const cheapest = prices.reduce((a, b) => (a[1] <= b[1] ? a : b));
|
|
if (cheapest[1] <= 0) continue;
|
|
choices.push({
|
|
id: model.id,
|
|
label: `${model.name ?? model.id} · $${cheapest[1].toFixed(6)}/image · ${cheapest[0]}`,
|
|
size: cheapest[0],
|
|
estimate: cheapest[1],
|
|
});
|
|
}
|
|
return choices;
|
|
}
|
|
|
|
function timestamp(): string {
|
|
const now = new Date();
|
|
const pad = (n: number) => String(n).padStart(2, "0");
|
|
return `${now.getFullYear()}${pad(now.getMonth() + 1)}${pad(now.getDate())}-${pad(now.getHours())}${pad(now.getMinutes())}${pad(now.getSeconds())}`;
|
|
}
|
|
|
|
function extension(bytes: Buffer): string {
|
|
if (bytes.subarray(0, 4).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47]))) return ".png";
|
|
if (bytes.subarray(0, 4).toString() === "RIFF" && bytes.subarray(8, 12).toString() === "WEBP")
|
|
return ".webp";
|
|
return ".jpg";
|
|
}
|
|
|
|
async function main() {
|
|
const { values } = parseArgs({
|
|
options: {
|
|
model: { type: "string" },
|
|
prompt: { type: "string" },
|
|
negative: { type: "string" },
|
|
steps: { type: "string" },
|
|
cfg: { type: "string" },
|
|
"list-models": { type: "boolean" },
|
|
help: { type: "boolean" },
|
|
},
|
|
});
|
|
if (values.help) {
|
|
console.log(
|
|
'node scripts/image-generation/nanogpt-quick.ts [--list-models]\nInteractive (default): fuzzy-search image model, then prompt/negative/steps/CFG.\nNon-interactive: --model ID --prompt "..." [--negative "..." --steps 4 --cfg 2]\nSaves to .cache/nanogpt-images/<timestamp>-<model>.<png|webp|jpg>.',
|
|
);
|
|
return;
|
|
}
|
|
const key = await apiKey();
|
|
const models = await fetchModels(key);
|
|
if (models.length === 0) throw new Error("No image-generation models with pricing found.");
|
|
if (values["list-models"]) {
|
|
for (const model of models) console.log(`${model.id}: ${model.label}`);
|
|
return;
|
|
}
|
|
|
|
let id = values.model;
|
|
let size = models.find((m) => m.id === id)?.size;
|
|
let estimate = models.find((m) => m.id === id)?.estimate;
|
|
if (!id || !size || estimate === undefined) {
|
|
if (id && !size) throw new Error(`Unknown model "${id}". Run with --list-models.`);
|
|
const picked = await search({
|
|
message: "Search image model",
|
|
source: async (term) => {
|
|
const query = (term ?? "").toLowerCase();
|
|
return models
|
|
.filter(
|
|
(m) => m.label.toLowerCase().includes(query) || m.id.toLowerCase().includes(query),
|
|
)
|
|
.map((m) => ({ name: m.label, value: m }));
|
|
},
|
|
});
|
|
id = picked.id;
|
|
size = picked.size;
|
|
estimate = picked.estimate;
|
|
}
|
|
|
|
const prompt =
|
|
values.prompt ??
|
|
(await input({
|
|
message: "Prompt",
|
|
validate: (v) => (v.trim() ? true : "Prompt is required."),
|
|
}));
|
|
const negative = values.negative ?? (await input({ message: "Negative prompt (optional)" }));
|
|
const stepsRaw =
|
|
values.steps ??
|
|
(await input({
|
|
message: "Steps",
|
|
default: "4",
|
|
validate: (v) =>
|
|
Number.isInteger(Number(v)) && Number(v) > 0 ? true : "Steps must be a positive integer.",
|
|
}));
|
|
const cfgRaw =
|
|
values.cfg ??
|
|
(await input({
|
|
message: "CFG scale",
|
|
default: "2",
|
|
validate: (v) =>
|
|
Number.isFinite(Number(v)) && Number(v) > 0 ? true : "CFG must be a positive number.",
|
|
}));
|
|
const steps = Number.parseInt(stepsRaw, 10);
|
|
const cfg = Number.parseFloat(cfgRaw);
|
|
if (!Number.isInteger(steps) || steps <= 0) throw new Error(`Invalid steps "${stepsRaw}".`);
|
|
if (!Number.isFinite(cfg) || cfg <= 0) throw new Error(`Invalid CFG "${cfgRaw}".`);
|
|
|
|
const response = await fetch(`${api}/v1/images/generations`, {
|
|
method: "POST",
|
|
headers: { Authorization: `Bearer ${key}`, "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
model: id,
|
|
prompt,
|
|
negative_prompt: negative,
|
|
num_inference_steps: steps,
|
|
guidance_scale: cfg,
|
|
size,
|
|
n: 1,
|
|
response_format: "b64_json",
|
|
}),
|
|
signal: AbortSignal.timeout(180_000),
|
|
});
|
|
if (!response.ok)
|
|
throw new Error(
|
|
`Generation returned HTTP ${response.status}: ${(await response.text()).slice(0, 300)}`,
|
|
);
|
|
const result = (await response.json()) as {
|
|
data?: Array<{ b64_json?: string; url?: string }>;
|
|
cost?: unknown;
|
|
};
|
|
const image = result.data?.[0];
|
|
if (!image) throw new Error("Generation returned no image.");
|
|
const blob = image.b64_json
|
|
? Buffer.from(image.b64_json, "base64")
|
|
: Buffer.from(
|
|
await (await fetch(image.url!, { signal: AbortSignal.timeout(60_000) })).arrayBuffer(),
|
|
);
|
|
const path = join(
|
|
root,
|
|
".cache/nanogpt-images",
|
|
`${timestamp()}-${id!.replace(/\//g, "-")}${extension(blob)}`,
|
|
);
|
|
await mkdir(join(root, ".cache/nanogpt-images"), { recursive: true });
|
|
await writeFile(path, blob);
|
|
console.log(
|
|
`Saved: ${path}\nModel: ${id} · ${size} · ${steps} steps · CFG ${cfg}\nListed: $${estimate!.toFixed(6)} · Charged: ${String(result.cost ?? "not reported")} USD`,
|
|
);
|
|
}
|
|
|
|
try {
|
|
await main();
|
|
} catch (error) {
|
|
console.error(error instanceof Error ? error.message : error);
|
|
process.exitCode ||= 1;
|
|
}
|