feat: improve references

This commit is contained in:
2026-02-25 00:20:33 +07:00
parent 1826c23826
commit c031c8e0be
2 changed files with 140 additions and 20 deletions
+78 -9
View File
@@ -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);
}