feat: 新增图像生成插件 + 记忆系统插件 + 7工具横评页面
- plugins/image_generator.py: 基于 ai.2x.nz ComfyUI 的图像生成插件 - 自然语言/Tag 双模式,支持中文描述 - WebSocket 实时进度,28 个角色工作流 - 分辨率可调(最大 1344x1344) - plugins/memory_plugin.py: 包装 xcli_core/memory.py 的记忆插件 - AI 可写日记、读日记、搜索日记 - 全局搜索(MEMORY + 日记 + 聊天记录) - 自动保存对话 + 上下文压缩 - 长期记忆 MEMORY.md 管理 - compare.html: 7 款 AI 编程助手横评页面 - Claude Code / Codex CLI / Gemini CLI / Aider / Copilot CLI / 小狸 / TG HELPER - 25+ 维度功能对比 + 评分 + 场景推荐 - .gitignore: 排除运行时生成的记忆数据和图片
This commit is contained in:
1 parent
62aa5ed128
commit
3fa1d56632
4 files changed
+1724
No files matched your search
@@ -0,0 +1,370 @@
|
||||
"""
|
||||
图像生成插件 - 基于 ai.2x.nz 的 ComfyUI 在线服务
|
||||
支持自然语言描述和英文 Tag 两种模式
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
import urllib.request
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class Plugin:
|
||||
"""图像生成工具 — 自然语言/Tag 生图,基于 ComfyUI 在线服务"""
|
||||
|
||||
def __init__(self):
|
||||
self.usage = """图像生成工具 — 基于 ai.2x.nz 的 ComfyUI 在线服务
|
||||
|
||||
操作:
|
||||
generate <描述> - 自然语言描述生成图片(中文/英文均可)
|
||||
tag <英文tags> - 直接 Tag 模式(逗号分隔的英文标签)
|
||||
list - 列出可用的工作流/角色
|
||||
status - 查看 GPU 状态
|
||||
|
||||
示例:
|
||||
image_generator generate 一个白发蓝眼的猫耳少女在写代码
|
||||
image_generator tag masterpiece, 1girl, white hair, blue eyes, cat ears, coding
|
||||
image_generator generate 画一个赛博朋克风格的城市夜景 --width 1024 --height 768
|
||||
|
||||
说明:
|
||||
- 自然语言模式会经过 LLM 翻译为 Tag,效果更好
|
||||
- 直接 Tag 模式适合熟悉 Danbooru 标签的用户
|
||||
- 默认分辨率 1024x1344(WAI 推荐)
|
||||
- 最大分辨率 1344x1344
|
||||
- 生成的图片保存在 screenshots/ 目录
|
||||
"""
|
||||
self.cli = None
|
||||
self.api_base = "https://ai.2x.nz"
|
||||
self.ws_url = "wss://ai.2x.nz/ws/run"
|
||||
self.default_workflow = "WAI - 无Lora.json"
|
||||
self.output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "screenshots")
|
||||
|
||||
def set_cli(self, cli):
|
||||
self.cli = cli
|
||||
|
||||
def get_tool_info(self):
|
||||
return {
|
||||
"name": "image_generator",
|
||||
"description": "图像生成 — 自然语言/Tag 生图,基于 ComfyUI 在线服务(ai.2x.nz),支持中文描述",
|
||||
"keywords": ["图像", "生成", "画图", "绘图", "AI画", "生图", "image", "generate", "draw", "paint",
|
||||
"diffusion", "comfyui", "stable", "sd", "文生图", "text2img"],
|
||||
"usage": self.usage
|
||||
}
|
||||
|
||||
def handle(self, args: str) -> str:
|
||||
try:
|
||||
parts = args.strip().split(maxsplit=1)
|
||||
if not parts:
|
||||
return "错误:请提供操作类型\n\n" + self.usage
|
||||
|
||||
op = parts[0].lower()
|
||||
rest = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
handlers = {
|
||||
"generate": self._generate,
|
||||
"tag": self._generate_tag,
|
||||
"gen": self._generate,
|
||||
"g": self._generate,
|
||||
"list": self._list_workflows,
|
||||
"ls": self._list_workflows,
|
||||
"status": self._gpu_status,
|
||||
}
|
||||
|
||||
handler = handlers.get(op)
|
||||
if not handler:
|
||||
# 如果没有指定操作,把整个 args 当作自然语言描述
|
||||
return self._generate(args)
|
||||
|
||||
return handler(rest)
|
||||
|
||||
except Exception as e:
|
||||
return f"图像生成错误: {e}"
|
||||
|
||||
def _generate(self, args: str) -> str:
|
||||
"""自然语言模式生成图片"""
|
||||
# 解析 --width 和 --height 参数
|
||||
width, height = 1024, 1344
|
||||
prompt = args.strip()
|
||||
|
||||
if "--width" in prompt or "--height" in prompt or "--w" in prompt or "--h" in prompt:
|
||||
import re
|
||||
w_match = re.search(r'--w(?:idth)?\s+(\d+)', prompt)
|
||||
h_match = re.search(r'--h(?:eight)?\s+(\d+)', prompt)
|
||||
if w_match:
|
||||
width = int(w_match.group(1))
|
||||
prompt = prompt[:w_match.start()] + prompt[w_match.end():]
|
||||
if h_match:
|
||||
height = int(h_match.group(1))
|
||||
prompt = prompt[:h_match.start()] + prompt[h_match.end():]
|
||||
prompt = prompt.strip()
|
||||
|
||||
if not prompt:
|
||||
return "错误:请提供图片描述\n示例: image_generator generate 一个白发猫耳少女"
|
||||
|
||||
# 限制分辨率
|
||||
width = min(width, 1344)
|
||||
height = min(height, 1344)
|
||||
|
||||
return self._run_generation(prompt, "", width, height, rewrite=True, override=False)
|
||||
|
||||
def _generate_tag(self, args: str) -> str:
|
||||
"""直接 Tag 模式"""
|
||||
width, height = 1024, 1344
|
||||
prompt = args.strip()
|
||||
|
||||
if "--width" in prompt or "--height" in prompt or "--w" in prompt or "--h" in prompt:
|
||||
import re
|
||||
w_match = re.search(r'--w(?:idth)?\s+(\d+)', prompt)
|
||||
h_match = re.search(r'--h(?:eight)?\s+(\d+)', prompt)
|
||||
if w_match:
|
||||
width = int(w_match.group(1))
|
||||
prompt = prompt[:w_match.start()] + prompt[w_match.end():]
|
||||
if h_match:
|
||||
height = int(h_match.group(1))
|
||||
prompt = prompt[:h_match.start()] + prompt[h_match.end():]
|
||||
prompt = prompt.strip()
|
||||
|
||||
if not prompt:
|
||||
return "错误:请提供 Tag\n示例: image_generator tag masterpiece, 1girl, white hair, blue eyes"
|
||||
|
||||
width = min(width, 1344)
|
||||
height = min(height, 1344)
|
||||
|
||||
return self._run_generation("", prompt, width, height, rewrite=False, override=False)
|
||||
|
||||
def _run_generation(self, nl_prompt: str, direct_prompt: str, width: int, height: int,
|
||||
rewrite: bool = False, override: bool = False) -> str:
|
||||
"""执行图像生成(同步包装)"""
|
||||
try:
|
||||
import websocket as ws_lib
|
||||
except ImportError:
|
||||
# 回退:用 urllib 的方式不行,WebSocket 必须用专用库
|
||||
# 尝试用 subprocess 调用 python 脚本
|
||||
return self._run_generation_subprocess(nl_prompt, direct_prompt, width, height, rewrite, override)
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
ws = ws_lib.create_connection(self.ws_url, timeout=180)
|
||||
|
||||
payload = {
|
||||
"workflow_path": self.default_workflow,
|
||||
"inline_workflow": None,
|
||||
"direct_prompt": direct_prompt,
|
||||
"nl_prompt": nl_prompt,
|
||||
"rewrite": rewrite,
|
||||
"override": override,
|
||||
"width": width,
|
||||
"height": height,
|
||||
}
|
||||
ws.send(json.dumps(payload))
|
||||
|
||||
image_path = None
|
||||
final_prompt = ""
|
||||
error_msg = ""
|
||||
|
||||
while True:
|
||||
try:
|
||||
raw = ws.recv()
|
||||
data = json.loads(raw)
|
||||
msg_type = data.get("type", "")
|
||||
|
||||
if msg_type == "image":
|
||||
url = data.get("url", "")
|
||||
filename = data.get("filename", "")
|
||||
if url:
|
||||
image_path = self._download_image(url, filename)
|
||||
elif msg_type == "final_prompt":
|
||||
final_prompt = data.get("final_prompt", "") if isinstance(data.get("final_prompt"), str) else data.get("prompt", "")
|
||||
elif msg_type == "done":
|
||||
break
|
||||
elif msg_type == "error":
|
||||
error_msg = data.get("message", "未知错误")
|
||||
break
|
||||
except Exception as e:
|
||||
error_msg = str(e)
|
||||
break
|
||||
|
||||
ws.close()
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
if error_msg:
|
||||
return f" 生成失败: {error_msg}"
|
||||
|
||||
if image_path:
|
||||
size = os.path.getsize(image_path)
|
||||
size_str = f"{size/1024:.0f}KB" if size < 1024*1024 else f"{size/1024/1024:.1f}MB"
|
||||
result = f" 图片生成成功!\n"
|
||||
result += f" 路径: {image_path}\n"
|
||||
result += f" 大小: {size_str}\n"
|
||||
result += f" 分辨率: {width}x{height}\n"
|
||||
result += f"⏱ 耗时: {elapsed:.1f}s\n"
|
||||
if final_prompt:
|
||||
result += f" Prompt: {final_prompt[:120]}...\n"
|
||||
result += f"\n💡 提示: 用 'send_image {image_path}' 可以推送到手机"
|
||||
return result
|
||||
else:
|
||||
return " 生成完成但未收到图片数据"
|
||||
|
||||
except Exception as e:
|
||||
return f" 连接失败: {e}\n💡 请检查网络连接或稍后重试"
|
||||
|
||||
def _run_generation_subprocess(self, nl_prompt: str, direct_prompt: str, width: int, height: int,
|
||||
rewrite: bool, override: bool) -> str:
|
||||
"""用子进程执行 WebSocket 生成(无需 websocket 库)"""
|
||||
import subprocess
|
||||
import tempfile
|
||||
|
||||
script = f'''
|
||||
import json, asyncio, urllib.request, os, sys
|
||||
|
||||
async def main():
|
||||
try:
|
||||
import websockets
|
||||
except ImportError:
|
||||
print("ERROR:需要安装 websockets: pip install websockets")
|
||||
sys.exit(1)
|
||||
|
||||
uri = "wss://ai.2x.nz/ws/run"
|
||||
async with websockets.connect(uri) as ws:
|
||||
payload = {{
|
||||
"workflow_path": "{self.default_workflow}",
|
||||
"inline_workflow": None,
|
||||
"direct_prompt": {json.dumps(direct_prompt)},
|
||||
"nl_prompt": {json.dumps(nl_prompt)},
|
||||
"rewrite": {"true" if rewrite else "false"},
|
||||
"override": {"true" if override else "false"},
|
||||
"width": {width},
|
||||
"height": {height},
|
||||
}}
|
||||
await ws.send(json.dumps(payload))
|
||||
|
||||
image_url = None
|
||||
while True:
|
||||
try:
|
||||
msg = await asyncio.wait_for(ws.recv(), timeout=180)
|
||||
data = json.loads(msg)
|
||||
t = data.get("type", "")
|
||||
if t == "image":
|
||||
image_url = data.get("url", "")
|
||||
elif t == "done":
|
||||
break
|
||||
elif t == "error":
|
||||
print(f"ERROR:{{data.get('message','')}}")
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
print("ERROR:超时")
|
||||
break
|
||||
|
||||
if image_url:
|
||||
full_url = f"https://ai.2x.nz{{image_url}}"
|
||||
out_path = os.path.join("{self.output_dir}", f"generated_{{os.path.basename(image_url)}}")
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
urllib.request.urlretrieve(full_url, out_path)
|
||||
print(f"OK:{{out_path}}")
|
||||
|
||||
asyncio.run(main())
|
||||
'''
|
||||
|
||||
script_path = os.path.join(self.output_dir, "_gen_temp.py")
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
with open(script_path, 'w') as f:
|
||||
f.write(script)
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["python3", script_path],
|
||||
capture_output=True, text=True, timeout=200
|
||||
)
|
||||
output = result.stdout.strip()
|
||||
|
||||
if output.startswith("OK:"):
|
||||
image_path = output[3:]
|
||||
size = os.path.getsize(image_path)
|
||||
size_str = f"{size/1024:.0f}KB" if size < 1024*1024 else f"{size/1024/1024:.1f}MB"
|
||||
res = f" 图片生成成功!\n"
|
||||
res += f" 路径: {image_path}\n"
|
||||
res += f" 大小: {size_str}\n"
|
||||
res += f" 分辨率: {width}x{height}\n"
|
||||
res += f"\n💡 提示: 用 'send_image {image_path}' 可以推送到手机"
|
||||
return res
|
||||
elif output.startswith("ERROR:"):
|
||||
return f" {output[6:]}"
|
||||
else:
|
||||
return f" 生成异常:\n{result.stdout}\n{result.stderr}"
|
||||
except subprocess.TimeoutExpired:
|
||||
return " 生成超时(>200s),请稍后重试"
|
||||
finally:
|
||||
try:
|
||||
os.remove(script_path)
|
||||
except:
|
||||
pass
|
||||
|
||||
def _download_image(self, url: str, filename: str) -> str:
|
||||
"""下载生成的图片"""
|
||||
full_url = f"{self.api_base}{url}"
|
||||
os.makedirs(self.output_dir, exist_ok=True)
|
||||
out_path = os.path.join(self.output_dir, f"gen_{filename}")
|
||||
urllib.request.urlretrieve(full_url, out_path)
|
||||
return out_path
|
||||
|
||||
def _list_workflows(self, args: str) -> str:
|
||||
"""列出可用的工作流"""
|
||||
try:
|
||||
import urllib.request
|
||||
req = urllib.request.Request(f"{self.api_base}/api/workflows")
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
data = json.loads(resp.read())
|
||||
|
||||
workflows = data.get("workflows", [])
|
||||
result = f" 可用工作流 ({len(workflows)} 个):\n\n"
|
||||
|
||||
lora_workflows = []
|
||||
no_lora = []
|
||||
for w in workflows:
|
||||
name = w.get("path", "")
|
||||
if "无Lora" in name or "no" in name.lower():
|
||||
no_lora.append(name)
|
||||
else:
|
||||
lora_workflows.append(name)
|
||||
|
||||
if no_lora:
|
||||
result += " 通用(无角色绑定):\n"
|
||||
for w in no_lora:
|
||||
result += f" • {w}\n"
|
||||
|
||||
if lora_workflows:
|
||||
result += "\n 角色工作流:\n"
|
||||
for w in lora_workflows:
|
||||
result += f" • {w}\n"
|
||||
|
||||
result += "\n💡 使用 'image_generator generate <描述>' 自然语言生图"
|
||||
result += "\n💡 使用 'image_generator tag <tags>' 直接 Tag 生图"
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return f" 获取工作流列表失败: {e}"
|
||||
|
||||
def _gpu_status(self, args: str) -> str:
|
||||
"""查看 GPU 状态"""
|
||||
try:
|
||||
import urllib.request
|
||||
req = urllib.request.Request(f"{self.api_base}/api/gpu")
|
||||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||||
data = json.loads(resp.read())
|
||||
|
||||
result = " 服务器 GPU 状态:\n"
|
||||
for gpu in data.get("gpus", []):
|
||||
name = gpu.get("name", "未知")
|
||||
util = gpu.get("utilization", 0)
|
||||
mem_used = gpu.get("memory_used", 0)
|
||||
mem_total = gpu.get("memory_total", 0)
|
||||
temp = gpu.get("temperature", 0)
|
||||
result += f" {name}\n"
|
||||
result += f" GPU 负载: {util}%\n"
|
||||
result += f" 显存: {mem_used}/{mem_total} MiB\n"
|
||||
result += f" 温度: {temp}°C\n"
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return f" 获取 GPU 状态失败: {e}"
|
||||
Reference in new issue
Block a user