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