remove-bg.mjs

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

const BASE = process.env.AI_API_BASE_URL;
const KEY  = process.env.AI_API_KEY;

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: {
        in:    { type: 'string' },
        out:   { type: 'string' },
        debug: { type: 'boolean', default: false },
    },
});

if (!values.in || !values.out) {
    console.error('Usage: node scripts/remove-bg.mjs --in input.png --out output.png');
    process.exit(1);
}

const FILEID_CACHE = '.file-id-cache.json';

async function getFileId(localPath) {
    const cache = existsSync(FILEID_CACHE)
    ? JSON.parse(readFileSync(FILEID_CACHE, 'utf8'))
    : {};

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

    console.log(`Uploading ${basename(localPath)}...`);
    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));
    return json.id;
}

async function removeBg() {
    const fileId = await getFileId(values.in);

    console.log(`Removing background via ideogram-remove-background...`);

    const res = await fetch(`${BASE}/images/edit`, {
        method: 'POST',
        headers: {
            Authorization: `Bearer ${KEY}`,
            'Content-Type': 'application/json',
        },
        body: JSON.stringify({
            model: 'ideogram-remove-background',
            image: fileId,
            prompt: 'remove background',
            async: true,
        }),
    });

    const json = await res.json();

    if (values.debug) {
        console.log(`Status: ${res.status}`);
        console.log(JSON.stringify(json, null, 2).slice(0, 2000));
    }

    if (!res.ok) throw new Error(`API error: ${JSON.stringify(json)}`);

    // Handle async polling
    const jobId = json.id;
    if (jobId && json.status !== 'completed') {
        console.log(`Job ${jobId} — polling...`);
        const maxWait = 120;
        const start = Date.now();

        while ((Date.now() - start) / 1000 < maxWait) {
            await new Promise(r => setTimeout(r, 3000));

            const pollRes = await fetch(`${BASE}/images/edits/${jobId}`, {
                headers: { Authorization: `Bearer ${KEY}` },
            });
            const pollJson = await pollRes.json();
            process.stdout.write(` [${pollJson.status}]`);

            if (pollJson.status === 'failed') throw new Error(`Job failed`);
            if (pollJson.status === 'completed') {
                console.log(' ✓');
                return extractImage(pollJson);
            }
        }
        throw new Error('Timed out');
    }

    return extractImage(json);
}

function extractImage(json) {
    const item = json.data?.[0];
    if (item?.b64_json) return Buffer.from(item.b64_json, 'base64');
    if (item?.url) return fetchBuf(item.url);
    if (json.b64_json) return Buffer.from(json.b64_json, 'base64');
    if (json.url) return fetchBuf(json.url);
    const out = json.output?.[0];
    if (out?.b64_json) return Buffer.from(out.b64_json, 'base64');
    if (out?.url) return fetchBuf(out.url);
    throw new Error(`No image in response: ${JSON.stringify(json).slice(0, 500)}`);
}

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

try {
    const buffer = await removeBg();
    writeFileSync(values.out, buffer);
    console.log(`\nSaved: ${values.out}`);
} catch (err) {
    console.error(`\nFailed: ${err.message}`);
    process.exit(1);
}