400 lines
12 KiB
JavaScript
400 lines
12 KiB
JavaScript
/**
|
||
* POC Workflow V4: 完整优化版
|
||
*
|
||
* 改进：
|
||
* 1. Vision超时重试机制
|
||
* 2. 可配置的套路模板系统
|
||
* 3. 基于真实交付成品的Prompt
|
||
* 4. 强化产品一致性约束
|
||
*
|
||
* 用法：
|
||
* node poc-workflow-v4.js --material-dir="./素材/素材/已有的素材" --sku-name="Touchdog冰蓝色伊丽莎白圈"
|
||
*/
|
||
|
||
require('dotenv').config();
|
||
const axios = require('axios');
|
||
const fs = require('fs');
|
||
const path = require('path');
|
||
const { S3Client, PutObjectCommand } = require('@aws-sdk/client-s3');
|
||
const { extractProductFeatures, buildGoldenDescription } = require('./lib/vision-extractor');
|
||
const { generateAllPrompts } = require('./lib/template-config');
|
||
|
||
// ============================================================
|
||
// 配置
|
||
// ============================================================
|
||
const API_KEY = 'G9rXx3Ag2Xfa7Gs8zou6t6HqeZ';
|
||
const API_BASE = 'https://api.wuyinkeji.com/api';
|
||
|
||
// R2客户端
|
||
const r2Client = new S3Client({
|
||
region: 'auto',
|
||
endpoint: `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com`,
|
||
credentials: {
|
||
accessKeyId: process.env.R2_ACCESS_KEY_ID,
|
||
secretAccessKey: process.env.R2_SECRET_ACCESS_KEY,
|
||
},
|
||
forcePathStyle: true,
|
||
});
|
||
|
||
// ============================================================
|
||
// 命令行参数解析
|
||
// ============================================================
|
||
function parseArgs() {
|
||
const args = process.argv.slice(2);
|
||
const config = {
|
||
materialDir: path.join(__dirname, '素材/素材/已有的素材'),
|
||
skuName: 'Touchdog 软质伊丽莎白圈',
|
||
brandName: 'TOUCHDOG',
|
||
productName: 'CAT SOFT CONE COLLAR',
|
||
petType: 'cat',
|
||
outputDir: null
|
||
};
|
||
|
||
for (const arg of args) {
|
||
if (arg.startsWith('--material-dir=')) {
|
||
config.materialDir = arg.split('=')[1];
|
||
} else if (arg.startsWith('--sku-name=')) {
|
||
config.skuName = arg.split('=')[1];
|
||
} else if (arg.startsWith('--brand-name=')) {
|
||
config.brandName = arg.split('=')[1];
|
||
} else if (arg.startsWith('--product-name=')) {
|
||
config.productName = arg.split('=')[1];
|
||
} else if (arg.startsWith('--output-dir=')) {
|
||
config.outputDir = arg.split('=')[1];
|
||
} else if (arg.startsWith('--pet-type=')) {
|
||
config.petType = arg.split('=')[1];
|
||
}
|
||
}
|
||
|
||
// 自动生成输出目录
|
||
if (!config.outputDir) {
|
||
const timestamp = new Date().toISOString().slice(0, 19).replace(/[T:]/g, '-');
|
||
const safeName = config.skuName.replace(/[^a-zA-Z0-9\u4e00-\u9fa5]/g, '_').slice(0, 20);
|
||
config.outputDir = path.join(__dirname, `output_v4_${safeName}_${timestamp}`);
|
||
}
|
||
|
||
return config;
|
||
}
|
||
|
||
// ============================================================
|
||
// 工具函数
|
||
// ============================================================
|
||
|
||
async function uploadToR2(filePath) {
|
||
const fileName = `v4-${Date.now()}-${path.basename(filePath)}`;
|
||
const fileBuffer = fs.readFileSync(filePath);
|
||
|
||
await r2Client.send(new PutObjectCommand({
|
||
Bucket: process.env.R2_BUCKET_NAME || 'ai-flow',
|
||
Key: fileName,
|
||
Body: fileBuffer,
|
||
ContentType: filePath.endsWith('.png') ? 'image/png' : 'image/jpeg',
|
||
}));
|
||
|
||
return process.env.R2_PUBLIC_DOMAIN
|
||
? `${process.env.R2_PUBLIC_DOMAIN}/${fileName}`
|
||
: `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com/${process.env.R2_BUCKET_NAME}/${fileName}`;
|
||
}
|
||
|
||
// 扫描素材目录
|
||
async function scanMaterials(materialDir) {
|
||
console.log('📁 扫描素材目录:', materialDir);
|
||
|
||
if (!fs.existsSync(materialDir)) {
|
||
throw new Error(`素材目录不存在: ${materialDir}`);
|
||
}
|
||
|
||
const files = fs.readdirSync(materialDir);
|
||
const images = files.filter(f => /\.(jpg|jpeg|png)$/i.test(f));
|
||
|
||
console.log(` 找到 ${images.length} 张图片: ${images.join(', ')}`);
|
||
|
||
// 识别图片类型
|
||
let flatImage = null; // 平铺图
|
||
let wornImage = null; // 佩戴图
|
||
|
||
for (const img of images) {
|
||
const imgPath = path.join(materialDir, img);
|
||
const imgLower = img.toLowerCase();
|
||
|
||
// 平铺图识别
|
||
if (img.includes('5683') || imgLower.includes('flat') || imgLower.includes('1.png')) {
|
||
flatImage = imgPath;
|
||
}
|
||
// 佩戴图识别
|
||
if (img.includes('6514') || imgLower.includes('worn') || imgLower.includes('wear') || imgLower.includes('3.png')) {
|
||
wornImage = imgPath;
|
||
}
|
||
}
|
||
|
||
// 如果没找到，用启发式规则
|
||
if (!flatImage) {
|
||
// 优先选png格式（通常更清晰）
|
||
const pngImages = images.filter(i => i.toLowerCase().endsWith('.png'));
|
||
flatImage = pngImages.length > 0
|
||
? path.join(materialDir, pngImages[0])
|
||
: path.join(materialDir, images[0]);
|
||
}
|
||
|
||
if (!wornImage) {
|
||
const jpgImages = images.filter(i => i.toLowerCase().endsWith('.jpg') || i.toLowerCase().endsWith('.jpeg'));
|
||
wornImage = jpgImages.length > 0
|
||
? path.join(materialDir, jpgImages[0])
|
||
: flatImage;
|
||
}
|
||
|
||
console.log(' ✓ 平铺图:', path.basename(flatImage));
|
||
console.log(' ✓ 佩戴图:', path.basename(wornImage));
|
||
|
||
return { flatImage, wornImage, allImages: images.map(i => path.join(materialDir, i)) };
|
||
}
|
||
|
||
// 生图任务
|
||
async function submitImageTask(prompt, refImageUrl, aspectRatio = '1:1') {
|
||
const payload = {
|
||
key: API_KEY,
|
||
prompt: prompt,
|
||
img_url: refImageUrl,
|
||
aspectRatio: aspectRatio,
|
||
imageSize: '1K'
|
||
};
|
||
|
||
const response = await axios.post(`${API_BASE}/img/nanoBanana-pro`, payload, {
|
||
timeout: 30000
|
||
});
|
||
|
||
const taskId = response.data.data?.id;
|
||
if (!taskId) {
|
||
console.error('Submit response:', response.data);
|
||
throw new Error('No task ID returned');
|
||
}
|
||
return taskId;
|
||
}
|
||
|
||
async function pollImageResult(taskId) {
|
||
let attempts = 0;
|
||
const maxAttempts = 90;
|
||
|
||
while (attempts < maxAttempts) {
|
||
try {
|
||
const response = await axios.get(`${API_BASE}/img/drawDetail`, {
|
||
params: { key: API_KEY, id: taskId },
|
||
timeout: 10000
|
||
});
|
||
|
||
const data = response.data.data;
|
||
|
||
if (data && data.status === 2 && data.image_url) {
|
||
return data.image_url;
|
||
} else if (data && data.status === 3) {
|
||
throw new Error('Generation failed: ' + (data.fail_reason || 'Unknown'));
|
||
}
|
||
|
||
process.stdout.write('.');
|
||
await new Promise(r => setTimeout(r, 2000));
|
||
attempts++;
|
||
} catch (error) {
|
||
if (error.message.includes('Generation failed')) throw error;
|
||
// 网络错误，继续重试
|
||
process.stdout.write('x');
|
||
await new Promise(r => setTimeout(r, 3000));
|
||
attempts++;
|
||
}
|
||
}
|
||
throw new Error('Timeout after ' + maxAttempts + ' attempts');
|
||
}
|
||
|
||
async function generateSingleImage(item, refUrl, outputDir) {
|
||
console.log(`\n🎨 [${item.id}] ${item.name}`);
|
||
|
||
let retryCount = 0;
|
||
const maxRetries = 2;
|
||
|
||
while (retryCount <= maxRetries) {
|
||
try {
|
||
if (retryCount > 0) console.log(` 重试 ${retryCount}/${maxRetries}...`);
|
||
|
||
console.log(' 提交任务...');
|
||
const taskId = await submitImageTask(item.prompt, refUrl, item.aspectRatio);
|
||
console.log(` Task ID: ${taskId}`);
|
||
|
||
const imageUrl = await pollImageResult(taskId);
|
||
console.log('\n ✓ 生成成功');
|
||
|
||
// 下载保存
|
||
const imgRes = await axios.get(imageUrl, { responseType: 'arraybuffer', timeout: 30000 });
|
||
const outputPath = path.join(outputDir, `${item.id}.jpg`);
|
||
fs.writeFileSync(outputPath, imgRes.data);
|
||
console.log(` ✓ 保存: ${item.id}.jpg`);
|
||
|
||
return { id: item.id, success: true, path: outputPath };
|
||
|
||
} catch (error) {
|
||
console.error(`\n ✗ 失败: ${error.message}`);
|
||
retryCount++;
|
||
if (retryCount <= maxRetries) {
|
||
await new Promise(r => setTimeout(r, 5000));
|
||
}
|
||
}
|
||
}
|
||
|
||
return { id: item.id, success: false, error: 'Max retries exceeded' };
|
||
}
|
||
|
||
// ============================================================
|
||
// 主流程
|
||
// ============================================================
|
||
|
||
async function main() {
|
||
const config = parseArgs();
|
||
|
||
console.log('\n' + '═'.repeat(70));
|
||
console.log('🚀 POC Workflow V4 - 完整优化版');
|
||
console.log('═'.repeat(70));
|
||
console.log('\n📋 配置:');
|
||
console.log(' 素材目录:', config.materialDir);
|
||
console.log(' SKU名称:', config.skuName);
|
||
console.log(' 品牌:', config.brandName);
|
||
console.log(' 输出目录:', config.outputDir);
|
||
|
||
// 创建输出目录
|
||
if (!fs.existsSync(config.outputDir)) {
|
||
fs.mkdirSync(config.outputDir, { recursive: true });
|
||
}
|
||
|
||
// ========== 阶段1: 扫描素材 ==========
|
||
console.log('\n' + '─'.repeat(70));
|
||
console.log('📁 阶段1: 扫描素材');
|
||
console.log('─'.repeat(70));
|
||
|
||
const materials = await scanMaterials(config.materialDir);
|
||
|
||
if (!materials.flatImage) {
|
||
console.error('❌ 未找到素材图片');
|
||
process.exit(1);
|
||
}
|
||
|
||
// ========== 阶段2: 上传素材 ==========
|
||
console.log('\n' + '─'.repeat(70));
|
||
console.log('📤 阶段2: 上传素材到云存储');
|
||
console.log('─'.repeat(70));
|
||
|
||
const flatImageUrl = await uploadToR2(materials.flatImage);
|
||
console.log(' 平铺图:', flatImageUrl);
|
||
|
||
let wornImageUrl = flatImageUrl;
|
||
if (materials.wornImage && materials.wornImage !== materials.flatImage) {
|
||
wornImageUrl = await uploadToR2(materials.wornImage);
|
||
console.log(' 佩戴图:', wornImageUrl);
|
||
}
|
||
|
||
// ========== 阶段3: Vision分析 ==========
|
||
console.log('\n' + '─'.repeat(70));
|
||
console.log('🔍 阶段3: Vision分析产品特征 (超时重试机制)');
|
||
console.log('─'.repeat(70));
|
||
|
||
const visionResult = await extractProductFeatures(flatImageUrl, {
|
||
maxRetries: 3,
|
||
timeout: 150000, // 150秒
|
||
retryDelay: 8000,
|
||
cacheDir: config.outputDir,
|
||
cacheKey: path.basename(materials.flatImage).replace(/\.[^.]+$/, '')
|
||
});
|
||
|
||
const goldenDescription = buildGoldenDescription(visionResult);
|
||
|
||
console.log('\n📝 Golden Description预览:');
|
||
console.log('─'.repeat(50));
|
||
console.log(goldenDescription.substring(0, 600) + '...');
|
||
|
||
// 保存分析结果
|
||
fs.writeFileSync(
|
||
path.join(config.outputDir, 'vision-analysis.json'),
|
||
JSON.stringify(visionResult, null, 2)
|
||
);
|
||
fs.writeFileSync(
|
||
path.join(config.outputDir, 'golden-description.txt'),
|
||
goldenDescription
|
||
);
|
||
|
||
// ========== 阶段4: 生成Prompts ==========
|
||
console.log('\n' + '─'.repeat(70));
|
||
console.log('📝 阶段4: 生成12张图的Prompts (可配置模板系统)');
|
||
console.log('─'.repeat(70));
|
||
|
||
const product = { goldenDescription };
|
||
const skuInfo = {
|
||
brandName: config.brandName,
|
||
productName: config.productName,
|
||
petType: config.petType
|
||
};
|
||
|
||
const prompts = generateAllPrompts(product, skuInfo);
|
||
console.log(` ✓ 生成了 ${prompts.length} 个Prompt`);
|
||
|
||
// 保存prompts供参考
|
||
fs.writeFileSync(
|
||
path.join(config.outputDir, 'all-prompts.json'),
|
||
JSON.stringify(prompts.map(p => ({
|
||
id: p.id,
|
||
name: p.name,
|
||
type: p.type,
|
||
aspectRatio: p.aspectRatio,
|
||
prompt: p.prompt
|
||
})), null, 2)
|
||
);
|
||
|
||
// ========== 阶段5: 批量生成 ==========
|
||
console.log('\n' + '─'.repeat(70));
|
||
console.log('🎨 阶段5: 批量生成图片');
|
||
console.log('─'.repeat(70));
|
||
|
||
const results = [];
|
||
for (const promptItem of prompts) {
|
||
// Main_02用平铺图参考，其他用佩戴图
|
||
const refUrl = promptItem.id === 'Main_02' ? flatImageUrl : wornImageUrl;
|
||
|
||
const result = await generateSingleImage(promptItem, refUrl, config.outputDir);
|
||
results.push(result);
|
||
}
|
||
|
||
// ========== 阶段6: 总结 ==========
|
||
console.log('\n' + '═'.repeat(70));
|
||
console.log('📊 生成完成');
|
||
console.log('═'.repeat(70));
|
||
|
||
const successCount = results.filter(r => r.success).length;
|
||
console.log(`\n✅ 成功: ${successCount}/${results.length}`);
|
||
|
||
if (successCount < results.length) {
|
||
console.log('\n❌ 失败的图片:');
|
||
results.filter(r => !r.success).forEach(r => {
|
||
console.log(` - ${r.id}: ${r.error}`);
|
||
});
|
||
}
|
||
|
||
console.log(`\n📁 输出目录: ${config.outputDir}`);
|
||
|
||
// 保存结果摘要
|
||
fs.writeFileSync(
|
||
path.join(config.outputDir, 'results-summary.json'),
|
||
JSON.stringify({
|
||
config,
|
||
visionSuccess: !!visionResult,
|
||
results,
|
||
summary: {
|
||
total: results.length,
|
||
success: successCount,
|
||
failed: results.length - successCount
|
||
},
|
||
timestamp: new Date().toISOString()
|
||
}, null, 2)
|
||
);
|
||
|
||
console.log('\n💡 提示: 查看 all-prompts.json 了解每张图的Prompt详情');
|
||
}
|
||
|
||
main().catch(console.error);
|
||
|
||
|