amz-pic-flow/upscale-to-4k-ai.js

/**
 * 4K AI 升级脚本 (1:1 还原版)
 * 逻辑：
 * 1. 读取 v2 版本已生成的图片（客户已对齐的版本）
 * 2. 将其作为 img_url (参考图) 传给 AI
 * 3. 使用原 prompt，并开启 imageSize: '4K'
 * 4. 这样 AI 会在 v2 的基础上进行高清化重绘，保持构图 1:1 还原
 */

require('dotenv').config();
const axios = require('axios');
const fs = require('fs');
const path = require('path');
const { S3Client, PutObjectCommand } = require('@aws-sdk/client-s3');
const imageProcessor = require('./lib/image-processor');

// ============================================================
// 配置
// ============================================================
const API_KEY = 'G9rXx3Ag2Xfa7Gs8zou6t6HqeZ';
const API_BASE = 'https://api.wuyinkeji.com/api';

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,
});

// ============================================================
// 工具函数
// ============================================================

async function uploadToR2(buffer, filename) {
  const fileName = `upscale-4k-${Date.now()}-${filename}`;
  await r2Client.send(new PutObjectCommand({
    Bucket: process.env.R2_BUCKET_NAME || 'ai-flow',
    Key: fileName,
    Body: buffer,
    ContentType: '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 submitUpscaleTask(prompt, refImageUrl, aspectRatio) {
  const payload = {
    key: API_KEY,
    prompt: prompt, // 使用原 prompt 保持语义一致
    img_url: refImageUrl, // 使用已生成的图作为底图
    aspectRatio: aspectRatio,
    imageSize: '4K' // 升级到 4K
  };

  const response = await axios.post(`${API_BASE}/img/nanoBanana-pro`, payload, { timeout: 60000 });
  return response.data.data?.id || response.data.id;
}

async function pollImageResult(taskId) {
  let attempts = 0;
  const maxAttempts = 120; // 4K 稍微给多点时间
  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;
      if (data && data.status === 3) throw new Error('Upscale failed: ' + (data.fail_reason || 'Unknown'));
      
      process.stdout.write('.');
      await new Promise(r => setTimeout(r, 2000));
      attempts++;
    } catch (error) {
      if (error.message.includes('Upscale failed')) throw error;
      process.stdout.write('x');
      await new Promise(r => setTimeout(r, 3000));
      attempts++;
    }
  }
  throw new Error('Timeout');
}

async function downloadImage(url) {
  const response = await axios.get(url, { responseType: 'arraybuffer', timeout: 30000 });
  return Buffer.from(response.data);
}

// ============================================================
// 主流程
// ============================================================

async function main() {
  const v2Dir = path.join(__dirname, 'output_carrier_v2_2025-12-27');
  const outputDir = path.join(__dirname, `output_carrier_v3_4K_Upscale_${new Date().toISOString().slice(0, 19).replace(/:/g, '-')}`);

  if (!fs.existsSync(v2Dir)) {
    console.error('❌ 找不到 v2 目录:', v2Dir);
    return;
  }

  if (!fs.existsSync(outputDir)) fs.mkdirSync(outputDir, { recursive: true });

  console.log('🚀 开始 4K AI 升级 (基于 v2 已对齐版本)...');
  console.log('📂 输入目录:', v2Dir);
  console.log('📂 输出目录:', outputDir);

  const files = fs.readdirSync(v2Dir).filter(f => f.endsWith('.jpg'));
  
  for (const file of files) {
    const id = file.replace('.jpg', '');
    const promptFile = path.join(v2Dir, `${id}_prompt.txt`);
    
    if (!fs.existsSync(promptFile)) {
      console.log(`⏭️  跳过 ${file} (找不到对应的 prompt 文件)`);
      continue;
    }

    console.log(`\n🎨 升级 [${id}]...`);
    
    try {
      const prompt = fs.readFileSync(promptFile, 'utf-8');
      const imgBuffer = fs.readFileSync(path.join(v2Dir, file));
      const aspectRatio = id.startsWith('APlus') ? '3:2' : '1:1';

      // 1. 上传 v2 图片到 R2
      const imgUrl = await uploadToR2(imgBuffer, file);
      
      // 2. 提交 4K 任务
      const taskId = await submitUpscaleTask(prompt, imgUrl, aspectRatio);
      console.log(`   Task ID: ${taskId}`);

      // 3. 轮询结果
      const resultUrl = await pollImageResult(taskId);
      console.log('\n   ✓ 4K 生成完成');

      // 4. 下载并保存
      const resultBuffer = await downloadImage(resultUrl);
      await imageProcessor.saveImage(resultBuffer, path.join(outputDir, file), 'jpeg', 95);
      
      // 拷贝 prompt 文件
      fs.copyFileSync(promptFile, path.join(outputDir, `${id}_prompt.txt`));
      
      console.log(`✅ ${file} 已升级到 4K`);
    } catch (e) {
      console.error(`\n❌ ${id} 失败:`, e.message);
    }
  }

  console.log('\n✨ 全部升级任务完成！');
}

main().catch(console.error);