OpenClaw
357cd0958f
18个插件: class Plugin -> class Liugin 10个插件: 测试块 Plugin() -> Liugin() 移除 ai_image.py 中多余的 Liugin = Plugin 别名 所有自研插件现统一使用 class Liugin 作为入口类名， 加载器保留 Plugin 兼容以支持第三方插件。
371 lines
14 KiB
Python
371 lines
14 KiB
Python
"""
|
||
图像生成插件 - 基于 ai.2x.nz 的 ComfyUI 在线服务
|
||
支持自然语言描述和英文 Tag 两种模式
|
||
"""
|
||
import os
|
||
import json
|
||
import asyncio
|
||
import urllib.request
|
||
import time
|
||
from typing import Optional
|
||
|
||
|
||
class Liugin:
|
||
"""图像生成工具 — 自然语言/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}"
|