Ubuntu
92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
baoyu-imagine wrapper for nanobot
|
|
简化信息图生成流程
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
# 配置
|
|
BAOYU_IMAGINE_DIR = "/home/ubuntu/forgithub/baoyu-skills/skills/baoyu-imagine"
|
|
DEFAULT_PROVIDER = "seedream"
|
|
|
|
def get_api_key():
|
|
"""获取 API Key"""
|
|
# 优先环境变量
|
|
if "ARK_API_KEY" in os.environ:
|
|
return os.environ["ARK_API_KEY"]
|
|
|
|
# 其次配置文件
|
|
config_path = Path.home() / ".config" / "baoyu" / "keys.conf"
|
|
if config_path.exists():
|
|
with open(config_path) as f:
|
|
for line in f:
|
|
if line.startswith("ARK_API_KEY="):
|
|
return line.strip().split("=", 1)[1]
|
|
|
|
raise ValueError("ARK_API_KEY not found. Please set it in ~/.config/baoyu/keys.conf")
|
|
|
|
def generate_infographic(system_prompt: str, content_prompt: str, output_path: str):
|
|
"""生成信息图"""
|
|
|
|
api_key = get_api_key()
|
|
|
|
# 创建临时 prompt 文件
|
|
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
|
|
system_file = f.name
|
|
f.write(system_prompt)
|
|
|
|
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
|
|
content_file = f.name
|
|
f.write(content_prompt)
|
|
|
|
try:
|
|
# 调用 baoyu-imagine
|
|
cmd = [
|
|
"npx", "-y", "bun",
|
|
f"{BAOYU_IMAGINE_DIR}/scripts/main.ts",
|
|
"--provider", DEFAULT_PROVIDER,
|
|
"--promptfiles", system_file, content_file,
|
|
"--image", output_path
|
|
]
|
|
|
|
env = os.environ.copy()
|
|
env["ARK_API_KEY"] = api_key
|
|
|
|
result = subprocess.run(
|
|
cmd,
|
|
env=env,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=300
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
print(f"Error: {result.stderr}", file=sys.stderr)
|
|
return False
|
|
|
|
return True
|
|
|
|
finally:
|
|
# 清理临时文件
|
|
os.unlink(system_file)
|
|
os.unlink(content_file)
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 4:
|
|
print("Usage: python generate.py <system-prompt.md> <content-prompt.md> <output.png>")
|
|
sys.exit(1)
|
|
|
|
with open(sys.argv[1]) as f:
|
|
system = f.read()
|
|
|
|
with open(sys.argv[2]) as f:
|
|
content = f.read()
|
|
|
|
success = generate_infographic(system, content, sys.argv[3])
|
|
sys.exit(0 if success else 1)
|