230 lines
7.5 KiB
JavaScript
230 lines
7.5 KiB
JavaScript
import { writeFileSync, mkdirSync, readFileSync, existsSync } from 'fs';
|
|
import { resolve, dirname, basename } from 'path';
|
|
|
|
const envPath = resolve(process.cwd(), '.env');
|
|
try {
|
|
const envContent = readFileSync(envPath, 'utf-8');
|
|
for (const line of envContent.split('\n')) {
|
|
const match = line.match(/^([^#=]+)=(.*)$/);
|
|
if (match && !process.env[match[1].trim()]) {
|
|
process.env[match[1].trim()] = match[2].trim();
|
|
}
|
|
}
|
|
} catch {}
|
|
|
|
|
|
const API_BASE = 'https://api.banatie.app/api/v1';
|
|
const API_KEY = process.env.BANATIE_KEY || '';
|
|
|
|
const POLL_INTERVAL_MS = 2000;
|
|
const POLL_MAX_ATTEMPTS = 60; // 2 minutes total
|
|
|
|
async function pollGeneration(generationId) {
|
|
for (let attempt = 1; attempt <= POLL_MAX_ATTEMPTS; attempt++) {
|
|
const response = await fetch(`${API_BASE}/generations/${generationId}`, {
|
|
headers: { 'X-API-Key': API_KEY },
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const text = await response.text();
|
|
console.error(`Poll error ${response.status}: ${text}`);
|
|
process.exit(1);
|
|
}
|
|
|
|
const result = await response.json();
|
|
if (!result.success) {
|
|
console.error(`Poll failed:`, result.error);
|
|
process.exit(1);
|
|
}
|
|
|
|
const { status } = result.data;
|
|
|
|
if (status === 'success') {
|
|
return result.data;
|
|
}
|
|
|
|
if (status === 'failed') {
|
|
console.error(`Generation failed: ${result.data.errorMessage || 'unknown error'}`);
|
|
console.error('Suggestions: try a simpler prompt, check rate limits, or verify your BANATIE_KEY.');
|
|
process.exit(1);
|
|
}
|
|
|
|
// Still pending/processing — wait and retry
|
|
console.log(`Status: ${status} (attempt ${attempt}/${POLL_MAX_ATTEMPTS})...`);
|
|
await new Promise(r => setTimeout(r, POLL_INTERVAL_MS));
|
|
}
|
|
|
|
console.error(`Generation timed out after ${POLL_MAX_ATTEMPTS * POLL_INTERVAL_MS / 1000}s. Try again or use a simpler prompt.`);
|
|
process.exit(1);
|
|
}
|
|
|
|
async function uploadImage(filePath, { flowId, alias }) {
|
|
const absolutePath = resolve(filePath);
|
|
const fileBuffer = readFileSync(absolutePath);
|
|
const fileName = basename(absolutePath);
|
|
|
|
const ext = fileName.split('.').pop().toLowerCase();
|
|
const mimeTypes = { png: 'image/png', jpg: 'image/jpeg', jpeg: 'image/jpeg', webp: 'image/webp' };
|
|
const mime = mimeTypes[ext] || 'application/octet-stream';
|
|
|
|
const form = new FormData();
|
|
form.append('file', new Blob([fileBuffer], { type: mime }), fileName);
|
|
form.append('alias', alias);
|
|
if (flowId) form.append('flowId', flowId);
|
|
|
|
console.log(`Uploading: ${filePath} as ${alias}${flowId ? ` (flow: ${flowId})` : ''}...`);
|
|
|
|
const response = await fetch(`${API_BASE}/images/upload`, {
|
|
method: 'POST',
|
|
headers: { 'X-API-Key': API_KEY },
|
|
body: form,
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const text = await response.text();
|
|
console.error(`Upload error ${response.status}: ${text}`);
|
|
if (response.status === 401) console.error('Check that BANATIE_KEY is set correctly in .env');
|
|
if (response.status === 429) console.error('Rate limit exceeded. Wait before retrying (limit: 100 req/hour).');
|
|
process.exit(1);
|
|
}
|
|
|
|
const result = await response.json();
|
|
if (!result.success) {
|
|
console.error(`Upload failed:`, result.error);
|
|
process.exit(1);
|
|
}
|
|
|
|
console.log(`Uploaded: ${alias} (${result.data.id})`);
|
|
return result.data;
|
|
}
|
|
|
|
function makeAlias(filePath, index) {
|
|
const name = basename(filePath).replace(/\.[^.]+$/, '').replace(/[^a-z0-9-]/gi, '-').toLowerCase();
|
|
return `@ref-${index + 1}-${name}`;
|
|
}
|
|
|
|
async function resolveRefs(refs) {
|
|
if (!refs || refs.length === 0) return undefined;
|
|
|
|
const aliases = [];
|
|
let flowId = null;
|
|
|
|
for (let i = 0; i < refs.length; i++) {
|
|
const ref = refs[i];
|
|
if (ref.startsWith('@')) {
|
|
aliases.push(ref);
|
|
} else if (existsSync(ref)) {
|
|
const alias = makeAlias(ref, i);
|
|
const data = await uploadImage(ref, { flowId, alias });
|
|
if (!flowId) flowId = data.flowId;
|
|
aliases.push(alias);
|
|
} else {
|
|
console.error(`Reference not found: ${ref}`);
|
|
process.exit(1);
|
|
}
|
|
}
|
|
|
|
return { referenceImages: aliases, flowId };
|
|
}
|
|
|
|
export async function generateImage({ prompt, output, aspectRatio = '1:1', refs, template, autoEnhance = true }) {
|
|
if (!API_KEY) {
|
|
console.error('BANATIE_KEY environment variable is not set.');
|
|
console.error('Add BANATIE_KEY=your_key to the .env file in the project root.');
|
|
process.exit(1);
|
|
}
|
|
|
|
const resolved = await resolveRefs(refs);
|
|
|
|
const body = { prompt, aspectRatio, autoEnhance };
|
|
|
|
if (template && autoEnhance) {
|
|
body.enhancementOptions = { template };
|
|
}
|
|
|
|
if (resolved) {
|
|
body.referenceImages = resolved.referenceImages;
|
|
if (resolved.flowId) body.flowId = resolved.flowId;
|
|
}
|
|
|
|
const enhanceInfo = autoEnhance ? ` [template: ${template || 'general'}]` : ' [no enhance]';
|
|
console.log(`Generating: "${prompt}" (${body.aspectRatio})${enhanceInfo}${resolved ? ` with ${resolved.referenceImages.length} ref(s)` : ''}...`);
|
|
|
|
const response = await fetch(`${API_BASE}/generations`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'X-API-Key': API_KEY,
|
|
},
|
|
body: JSON.stringify(body),
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const text = await response.text();
|
|
console.error(`API error ${response.status}: ${text}`);
|
|
if (response.status === 401) console.error('Check that BANATIE_KEY is set correctly in .env');
|
|
if (response.status === 429) console.error('Rate limit exceeded. Wait before retrying (limit: 100 req/hour).');
|
|
process.exit(1);
|
|
}
|
|
|
|
const result = await response.json();
|
|
if (!result.success) {
|
|
console.error(`Generation failed:`, result.error);
|
|
process.exit(1);
|
|
}
|
|
|
|
// Handle async generation: poll if not yet complete
|
|
let data = result.data;
|
|
if (data.status === 'pending' || data.status === 'processing') {
|
|
console.log(`Generation queued (id: ${data.id}), waiting for completion...`);
|
|
data = await pollGeneration(data.id);
|
|
}
|
|
|
|
if (!data.outputImage?.storageUrl) {
|
|
console.error('Generation completed but no output image found.');
|
|
console.error('Response data:', JSON.stringify(data, null, 2));
|
|
process.exit(1);
|
|
}
|
|
|
|
const imageUrl = data.outputImage.storageUrl;
|
|
console.log(`Downloading from ${imageUrl}...`);
|
|
|
|
const imageResponse = await fetch(imageUrl);
|
|
if (!imageResponse.ok) {
|
|
console.error(`Failed to download image: ${imageResponse.status}`);
|
|
process.exit(1);
|
|
}
|
|
|
|
const buffer = Buffer.from(await imageResponse.arrayBuffer());
|
|
const outputPath = resolve(output);
|
|
mkdirSync(dirname(outputPath), { recursive: true });
|
|
writeFileSync(outputPath, buffer);
|
|
|
|
console.log(`Image saved: ${outputPath} (${data.outputImage.width}x${data.outputImage.height})`);
|
|
return { path: outputPath, generation: data };
|
|
}
|
|
|
|
function parseArgs(args) {
|
|
const result = { autoEnhance: true };
|
|
for (let i = 0; i < args.length; i++) {
|
|
if (args[i] === '--prompt') result.prompt = args[++i];
|
|
else if (args[i] === '--output') result.output = args[++i];
|
|
else if (args[i] === '--aspect-ratio') result.aspectRatio = args[++i];
|
|
else if (args[i] === '--template') result.template = args[++i];
|
|
else if (args[i] === '--no-enhance') result.autoEnhance = false;
|
|
else if (args[i] === '--ref') {
|
|
if (!result.refs) result.refs = [];
|
|
result.refs.push(args[++i]);
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
const args = parseArgs(process.argv.slice(2));
|
|
if (args.prompt && args.output) {
|
|
generateImage(args);
|
|
} else if (process.argv.length > 2) {
|
|
console.error('Usage: node banatie-gen.mjs --prompt "<description>" --output <path> [--aspect-ratio <ratio>] [--template <template>] [--no-enhance] [--ref <file|@alias>]...');
|
|
process.exit(1);
|
|
}
|