新增AI绘图功能
This commit is contained in:
2
plugins/ai_gen_image/__init__.py
Normal file
2
plugins/ai_gen_image/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
|||||||
|
# AI绘图插件
|
||||||
|
# 基于pollinations.ai API生成图片
|
||||||
14
plugins/ai_gen_image/config.toml
Normal file
14
plugins/ai_gen_image/config.toml
Normal file
@@ -0,0 +1,14 @@
|
|||||||
|
[AIGenImage]
|
||||||
|
enable = true
|
||||||
|
command = ["AI绘图", "绘图", "画图", "生成图片"]
|
||||||
|
command-format = """
|
||||||
|
🎨AI绘图指令:
|
||||||
|
AI绘图 描述文字
|
||||||
|
"""
|
||||||
|
|
||||||
|
# 图片生成API配置
|
||||||
|
image_api_url = "https://image.pollinations.ai/prompt/{prompt}"
|
||||||
|
default_width = 1024
|
||||||
|
default_height = 1024
|
||||||
|
default_model = "turbo"
|
||||||
|
default_timeout = 300
|
||||||
188
plugins/ai_gen_image/main.py
Normal file
188
plugins/ai_gen_image/main.py
Normal file
@@ -0,0 +1,188 @@
|
|||||||
|
import os
|
||||||
|
import time
|
||||||
|
import urllib.parse
|
||||||
|
import uuid
|
||||||
|
from typing import Dict, Any, List, Optional, Tuple
|
||||||
|
|
||||||
|
import requests
|
||||||
|
from loguru import logger
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from base.plugin_common.message_plugin_interface import MessagePluginInterface
|
||||||
|
from base.plugin_common.plugin_interface import PluginStatus
|
||||||
|
from utils.decorator.plugin_decorators import plugin_stats_decorator
|
||||||
|
from utils.robot_cmd.robot_command import Feature, PermissionStatus, GroupBotManager
|
||||||
|
from utils.decorator.points_decorator import plugin_points_cost
|
||||||
|
from wechat_ipad import WechatAPIClient
|
||||||
|
|
||||||
|
|
||||||
|
class AIGenImagePlugin(MessagePluginInterface):
|
||||||
|
"""AI绘图插件"""
|
||||||
|
|
||||||
|
# 功能权限常量
|
||||||
|
FEATURE_KEY = "AI_GEN_IMAGE"
|
||||||
|
FEATURE_DESCRIPTION = "🎨 AI绘图功能 [AI绘图, 绘图, 画图, 生成图片]"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def name(self) -> str:
|
||||||
|
return "AI绘图"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def version(self) -> str:
|
||||||
|
return "1.0.0"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def description(self) -> str:
|
||||||
|
return "提供AI绘图功能,基于pollinations.ai生成图片"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def author(self) -> str:
|
||||||
|
return "liu.wei"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def command_prefix(self) -> Optional[str]:
|
||||||
|
return "" # 不需要前缀,直接匹配命令
|
||||||
|
|
||||||
|
@property
|
||||||
|
def commands(self) -> List[str]:
|
||||||
|
return self._commands
|
||||||
|
|
||||||
|
@property
|
||||||
|
def feature_key(self) -> Optional[str]:
|
||||||
|
return self.FEATURE_KEY
|
||||||
|
|
||||||
|
@property
|
||||||
|
def feature_description(self) -> Optional[str]:
|
||||||
|
return self.FEATURE_DESCRIPTION
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self.feature = self.register_feature()
|
||||||
|
|
||||||
|
def initialize(self, context: Dict[str, Any]) -> bool:
|
||||||
|
"""初始化插件"""
|
||||||
|
self.LOG = logger
|
||||||
|
self.LOG.info(f"正在初始化 {self.name} 插件...")
|
||||||
|
|
||||||
|
# 保存上下文对象
|
||||||
|
self.event_system = context.get("event_system")
|
||||||
|
|
||||||
|
# 从配置文件加载配置
|
||||||
|
self._commands = self._config.get("AIGenImage", {}).get("command", ["AI绘图", "绘图", "画图", "生成图片"])
|
||||||
|
self.command_format = self._config.get("AIGenImage", {}).get("command-format", "AI绘图 描述文字")
|
||||||
|
self.enable = self._config.get("AIGenImage", {}).get("enable", True)
|
||||||
|
|
||||||
|
# API配置
|
||||||
|
self.image_api_url = self._config.get("AIGenImage", {}).get("image_api_url", "https://image.pollinations.ai/prompt/{prompt}")
|
||||||
|
self.default_width = self._config.get("AIGenImage", {}).get("default_width", 1024)
|
||||||
|
self.default_height = self._config.get("AIGenImage", {}).get("default_height", 1024)
|
||||||
|
self.default_model = self._config.get("AIGenImage", {}).get("default_model", "turbo")
|
||||||
|
self.default_timeout = self._config.get("AIGenImage", {}).get("default_timeout", 300)
|
||||||
|
|
||||||
|
# 确保临时目录存在
|
||||||
|
self.temp_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'temp')
|
||||||
|
os.makedirs(self.temp_dir, exist_ok=True)
|
||||||
|
|
||||||
|
self.LOG.info(f"[{self.name}] 插件初始化完成,指令:{self._commands}")
|
||||||
|
return True
|
||||||
|
|
||||||
|
def start(self) -> bool:
|
||||||
|
"""启动插件"""
|
||||||
|
self.LOG.info(f"[{self.name}] 插件已启动")
|
||||||
|
self.status = PluginStatus.RUNNING
|
||||||
|
return True
|
||||||
|
|
||||||
|
def stop(self) -> bool:
|
||||||
|
"""停止插件"""
|
||||||
|
self.LOG.info(f"[{self.name}] 插件已停止")
|
||||||
|
self.status = PluginStatus.STOPPED
|
||||||
|
return True
|
||||||
|
|
||||||
|
def can_process(self, message: Dict[str, Any]) -> bool:
|
||||||
|
"""检查是否可以处理该消息"""
|
||||||
|
if not self.enable:
|
||||||
|
return False
|
||||||
|
|
||||||
|
content = str(message.get("content", "")).strip()
|
||||||
|
command = content.split(" ")[0]
|
||||||
|
|
||||||
|
return command in self._commands
|
||||||
|
|
||||||
|
@plugin_stats_decorator(plugin_name="AI绘图")
|
||||||
|
@plugin_points_cost(5, "AI绘图消耗积分", FEATURE_KEY)
|
||||||
|
async def process_message(self, message: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
|
||||||
|
"""处理消息"""
|
||||||
|
content = str(message.get("content", "")).strip()
|
||||||
|
self.LOG.debug(f"插件执行: {self.name}:{content}")
|
||||||
|
command = content.split(" ")[0]
|
||||||
|
sender = message.get("sender")
|
||||||
|
roomid = message.get("roomid", "")
|
||||||
|
gbm: GroupBotManager = message.get("gbm")
|
||||||
|
bot: WechatAPIClient = message.get("bot")
|
||||||
|
|
||||||
|
# 检查命令格式
|
||||||
|
if len(content.split(" ")) == 1:
|
||||||
|
await bot.send_text_message((roomid if roomid else sender), f"❌命令格式错误!\n{self.command_format}"
|
||||||
|
, sender)
|
||||||
|
return False, "命令格式错误"
|
||||||
|
|
||||||
|
# 检查权限
|
||||||
|
if roomid and gbm.get_group_permission(roomid, self.feature) == PermissionStatus.DISABLED:
|
||||||
|
return False, "没有权限"
|
||||||
|
|
||||||
|
# 提取描述文字
|
||||||
|
prompt = content[len(command):].strip()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# 发送提示消息
|
||||||
|
await bot.send_text_message((roomid if roomid else sender), f"🎨正在生成图片,请稍候...", sender)
|
||||||
|
|
||||||
|
# 生成图片
|
||||||
|
image_path = self._generate_image(prompt)
|
||||||
|
if not image_path or not os.path.exists(image_path):
|
||||||
|
await bot.send_text_message((roomid if roomid else sender), f"❌生成图片失败,请重试", sender)
|
||||||
|
return False, "生成图片失败"
|
||||||
|
|
||||||
|
# 发送图片
|
||||||
|
await bot.send_image_message((roomid if roomid else sender), Path(image_path))
|
||||||
|
return True, "发送成功"
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.LOG.error(f"处理AI绘图请求出错: {e}")
|
||||||
|
await bot.send_text_message((roomid if roomid else sender), f"❌生成图片出错: {str(e)}", sender)
|
||||||
|
return False, f"处理出错: {e}"
|
||||||
|
|
||||||
|
def _generate_image(self, prompt: str) -> str:
|
||||||
|
"""生成图片并返回图片路径"""
|
||||||
|
try:
|
||||||
|
# 准备API参数
|
||||||
|
params = {
|
||||||
|
"width": self.default_width,
|
||||||
|
"height": self.default_height,
|
||||||
|
"model": self.default_model,
|
||||||
|
"seed": int(time.time()) % 1000000 # 使用时间戳作为随机种子
|
||||||
|
}
|
||||||
|
|
||||||
|
# 编码提示词
|
||||||
|
encoded_prompt = urllib.parse.quote(prompt)
|
||||||
|
url = self.image_api_url.format(prompt=encoded_prompt)
|
||||||
|
|
||||||
|
self.LOG.info(f"正在生成图片,提示词: {prompt[:30]}...")
|
||||||
|
|
||||||
|
# 发送请求
|
||||||
|
response = requests.get(url, params=params, timeout=self.default_timeout)
|
||||||
|
response.raise_for_status()
|
||||||
|
|
||||||
|
# 保存图片
|
||||||
|
image_filename = f"ai_image_{uuid.uuid4().hex[:8]}.jpg"
|
||||||
|
image_path = os.path.join(self.temp_dir, image_filename)
|
||||||
|
|
||||||
|
with open(image_path, 'wb') as f:
|
||||||
|
f.write(response.content)
|
||||||
|
|
||||||
|
self.LOG.info(f"图片生成成功,保存至: {image_path}")
|
||||||
|
return image_path
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.LOG.error(f"生成图片出错: {e}")
|
||||||
|
return ""
|
||||||
Reference in New Issue
Block a user