feat: improve references
This commit is contained in:
@@ -16,6 +16,48 @@ try {
|
||||
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);
|
||||
@@ -41,6 +83,8 @@ async function uploadImage(filePath, { flowId, alias }) {
|
||||
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);
|
||||
}
|
||||
|
||||
@@ -83,20 +127,28 @@ async function resolveRefs(refs) {
|
||||
return { referenceImages: aliases, flowId };
|
||||
}
|
||||
|
||||
export async function generateImage({ prompt, output, aspectRatio = '1:1', refs }) {
|
||||
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('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 };
|
||||
|
||||
const body = { prompt, aspectRatio, autoEnhance };
|
||||
|
||||
if (template && autoEnhance) {
|
||||
body.enhancementOptions = { template };
|
||||
}
|
||||
|
||||
if (resolved) {
|
||||
body.referenceImages = resolved.referenceImages;
|
||||
if (resolved.flowId) body.flowId = resolved.flowId;
|
||||
}
|
||||
|
||||
console.log(`Generating: "${prompt}" (${body.aspectRatio})${resolved ? ` with ${resolved.referenceImages.length} ref(s)` : ''}...`);
|
||||
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',
|
||||
@@ -110,6 +162,8 @@ export async function generateImage({ prompt, output, aspectRatio = '1:1', refs
|
||||
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);
|
||||
}
|
||||
|
||||
@@ -119,7 +173,20 @@ export async function generateImage({ prompt, output, aspectRatio = '1:1', refs
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const imageUrl = result.data.outputImage.storageUrl;
|
||||
// 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);
|
||||
@@ -133,16 +200,18 @@ export async function generateImage({ prompt, output, aspectRatio = '1:1', refs
|
||||
mkdirSync(dirname(outputPath), { recursive: true });
|
||||
writeFileSync(outputPath, buffer);
|
||||
|
||||
console.log(`Image saved: ${outputPath} (${result.data.outputImage.width}x${result.data.outputImage.height})`);
|
||||
return { path: outputPath, generation: result.data };
|
||||
console.log(`Image saved: ${outputPath} (${data.outputImage.width}x${data.outputImage.height})`);
|
||||
return { path: outputPath, generation: data };
|
||||
}
|
||||
|
||||
function parseArgs(args) {
|
||||
const result = {};
|
||||
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]);
|
||||
@@ -155,6 +224,6 @@ 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>] [--ref <file|@alias>]...');
|
||||
console.error('Usage: node banatie-gen.mjs --prompt "<description>" --output <path> [--aspect-ratio <ratio>] [--template <template>] [--no-enhance] [--ref <file|@alias>]...');
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user