feat: 新增 AI 图片生成功能(通义万相 wanx-v1)
新增内容: - 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
This commit is contained in:
parent
9a5b0828f7
commit
a0362e0ba8
253
feishu_bot.py
253
feishu_bot.py
@ -3,6 +3,7 @@
|
||||
"""
|
||||
飞书机器人消息接收与处理脚本
|
||||
接收飞书消息,解析内容,调用 WordPress 发布脚本
|
||||
支持 AI 图片生成
|
||||
"""
|
||||
|
||||
import os
|
||||
@ -56,6 +57,7 @@ class FeishuBot:
|
||||
self.config = self._load_config()
|
||||
self.wp_script = os.path.join(BASE_DIR, 'scripts', 'wp_publish_text.py')
|
||||
self.word_script = os.path.join(BASE_DIR, 'scripts', 'wp_publish.py')
|
||||
self.image_gen_script = os.path.join(BASE_DIR, 'scripts', 'wp_generate_image.py')
|
||||
|
||||
# 初始化飞书 API 客户端
|
||||
self.feishu_client = create_feishu_client(
|
||||
@ -66,6 +68,7 @@ class FeishuBot:
|
||||
logger.info("🤖 飞书机器人初始化完成")
|
||||
logger.info(f" WordPress 发布脚本:{self.wp_script}")
|
||||
logger.info(f" Word 发布脚本:{self.word_script}")
|
||||
logger.info(f" AI 生图脚本:{self.image_gen_script}")
|
||||
logger.info(f" 飞书 API 客户端:已初始化")
|
||||
|
||||
def _load_config(self):
|
||||
@ -155,8 +158,30 @@ class FeishuBot:
|
||||
if message_id:
|
||||
images = self._get_message_images(message_id)
|
||||
|
||||
# 检查是否为生图指令
|
||||
if instruction.get('action') == 'generate_image':
|
||||
return self._generate_image(
|
||||
prompt=instruction.get('prompt', ''),
|
||||
count=instruction.get('image_count', 1),
|
||||
model=instruction.get('image_model', None),
|
||||
size=instruction.get('image_size', None),
|
||||
style=instruction.get('image_style', None)
|
||||
)
|
||||
# 检查是否为发布+生图指令
|
||||
elif instruction.get('action') == 'publish_with_image':
|
||||
return self._publish_with_generated_image(
|
||||
text=instruction.get('text', ''),
|
||||
title=instruction.get('title', ''),
|
||||
category=instruction.get('category', ''),
|
||||
tags=instruction.get('tags', ''),
|
||||
status=instruction.get('status', 'publish'),
|
||||
image_count=instruction.get('image_count', 1),
|
||||
model=instruction.get('image_model', None),
|
||||
size=instruction.get('image_size', None),
|
||||
style=instruction.get('image_style', None)
|
||||
)
|
||||
# 检查是否为发布指令
|
||||
if instruction.get('action') == 'publish':
|
||||
elif instruction.get('action') == 'publish':
|
||||
return self._publish_article(
|
||||
text=instruction.get('text', ''),
|
||||
title=instruction.get('title', ''),
|
||||
@ -343,7 +368,12 @@ class FeishuBot:
|
||||
'title': '',
|
||||
'category': '',
|
||||
'tags': '',
|
||||
'status': 'publish'
|
||||
'status': 'publish',
|
||||
'prompt': '',
|
||||
'image_count': 1,
|
||||
'image_model': None,
|
||||
'image_size': None,
|
||||
'image_style': None
|
||||
}
|
||||
|
||||
lines = content.strip().split('\n')
|
||||
@ -375,6 +405,29 @@ class FeishuBot:
|
||||
elif line.startswith('#更新') or line.startswith('#update'):
|
||||
instruction['action'] = 'update'
|
||||
instruction['target'] = line.replace('#更新', '').replace('#update', '').strip()
|
||||
# AI 生图指令
|
||||
elif line.startswith('#生图') or line.startswith('#生成图片'):
|
||||
instruction['action'] = 'generate_image'
|
||||
prompt_text = line.replace('#生图', '').replace('#生成图片', '').strip()
|
||||
instruction['prompt'] = prompt_text
|
||||
elif line.startswith('#生图数量') or line.startswith('#图片数量'):
|
||||
try:
|
||||
count = int(line.replace('#生图数量', '').replace('#图片数量', '').strip())
|
||||
instruction['image_count'] = min(max(count, 1), 4)
|
||||
except ValueError:
|
||||
pass
|
||||
elif line.startswith('#生图模型') or line.startswith('#图片模型'):
|
||||
model = line.replace('#生图模型', '').replace('#图片模型', '').strip()
|
||||
instruction['image_model'] = model
|
||||
elif line.startswith('#生图尺寸') or line.startswith('#图片尺寸'):
|
||||
size = line.replace('#生图尺寸', '').replace('#图片尺寸', '').strip()
|
||||
instruction['image_size'] = size
|
||||
elif line.startswith('#生图风格') or line.startswith('#图片风格'):
|
||||
style = line.replace('#生图风格', '').replace('#图片风格', '').strip()
|
||||
instruction['image_style'] = style
|
||||
# 发布+生图指令
|
||||
elif line.startswith('#发布并生图') or line.startswith('#发布+生图'):
|
||||
instruction['action'] = 'publish_with_image'
|
||||
else:
|
||||
text_lines.append(line)
|
||||
|
||||
@ -476,6 +529,165 @@ class FeishuBot:
|
||||
logger.error(f"发布文章失败:{str(e)}", exc_info=True)
|
||||
return f"❌ 发布失败:{str(e)}"
|
||||
|
||||
def _publish_with_generated_image(self, text='', title='', category='', tags='', status='publish',
|
||||
image_count=1, model=None, size=None, style=None):
|
||||
"""
|
||||
发布文章并自动生成配图
|
||||
|
||||
Args:
|
||||
text: 文章正文
|
||||
title: 文章标题
|
||||
category: 分类
|
||||
tags: 标签
|
||||
status: 发布状态
|
||||
image_count: 生成图片数量
|
||||
model: 生图模型
|
||||
size: 图片尺寸
|
||||
style: 图片风格
|
||||
|
||||
Returns:
|
||||
str: 回复消息
|
||||
"""
|
||||
if not text and not title:
|
||||
return "⚠️ 请提供文章标题或内容"
|
||||
|
||||
logger.info(f"📝 准备发布文章(含 AI 配图)- 标题:{title}")
|
||||
|
||||
try:
|
||||
# 加载配置获取 API Key
|
||||
config = self._load_image_gen_config()
|
||||
|
||||
if not config.get('dashscope_api_key'):
|
||||
return "⚠️ 未配置 DashScope API Key,无法生成图片\n\n请在 config.py 中配置 dashscope_api_key"
|
||||
|
||||
# 生成图片
|
||||
from scripts.wp_generate_image import generate_images_for_article
|
||||
|
||||
pl = logger
|
||||
pl.info(f"🎨 开始生成 AI 配图...")
|
||||
|
||||
image_result = generate_images_for_article(
|
||||
title=title or '自动提取',
|
||||
content=text,
|
||||
count=image_count,
|
||||
api_key=config.get('dashscope_api_key')
|
||||
)
|
||||
|
||||
if not image_result.get('success'):
|
||||
return f"⚠️ AI 配图生成失败:{image_result.get('error', '未知错误')}\n\n文章仍可发布,但不含配图。"
|
||||
|
||||
image_paths = image_result.get('paths', [])
|
||||
pl.info(f"✅ 已生成 {len(image_paths)} 张配图")
|
||||
|
||||
# 发布文章(带图片)
|
||||
from scripts.wp_publish_text import publish_text_with_images
|
||||
|
||||
category_id = self._resolve_category(category) if category else None
|
||||
|
||||
publish_result = publish_text_with_images(
|
||||
text=text,
|
||||
images=image_paths,
|
||||
instruction=f"#分类 {category}" if category else None,
|
||||
status=status,
|
||||
category_id=category_id,
|
||||
title=title if title else None
|
||||
)
|
||||
|
||||
# 处理结果
|
||||
if publish_result.get('success'):
|
||||
post_url = publish_result.get('post_url', '')
|
||||
post_id = publish_result.get('post_id', '')
|
||||
|
||||
reply = "✅ 文章发布成功(含 AI 配图)!\n"
|
||||
reply += f"📝 标题:{publish_result.get('title', title or '自动提取')}\n"
|
||||
reply += f"🔗 链接:{post_url}\n"
|
||||
reply += f"📊 文章 ID:{post_id}\n"
|
||||
reply += f"🎨 AI 配图:{len(image_paths)} 张"
|
||||
|
||||
return reply
|
||||
else:
|
||||
error_msg = publish_result.get('error', '未知错误')
|
||||
return f"❌ 发布失败:{error_msg}"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"发布文章(含 AI 配图)失败:{str(e)}", exc_info=True)
|
||||
return f"❌ 发布失败:{str(e)}"
|
||||
|
||||
def _generate_image(self, prompt='', count=1, model=None, size=None, style=None):
|
||||
"""
|
||||
生成 AI 图片
|
||||
|
||||
Args:
|
||||
prompt: 图片描述
|
||||
count: 生成数量
|
||||
model: 模型名称
|
||||
size: 图片尺寸
|
||||
style: 图片风格
|
||||
|
||||
Returns:
|
||||
str: 回复消息
|
||||
"""
|
||||
if not prompt:
|
||||
return "⚠️ 请提供图片描述\n\n示例:`#生图 一只可爱的猫咪在草地上晒太阳`"
|
||||
|
||||
logger.info(f"🎨 开始 AI 生图 - 描述:{prompt}")
|
||||
|
||||
try:
|
||||
# 加载配置获取 API Key
|
||||
config = self._load_image_gen_config()
|
||||
|
||||
if not config.get('dashscope_api_key'):
|
||||
return "⚠️ 未配置 DashScope API Key,无法生成图片\n\n请在 config.py 中配置 dashscope_api_key"
|
||||
|
||||
# 调用生图脚本
|
||||
from scripts.wp_generate_image import generate_image
|
||||
|
||||
result = generate_image(
|
||||
prompt=prompt,
|
||||
api_key=config.get('dashscope_api_key'),
|
||||
model=model,
|
||||
size=size,
|
||||
count=count,
|
||||
style=style
|
||||
)
|
||||
|
||||
if result.get('success'):
|
||||
paths = result.get('paths', [])
|
||||
reply = f"✅ AI 生图完成!成功生成 {len(paths)} 张图片\n"
|
||||
reply += f"📊 模型:{result.get('model', 'wanx-v1')}\n"
|
||||
reply += f"📐 尺寸:{result.get('size', '1024*1024')}\n\n"
|
||||
reply += "📁 本地路径:\n"
|
||||
for i, path in enumerate(paths):
|
||||
reply += f" {i+1}. {path}\n"
|
||||
|
||||
return reply
|
||||
else:
|
||||
error_msg = result.get('error', '未知错误')
|
||||
return f"❌ AI 生图失败:{error_msg}"
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"AI 生图失败:{str(e)}", exc_info=True)
|
||||
return f"❌ AI 生图失败:{str(e)}"
|
||||
|
||||
def _load_image_gen_config(self):
|
||||
"""加载图片生成配置"""
|
||||
config = {
|
||||
'dashscope_api_key': '',
|
||||
'image_model': 'wanx-v1',
|
||||
'image_size': '1024*1024',
|
||||
'image_count': 1
|
||||
}
|
||||
|
||||
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:
|
||||
logger.warning(f"加载配置文件失败:{str(e)}")
|
||||
|
||||
return config
|
||||
|
||||
def _update_article(self, target='', text='', title='', status='publish'):
|
||||
"""
|
||||
更新已有文章(追加内容模式)
|
||||
@ -579,9 +791,12 @@ class FeishuBot:
|
||||
🖼️ **发布图片文章**
|
||||
发送图片即可发布
|
||||
|
||||
🎨 **AI 图片生成**(新增)
|
||||
使用 AI 生成文章配图或独立图片
|
||||
|
||||
---
|
||||
|
||||
**指令说明**:
|
||||
**基础指令**:
|
||||
|
||||
`#标题 文章标题` - 指定文章标题
|
||||
`#分类 分类名` - 指定分类(如:#分类 ai)
|
||||
@ -592,6 +807,21 @@ class FeishuBot:
|
||||
|
||||
---
|
||||
|
||||
**AI 生图指令**(新增):
|
||||
|
||||
`#生图 图片描述` - 生成指定描述的图片
|
||||
`#生图数量 2` - 设置生成数量(1-4)
|
||||
`#生图模型 wanx-v1` - 选择模型
|
||||
`#生图尺寸 1024*1024` - 设置尺寸
|
||||
`#生图风格 写实` - 设置风格
|
||||
|
||||
**发布+生图指令**(新增):
|
||||
|
||||
`#发布并生图` - 发布文章并自动生成配图
|
||||
(与 #标题、#分类 等指令配合使用)
|
||||
|
||||
---
|
||||
|
||||
**示例**:
|
||||
|
||||
```
|
||||
@ -600,17 +830,21 @@ class FeishuBot:
|
||||
人工智能正在改变世界...
|
||||
```
|
||||
|
||||
**更新示例**:
|
||||
**AI 生图示例**:
|
||||
|
||||
```
|
||||
#更新 12345
|
||||
新增的内容追加到原文末尾...
|
||||
#生图 一只可爱的猫咪在草地上晒太阳,高质量插画
|
||||
#生图数量 2
|
||||
#生图尺寸 1024*1024
|
||||
```
|
||||
|
||||
**发布+生图示例**:
|
||||
|
||||
```
|
||||
#更新 AI发展趋势
|
||||
#标题 AI 发展趋势 2026 更新版
|
||||
补充一段关于 Agent 的最新进展...
|
||||
#发布并生图
|
||||
#标题 AI 未来展望
|
||||
#分类 ai
|
||||
人工智能正在改变世界...
|
||||
```
|
||||
|
||||
**可用分类**:
|
||||
@ -641,6 +875,7 @@ class FeishuBot:
|
||||
- **状态**:✅ 运行中
|
||||
- **WordPress**:https://www.nanlou.net
|
||||
- **发布脚本**:已就绪
|
||||
- **AI 生图**:✅ 支持
|
||||
|
||||
发送 `#帮助` 查看使用说明"""
|
||||
|
||||
|
||||
323
modules/wp_image_generator.py
Normal file
323
modules/wp_image_generator.py
Normal file
@ -0,0 +1,323 @@
|
||||
#!/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}")
|
||||
192
scripts/wp_generate_image.py
Normal file
192
scripts/wp_generate_image.py
Normal file
@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
WordPress 发布系统 - AI 图片生成脚本
|
||||
支持命令行调用和模块调用
|
||||
"""
|
||||
|
||||
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_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):
|
||||
"""
|
||||
生成图片(主入口函数)
|
||||
|
||||
Args:
|
||||
prompt: 图片描述
|
||||
api_key: DashScope API Key(可选)
|
||||
model: 模型名称(可选)
|
||||
size: 图片尺寸(可选)
|
||||
count: 生成数量(可选)
|
||||
negative_prompt: 反向提示词(可选)
|
||||
style: 风格(可选)
|
||||
|
||||
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)
|
||||
|
||||
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_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
|
||||
}
|
||||
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):
|
||||
"""
|
||||
根据文章标题和内容生成配图
|
||||
|
||||
Args:
|
||||
title: 文章标题
|
||||
content: 文章内容
|
||||
count: 生成数量
|
||||
api_key: DashScope API Key(可选)
|
||||
|
||||
Returns:
|
||||
dict: 生成结果
|
||||
"""
|
||||
pl = get_publish_logger()
|
||||
dl = get_debug_logger()
|
||||
|
||||
# 加载配置
|
||||
config = load_config()
|
||||
api_key = api_key or config.get('dashscope_api_key', '')
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
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('--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
|
||||
)
|
||||
|
||||
# 输出 JSON 结果
|
||||
print("\n" + json.dumps(result, ensure_ascii=False, indent=2))
|
||||
|
||||
# 返回状态码
|
||||
sys.exit(0 if result.get('success') else 1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Loading…
Reference in New Issue
Block a user