宜博 / lemonboss

#!/usr/bin/env node
/**
 * generate-images.js — 为剧本生成角色图、背景图、封面图，并上传到 OSS
 *
 * 用法:
 *   node generate-images.js <drama-json文件>
 *
 * 输入 JSON 中需要包含:
 *   - characters[].imagePrompt (角色生图 prompt，由 SKILL.md 指导 Claude 生成)
 *   - backgroundPrompt (背景图 prompt)
 *   - coverPrompt (封面图 prompt)
 *   - dramaKey (用于 OSS 路径)
 *
 * 输出: 在原 JSON 基础上补充 URL 字段，写回文件
 *   - characters[].bodyUrl
 *   - coverUrl
 *   - backgroundUrl
 *
 * 依赖: ali-oss, sharp (可选), rembg (Python, 可选)
 */

const path = require('path');
const fs = require('fs');
const https = require('https');
const http = require('http');
const { execSync } = require('child_process');
const os = require('os');

// ---------- 加载 .env ----------
function loadEnv() {
    const envArgIdx = process.argv.indexOf('--env');
    const candidates = [
        envArgIdx !== -1 ? path.resolve(process.argv[envArgIdx + 1]) : null,
        process.env.DRAMA_ENV_PATH ? path.resolve(process.env.DRAMA_ENV_PATH) : null,
        path.resolve(__dirname, '.env'),
    ].filter(Boolean);

    for (const envPath of candidates) {
        if (fs.existsSync(envPath)) {
            const envContent = fs.readFileSync(envPath, 'utf-8');
            for (const line of envContent.split('\n')) {
                const trimmed = line.trim();
                if (!trimmed || trimmed.startsWith('#')) continue;
                const eqIdx = trimmed.indexOf('=');
                if (eqIdx === -1) continue;
                const key = trimmed.slice(0, eqIdx).trim();
                const val = trimmed.slice(eqIdx + 1).trim();
                if (!process.env[key]) process.env[key] = val;
            }
            console.log(`已加载配置: ${envPath}`);
            return;
        }
    }
    console.log('未找到 .env 文件，使用环境变量');
}
loadEnv();

// ========== Wanjie 图片生成 API ==========

const WANJIE_BASE = 'maas-openapi.wanjiedata.com';

function wanjieRequest(action, body) {
    const apiKey = process.env.WANJIE_API_KEY;
    if (!apiKey) throw new Error('缺少 WANJIE_API_KEY 环境变量');

    const jsonBody = JSON.stringify(body);
    return new Promise((resolve, reject) => {
        const req = https.request({
            hostname: WANJIE_BASE,
            path: `/api/jimeng/v1?Action=${action}&Version=2022-08-31`,
            method: 'POST',
            headers: {
                'Authorization': apiKey,
                'Content-Type': 'application/json',
                'Content-Length': Buffer.byteLength(jsonBody),
            },
        }, (res) => {
            let data = '';
            res.on('data', chunk => data += chunk);
            res.on('end', () => {
                try { resolve(JSON.parse(data)); }
                catch (err) { reject(new Error(`解析响应失败: ${data.slice(0, 300)}`)); }
            });
        });
        req.on('error', reject);
        req.write(jsonBody);
        req.end();
    });
}

/** Jimeng 文生图（异步模式） */
async function generateImageJimeng(prompt, width = 1024, height = 1024) {
    const res = await wanjieRequest('CVSync2AsyncSubmitTask', {
        req_key: 'jimeng_t2i_v40',
        prompt, width, height,
    });

    if (res.code !== 10000 || !res.data?.task_id) {
        throw new Error(`Jimeng 提交失败: ${JSON.stringify(res).slice(0, 300)}`);
    }

    const taskId = res.data.task_id;
    console.log(`  [Jimeng] 任务提交: task_id=${taskId}`);

    // 轮询等待
    const start = Date.now();
    while (Date.now() - start < 120000) {
        const poll = await wanjieRequest('CVSync2AsyncGetResult', {
            req_key: 'jimeng_t2i_v40', task_id: taskId,
        });

        if (poll.code === 10000 && poll.data?.image_urls?.length > 0) {
            const buffer = await downloadUrl(poll.data.image_urls[0]);
            return { buffer, mimeType: 'image/png' };
        }
        if (poll.code === 10000 && poll.data?.binary_data_base64?.length > 0) {
            return { buffer: Buffer.from(poll.data.binary_data_base64[0], 'base64'), mimeType: 'image/png' };
        }
        if (poll.code && poll.code !== 10000 && !poll.message?.includes('processing') && !poll.message?.includes('pending')) {
            throw new Error(`Jimeng 任务失败: ${JSON.stringify(poll).slice(0, 300)}`);
        }

        await sleep(3000);
    }
    throw new Error(`Jimeng 超时: task_id=${taskId}`);
}

/** Nano Banana Pro 文生图（同步模式，fallback） */
async function generateImageNano(prompt) {
    const apiKey = process.env.WANJIE_API_KEY;
    if (!apiKey) throw new Error('缺少 WANJIE_API_KEY');

    const url = `https://${WANJIE_BASE}/api/v1beta/models/gemini-3-pro-image-preview:generateContent`;
    const body = JSON.stringify({ contents: [{ parts: [{ text: prompt }] }] });

    for (let attempt = 1; attempt <= 3; attempt++) {
        try {
            const data = await httpsPost(url, body, { 'Authorization': `Bearer ${apiKey}` });
            const response = JSON.parse(data);
            const img = extractImageFromResponse(response);
            if (img) return img;
            throw new Error('无图片数据');
        } catch (err) {
            if (err.statusCode === 429 && attempt < 3) {
                console.warn(`  [Nano] 429 限流，${3 * attempt}s 后重试...`);
                await sleep(3000 * attempt);
                continue;
            }
            throw err;
        }
    }
}

function extractImageFromResponse(response) {
    const parts = response?.candidates?.[0]?.content?.parts;
    if (!parts) return null;
    for (const part of parts) {
        if (part.inlineData) {
            return { buffer: Buffer.from(part.inlineData.data, 'base64'), mimeType: part.inlineData.mimeType || 'image/png' };
        }
        if (part.text && part.text.includes('data:image')) {
            const match = part.text.match(/data:image\/(png|jpeg);base64,([A-Za-z0-9+/=]+)/);
            if (match) return { buffer: Buffer.from(match[2], 'base64'), mimeType: `image/${match[1]}` };
        }
    }
    return null;
}

/** 生图：优先 Jimeng，fallback Nano */
async function generateImage(prompt) {
    try {
        return await generateImageJimeng(prompt);
    } catch (jimengErr) {
        console.warn(`  [Jimeng] 失败: ${jimengErr.message.slice(0, 80)}，fallback Nano...`);
        return await generateImageNano(prompt);
    }
}

// ========== 图片处理 ==========

/** 去背景（rembg） */
function removeBg(inputBuffer) {
    const tmpDir = os.tmpdir();
    const ts = Date.now();
    const inputPath = path.join(tmpDir, `rembg_in_${ts}.png`);
    const outputPath = path.join(tmpDir, `rembg_out_${ts}.png`);

    try {
        fs.writeFileSync(inputPath, inputBuffer);
        execSync(`rembg i "${inputPath}" "${outputPath}"`, { timeout: 60000 });
        const result = fs.readFileSync(outputPath);
        console.log('  [rembg] 去背景成功');

        // 尝试 sharp 裁剪
        try {
            const sharp = require('sharp');
            // sharp 操作是 async，但我们在 sync context 用 execSync 调 sharp
            // 改用 async
            return trimImage(result);
        } catch {
            return result;
        }
    } catch (err) {
        console.warn(`  [rembg] 失败(${err.message.slice(0, 60)})，返回原图`);
        return inputBuffer;
    } finally {
        try { fs.unlinkSync(inputPath); } catch { }
        try { fs.unlinkSync(outputPath); } catch { }
    }
}

async function trimImage(buffer) {
    try {
        const sharp = require('sharp');
        const trimmed = await sharp(buffer).trim().png().toBuffer();
        const meta = await sharp(trimmed).metadata();
        const padX = Math.round(meta.width * 0.05);
        const padY = Math.round(meta.height * 0.05);
        return await sharp(trimmed)
            .extend({ top: padY, bottom: padY, left: padX, right: padX, background: { r: 0, g: 0, b: 0, alpha: 0 } })
            .png().toBuffer();
    } catch {
        return buffer;
    }
}

// ========== OSS 上传 ==========

function getOSSClient() {
    const OSS = require('ali-oss');
    return new OSS({
        region: process.env.OSS_REGION,
        accessKeyId: process.env.OSS_ACCESS_KEY_ID,
        accessKeySecret: process.env.OSS_ACCESS_KEY_SECRET,
        bucket: process.env.OSS_BUCKET,
    });
}

function getPublicUrl(key) {
    const endpoint = process.env.OSS_ENDPOINT || `https://${process.env.OSS_REGION}.aliyuncs.com`;
    const bucket = process.env.OSS_BUCKET;
    return `https://${bucket}.${endpoint.replace('https://', '')}/${key}`;
}

async function uploadToOSS(ossKey, buffer, contentType = 'image/png') {
    const client = getOSSClient();
    const result = await client.put(ossKey, buffer, { headers: { 'Content-Type': contentType } });
    return getPublicUrl(result.name);
}

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

function sleep(ms) { return new Promise(r => setTimeout(r, ms)); }

function httpsPost(url, body, extraHeaders = {}) {
    return new Promise((resolve, reject) => {
        const urlObj = new URL(url);
        const req = https.request({
            hostname: urlObj.hostname,
            path: urlObj.pathname + urlObj.search,
            method: 'POST',
            headers: { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(body), ...extraHeaders },
        }, (res) => {
            let data = '';
            res.on('data', chunk => data += chunk);
            res.on('end', () => {
                if (res.statusCode === 429) { const e = new Error('429'); e.statusCode = 429; reject(e); return; }
                resolve(data);
            });
        });
        req.on('error', reject);
        req.write(body);
        req.end();
    });
}

function downloadUrl(url) {
    return new Promise((resolve, reject) => {
        const protocol = url.startsWith('https') ? https : http;
        protocol.get(url, (res) => {
            if (res.statusCode >= 300 && res.statusCode < 400 && res.headers.location) {
                return downloadUrl(res.headers.location).then(resolve).catch(reject);
            }
            const chunks = [];
            res.on('data', chunk => chunks.push(chunk));
            res.on('end', () => resolve(Buffer.concat(chunks)));
            res.on('error', reject);
        }).on('error', reject);
    });
}

// ========== 主逻辑 ==========

async function main() {
    const filePath = process.argv.find(a => a.endsWith('.json'));
    if (!filePath) {
        console.error('用法: node generate-images.js <drama-json文件>');
        process.exit(1);
    }

    const absPath = path.resolve(filePath);
    const data = JSON.parse(fs.readFileSync(absPath, 'utf-8'));
    const dramaKey = data.dramaKey || `drama_${Date.now().toString(36)}`;
    data.dramaKey = dramaKey;

    // 检查 OSS 配置
    const ossOk = process.env.OSS_REGION && process.env.OSS_ACCESS_KEY_ID && process.env.OSS_ACCESS_KEY_SECRET && process.env.OSS_BUCKET;
    if (!ossOk) {
        console.error('OSS 配置不完整，请检查 .env 中的 OSS_REGION, OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET, OSS_BUCKET');
        process.exit(1);
    }

    // 检查生图 API 配置
    if (!process.env.WANJIE_API_KEY) {
        console.error('缺少 WANJIE_API_KEY');
        process.exit(1);
    }

    const characters = data.characters || [];
    const total = characters.length + 2; // chars + background + cover
    let completed = 0;
    const tick = (label) => { completed++; console.log(`[${completed}/${total}] ${label} 完成`); };

    // 并行处理所有图片
    const tasks = [];

    // 角色图
    for (const char of characters) {
        tasks.push((async () => {
            const prompt = char.imagePrompt || `全身立绘，纯色背景，沙雕动漫风格。角色：${char.name}。脸部空白无五官。`;
            console.log(`\n生成角色图: ${char.name}`);
            console.log(`  prompt: ${prompt.slice(0, 80)}...`);

            const { buffer } = await generateImage(prompt);

            // 去背景
            console.log(`  正在去背景...`);
            const processed = await removeBg(buffer);

            // 上传 OSS
            const safeName = Buffer.from(char.name).toString('hex').slice(0, 16);
            const ossKey = `dramas/${dramaKey}/characters/${safeName}.png`;
            const url = await uploadToOSS(ossKey, processed);
            char.bodyUrl = url;
            console.log(`  上传完成: ${url}`);
            tick(char.name);
        })());
    }

    // 背景图
    tasks.push((async () => {
        const prompt = data.backgroundPrompt || `横版场景背景图，宽屏比例，场景：${data.description || data.name}，卡通风格，无人物。`;
        console.log(`\n生成背景图`);
        console.log(`  prompt: ${prompt.slice(0, 80)}...`);

        const { buffer } = await generateImage(prompt);
        const ossKey = `dramas/${dramaKey}/background.png`;
        const url = await uploadToOSS(ossKey, buffer);
        data.backgroundUrl = url;
        console.log(`  上传完成: ${url}`);
        tick('背景图');
    })());

    // 封面图
    tasks.push((async () => {
        const prompt = data.coverPrompt || `剧目封面海报，标题：${data.title || data.name}，色彩鲜明。`;
        console.log(`\n生成封面图`);
        console.log(`  prompt: ${prompt.slice(0, 80)}...`);

        const { buffer } = await generateImage(prompt);
        const ossKey = `dramas/${dramaKey}/cover.png`;
        const url = await uploadToOSS(ossKey, buffer);
        data.coverUrl = url;
        console.log(`  上传完成: ${url}`);
        tick('封面图');
    })());

    await Promise.allSettled(tasks);

    // 写回文件
    fs.writeFileSync(absPath, JSON.stringify(data, null, 2), 'utf-8');
    console.log(`\n图片生成完成，已更新 JSON 文件: ${absPath}`);
    console.log(`  角色图: ${characters.filter(c => c.bodyUrl).length}/${characters.length}`);
    console.log(`  背景图: ${data.backgroundUrl ? '✓' : '✗'}`);
    console.log(`  封面图: ${data.coverUrl ? '✓' : '✗'}`);
}

main().catch(err => {
    console.error('执行失败:', err);
    process.exit(1);
});