625 lines
18 KiB
JavaScript
625 lines
18 KiB
JavaScript
require('dotenv').config();
|
||
const express = require('express');
|
||
const multer = require('multer');
|
||
const { S3Client, PutObjectCommand, DeleteObjectCommand } = require('@aws-sdk/client-s3');
|
||
const axios = require('axios');
|
||
const fs = require('fs');
|
||
const path = require('path');
|
||
const cors = require('cors');
|
||
|
||
// 导入新模块
|
||
const { planImages } = require('./lib/brain');
|
||
const { injectConstraints, getAspectRatio } = require('./lib/constraint-injector');
|
||
|
||
const app = express();
|
||
const port = process.env.PORT || 3000;
|
||
|
||
app.use(cors());
|
||
app.use(express.json({ limit: '50mb' }));
|
||
app.use(express.static('public'));
|
||
|
||
// Configure R2 Client
|
||
const bucketName = process.env.R2_BUCKET_NAME || 'ai-flow';
|
||
console.log('Using R2 Bucket:', bucketName);
|
||
|
||
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,
|
||
});
|
||
|
||
// Multer for memory storage (upload directly to R2)
|
||
const upload = multer({ storage: multer.memoryStorage() });
|
||
|
||
// ============================================
|
||
// 原有API端点
|
||
// ============================================
|
||
|
||
// Endpoint to get prompt template
|
||
app.get('/api/prompt-template', (req, res) => {
|
||
try {
|
||
const promptPath = path.join(__dirname, 'prompt.md');
|
||
const content = fs.readFileSync(promptPath, 'utf-8');
|
||
res.json({ content });
|
||
} catch (error) {
|
||
res.status(500).json({ error: 'Failed to read prompt template' });
|
||
}
|
||
});
|
||
|
||
// Endpoint to upload to R2
|
||
app.post('/api/upload-r2', upload.array('files'), async (req, res) => {
|
||
try {
|
||
const files = req.files;
|
||
if (!files || files.length === 0) {
|
||
return res.status(400).json({ error: 'No files uploaded' });
|
||
}
|
||
|
||
const uploadedUrls = [];
|
||
|
||
for (const file of files) {
|
||
const fileName = `${Date.now()}-${file.originalname}`;
|
||
const command = new PutObjectCommand({
|
||
Bucket: bucketName,
|
||
Key: fileName,
|
||
Body: file.buffer,
|
||
ContentType: file.mimetype,
|
||
});
|
||
|
||
await r2Client.send(command);
|
||
|
||
const publicUrl = process.env.R2_PUBLIC_DOMAIN
|
||
? `${process.env.R2_PUBLIC_DOMAIN}/${fileName}`
|
||
: `https://${process.env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com/${bucketName}/${fileName}`;
|
||
|
||
uploadedUrls.push(publicUrl);
|
||
}
|
||
|
||
res.json({ urls: uploadedUrls });
|
||
} catch (error) {
|
||
console.error('R2 Upload Error:', error);
|
||
res.status(500).json({ error: 'Upload failed: ' + error.message });
|
||
}
|
||
});
|
||
|
||
// Endpoint to delete from R2
|
||
app.delete('/api/delete-r2', async (req, res) => {
|
||
try {
|
||
const { url } = req.body;
|
||
if (!url) return res.status(400).json({ error: 'URL is required' });
|
||
|
||
const urlObj = new URL(url);
|
||
const key = decodeURIComponent(urlObj.pathname.substring(1));
|
||
|
||
const bucketName = process.env.R2_BUCKET_NAME || 'ai-flow';
|
||
|
||
await r2Client.send(new DeleteObjectCommand({
|
||
Bucket: bucketName,
|
||
Key: key,
|
||
}));
|
||
|
||
res.json({ success: true });
|
||
} catch (error) {
|
||
console.error('R2 Delete Error:', error);
|
||
res.status(500).json({ error: 'Delete failed: ' + error.message });
|
||
}
|
||
});
|
||
|
||
// Endpoint to generate prompts (LLM) - 保留原有功能
|
||
app.post('/api/generate-prompts', async (req, res) => {
|
||
try {
|
||
const { prompt } = req.body;
|
||
|
||
const response = await axios.post('https://api2img.shubiaobiao.com/v1/chat/completions', {
|
||
model: 'gemini-3-pro-preview',
|
||
messages: [
|
||
{ role: 'user', content: prompt }
|
||
],
|
||
stream: false
|
||
}, {
|
||
headers: {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': `Bearer ${process.env.API_KEY}`
|
||
}
|
||
});
|
||
|
||
res.json({ data: response.data });
|
||
} catch (error) {
|
||
console.error('LLM API Error:', error.response?.data || error.message);
|
||
res.status(500).json({ error: 'Prompt generation failed' });
|
||
}
|
||
});
|
||
|
||
// Endpoint to generate image (Sync) - 保留原有功能
|
||
app.post('/api/generate-image', async (req, res) => {
|
||
try {
|
||
const { prompt, aspectRatio, refImages } = req.body;
|
||
console.log('Generate Image Request:', { promptLength: prompt?.length, refImagesCount: refImages?.length });
|
||
|
||
let finalPrompt = prompt;
|
||
let ar = aspectRatio || '1:1';
|
||
|
||
const arMatch = prompt.match(/--ar\s+(\d+:\d+)/i);
|
||
if (arMatch) {
|
||
ar = arMatch[1];
|
||
}
|
||
|
||
const payload = {
|
||
model: 'gemini-3-pro-image-preview',
|
||
prompt: finalPrompt,
|
||
n: 1,
|
||
size: "1K",
|
||
aspect_ratio: ar
|
||
};
|
||
|
||
if (refImages && Array.isArray(refImages) && refImages.length > 0) {
|
||
payload.image_urls = refImages;
|
||
}
|
||
|
||
const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
|
||
headers: {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': `Bearer ${process.env.API_KEY}`
|
||
},
|
||
timeout: 120000
|
||
});
|
||
|
||
res.json(response.data);
|
||
} catch (error) {
|
||
console.error('Image Gen API Error:', error.response?.data || error.message);
|
||
res.status(500).json({ error: 'Image generation failed' });
|
||
}
|
||
});
|
||
|
||
// Endpoint to get image result
|
||
app.post('/api/get-image-result', async (req, res) => {
|
||
try {
|
||
const { id } = req.body;
|
||
if (!id) return res.status(400).json({ error: 'Task ID is required' });
|
||
|
||
const queryUrl = process.env.IMAGE_QUERY_API_URL || 'https://api.wuyinkeji.com/api/img/drawDetail';
|
||
|
||
const response = await axios.get(queryUrl, {
|
||
params: {
|
||
key: process.env.API_KEY,
|
||
id: id
|
||
},
|
||
headers: { 'Content-Type': 'application/json' }
|
||
});
|
||
|
||
res.json(response.data);
|
||
} catch (error) {
|
||
console.error('Image Query API Error:', error.response?.data || error.message);
|
||
res.status(500).json({ error: 'Failed to query image result' });
|
||
}
|
||
});
|
||
|
||
// ============================================
|
||
// 新增API端点 - SKU智能生图
|
||
// ============================================
|
||
|
||
/**
|
||
* 获取Brain System Prompt
|
||
*/
|
||
app.get('/api/brain-prompt', (req, res) => {
|
||
try {
|
||
const promptPath = path.join(__dirname, 'prompts/brain-system.md');
|
||
const content = fs.readFileSync(promptPath, 'utf-8');
|
||
res.json({ content });
|
||
} catch (error) {
|
||
res.status(500).json({ error: 'Failed to read brain prompt' });
|
||
}
|
||
});
|
||
|
||
/**
|
||
* 获取示例SKU数据
|
||
*/
|
||
app.get('/api/sample-sku', (req, res) => {
|
||
try {
|
||
const skuPath = path.join(__dirname, 'data/sample-sku.json');
|
||
const content = fs.readFileSync(skuPath, 'utf-8');
|
||
res.json(JSON.parse(content));
|
||
} catch (error) {
|
||
res.status(500).json({ error: 'Failed to read sample SKU' });
|
||
}
|
||
});
|
||
|
||
/**
|
||
* Stage 1: Brain决策 - 生成图片规划
|
||
* POST /api/plan-images
|
||
* Body: { sku: SKU数据对象 }
|
||
* Returns: { plan: Brain输出的图片规划 }
|
||
*/
|
||
app.post('/api/plan-images', async (req, res) => {
|
||
try {
|
||
const { sku } = req.body;
|
||
|
||
if (!sku || !sku.sku_id) {
|
||
return res.status(400).json({ error: 'SKU data is required' });
|
||
}
|
||
|
||
console.log(`Planning images for SKU: ${sku.sku_id}`);
|
||
|
||
const plan = await planImages(sku, {
|
||
apiKey: process.env.API_KEY
|
||
});
|
||
|
||
res.json({
|
||
success: true,
|
||
sku_id: sku.sku_id,
|
||
plan
|
||
});
|
||
} catch (error) {
|
||
console.error('Plan Images Error:', error.message);
|
||
res.status(500).json({ error: 'Failed to plan images: ' + error.message });
|
||
}
|
||
});
|
||
|
||
/**
|
||
* Stage 2: 生成单张图片（带约束注入）
|
||
* POST /api/generate-image-with-constraints
|
||
* Body: { sku, imageSpec, plan }
|
||
* Returns: { imageUrl }
|
||
*/
|
||
app.post('/api/generate-image-with-constraints', async (req, res) => {
|
||
try {
|
||
const { sku, imageSpec } = req.body;
|
||
|
||
if (!sku || !imageSpec) {
|
||
return res.status(400).json({ error: 'SKU and imageSpec are required' });
|
||
}
|
||
|
||
// 注入约束
|
||
const isComparison = imageSpec.is_comparison || imageSpec.id === 'APlus_02';
|
||
const fullPrompt = injectConstraints(
|
||
imageSpec.ai_prompt,
|
||
sku,
|
||
{
|
||
isComparison,
|
||
logoPlacement: imageSpec.logo_placement || {}
|
||
}
|
||
);
|
||
|
||
// 确定宽高比
|
||
const aspectRatio = getAspectRatio(imageSpec.id);
|
||
|
||
console.log(`Generating image ${imageSpec.id} for SKU ${sku.sku_id}`);
|
||
|
||
// 调用图像生成API
|
||
const payload = {
|
||
model: 'gemini-3-pro-image-preview',
|
||
prompt: fullPrompt,
|
||
n: 1,
|
||
size: "1K",
|
||
aspect_ratio: aspectRatio
|
||
};
|
||
|
||
// 添加参考图
|
||
if (sku.ref_images && sku.ref_images.length > 0) {
|
||
payload.image_urls = sku.ref_images;
|
||
}
|
||
|
||
const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
|
||
headers: {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': `Bearer ${process.env.API_KEY}`
|
||
},
|
||
timeout: 180000 // 3分钟超时
|
||
});
|
||
|
||
// 提取图片URL
|
||
let imageUrl = null;
|
||
if (response.data.data && Array.isArray(response.data.data) && response.data.data.length > 0) {
|
||
imageUrl = response.data.data[0].url;
|
||
}
|
||
|
||
res.json({
|
||
success: true,
|
||
image_id: imageSpec.id,
|
||
image_url: imageUrl,
|
||
aspect_ratio: aspectRatio
|
||
});
|
||
|
||
} catch (error) {
|
||
console.error('Generate Image With Constraints Error:', error.response?.data || error.message);
|
||
res.status(500).json({ error: 'Image generation failed: ' + error.message });
|
||
}
|
||
});
|
||
|
||
/**
|
||
* 完整流程：为单个SKU生成全部12张图
|
||
* POST /api/generate-sku
|
||
* Body: { sku: SKU数据对象 }
|
||
* Returns: { results: [{ id, url, status }] }
|
||
*/
|
||
app.post('/api/generate-sku', async (req, res) => {
|
||
try {
|
||
const { sku } = req.body;
|
||
|
||
if (!sku || !sku.sku_id) {
|
||
return res.status(400).json({ error: 'SKU data is required' });
|
||
}
|
||
|
||
console.log(`\n========== Starting SKU Generation: ${sku.sku_id} ==========`);
|
||
|
||
// Stage 1: Brain规划
|
||
console.log('Stage 1: Planning images...');
|
||
const plan = await planImages(sku, { apiKey: process.env.API_KEY });
|
||
console.log(`Plan generated with ${plan.images.length} images`);
|
||
|
||
// Stage 2: 逐张生成图片
|
||
console.log('Stage 2: Generating images...');
|
||
const results = [];
|
||
|
||
for (const imageSpec of plan.images) {
|
||
try {
|
||
console.log(` Generating ${imageSpec.id}...`);
|
||
|
||
// 注入约束
|
||
const isComparison = imageSpec.is_comparison || imageSpec.id === 'APlus_02';
|
||
const fullPrompt = injectConstraints(
|
||
imageSpec.ai_prompt,
|
||
sku,
|
||
{
|
||
isComparison,
|
||
logoPlacement: imageSpec.logo_placement || {}
|
||
}
|
||
);
|
||
|
||
const aspectRatio = getAspectRatio(imageSpec.id);
|
||
|
||
const payload = {
|
||
model: 'gemini-3-pro-image-preview',
|
||
prompt: fullPrompt,
|
||
n: 1,
|
||
size: "1K",
|
||
aspect_ratio: aspectRatio
|
||
};
|
||
|
||
if (sku.ref_images && sku.ref_images.length > 0) {
|
||
payload.image_urls = sku.ref_images;
|
||
}
|
||
|
||
const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
|
||
headers: {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': `Bearer ${process.env.API_KEY}`
|
||
},
|
||
timeout: 180000
|
||
});
|
||
|
||
let imageUrl = null;
|
||
if (response.data.data && Array.isArray(response.data.data) && response.data.data.length > 0) {
|
||
imageUrl = response.data.data[0].url;
|
||
}
|
||
|
||
results.push({
|
||
id: imageSpec.id,
|
||
type: imageSpec.type,
|
||
purpose: imageSpec.purpose,
|
||
url: imageUrl,
|
||
status: imageUrl ? 'success' : 'failed',
|
||
aspect_ratio: aspectRatio
|
||
});
|
||
|
||
console.log(` ✓ ${imageSpec.id} generated`);
|
||
|
||
// 添加小延迟避免API限流
|
||
await new Promise(resolve => setTimeout(resolve, 1000));
|
||
|
||
} catch (imgError) {
|
||
console.error(` ✗ ${imageSpec.id} failed:`, imgError.message);
|
||
results.push({
|
||
id: imageSpec.id,
|
||
type: imageSpec.type,
|
||
purpose: imageSpec.purpose,
|
||
url: null,
|
||
status: 'failed',
|
||
error: imgError.message
|
||
});
|
||
}
|
||
}
|
||
|
||
const successCount = results.filter(r => r.status === 'success').length;
|
||
console.log(`\n========== SKU Generation Complete: ${successCount}/${results.length} ==========\n`);
|
||
|
||
res.json({
|
||
success: true,
|
||
sku_id: sku.sku_id,
|
||
plan: {
|
||
analysis: plan.analysis,
|
||
image_count: plan.images.length
|
||
},
|
||
results,
|
||
summary: {
|
||
total: results.length,
|
||
success: successCount,
|
||
failed: results.length - successCount
|
||
}
|
||
});
|
||
|
||
} catch (error) {
|
||
console.error('Generate SKU Error:', error.message);
|
||
res.status(500).json({ error: 'SKU generation failed: ' + error.message });
|
||
}
|
||
});
|
||
|
||
/**
|
||
* 批量生成：为多个SKU生成图片
|
||
* POST /api/batch-generate
|
||
* Body: { skus: [SKU数据数组], concurrency: 并发数 }
|
||
* Returns: { results: [{ sku_id, status, images }] }
|
||
*/
|
||
app.post('/api/batch-generate', async (req, res) => {
|
||
try {
|
||
const { skus, concurrency = 1 } = req.body;
|
||
|
||
if (!skus || !Array.isArray(skus) || skus.length === 0) {
|
||
return res.status(400).json({ error: 'SKUs array is required' });
|
||
}
|
||
|
||
console.log(`\n========== Batch Generation Started: ${skus.length} SKUs ==========`);
|
||
|
||
const allResults = [];
|
||
|
||
// 简单的串行处理（避免API限流）
|
||
for (let i = 0; i < skus.length; i++) {
|
||
const sku = skus[i];
|
||
console.log(`\nProcessing SKU ${i + 1}/${skus.length}: ${sku.sku_id}`);
|
||
|
||
try {
|
||
// 调用单个SKU生成
|
||
const plan = await planImages(sku, { apiKey: process.env.API_KEY });
|
||
|
||
const skuResults = {
|
||
sku_id: sku.sku_id,
|
||
status: 'processing',
|
||
images: [],
|
||
plan_analysis: plan.analysis
|
||
};
|
||
|
||
for (const imageSpec of plan.images) {
|
||
try {
|
||
const isComparison = imageSpec.is_comparison || imageSpec.id === 'APlus_02';
|
||
const fullPrompt = injectConstraints(
|
||
imageSpec.ai_prompt,
|
||
sku,
|
||
{ isComparison, logoPlacement: imageSpec.logo_placement || {} }
|
||
);
|
||
|
||
const aspectRatio = getAspectRatio(imageSpec.id);
|
||
|
||
const payload = {
|
||
model: 'gemini-3-pro-image-preview',
|
||
prompt: fullPrompt,
|
||
n: 1,
|
||
size: "1K",
|
||
aspect_ratio: aspectRatio
|
||
};
|
||
|
||
if (sku.ref_images && sku.ref_images.length > 0) {
|
||
payload.image_urls = sku.ref_images;
|
||
}
|
||
|
||
const response = await axios.post('https://api2img.shubiaobiao.com/v1/images/generations', payload, {
|
||
headers: {
|
||
'Content-Type': 'application/json',
|
||
'Authorization': `Bearer ${process.env.API_KEY}`
|
||
},
|
||
timeout: 180000
|
||
});
|
||
|
||
let imageUrl = null;
|
||
if (response.data.data && Array.isArray(response.data.data) && response.data.data.length > 0) {
|
||
imageUrl = response.data.data[0].url;
|
||
}
|
||
|
||
skuResults.images.push({
|
||
id: imageSpec.id,
|
||
url: imageUrl,
|
||
status: imageUrl ? 'success' : 'failed'
|
||
});
|
||
|
||
// 延迟避免限流
|
||
await new Promise(resolve => setTimeout(resolve, 1500));
|
||
|
||
} catch (imgErr) {
|
||
skuResults.images.push({
|
||
id: imageSpec.id,
|
||
url: null,
|
||
status: 'failed',
|
||
error: imgErr.message
|
||
});
|
||
}
|
||
}
|
||
|
||
const successCount = skuResults.images.filter(img => img.status === 'success').length;
|
||
skuResults.status = successCount === skuResults.images.length ? 'complete' : 'partial';
|
||
skuResults.summary = {
|
||
total: skuResults.images.length,
|
||
success: successCount,
|
||
failed: skuResults.images.length - successCount
|
||
};
|
||
|
||
allResults.push(skuResults);
|
||
console.log(` SKU ${sku.sku_id} complete: ${successCount}/${skuResults.images.length} images`);
|
||
|
||
} catch (skuErr) {
|
||
console.error(` SKU ${sku.sku_id} failed:`, skuErr.message);
|
||
allResults.push({
|
||
sku_id: sku.sku_id,
|
||
status: 'failed',
|
||
error: skuErr.message,
|
||
images: []
|
||
});
|
||
}
|
||
}
|
||
|
||
const completedCount = allResults.filter(r => r.status === 'complete').length;
|
||
console.log(`\n========== Batch Complete: ${completedCount}/${skus.length} SKUs ==========\n`);
|
||
|
||
res.json({
|
||
success: true,
|
||
total_skus: skus.length,
|
||
completed: completedCount,
|
||
results: allResults
|
||
});
|
||
|
||
} catch (error) {
|
||
console.error('Batch Generate Error:', error.message);
|
||
res.status(500).json({ error: 'Batch generation failed: ' + error.message });
|
||
}
|
||
});
|
||
|
||
/**
|
||
* 获取约束预览（调试用）
|
||
* POST /api/preview-constraints
|
||
*/
|
||
app.post('/api/preview-constraints', (req, res) => {
|
||
try {
|
||
const { sku, imageSpec } = req.body;
|
||
|
||
if (!sku) {
|
||
return res.status(400).json({ error: 'SKU data is required' });
|
||
}
|
||
|
||
const isComparison = imageSpec?.is_comparison || imageSpec?.id === 'APlus_02';
|
||
const fullPrompt = injectConstraints(
|
||
imageSpec?.ai_prompt || 'Sample prompt for preview',
|
||
sku,
|
||
{
|
||
isComparison,
|
||
logoPlacement: imageSpec?.logo_placement || {}
|
||
}
|
||
);
|
||
|
||
res.json({
|
||
original_prompt: imageSpec?.ai_prompt || 'Sample prompt for preview',
|
||
full_prompt_with_constraints: fullPrompt,
|
||
prompt_length: fullPrompt.length
|
||
});
|
||
} catch (error) {
|
||
res.status(500).json({ error: error.message });
|
||
}
|
||
});
|
||
|
||
// ============================================
|
||
// 启动服务器
|
||
// ============================================
|
||
|
||
app.listen(port, () => {
|
||
console.log(`\n🚀 Server running at http://localhost:${port}`);
|
||
console.log('📁 Available endpoints:');
|
||
console.log(' GET /api/prompt-template - Get prompt template');
|
||
console.log(' GET /api/brain-prompt - Get brain system prompt');
|
||
console.log(' GET /api/sample-sku - Get sample SKU data');
|
||
console.log(' POST /api/plan-images - Stage 1: Brain planning');
|
||
console.log(' POST /api/generate-sku - Full SKU generation (12 images)');
|
||
console.log(' POST /api/batch-generate - Batch SKU generation');
|
||
console.log(' POST /api/preview-constraints - Preview constraints injection');
|
||
console.log('');
|
||
});
|