gen-image.mjs

#!/usr/bin/env node
import { writeFileSync, readFileSync, existsSync } from 'node:fs';
import { parseArgs } from 'node:util';
import { basename } from 'node:path';

// --- config ---
const BASE = process.env.AI_API_BASE_URL;
const KEY  = process.env.AI_API_KEY;
const BUDGET_FILE   = '.image-budget.json';
const FILEID_CACHE  = '.file-id-cache.json';
const MAX_IMAGES    = Number(process.env.MAX_IMAGES ?? 50);
const REF_DEFAULT   = 'sprites/reference/companion-base.png';

if (!BASE || !KEY) {
    console.error('Missing AI_API_BASE_URL / AI_API_KEY. Set them in .env (see .env.example).');
    process.exit(1);
}

const { values } = parseArgs({
    options: {
        prompt:  { type: 'string' },
        out:     { type: 'string' },
        model:   { type: 'string', default: 'flux-2-pro' },
        size:    { type: 'string', default: '1024x1024' },
        ref:     { type: 'string', default: REF_DEFAULT },
        noref:   { type: 'boolean', default: false },
        debug:   { type: 'boolean', default: false },
        poll:    { type: 'string', default: '4' },
        timeout: { type: 'string', default: '180' },
    },
});

if (!values.prompt || !values.out) {
    console.error(
        'Usage: node scripts/gen-image.mjs \\\n' +
        '  --prompt "..." --out path.png \\\n' +
        '  [--model flux-2-pro] [--ref path.png] [--noref] [--debug]'
    );
    process.exit(1);
}

// --- spend guard ---
const budget = existsSync(BUDGET_FILE)
? JSON.parse(readFileSync(BUDGET_FILE, 'utf8'))
: { count: 0 };

if (budget.count >= MAX_IMAGES) {
    console.error(`Budget cap reached (${MAX_IMAGES}). Delete ${BUDGET_FILE} to reset.`);
    process.exit(1);
}

// --- helpers ---
function sleep(s) { return new Promise(r => setTimeout(r, s * 1000)); }

async function api(path, options = {}) {
    const res = await fetch(`${BASE}${path}`, {
        ...options,
        headers: {
            Authorization: `Bearer ${KEY}`,
            ...options.headers,
        },
    });
    const json = await res.json();
    if (values.debug) {
        console.log(`\n${options.method ?? 'GET'} ${path} → ${res.status}`);
        console.log(JSON.stringify(json, null, 2).slice(0, 3000));
    }
    return { ok: res.ok, status: res.status, json };
}

// --- upload file & cache the id ---
async function getFileId(localPath) {
    const cache = existsSync(FILEID_CACHE)
    ? JSON.parse(readFileSync(FILEID_CACHE, 'utf8'))
    : {};

    if (cache[localPath]) {
        console.log(`Cached file_id for ${basename(localPath)}: ${cache[localPath]}`);
        return cache[localPath];
    }

    console.log(`Uploading ${basename(localPath)} to /v1/files...`);
    const blob = new Blob([readFileSync(localPath)], { type: 'image/png' });
    const form = new FormData();
    form.append('file', blob, basename(localPath));
    form.append('purpose', 'vision');

    const res = await fetch(`${BASE}/files`, {
        method: 'POST',
        headers: { Authorization: `Bearer ${KEY}` },
        body: form,
    });
    const json = await res.json();

    if (!json.id) throw new Error(`Upload failed: ${JSON.stringify(json)}`);

    cache[localPath] = json.id;
    writeFileSync(FILEID_CACHE, JSON.stringify(cache, null, 2));
    console.log(`Uploaded → file_id: ${json.id}`);
    return json.id;
}

// --- main generation flow ---
async function generate() {
    const useRef = !values.noref && existsSync(values.ref);

    // ---- TEXT-TO-IMAGE (no reference) ----
    if (!useRef) {
        console.log(`Generating via /images/generations (no reference)...`);
        const { ok, json } = await api('/images/generations', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({
                model: values.model,
                prompt: values.prompt,
                n: 1,
                size: values.size,
            }),
        });
        if (!ok) throw new Error(`generations failed: ${JSON.stringify(json)}`);
        return extractImage(json);
    }

    // ---- IMAGE EDIT (with reference) ----
    const fileId = await getFileId(values.ref);
    console.log(`Generating via /images/edit with reference...`);

    const { ok, json } = await api('/images/edit', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
            model: values.model,
            prompt: values.prompt,
            image: fileId,
            size: values.size,
            async: true,
        }),
    });

    if (!ok) throw new Error(`edit failed: ${JSON.stringify(json)}`);

    // --- async: poll for result ---
    const jobId = json.id;
    if (!jobId) {
        // maybe it returned synchronously
        const buf = extractImage(json);
        if (buf) return buf;
        throw new Error(`No job id and no image: ${JSON.stringify(json).slice(0, 500)}`);
    }

    console.log(`Job ${jobId} — polling...`);
    const pollInterval = Number(values.poll);
    const maxWait = Number(values.timeout);
    const start = Date.now();

    while ((Date.now() - start) / 1000 < maxWait) {
        await sleep(pollInterval);

        const { json: pollJson } = await api(`/images/edits/${jobId}`, {
            method: 'GET',
        });

        const status = pollJson.status;
        process.stdout.write(` [${status}]`);

        if (status === 'failed') {
            throw new Error(`Job failed: ${JSON.stringify(pollJson)}`);
        }

        if (status === 'completed') {
            console.log(' ✓');
            const buf = extractImage(pollJson);
            if (buf) return buf;
            throw new Error(`Completed but no image: ${JSON.stringify(pollJson).slice(0, 500)}`);
        }
    }

    throw new Error(`Timed out after ${maxWait}s`);
}

// --- extract image buffer from any response shape ---
function extractImage(json) {
    // data array (OpenAI format)
    const item = json.data?.[0];
    if (item?.b64_json) return Buffer.from(item.b64_json, 'base64');
    if (item?.url) return fetchBuffer(item.url);

    // direct fields
    if (json.b64_json) return Buffer.from(json.b64_json, 'base64');
    if (json.url) return fetchBuffer(json.url);

    // output array
    const out = json.output?.[0];
    if (out?.b64_json) return Buffer.from(out.b64_json, 'base64');
    if (out?.url) return fetchBuffer(out.url);

    return null;
}

async function fetchBuffer(url) {
    return Buffer.from(await (await fetch(url)).arrayBuffer());
}

// --- run ---
try {
    const buffer = await generate();
    if (!buffer) throw new Error('No image data extracted');

    writeFileSync(values.out, buffer);
    budget.count += 1;
    writeFileSync(BUDGET_FILE, JSON.stringify(budget));

    console.log(`\nSaved: ${values.out}`);
    console.log(`Budget: ${budget.count}/${MAX_IMAGES} images used`);

} catch (err) {
    console.error(`\nFailed: ${err.message}`);
    if (values.debug) console.error(err.stack);
    process.exit(1);
}