feishu_fabu/scripts/wp_generate_image.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WordPress 发布系统 - AI 图片生成脚本
支持命令行调用和模块调用
支持格式控制：jpg, png, webp
"""

import os
import sys
import json
import argparse

# 添加项目根目录到 Python 路径
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BASE_DIR)

from modules.wp_image_generator import create_image_generator
from modules.wp_logger import get_publish_logger, get_debug_logger


def load_config():
    """加载配置文件"""
    config = {
        'dashscope_api_key': '',
        'image_model': 'wanx-v1',
        'image_size': '1024*1024',
        'image_format': 'webp',
        'image_count': 1,
        'image_style': None
    }
    
    config_file = os.path.join(BASE_DIR, 'config.py')
    if os.path.exists(config_file):
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                exec(f.read(), config)
        except Exception as e:
            print(f"加载配置文件失败：{str(e)}，使用默认配置")
    
    return config


def generate_image(prompt, api_key=None, model=None, size=None, count=None, 
                   negative_prompt=None, style=None, image_format=None):
    """
    生成图片（主入口函数）
    
    Args:
        prompt: 图片描述
        api_key: DashScope API Key（可选）
        model: 模型名称（可选）
        size: 图片尺寸（可选）
        count: 生成数量（可选）
        negative_prompt: 反向提示词（可选）
        style: 风格（可选）
        image_format: 图片格式 (jpg, png, webp)（可选）
        
    Returns:
        dict: 生成结果
    """
    pl = get_publish_logger()
    dl = get_debug_logger()
    
    # 加载配置
    config = load_config()
    
    # 使用传入参数或配置
    api_key = api_key or config.get('dashscope_api_key', '')
    model = model or config.get('image_model', 'wanx-v1')
    size = size or config.get('image_size', '1024*1024')
    count = count or config.get('image_count', 1)
    style = style or config.get('image_style', None)
    image_format = image_format or config.get('image_format', 'webp')
    
    if not api_key:
        pl.error("未配置 DashScope API Key")
        return {'success': False, 'error': '未配置 DashScope API Key'}
    
    pl.start_publish('AI 图片生成', prompt)
    
    try:
        # 创建生成器
        generator = create_image_generator(
            api_key=api_key,
            model=model,
            size=size,
            image_format=image_format
        )
        
        # 生成图片
        image_paths = generator.generate_image(
            prompt=prompt,
            negative_prompt=negative_prompt,
            n=count,
            style=style
        )
        
        if image_paths:
            pl.end_publish(True, post_id=None, post_url=f"生成 {len(image_paths)} 张图片")
            return {
                'success': True,
                'count': len(image_paths),
                'paths': image_paths,
                'model': model,
                'size': size,
                'format': image_format
            }
        else:
            pl.end_publish(False, error_msg='生成图片失败')
            return {
                'success': False,
                'error': '生成图片失败'
            }
    
    except Exception as e:
        pl.end_publish(False, error_msg=str(e))
        dl.error(f"生图异常：{str(e)}", exc_info=True)
        return {
            'success': False,
            'error': str(e)
        }


def generate_images_for_article(title, content, count=1, api_key=None, image_format='webp'):
    """
    根据文章标题和内容生成配图
    
    Args:
        title: 文章标题
        content: 文章内容
        count: 生成数量
        api_key: DashScope API Key（可选）
        image_format: 图片格式 (jpg, png, webp)
        
    Returns:
        dict: 生成结果
    """
    pl = get_publish_logger()
    dl = get_debug_logger()
    
    # 加载配置
    config = load_config()
    api_key = api_key or config.get('dashscope_api_key', '')
    image_format = image_format or config.get('image_format', 'webp')
    
    if not api_key:
        pl.error("未配置 DashScope API Key")
        return {'success': False, 'error': '未配置 DashScope API Key'}
    
    pl.info(f"🎨 为文章生成配图 - 标题：{title}")
    
    # 提取内容前 500 字符作为提示
    content_preview = content[:500] if len(content) > 500 else content
    
    # 构建提示词
    prompt = f"高质量插画，主题：{title}。内容描述：{content_preview}"
    
    # 调用生图
    result = generate_image(
        prompt=prompt,
        api_key=api_key,
        count=count,
        image_format=image_format
    )
    
    return result


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(description='WordPress AI 图片生成工具')
    parser.add_argument('prompt', help='图片描述')
    parser.add_argument('--model', '-m', default=None, help='模型名称 (wanx-v1, wanx2.1-t2i-turbo, wanx2.1-t2i-plus)')
    parser.add_argument('--size', '-s', default=None, help='图片尺寸 (1024*1024, 720*1280, 1280*720 等)')
    parser.add_argument('--format', '-f', default=None, choices=['jpg', 'png', 'webp'], help='图片格式 (jpg, png, webp)')
    parser.add_argument('--count', '-n', type=int, default=None, help='生成数量 (1-4)')
    parser.add_argument('--style', default=None, help='图片风格 (写实、动漫、水彩等)')
    parser.add_argument('--negative', default=None, help='反向提示词')
    parser.add_argument('--api-key', default=None, help='DashScope API Key')
    
    args = parser.parse_args()
    
    result = generate_image(
        prompt=args.prompt,
        api_key=args.api_key,
        model=args.model,
        size=args.size,
        count=args.count,
        negative_prompt=args.negative,
        style=args.style,
        image_format=args.format
    )
    
    # 输出 JSON 结果
    print("\n" + json.dumps(result, ensure_ascii=False, indent=2))
    
    # 返回状态码
    sys.exit(0 if result.get('success') else 1)


if __name__ == '__main__':
    main()