Upload latest code and optimized prompts (v6)
This commit is contained in:
tony
399
poc-workflow-v4.js
Normal file
399
poc-workflow-v4.js
Normal file
@@ -0,0 +1,399 @@
|
||||
/**
|
||||
* 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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user