diff --git a/plugins/ai_image.py b/plugins/ai_image.py new file mode 100644 index 0000000..945f318 --- /dev/null +++ b/plugins/ai_image.py @@ -0,0 +1,590 @@ +""" +AI 生图插件 - 基于 ai.2x.nz 自然语言生图接口 +支持自然语言描述、直接 Tag、工作流切换、画风选择等 +""" + +import os +import json +import time +import threading +import hashlib +import re +from datetime import datetime +from urllib.parse import quote + +try: + import requests + REQUESTS_AVAILABLE = True +except ImportError: + REQUESTS_AVAILABLE = False + +try: + import websocket + WEBSOCKET_AVAILABLE = True +except ImportError: + WEBSOCKET_AVAILABLE = False + + +class Plugin: + """AI 生图插件 - 调用 ai.2x.nz 接口生成图片""" + + BASE_URL = "https://ai.2x.nz" + WS_URL = "wss://ai.2x.nz/ws/run" + API_URL = "https://ai.2x.nz/api" + + # 保存目录 + DEFAULT_SAVE_DIR = "generated_images" + + def __init__(self): + self.usage = """AI 生图工具 - 基于自然语言描述生成图片 + +基本用法: + generate <描述> - 用自然语言描述生成图片 + tag <直接Tag> - 用 Danbooru Tag 生成图片 + both <<<>>> <描述> - 同时使用 Tag 和自然语言 + +高级用法: + workflows - 列出所有可用工作流 + switch <工作流名> - 切换工作流 + width <数值> - 设置图片宽度 (最大1344) + height <数值> - 设置图片高度 (最大1344) + size <宽> <高> - 同时设置宽高 + rewrite on|off - 改写模式开关 + override on|off - 覆写模式开关 + status - 查看当前状态和配置 + +示例: + ai_image generate 一个可爱的猫耳少女,樱花背景 + ai_image tag 1girl, cat_ears, cherry_blossoms, smile + ai_image both 1girl, cat_ears <<<>>> 一个可爱的猫耳少女在樱花树下微笑 + ai_image workflows + ai_image switch WAI - 37.json + ai_image size 1024 1024 + ai_image rewrite on +""" + self.cli = None + self.save_dir = self.DEFAULT_SAVE_DIR + + # 配置状态 + self.current_workflow = None + self.direct_prompt = "" + self.nl_prompt = "" + self.rewrite = False + self.override = False + self.width = None + self.height = None + self.workflows_cache = None + self.is_generating = False + + def set_cli(self, cli): + """设置 CLI 实例引用""" + self.cli = cli + # 注册插件命令 + if hasattr(cli, 'register_liugin_command'): + cli.register_liugin_command('ai_image', self.command_handler) + + def get_tool_info(self): + return { + "name": "ai_image", + "description": "AI 生图工具 — 基于 ai.2x.nz 接口,支持自然语言描述和 Danbooru Tag 生成图片。当用户想要生成图片、画图、AI绘画、生图时使用此工具。", + "keywords": [ + "生图", "画图", "图片", "AI绘画", "生成图片", "image", "generate", + "draw", "绘画", "插画", "二次元", "动漫", "Danbooru", "tag", + "工作流", "workflow", "ComfyUI" + ], + "usage": self.usage + } + + def get_mcp_definition(self): + return { + "name": "ai_image", + "description": "AI 生图工具,基于 ai.2x.nz 接口生成图片", + "inputSchema": { + "type": "object", + "properties": { + "operation": { + "type": "string", + "enum": ["generate", "tag", "both", "workflows", "switch", "size", "status"], + "description": "操作类型" + }, + "prompt": { + "type": "string", + "description": "生成提示词" + } + }, + "required": ["operation"] + } + } + + def convert_mcp_args(self, arguments): + op = arguments.get("operation", "") + prompt = arguments.get("prompt", "") + return f"{op} {prompt}" + + def command_handler(self, args): + """处理 /ai_image 命令""" + return self.handle(args) + + def handle(self, args): + """处理 AI 生图请求""" + if not REQUESTS_AVAILABLE: + return "❌ 缺少 requests 库,请运行: pip install requests" + + args = args.strip() + if not args: + return self.usage + + # 解析操作 + parts = args.split(maxsplit=1) + operation = parts[0].lower() + remaining = parts[1] if len(parts) > 1 else "" + + try: + if operation in ["generate", "gen", "g"]: + return self._cmd_generate(remaining) + elif operation in ["tag", "t", "direct", "d"]: + return self._cmd_tag(remaining) + elif operation in ["both", "b"]: + return self._cmd_both(remaining) + elif operation in ["workflows", "wf", "list"]: + return self._cmd_workflows() + elif operation in ["switch", "sw"]: + return self._cmd_switch(remaining) + elif operation in ["width", "w"]: + return self._cmd_set_width(remaining) + elif operation in ["height", "h"]: + return self._cmd_set_height(remaining) + elif operation in ["size", "s"]: + return self._cmd_set_size(remaining) + elif operation in ["rewrite", "rw"]: + return self._cmd_set_rewrite(remaining) + elif operation in ["override", "ov"]: + return self._cmd_set_override(remaining) + elif operation in ["status", "st"]: + return self._cmd_status() + elif operation in ["help", "?"]: + return self.usage + else: + # 尝试当作自然语言描述处理 + return self._cmd_generate(args) + except Exception as e: + return f"❌ 操作失败: {str(e)}" + + # ── 命令实现 ── + + def _cmd_generate(self, prompt): + """自然语言生成图片""" + if not prompt.strip(): + return "❌ 请提供描述内容。示例: ai_image generate 一个可爱的猫耳少女" + + self.nl_prompt = prompt.strip() + self.direct_prompt = "" + self.override = False + return self._run_generation() + + def _cmd_tag(self, tags): + """使用直接 Tag 生成图片""" + if not tags.strip(): + return "❌ 请提供 Tag 内容。示例: ai_image tag 1girl, cat_ears, smile" + + self.direct_prompt = tags.strip() + self.nl_prompt = "" + self.override = True + return self._run_generation() + + def _cmd_both(self, args): + """同时使用 Tag 和自然语言""" + if "<<<>>>" in args: + parts = args.split("<<<>>>", 1) + self.direct_prompt = parts[0].strip() + self.nl_prompt = parts[1].strip() + else: + # 尝试用 || 分隔 + if "||" in args: + parts = args.split("||", 1) + self.direct_prompt = parts[0].strip() + self.nl_prompt = parts[1].strip() + else: + return "❌ 格式错误。用法: ai_image both <<<>>> <描述>" + + self.override = False + return self._run_generation() + + def _cmd_workflows(self): + """列出所有工作流""" + workflows = self._fetch_workflows() + if not workflows: + return "❌ 获取工作流列表失败" + + result = f"📂 共 {len(workflows)} 个工作流:\n\n" + for i, wf in enumerate(workflows, 1): + name = wf.get("path", "未知") + current = " ← [当前]" if name == self.current_workflow else "" + result += f" {i:2d}. {name}{current}\n" + + result += f"\n💡 使用 'ai_image switch <工作流名>' 切换工作流" + return result + + def _cmd_switch(self, name): + """切换工作流""" + if not name.strip(): + return "❌ 请提供工作流名称。先用 'ai_image workflows' 查看列表" + + # 模糊匹配 + workflows = self._fetch_workflows() + if not workflows: + return "❌ 获取工作流列表失败" + + matched = None + name_lower = name.strip().lower() + for wf in workflows: + wf_path = wf.get("path", "") + if name_lower == wf_path.lower(): + matched = wf_path + break + if name_lower in wf_path.lower(): + matched = wf_path + + if not matched: + # 尝试数字索引 + try: + idx = int(name.strip()) - 1 + if 0 <= idx < len(workflows): + matched = workflows[idx].get("path", "") + except ValueError: + pass + + if not matched: + return f"❌ 未找到匹配的工作流: {name}\n💡 先用 'ai_image workflows' 查看列表" + + self.current_workflow = matched + return f"✅ 已切换工作流: {matched}" + + def _cmd_set_width(self, val): + """设置宽度""" + try: + w = int(val.strip()) + if w < 64 or w > 1344: + return "❌ 宽度范围: 64 ~ 1344" + self.width = w + return f"✅ 宽度已设置为: {w}" + except ValueError: + return "❌ 请输入有效数字" + + def _cmd_set_height(self, val): + """设置高度""" + try: + h = int(val.strip()) + if h < 64 or h > 1344: + return "❌ 高度范围: 64 ~ 1344" + self.height = h + return f"✅ 高度已设置为: {h}" + except ValueError: + return "❌ 请输入有效数字" + + def _cmd_set_size(self, val): + """同时设置宽高""" + parts = val.strip().split() + if len(parts) < 2: + return "❌ 用法: ai_image size <宽> <高>" + try: + w, h = int(parts[0]), int(parts[1]) + if w < 64 or w > 1344 or h < 64 or h > 1344: + return "❌ 宽高范围: 64 ~ 1344" + self.width = w + self.height = h + return f"✅ 尺寸已设置为: {w} x {h}" + except ValueError: + return "❌ 请输入有效数字" + + def _cmd_set_rewrite(self, val): + """设置改写模式""" + val = val.strip().lower() + if val in ["on", "开", "1", "true", "yes"]: + self.rewrite = True + return "✅ 改写模式已开启 — LLM 将智能重写 prompt" + elif val in ["off", "关", "0", "false", "no"]: + self.rewrite = False + return "✅ 改写模式已关闭 — 普通翻译/追加模式" + else: + self.rewrite = not self.rewrite + status = "开启" if self.rewrite else "关闭" + return f"✅ 改写模式已{status}" + + def _cmd_set_override(self, val): + """设置覆写模式""" + val = val.strip().lower() + if val in ["on", "开", "1", "true", "yes"]: + self.override = True + return "✅ 覆写模式已开启 — 忽略工作流内置 prompt" + elif val in ["off", "关", "0", "false", "no"]: + self.override = False + return "✅ 覆写模式已关闭 — 保留工作流内置 prompt" + else: + self.override = not self.override + status = "开启" if self.override else "关闭" + return f"✅ 覆写模式已{status}" + + def _cmd_status(self): + """查看当前状态""" + result = " AI 生图状态:\n\n" + result += f" 工作流: {self.current_workflow or '(默认)'}\n" + result += f" 直接 Tag: {self.direct_prompt or '(未设置)'}\n" + result += f" 自然语言: {self.nl_prompt or '(未设置)'}\n" + result += f" 改写模式: {'✅ 开启' if self.rewrite else '❌ 关闭'}\n" + result += f" 覆写模式: {'✅ 开启' if self.override else '❌ 关闭'}\n" + result += f" 宽度: {self.width or '(默认)'}\n" + result += f" 高度: {self.height or '(默认)'}\n" + result += f" 保存目录: {self.save_dir}\n" + result += f" 生成中: {'是' if self.is_generating else '否'}\n" + return result + + # ── 核心生图逻辑 ── + + def _fetch_workflows(self): + """获取工作流列表""" + if self.workflows_cache: + return self.workflows_cache + + try: + resp = requests.get(f"{self.API_URL}/workflows", timeout=15) + if resp.status_code == 200: + data = resp.json() + self.workflows_cache = data.get("workflows", []) + return self.workflows_cache + except Exception as e: + pass + return [] + + def _run_generation(self): + """执行图片生成""" + if self.is_generating: + return "⚠️ 已有生成任务在进行中,请等待完成" + + if not self.direct_prompt and not self.nl_prompt: + return "❌ 请提供至少一种 prompt(直接 Tag 或自然语言描述)" + + self.is_generating = True + result_lines = [] + image_paths = [] + + try: + # 获取默认工作流(如果未指定则使用列表中第一个) + workflow = self.current_workflow + if not workflow: + workflows = self._fetch_workflows() + if workflows: + # 优先选 "无Lora" 工作流,否则用第一个 + for wf in workflows: + if "无lora" in wf.get("path", "").lower(): + workflow = wf["path"] + break + if not workflow: + workflow = workflows[0].get("path", "") + + # 构建 payload + payload = { + "workflow_path": workflow or "", + "inline_workflow": None, + "direct_prompt": self.direct_prompt, + "nl_prompt": self.nl_prompt, + "rewrite": self.rewrite, + "override": self.override, + "width": self.width, + "height": self.height, + } + + result_lines.append(" 正在连接 ai.2x.nz ...") + result_lines.append(f" 直接 Tag: {self.direct_prompt[:80] or '(无)'}") + result_lines.append(f" 自然语言: {self.nl_prompt[:80] or '(无)'}") + result_lines.append(f" 工作流: {workflow or '(默认)'}") + result_lines.append("") + + # 使用同步方式通过 REST + 轮询获取结果 + # 由于 websocket-client 可能未安装,提供降级方案 + if WEBSOCKET_AVAILABLE: + images = self._generate_via_ws(payload, result_lines) + else: + # 降级:尝试使用 requests + SSE 或提示安装 + result_lines.append("⚠️ websocket-client 未安装,尝试安装...") + images = self._generate_via_ws_fallback(payload, result_lines) + + if images: + image_paths = images + result_lines.append(f"\n✅ 生成完成!共 {len(images)} 张图片") + for i, path in enumerate(images, 1): + result_lines.append(f" {i}. {path}") + else: + result_lines.append("\n❌ 未生成任何图片") + + except Exception as e: + result_lines.append(f"\n❌ 生成失败: {str(e)}") + finally: + self.is_generating = False + + return "\n".join(result_lines) + + def _generate_via_ws(self, payload, result_lines): + """通过 WebSocket 生成图片""" + images = [] + final_prompt = "" + error_msg = None + + ws_result = {"images": [], "error": None, "final_prompt": "", "done": False} + + def on_message(ws, message): + try: + m = json.loads(message) + msg_type = m.get("type", "") + + if msg_type == "log": + result_lines.append(f" {m.get('message', '')}") + elif msg_type == "llm_start": + result_lines.append(" LLM 正在处理 prompt...") + elif msg_type == "llm_chunk": + pass # 流式输出,不逐块显示 + elif msg_type == "llm_done": + result_lines.append(" LLM 处理完成") + elif msg_type == "progress": + node = m.get("node", "") + value = m.get("value", 0) + max_val = m.get("max", 0) + if max_val and max_val > 1: + pct = int(value * 100 / max_val) + result_lines.append(f"⏳ 进度: {node} {value}/{max_val} ({pct}%)") + else: + result_lines.append(f"⏳ 执行: {node}") + elif msg_type == "prompt_id": + fp = m.get("final_prompt", "") + if fp: + ws_result["final_prompt"] = fp + result_lines.append(f"\n 最终 Prompt:\n{fp[:300]}{'...' if len(fp) > 300 else ''}\n") + elif msg_type == "image": + url = m.get("url", "") + filename = m.get("filename", "") + if url: + # 下载图片 + saved = self._download_image(url, filename) + if saved: + ws_result["images"].append(saved) + result_lines.append(f" 收到图片: {filename}") + elif msg_type == "done": + count = m.get("count", 0) + result_lines.append(f"✅ 服务端完成,共 {count} 张") + ws_result["done"] = True + ws.close() + elif msg_type == "error": + ws_result["error"] = m.get("message", "未知错误") + result_lines.append(f"❌ 错误: {ws_result['error']}") + ws.close() + except Exception as e: + result_lines.append(f"⚠️ 解析消息出错: {e}") + + def on_error(ws, error): + ws_result["error"] = str(error) + result_lines.append(f"❌ WebSocket 错误: {error}") + + def on_close(ws, close_status_code, close_msg): + ws_result["done"] = True + + def on_open(ws): + ws.send(json.dumps(payload)) + + try: + ws = websocket.WebSocketApp( + self.WS_URL, + on_message=on_message, + on_error=on_error, + on_close=on_close, + on_open=on_open + ) + + # 在独立线程中运行 WebSocket + ws_thread = threading.Thread(target=ws.run_forever, kwargs={"ping_timeout": 10}) + ws_thread.daemon = True + ws_thread.start() + + # 等待完成(最多 5 分钟) + timeout = 300 + start = time.time() + while not ws_result["done"] and (time.time() - start) < timeout: + time.sleep(1) + + if ws_result["error"]: + raise Exception(ws_result["error"]) + + if not ws_result["done"]: + ws.close() + raise Exception("生成超时(超过 5 分钟)") + + return ws_result["images"] + + except Exception as e: + raise Exception(f"WebSocket 连接失败: {e}") + + def _generate_via_ws_fallback(self, payload, result_lines): + """降级方案:尝试自动安装 websocket-client""" + try: + import subprocess + result_lines.append(" 正在安装 websocket-client ...") + subprocess.run( + ["pip", "install", "websocket-client"], + capture_output=True, timeout=60 + ) + import importlib + global websocket + websocket = importlib.import_module("websocket") + global WEBSOCKET_AVAILABLE + WEBSOCKET_AVAILABLE = True + result_lines.append("✅ websocket-client 安装成功") + return self._generate_via_ws(payload, result_lines) + except Exception as e: + result_lines.append(f"❌ 无法安装 websocket-client: {e}") + result_lines.append("💡 请手动运行: pip install websocket-client") + return [] + + def _download_image(self, url, filename): + """下载图片到本地""" + try: + # 确保 URL 是完整的 + if url.startswith("/"): + url = f"{self.BASE_URL}{url}" + elif not url.startswith("http"): + url = f"{self.BASE_URL}/api/output/file?path={quote(url)}" + + # 创建保存目录 + os.makedirs(self.save_dir, exist_ok=True) + + # 生成文件名 + if not filename: + filename = f"img_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" + + # 清理文件名 + filename = re.sub(r'[<>:"/\\|?*]', '_', filename) + save_path = os.path.join(self.save_dir, filename) + + # 如果文件已存在,添加序号 + if os.path.exists(save_path): + name, ext = os.path.splitext(filename) + for i in range(1, 100): + new_path = os.path.join(self.save_dir, f"{name}_{i}{ext}") + if not os.path.exists(new_path): + save_path = new_path + break + + # 下载 + resp = requests.get(url, timeout=60, stream=True) + if resp.status_code == 200: + with open(save_path, 'wb') as f: + for chunk in resp.iter_content(chunk_size=8192): + f.write(chunk) + return save_path + else: + return None + + except Exception as e: + return None + + +# 兼容 Liugin 类名 +Liugin = Plugin