xiaohongshufabu/server/lib/generate-image.js

const { imageToPHash, isSimilar } = require('./phash');
const { v4: uuidv4 } = require('uuid');
const fs = require('fs');
const path = require('path');

async function generateImage(db, noteId, productTitle, options = {}) {
  const imgConfig = {};
  const keys = ['image_provider', 'image_api_key', 'image_base_url', 'image_prompt_template', 'image_style_options', 'image_count', 'image_size', 'image_save_path'];
  for (const k of keys) {
    const row = db.prepare('SELECT value FROM configs WHERE key = ?').get(k);
    imgConfig[k] = row?.value || '';
  }
  if (!imgConfig.image_api_key) return { imagePaths: [], imagePrompt: '' };

  const note = db.prepare('SELECT * FROM notes WHERE id = ?').get(noteId);
  if (!note) return { imagePaths: [], imagePrompt: '' };

  const imageCount = parseInt(options.image_count || imgConfig.image_count || '3', 10);
  const imageSize = options.image_size || imgConfig.image_size || '1K';

  const styles = JSON.parse(imgConfig.image_style_options || '["默认"]');
  const style = styles[Math.floor(Math.random() * styles.length)];
  const shop = db.prepare('SELECT shop_name FROM shops WHERE id = ?').get(note.shop_id);
  const prompt = (options.image_prompt_template || imgConfig.image_prompt_template || '{{title}}')
    .replace(/\{\{title\}\}/g, productTitle)
    .replace(/\{商品标题\}/g, productTitle)
    .replace(/\{\{shopName\}\}/g, shop?.shop_name || '')
    .replace(/\{店铺名\}\}/g, shop?.shop_name || '')
    .replace(/\{\{style\}\}/g, style);

  const finalPrompt = prompt.includes(productTitle) ? prompt : `${productTitle}，${prompt}`;

  const imagesRoot = imgConfig.image_save_path || path.join(__dirname, '..', 'images');
  const imagesDir = path.join(imagesRoot, note.product_id);
  if (!fs.existsSync(imagesDir)) fs.mkdirSync(imagesDir, { recursive: true });

  // Load product reference image if available
  let referenceImage = null;
  const product = db.prepare('SELECT reference_image FROM products WHERE id = ?').get(note.product_id);
  if (product?.reference_image) {
    const refPath = path.join(imagesRoot, note.product_id, 'ref.png');
    if (fs.existsSync(refPath)) {
      const refBuffer = fs.readFileSync(refPath);
      referenceImage = `data:image/png;base64,${refBuffer.toString('base64')}`;
      console.log(`[generate] Using reference image for product ${note.product_id}`);
    }
  }

  const imagePaths = [];
  for (let i = 0; i < imageCount; i++) {
    try {
      let baseUrl = (imgConfig.image_base_url || '').replace(/\/+$/, '');
      if (!baseUrl) { console.error('[generate] image_base_url is empty, skipping'); continue; }
      const apiUrl = baseUrl.endsWith('/v1') ? baseUrl + '/images/generations' : baseUrl + '/v1/images/generations';
      console.log(`[generate] Image API request: ${apiUrl} model=${imgConfig.image_provider || 'agnes-image-2.1-flash'} size=${imageSize}`);

      const requestBody = {
        model: imgConfig.image_provider || 'agnes-image-2.1-flash',
        prompt: finalPrompt,
        size: imageSize,
      };
      if (referenceImage) {
        requestBody.image = referenceImage;
      }

      // Retry logic for transient errors (502, 503, 429)
      let data = null;
      for (let attempt = 0; attempt < 3; attempt++) {
        if (attempt > 0) {
          const wait = attempt * 5000;
          console.log(`[generate] Retry ${attempt}/2 after ${wait / 1000}s...`);
          await new Promise(r => setTimeout(r, wait));
        }
        const response = await fetch(apiUrl, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${imgConfig.image_api_key}`,
          },
          body: JSON.stringify(requestBody),
        });

        if (!response.ok) {
          const errBody = await response.text().catch(() => '');
          console.error(`[generate] Image API error: ${response.status} ${errBody}`);
          if ([502, 503, 429].includes(response.status) && attempt < 2) continue;
          break;
        }
        data = await response.json();
        break;
      }

      const imageUrl = data?.data?.[0]?.url || data?.data?.[0]?.b64_json;
      if (!imageUrl) continue;

      const imageFilename = `${noteId}_${i + 1}.png`;
      const imagePath = path.join(imagesDir, imageFilename);

      if (imageUrl.startsWith('http')) {
        const imgResp = await fetch(imageUrl);
        const buffer = Buffer.from(await imgResp.arrayBuffer());
        fs.writeFileSync(imagePath, buffer);
      } else {
        fs.writeFileSync(imagePath, Buffer.from(imageUrl, 'base64'));
      }

      // pHash dedup
      const fileBuffer = fs.readFileSync(imagePath);
      const newHash = await imageToPHash(fileBuffer);
      const existingHashes = db.prepare('SELECT phash FROM image_hashes ORDER BY created_at DESC LIMIT 50').all();
      const isDuplicate = existingHashes.some(h => isSimilar(h.phash, newHash));
      if (isDuplicate) {
        console.log(`[pHash] Similar image detected, skipping: ${imageFilename}`);
        fs.unlinkSync(imagePath);
        continue;
      }

      const publicPath = `/api/images/${note.product_id}/${imageFilename}`;
      db.prepare('INSERT INTO image_hashes (id, note_id, image_path, phash) VALUES (?, ?, ?, ?)')
        .run(uuidv4(), noteId, publicPath, newHash);
      imagePaths.push(publicPath);
    } catch (e) {
      console.error(`[generate] Image generation failed:`, e.message);
    }
  }
  return { imagePaths, imagePrompt: finalPrompt };
}

module.exports = { generateImage };