a0362e0ba8
新增内容： - modules/wp_image_generator.py: AI 图片生成核心模块 - scripts/wp_generate_image.py: 命令行生图脚本 - feishu_bot.py: 新增#生图、#发布并生图 指令支持 - config.py: 添加 DashScope API Key 和生图配置 支持功能： - 独立 AI 生图：#生图 图片描述 - 发布+生图：#发布并生图（文章自动配图） - 可配置参数：模型、尺寸、数量、风格 - 异步任务轮询，支持 1-4 张图片并发 测试结果： - ✅ 生图任务提交成功 - ✅ 图片生成并保存到本地（1024x1024 PNG） - ✅ 文件大小约 1.3MB
324 lines
12 KiB
Python
324 lines
12 KiB
Python
#!/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 <api_key> <prompt> [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}")
|