#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ WordPress 发布系统 - AI 图片生成模块 基于阿里云 DashScope API 实现文生图功能 支持：通义万相 (wanx-v1) """ import os import json import time import requests from datetime import datetime from modules.wp_logger import get_publish_logger, get_debug_logger # 基础目录 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) IMAGE_DIR = os.path.join(BASE_DIR, 'temp') os.makedirs(IMAGE_DIR, exist_ok=True) class ImageGenerator: """AI 图片生成器""" # 支持的模型列表 SUPPORTED_MODELS = { 'wanx-v1': { 'name': '通义万相 v1', 'base_url': 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis', 'sizes': ['1024*1024', '720*1280', '1280*720', '512*512', '512*1024', '1024*512'], 'max_images': 4 }, 'wanx2.1-t2i-turbo': { 'name': '通义万相 v2 Turbo', 'base_url': 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis', 'sizes': ['1024*1024', '720*1280', '1280*720', '512*512'], 'max_images': 4 }, 'wanx2.1-t2i-plus': { 'name': '通义万相 v2 Plus', 'base_url': 'https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis', 'sizes': ['1024*1024', '720*1280', '1280*720', '512*512'], 'max_images': 4 } } # 默认尺寸 DEFAULT_SIZE = '1024*1024' def __init__(self, api_key=None, model='wanx-v1', size=None): """ 初始化图片生成器 Args: api_key: DashScope API Key model: 模型名称 size: 图片尺寸 """ self.api_key = api_key or self._get_api_key_from_env() self.model = model if model in self.SUPPORTED_MODELS else 'wanx-v1' self.size = size or self.DEFAULT_SIZE # 验证尺寸 model_config = self.SUPPORTED_MODELS[self.model] if self.size not in model_config['sizes']: self.size = self.DEFAULT_SIZE self.base_url = model_config['base_url'] self.max_images = model_config['max_images'] self.pl = get_publish_logger() self.dl = get_debug_logger() def generate_image(self, prompt, negative_prompt=None, n=1, style=None): """ 生成图片（异步提交 + 轮询结果） Args: prompt: 图片描述（中文） negative_prompt: 反向提示词 n: 生成图片数量 (1-4) style: 风格（如：写实、动漫、水彩等） Returns: list: 生成的图片本地路径列表 """ if n < 1 or n > self.max_images: n = 1 self.pl.info(f"🎨 AI 生图开始 - 模型：{self.model}, 尺寸：{self.size}, 数量：{n}") self.dl.log_step("AI 生图", f"模型：{self.model}, 尺寸：{self.size}") self.dl.debug(f"提示词：{prompt}") if negative_prompt: self.dl.debug(f"反向提示词：{negative_prompt}") try: # 提交生图任务 task_id = self._submit_task(prompt, negative_prompt, n, style) if not task_id: return [] # 轮询等待结果 image_urls = self._poll_task(task_id) if not image_urls: self.pl.error("生图任务失败：未获取到图片 URL") return [] # 下载图片到本地 local_paths = self._download_images(image_urls, prompt) self.pl.success(f"AI 生图完成 - 成功 {len(local_paths)} 张图片") self.dl.log_result("生图结果", { 'model': self.model, 'count': len(local_paths), 'paths': local_paths }) return local_paths except Exception as e: self.pl.error(f"AI 生图异常：{str(e)}") self.dl.error(f"生图异常：{str(e)}", exc_info=True) return [] def _submit_task(self, prompt, negative_prompt=None, n=1, style=None): """提交生图任务""" headers = { 'Authorization': f'Bearer {self.api_key}', 'Content-Type': 'application/json', 'X-DashScope-Async': 'enable' } payload = { 'model': self.model, 'input': { 'prompt': prompt }, 'parameters': { 'size': self.size, 'n': n } } if negative_prompt: payload['parameters']['negative_prompt'] = negative_prompt if style: payload['parameters']['style'] = style try: response = requests.post( self.base_url, headers=headers, json=payload, timeout=30 ) self.dl.debug(f"提交任务响应：{response.text[:500]}") if response.status_code == 200: result = response.json() task_id = result.get('output', {}).get('task_id') if task_id: self.pl.info(f"📋 生图任务已提交 - Task ID: {task_id}") self.dl.debug(f"任务 ID: {task_id}") return task_id else: self.pl.error(f"未获取到任务 ID：{result}") return None else: error_msg = response.text self.pl.error(f"提交生图任务失败 - 状态码：{response.status_code}") self.pl.error(f"错误详情：{error_msg}") self.dl.error(f"提交失败：{error_msg}") return None except requests.exceptions.Timeout: self.pl.error("提交生图任务超时") return None except Exception as e: self.pl.error(f"提交生图任务异常：{str(e)}") return None def _poll_task(self, task_id, max_retries=60, interval=5): """轮询任务状态""" url = f'https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}' headers = { 'Authorization': f'Bearer {self.api_key}' } self.pl.info(f"⏳ 等待生图完成...") for attempt in range(max_retries): try: response = requests.get(url, headers=headers, timeout=30) if response.status_code != 200: self.dl.warning(f"查询任务状态失败：{response.status_code}") time.sleep(interval) continue result = response.json() status = result.get('output', {}).get('task_status', '') self.dl.debug(f"任务状态：{status} (尝试 {attempt + 1}/{max_retries})") if status == 'SUCCEEDED': image_urls = [] results = result.get('output', {}).get('results', []) for item in results: url = item.get('url', '') if url: image_urls.append(url) if image_urls: self.pl.success(f"✅ 生图任务完成 - 获取到 {len(image_urls)} 张图片 URL") return image_urls else: self.pl.error("任务完成但未获取到图片 URL") return [] elif status == 'FAILED': message = result.get('output', {}).get('message', '未知错误') self.pl.error(f"生图任务失败：{message}") self.dl.error(f"任务失败：{message}") return [] elif status in ['PENDING', 'RUNNING']: # 继续等待 if (attempt + 1) % 6 == 0: # 每 30 秒输出一次日志 self.pl.info(f"⏳ 生图进行中... ({attempt * interval} 秒)") time.sleep(interval) continue else: self.dl.warning(f"未知任务状态：{status}") time.sleep(interval) continue except requests.exceptions.Timeout: self.dl.warning("查询任务状态超时") time.sleep(interval) continue except Exception as e: self.pl.error(f"查询任务状态异常：{str(e)}") time.sleep(interval) continue self.pl.error(f"生图任务超时（{max_retries * interval} 秒）") return [] def _download_images(self, image_urls, prompt): """下载图片到本地""" local_paths = [] for i, url in enumerate(image_urls): try: response = requests.get(url, timeout=30) if response.status_code == 200: # 生成文件名 timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') prompt_hash = abs(hash(prompt)) % 10000 filename = f"ai_image_{timestamp}_{i+1}_{prompt_hash}.png" filepath = os.path.join(IMAGE_DIR, filename) with open(filepath, 'wb') as f: f.write(response.content) file_size = os.path.getsize(filepath) local_paths.append(filepath) self.dl.debug(f"图片已保存：{filepath} ({file_size} 字节)") else: self.pl.error(f"下载图片失败：{url} (状态码：{response.status_code})") except Exception as e: self.pl.error(f"下载图片异常：{url} - {str(e)}") return local_paths def _get_api_key_from_env(self): """从环境变量获取 API Key""" return os.environ.get('DASHSCOPE_API_KEY', '') def get_supported_models(self): """获取支持的模型列表""" return list(self.SUPPORTED_MODELS.keys()) def get_model_info(self, model=None): """获取模型信息""" if model: return self.SUPPORTED_MODELS.get(model, {}) return self.SUPPORTED_MODELS def create_image_generator(api_key=None, model='wanx-v1', size=None): """创建图片生成器实例""" return ImageGenerator(api_key=api_key, model=model, size=size) if __name__ == '__main__': import sys if len(sys.argv) < 3: print("用法：python wp_image_generator.py [model] [size] [count]") print("示例：python wp_image_generator.py sk-xxx '一只可爱的猫咪' wanx-v1 1024*1024 1") sys.exit(1) api_key = sys.argv[1] prompt = sys.argv[2] model = sys.argv[3] if len(sys.argv) > 3 else 'wanx-v1' size = sys.argv[4] if len(sys.argv) > 4 else '1024*1024' count = int(sys.argv[5]) if len(sys.argv) > 5 else 1 generator = create_image_generator(api_key=api_key, model=model, size=size) paths = generator.generate_image(prompt, n=count) print(f"\n生成结果：") for path in paths: print(f" {path}")