Publish current ComfyUI O1Key code baseline

Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
2026-09-24 19:56:48 +08:00
parent 3e337722ab
commit ba920f2b66
183 changed files with 49496 additions and 9934 deletions
+698 -108
View File
@@ -1,6 +1,6 @@
"""
Comfyui_o1key - ComfyUI 自定义节点集合
通过 api.o1key.com 调用 AI 模型进行图像生成和文本生成
通过 api.o1key.cn 调用 AI 模型进行图像生成和文本生成
项目结构:
├── nodes/ # 节点实现
@@ -71,12 +71,20 @@ if not getattr(asyncio, "_o1key_new_event_loop_patched", False):
asyncio.new_event_loop = _o1key_new_event_loop
asyncio._o1key_new_event_loop_patched = True
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, Google31Video, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, O1keyGrokVideo, KVideoFirstLast, KVideoImage2Video
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch, SaveImageFormat
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, LoadImagesFromFolder, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, Google31Video, MiniMaxH3Video, FluxImageEdit, UniversalLLMChat, BatchImagesO1key, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, O1keyGrokVideo, O1keyGrokVideoEdit
from .nodes import K3Video, K3MotionControl, SaveImageFormat
from .nodes import O1keySavePSD
from .nodes import O1keyRemoveBackground
from .nodes import O1keyColorRemoveBG
from .nodes import O1keyGridSplitter
from .nodes import O1keyPromptMultiFunction
from .nodes import O1keyVideoTrim
from .nodes import SeedanceElementCreate
from .nodes import SeedanceAutoPass
from .nodes import SeedanceAutoPassBatch
from .nodes import O1keyAutoRedCast
from .nodes import O1keyImageGenerator, O1keyImageSave
from .nodes import O1keyVideoGenerator, O1keyVideoResult
from .nodes import O1keyOmniFlashVideo
# 报错弹框友好文案(不修改原节点代码,仅在外层统一处理)
_MSG_TIMEOUT = "API 请求超时,请稍后重试或检查网络。"
@@ -110,8 +118,6 @@ def _wrap_generate_for_error_display(cls, attr="generate"):
_wrap_generate_for_error_display(NanoBananaPro)
_wrap_generate_for_error_display(BatchNanoBananaPro)
_wrap_generate_for_error_display(NanoBananaV2)
_wrap_generate_for_error_display(NanoBananaV2Batch)
# ComfyUI 节点注册
NODE_CLASS_MAPPINGS = {
@@ -119,21 +125,18 @@ NODE_CLASS_MAPPINGS = {
"BatchNanoBananaPro": BatchNanoBananaPro,
"GoogleGemini": GoogleGemini,
"LoadFile": LoadFile,
"O1keyLoadImagesFromFolder": LoadImagesFromFolder,
"ImageStitchPro": ImageStitchPro,
"BatchCleanMetadata": BatchCleanMetadata,
"VideoPreview": VideoPreview,
"GoogleVeo": GoogleVeo,
"Google31Video": Google31Video,
"MiniMaxH3Video": MiniMaxH3Video,
"FluxImageEdit": FluxImageEdit,
"UniversalLLMChat": UniversalLLMChat,
"KlingVideo": KlingVideo,
"KlingFirstLastFrame": KlingFirstLastFrame,
"KlingMotionControlTest": KlingMotionControlTest,
"AspectRatioPreset": AspectRatioPreset,
"BatchImagesO1key": BatchImagesO1key,
"Seedance": Seedance,
"SeedanceMultiModal": SeedanceMultiModal,
"StreamPreview": StreamPreview,
"DoubaoImage": DoubaoImage,
@@ -141,61 +144,68 @@ NODE_CLASS_MAPPINGS = {
"O1keyGPTImageBatch": O1keyGPTImageBatch,
"O1keyGrokImage": O1keyGrokImage,
"O1keyGrokVideo": O1keyGrokVideo,
"KVideoFirstLast": KVideoFirstLast,
"KVideoImage2Video": KVideoImage2Video,
"O1keyGrokVideoEdit": O1keyGrokVideoEdit,
"K3Video": K3Video,
"K3VideoFirstLast": K3VideoFirstLast,
"K3MotionControl": K3MotionControl,
"K3MotionVideoCheck": K3MotionVideoCheck,
"NanoBananaV2": NanoBananaV2,
"NanoBananaV2Batch": NanoBananaV2Batch,
"SaveImageFormat": SaveImageFormat,
"O1keySavePSD": O1keySavePSD,
"O1keyRemoveBackground": O1keyRemoveBackground,
"O1keyColorRemoveBG": O1keyColorRemoveBG,
"O1keyGridSplitter": O1keyGridSplitter,
"O1keyPromptMultiFunction": O1keyPromptMultiFunction,
"O1keyVideoTrim": O1keyVideoTrim,
"SeedanceElementCreate": SeedanceElementCreate,
"SeedanceAutoPass": SeedanceAutoPass,
"SeedanceAutoPassBatch": SeedanceAutoPassBatch,
"O1keyAutoRedCast": O1keyAutoRedCast,
"O1keyImageGenerator": O1keyImageGenerator,
"O1keyImageSave": O1keyImageSave,
"O1keyVideoGenerator": O1keyVideoGenerator,
"O1keyVideoResult": O1keyVideoResult,
"O1keyOmniFlashVideo": O1keyOmniFlashVideo,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"O1keyVideoGenerator": "o1key 视频生成",
"O1keyVideoResult": "o1key 视频结果",
"O1keyOmniFlashVideo": "Omni Flash 视频生成",
"NanoBanana": "Nano Banana",
"BatchNanoBananaPro": "批量 Nano Banana",
"BatchNanoBananaPro": "Nano Banana 批量跑图",
"GoogleGemini": "Google Gemini",
"LoadFile": "加载文件",
"O1keyLoadImagesFromFolder": "加载图像(文件夹)",
"ImageStitchPro": "图像拼接 Pro",
"BatchCleanMetadata": "批量任务(防AI识别)",
"VideoPreview": "预览视频",
"GoogleVeo": "Google Veo - ab",
"Google31Video": "Google 3.1 Video",
"MiniMaxH3Video": "MiniMax H3 / H3 Max 视频生成",
"FluxImageEdit": "Flux2 图像编辑",
"UniversalLLMChat": "全能LLM对话助手",
"KlingVideo": "文/图生视频 自研模型",
"KlingFirstLastFrame": "首尾帧生视频 自研模型",
"KlingMotionControlTest": "动作控制 自研模型",
"AspectRatioPreset": "图片宽高比预设",
"UniversalLLMChat": "提示词专家",
"BatchImagesO1key": "加载图像(批量)",
"Seedance": "Seedance 视频生成",
"SeedanceMultiModal": "Seedance 多模态参考生视频",
"StreamPreview": "流式文本预览",
"DoubaoImage": "豆包生图",
"O1keyGPTImage": "o1key GPT Image",
"O1keyGPTImage": "gpt image",
"O1keyGPTImageBatch": "o1key GPT Image(批量)",
"O1keyGrokImage": "Grok Image",
"O1keyGrokVideo": "Grok Video",
"KVideoFirstLast": "K26 图生视频(首尾帧)",
"KVideoImage2Video": "K26 图生视频",
"K3Video": "K3 图生视频 自研",
"K3VideoFirstLast": "首尾帧 K3 自研",
"K3MotionControl": "动作控制 K3 自研",
"K3MotionVideoCheck": "视频时长检测 K3",
"NanoBananaV2": "Nano Banana V2",
"NanoBananaV2Batch": "Nano Banana V2(批量)",
"O1keyGrokVideoEdit": "Grok Video Edit",
"K3Video": "K 视频生成",
"K3MotionControl": "K 动作模仿",
"SaveImageFormat": "保存图像(格式转换)",
"O1keySavePSD": "保存 PSD(分层)",
"O1keyRemoveBackground": "去背景(rembg",
"O1keyColorRemoveBG": "颜色去背景",
"O1keyGridSplitter": "合并图智能切割",
"O1keyPromptMultiFunction": "提示词(多功能)",
"O1keyVideoTrim": "视频裁剪",
"SeedanceElementCreate": "Seedance 创建素材",
"SeedanceAutoPass": "Seedance 全能生成视频",
"SeedanceAutoPassBatch": "Seedance 全能生成视频(批量)",
"O1keyAutoRedCast": "自动红偏校正",
"O1keyImageGenerator": "o1key 图片生成",
"O1keyImageSave": "o1key 保存图像",
}
WEB_DIRECTORY = "./web"
@@ -207,11 +217,149 @@ try:
from aiohttp import web
from server import PromptServer
import folder_paths
from .utils.config import CONFIG_FILE, load_config, NETWORK_ROUTES
from .utils.config import (
DEFAULT_NETWORK_ROUTE,
NETWORK_ROUTE_CONFIG_KEY,
NETWORK_ROUTE_OPTIONS,
NETWORK_ROUTES,
get_api_key,
get_network_route,
load_config,
update_config,
)
from .utils.chat_support import (
PROMPT_OPTIMIZER_TIMEOUT_SECONDS,
build_search_context,
expand_xlsx_attachments,
extract_search_query,
optimize_image_prompt,
rewrite_search_query,
web_search,
write_video_prompt,
)
from .utils.o1key_image_jobs import register_o1key_image_job_routes
from .utils.o1key_image_thumbnail import register_o1key_image_thumbnail_route
from .utils.o1key_video_jobs import register_o1key_video_job_routes
from .utils.updater import UpdateError, update_package
import threading as _update_threading
_update_lock = _update_threading.Lock()
_O1KEY_IMAGE_JOB_MANAGER = register_o1key_image_job_routes(
PromptServer,
web,
folder_paths,
)
register_o1key_image_thumbnail_route(PromptServer, web, folder_paths)
_O1KEY_VIDEO_JOB_MANAGER = register_o1key_video_job_routes(
PromptServer,
web,
folder_paths,
)
# 每次 ComfyUI 进程启动都会生成新的标识。前端据此确认后端确实完成了
# 重启,而不是仅仅重新加载了浏览器页面。
import os as _restart_os
import sys as _restart_sys
import threading as _restart_threading
import time as _restart_time
import uuid as _restart_uuid
_O1KEY_BOOT_ID = _restart_uuid.uuid4().hex
_o1key_restart_pending = _restart_threading.Event()
_o1key_update_lock = _restart_threading.Lock()
@PromptServer.instance.routes.post("/o1key/update")
async def update_o1key_package(request):
if request.headers.get("X-O1Key-Update") != "1":
return web.json_response(
{"code": "invalid_request", "error": "无效的更新请求。", "suggestion": "请从 O1Key 更新面板重新操作。"},
status=403,
)
if not _o1key_update_lock.acquire(blocking=False):
return web.json_response(
{"code": "update_in_progress", "error": "更新正在进行。", "suggestion": "请等待当前操作完成,不要重复点击。"},
status=409,
)
try:
result = await asyncio.to_thread(update_package)
return web.json_response(result)
except UpdateError as exc:
return web.json_response(exc.as_dict(), status=exc.status)
except Exception:
logging.exception("O1Key 更新失败")
return web.json_response(
{"code": "internal_error", "error": "更新失败。", "suggestion": "请查看 ComfyUI 日志,并在确认本地文件安全后重试。"},
status=500,
)
finally:
_o1key_update_lock.release()
def _o1key_restart_command():
"""复用当前解释器和启动参数,并禁止重启时额外打开浏览器。"""
auto_launch_flags = {"--auto-launch", "--auto_launch", "--launch"}
# orig_argv 包含嵌入式 Python 的 -s 等解释器参数;普通 sys.argv 不包含。
# 保留这些参数可确保便携版重启前后的运行环境完全一致。
original = getattr(_restart_sys, "orig_argv", None)
source_arguments = original[1:] if original else _restart_sys.argv
arguments = [arg for arg in source_arguments if arg not in auto_launch_flags]
if "--disable-auto-launch" not in arguments:
arguments.append("--disable-auto-launch")
return [_restart_sys.executable, *arguments]
def _restart_o1key_comfyui_process(delay=1.25):
"""在响应发送完成后,用相同终端进程重新启动 ComfyUI。"""
try:
_restart_time.sleep(delay)
try:
_restart_sys.stdout.flush()
_restart_sys.stderr.flush()
except Exception:
pass
command = _o1key_restart_command()
print("[O1Key] 正在重启 ComfyUI...", flush=True)
_restart_os.execv(command[0], command)
except Exception:
_o1key_restart_pending.clear()
logging.exception("O1Key 无法重启 ComfyUI 进程")
@PromptServer.instance.routes.get("/o1key/restart/status")
async def get_o1key_restart_status(request):
return web.json_response(
{
"ready": True,
"boot_id": _O1KEY_BOOT_ID,
"pid": _restart_os.getpid(),
},
headers={"Cache-Control": "no-store"},
)
@PromptServer.instance.routes.post("/o1key/restart")
async def restart_o1key_comfyui(request):
if _o1key_restart_pending.is_set():
return web.json_response(
{
"success": False,
"error": "ComfyUI 正在重启,请稍候。",
"boot_id": _O1KEY_BOOT_ID,
},
status=409,
)
_o1key_restart_pending.set()
worker = _restart_threading.Thread(
target=_restart_o1key_comfyui_process,
name="o1key-comfyui-restart",
daemon=True,
)
worker.start()
return web.json_response(
{
"success": True,
"message": "ComfyUI 正在重启。",
"boot_id": _O1KEY_BOOT_ID,
"pid": _restart_os.getpid(),
},
headers={"Cache-Control": "no-store"},
)
def _get_o1key_server_port():
try:
@@ -243,10 +391,30 @@ try:
)
def _get_o1key_notes_file():
# 笔记固定存 ComfyUI input 目录,插件更新/替换不会清空笔记
import os as _os_notes
input_dir = _os_notes.path.abspath(folder_paths.get_input_directory())
_os_notes.makedirs(input_dir, exist_ok=True)
return _os_notes.path.join(input_dir, "o1key-notes.json")
notes_file = _os_notes.path.join(input_dir, "o1key-notes.json")
return notes_file
def _get_o1key_cases_dir():
import os as _os_cases
cases_dir = _os_cases.path.join(_os_cases.path.dirname(__file__), "cases")
_os_cases.makedirs(cases_dir, exist_ok=True)
return cases_dir
def _get_o1key_case_file(filename):
import os as _os_cases
safe_name = _os_cases.path.basename(filename or "")
if not safe_name.lower().endswith(".json"):
return None
cases_dir = _get_o1key_cases_dir()
path = _os_cases.path.abspath(_os_cases.path.join(cases_dir, safe_name))
if not path.startswith(_os_cases.path.abspath(cases_dir) + _os_cases.sep):
return None
return path
def _extract_o1key_notes(payload):
if isinstance(payload, list):
@@ -255,6 +423,46 @@ try:
return payload["notes"]
return None
@PromptServer.instance.routes.get("/o1key/cases")
async def get_o1key_cases(request):
import os as _os_cases
import json as _json_cases
cases_dir = _get_o1key_cases_dir()
cases = []
for filename in sorted(_os_cases.listdir(cases_dir), key=str.lower):
if not filename.lower().endswith(".json"):
continue
path = _get_o1key_case_file(filename)
if not path or not _os_cases.path.isfile(path):
continue
title = _os_cases.path.splitext(filename)[0]
try:
with open(path, "r", encoding="utf-8") as cf:
data = _json_cases.load(cf)
if isinstance(data, dict):
title = str(data.get("title") or data.get("name") or title)
except Exception:
pass
cases.append({"id": filename, "filename": filename, "title": title})
return web.json_response({"cases": cases, "path": str(cases_dir)})
@PromptServer.instance.routes.get("/o1key/case")
async def get_o1key_case(request):
import os as _os_cases
import json as _json_cases
filename = request.query.get("file", "")
path = _get_o1key_case_file(filename)
if not path or not _os_cases.path.isfile(path):
return web.json_response({"error": "case not found"}, status=404)
try:
with open(path, "r", encoding="utf-8") as cf:
data = _json_cases.load(cf)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
return web.json_response({"filename": _os_cases.path.basename(path), "case": data})
@PromptServer.instance.routes.get("/o1key/notes")
async def get_o1key_notes(request):
import os as _os_notes
@@ -327,32 +535,69 @@ try:
masked = key[:3] + "****" + key[-4:]
else:
masked = "****"
return web.json_response({"has_key": bool(key), "masked": masked})
return web.json_response({
"has_key": bool(key),
"masked": masked,
"network_route": get_network_route(),
"network_route_options": NETWORK_ROUTE_OPTIONS,
})
@PromptServer.instance.routes.post("/o1key/config")
async def set_o1key_config(request):
data = await request.json()
route = str(data.get("network_route", "")).strip()
if route not in NETWORK_ROUTES:
return web.json_response({"error": "网络线路无效"}, status=400)
updates = {NETWORK_ROUTE_CONFIG_KEY: route}
if data.get("api_key") is not None:
new_key = str(data.get("api_key", "")).strip()
if not new_key:
return web.json_response({"error": "API Key 不能为空"}, status=400)
if "\n" in new_key or "\r" in new_key:
return web.json_response({"error": "API Key 格式无效"}, status=400)
updates["O1KEY_API_KEY"] = new_key
config = update_config(updates=updates)
key = config.get("O1KEY_API_KEY", "")
masked = key[:3] + "****" + key[-4:] if len(key) > 8 else ("****" if key else "")
return web.json_response({
"success": True,
"has_key": bool(key),
"masked": masked,
"network_route": route,
})
@PromptServer.instance.routes.post("/o1key/network_route")
async def set_network_route(request):
data = await request.json()
route = str(data.get("network_route", "")).strip()
if route not in NETWORK_ROUTES:
return web.json_response({"error": "网络线路无效"}, status=400)
update_config(updates={NETWORK_ROUTE_CONFIG_KEY: route})
return web.json_response({"success": True, "network_route": route})
@PromptServer.instance.routes.post("/o1key/api_key")
async def set_api_key_route(request):
import os
data = await request.json()
new_key = data.get("api_key", "").strip()
if not new_key:
return web.json_response({"error": "API Key 不能为空"}, status=400)
config = load_config()
config["O1KEY_API_KEY"] = new_key
lines = []
for k, v in config.items():
lines.append(f"{k}={v}")
with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
f.write("\n".join(lines) + "\n")
if "\n" in new_key or "\r" in new_key:
return web.json_response({"error": "API Key 格式无效"}, status=400)
update_config(updates={"O1KEY_API_KEY": new_key})
return web.json_response({"success": True})
@PromptServer.instance.routes.post("/o1key/test_key")
async def test_api_key_route(request):
import aiohttp as _aiohttp
data = await request.json()
test_key = data.get("api_key", "").strip()
test_key = str(data.get("api_key") or "").strip() or get_api_key()
if not test_key:
return web.json_response({"valid": False, "error": "密钥不能为空"})
base_url = NETWORK_ROUTES.get("CF加速", "https://cf-api.o1key.com")
return web.json_response({"valid": False, "error": "请先输入或保存 API Key"})
requested_route = str(data.get("network_route", "")).strip()
route = requested_route if requested_route in NETWORK_ROUTES else get_network_route()
base_url = NETWORK_ROUTES.get(route, NETWORK_ROUTES[DEFAULT_NETWORK_ROUTE])
url = f"{base_url}/v1/models"
headers = {"Authorization": f"Bearer {test_key}"}
try:
@@ -370,15 +615,261 @@ try:
@PromptServer.instance.routes.delete("/o1key/api_key")
async def delete_api_key_route(request):
config = load_config()
config.pop("O1KEY_API_KEY", None)
lines = []
for k, v in config.items():
lines.append(f"{k}={v}")
with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
f.write("\n".join(lines) + "\n")
update_config(remove=["O1KEY_API_KEY"])
return web.json_response({"success": True})
# === 主体(Element)代理:转发到 {base}/kling/v1/general/*,后端注入令牌 ===
_ELEMENT_PREFIX = "/kling/v1/general"
def _element_base_url(_route=None):
return NETWORK_ROUTES[get_network_route()].rstrip("/")
def _element_headers():
config = load_config()
key = config.get("O1KEY_API_KEY", "")
if not key:
return None
return {"Authorization": f"Bearer {key}"}
@PromptServer.instance.routes.get("/o1key/element/mine")
async def o1key_element_mine(request):
"""列表接口:GET /kling/v1/general/advanced-custom-elements"""
import aiohttp as _aiohttp
headers = _element_headers()
if not headers:
return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
base = _element_base_url(request.query.get("route"))
url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements"
params = {}
# 支持分页参数
page_num = request.query.get("pageNum", "1")
page_size = request.query.get("pageSize", "100")
params["pageNum"] = page_num
params["pageSize"] = page_size
try:
async with _aiohttp.ClientSession() as session:
async with session.get(url, headers=headers, params=params,
timeout=_aiohttp.ClientTimeout(total=30)) as up:
result = await up.json()
# 新API返回: {"success": true, "data": {"code": 0, "data": [...], "total": N}, "message": ""}
# 转换为前端期望的格式: {"success": true, "data": [...]}
if result.get("success") and isinstance(result.get("data"), dict):
elements = result["data"].get("data", [])
return web.json_response({"success": True, "data": elements, "message": ""})
return web.json_response(result, status=up.status)
except Exception as e:
return web.json_response({"success": False, "message": str(e)}, status=502)
@PromptServer.instance.routes.post("/o1key/element/upload")
async def o1key_element_upload(request):
"""转发 multipart 文件上传:POST /kling/v1/general/upload
视频可达 200MB,固定总超时会截断大文件上传,改用:不限总时长 +
读空闲 120s 超时(连接卡死才超时,慢速大文件不会被一刀切断)。"""
import aiohttp as _aiohttp
headers = _element_headers()
if not headers:
return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
base = _element_base_url(request.query.get("route"))
url = f"{base}{_ELEMENT_PREFIX}/upload"
try:
reader = await request.multipart()
field = await reader.next()
if field is None or field.name != "file":
return web.json_response({"success": False, "message": "缺少 file 字段"}, status=400)
file_bytes = await field.read(decode=False)
filename = field.filename or "image.png"
form = _aiohttp.FormData()
form.add_field("file", file_bytes, filename=filename,
content_type=field.headers.get("Content-Type", "application/octet-stream"))
timeout = _aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=120)
async with _aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, data=form,
timeout=timeout) as up:
data = await up.json()
return web.json_response(data, status=up.status)
except Exception as e:
return web.json_response({"success": False, "message": str(e)}, status=502)
@PromptServer.instance.routes.post("/o1key/element/create")
async def o1key_element_create(request):
"""创建主体:POST /kling/v1/general/advanced-custom-elements
新API字段映射:
- name -> element_name
- description -> element_description
- reference_type -> reference_type (image_refer / video_refer)
- frontal_image -> frontal_image
- refer_images -> refer_images
- video_list -> video_list
- element_voice_id, tag_ids, channel_id 保持不变
"""
import aiohttp as _aiohttp
import json as _json_element
headers = _element_headers()
if not headers:
return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
try:
payload = await request.json()
except Exception as e:
return web.json_response({"success": False, "message": f"请求体无效: {e}"}, status=400)
route = payload.pop("route", None)
base = _element_base_url(route)
url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements"
# 字段映射:前端使用旧字段名,转换为新API字段名
api_payload = {}
if "name" in payload:
api_payload["element_name"] = payload["name"]
if "description" in payload:
api_payload["element_description"] = payload["description"]
# 其他字段直接透传
for key in ["reference_type", "frontal_image", "refer_images", "video_list",
"element_voice_id", "tag_ids", "channel_id"]:
if key in payload:
api_payload[key] = payload[key]
send_headers = {**headers, "Content-Type": "application/json"}
# 打印创建主体的请求信息
try:
print(f"[主体创建] 请求 URL: {url}")
print("[主体创建] 请求体: " + _json_element.dumps(api_payload, ensure_ascii=False, indent=2))
except Exception:
pass
try:
async with _aiohttp.ClientSession() as session:
async with session.post(url, headers=send_headers, json=api_payload,
timeout=_aiohttp.ClientTimeout(total=60)) as up:
data = await up.json()
# 打印创建主体的响应信息
try:
print("[主体创建] 响应体: " + _json_element.dumps(data, ensure_ascii=False, indent=2))
except Exception:
pass
return web.json_response(data, status=up.status)
except Exception as e:
return web.json_response({"success": False, "message": str(e)}, status=502)
@PromptServer.instance.routes.post("/o1key/element/refresh")
async def o1key_element_refresh(request):
"""查询主体:GET /kling/v1/general/advanced-custom-elements/{task_id}
前端传 id(数据库主键),需要先查本地库拿到 job_id(即 task_id),再查询上游。
为了简化,这里改为前端直接传 task_id(即创建时返回的 job_id)。
"""
import aiohttp as _aiohttp
import json as _json_refresh
headers = _element_headers()
if not headers:
return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
try:
payload = await request.json()
except Exception:
payload = {}
task_id = payload.get("task_id") or payload.get("id")
if not task_id:
return web.json_response({"success": False, "message": "缺少 task_id"}, status=400)
base = _element_base_url(payload.get("route"))
url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements/{task_id}"
try:
async with _aiohttp.ClientSession() as session:
async with session.get(url, headers=headers,
timeout=_aiohttp.ClientTimeout(total=60)) as up:
data = await up.json()
# 打印查询响应
try:
print(f"[主体查询] task_id={task_id}")
print("[主体查询] 响应体: " + _json_refresh.dumps(data, ensure_ascii=False, indent=2))
except Exception:
pass
# 新API返回嵌套结构,需要提取 task_status 和 element_id
# 响应: {"success": true, "data": {"code": 0, "data": {"task_status": "succeed", "task_result": {"elements": [...]}}}}
if data.get("success") and isinstance(data.get("data"), dict):
inner = data["data"].get("data", {})
task_status = inner.get("task_status", "")
# 转换为前端期望的格式
element = {
"id": task_id,
"job_id": task_id,
"status": task_status,
"task_status": task_status,
}
if task_status == "succeed":
elements = inner.get("task_result", {}).get("elements", [])
if elements:
first = elements[0]
element["element_id"] = str(first.get("element_id", ""))
element["name"] = first.get("element_name", "")
element["frontal_image"] = first.get("element_image_list", {}).get("frontal_image", "")
elif task_status == "failed":
element["fail_reason"] = inner.get("task_status_msg", "")
return web.json_response({"success": True, "data": {"element": element, "detail": data["data"]}})
return web.json_response(data, status=up.status)
except Exception as e:
return web.json_response({"success": False, "message": str(e)}, status=502)
@PromptServer.instance.routes.post("/o1key/element/delete")
async def o1key_element_delete(request):
"""删除主体:POST /kling/v1/general/delete-advanced-elements
请求体: {"element_id": "315320838184520"}
"""
import aiohttp as _aiohttp
headers = _element_headers()
if not headers:
return web.json_response({"success": False, "message": "未配置 API Key"}, status=401)
try:
payload = await request.json()
except Exception:
payload = {}
element_id = payload.get("element_id") or payload.get("id")
if not element_id:
return web.json_response({"success": False, "message": "缺少 element_id"}, status=400)
base = _element_base_url(payload.get("route"))
url = f"{base}{_ELEMENT_PREFIX}/delete-advanced-elements"
delete_payload = {"element_id": str(element_id)}
send_headers = {**headers, "Content-Type": "application/json"}
try:
async with _aiohttp.ClientSession() as session:
async with session.post(url, headers=send_headers, json=delete_payload,
timeout=_aiohttp.ClientTimeout(total=60)) as up:
data = await up.json()
return web.json_response(data, status=up.status)
except Exception as e:
return web.json_response({"success": False, "message": str(e)}, status=502)
@PromptServer.instance.routes.get("/o1key/element/image")
async def o1key_element_image(request):
"""图片同源代理:ComfyUI 的 CSP 限制 img-src 'self',外链缩略图无法直接显示。
前端把缩略图 src 指向本路由,后端取回字节再吐给浏览器,对浏览器即同源。
仅允许 o1key 资源域,避免被当成任意 URL 抓取的 SSRF 跳板。"""
import aiohttp as _aiohttp
from urllib.parse import urlparse, unquote
raw = request.query.get("url", "")
if not raw:
return web.json_response({"success": False, "message": "缺少 url"}, status=400)
target = unquote(raw)
try:
parsed = urlparse(target)
except Exception:
parsed = None
if not parsed or parsed.scheme not in ("http", "https"):
return web.json_response({"success": False, "message": "非法 url"}, status=400)
host = (parsed.hostname or "").lower()
if not (host.endswith(".o1key.com") or host.endswith(".o1key.cn")
or host in ("o1key.com", "o1key.cn")):
return web.json_response({"success": False, "message": "不允许的图片来源"}, status=403)
try:
async with _aiohttp.ClientSession() as session:
async with session.get(target, timeout=_aiohttp.ClientTimeout(total=30)) as up:
if up.status != 200:
return web.Response(status=up.status)
body = await up.read()
ctype = up.headers.get("Content-Type", "image/jpeg").split(";")[0].strip()
return web.Response(body=body, content_type=ctype or "image/jpeg",
headers={"Cache-Control": "max-age=3600"})
except Exception as e:
return web.json_response({"success": False, "message": str(e)}, status=502)
@PromptServer.instance.routes.get("/o1key/output_history")
async def get_output_history(request):
"""读取 output 目录文件,按执行分组返回 /api/jobs 兼容格式"""
@@ -616,37 +1107,93 @@ try:
pass
return web.json_response({"success": True, "deleted": deleted_files})
@PromptServer.instance.routes.post("/o1key/update")
async def update_node_package(request):
if request.headers.get("X-O1Key-Update") != "1":
return web.json_response({"error": "无效的更新请求。"}, status=403)
if not _update_lock.acquire(blocking=False):
return web.json_response({"error": "更新正在进行,请稍候。"}, status=409)
try:
result = await asyncio.to_thread(update_package)
return web.json_response(result)
except UpdateError as exc:
return web.json_response({"error": str(exc)}, status=409)
except Exception:
logging.exception("o1key update failed")
return web.json_response({"error": "更新失败,请查看 ComfyUI 日志。"}, status=500)
finally:
_update_lock.release()
# === 图片生成提示词优化(服务端读取参考图,避免前端接触 API Key) ===
@PromptServer.instance.routes.post("/o1key/image/prompt-optimize")
async def optimize_o1key_image_prompt(request):
import aiohttp as _aiohttp
# === AI 聊天代理(流式 SSE 透传) ===
@PromptServer.instance.routes.post("/o1key/restart")
async def restart_server(request):
import sys, os as _ros, subprocess, threading
def _do_restart():
import time
time.sleep(1.5)
skip = {"--auto-launch", "--auto_launch", "--launch", "--windows-standalone-build"}
args = [a for a in sys.argv if a not in skip]
args.append("--disable-auto-launch")
subprocess.Popen([sys.executable] + args, cwd=_ros.getcwd())
_ros._exit(0)
threading.Thread(target=_do_restart, daemon=True).start()
return web.json_response({"success": True, "message": "正在重启..."})
try:
data = await request.json()
if not isinstance(data, dict):
raise ValueError("请求体必须是对象")
prompt = data.get("prompt", "")
references = data.get("references", [])
api_key = get_api_key() or ""
if not api_key:
return web.json_response({"error": "未配置 API Key"}, status=401)
timeout = _aiohttp.ClientTimeout(total=PROMPT_OPTIMIZER_TIMEOUT_SECONDS + 15)
async with _aiohttp.ClientSession(timeout=timeout) as session:
optimized = await optimize_image_prompt(
session,
NETWORK_ROUTES[get_network_route()],
api_key,
prompt,
references,
folder_paths.get_input_directory(),
)
return web.json_response(
{
"prompt": optimized,
"model": "gpt-5.6-sol",
"reasoning_effort": "high",
},
headers={"Cache-Control": "no-store"},
)
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
except RuntimeError as exc:
return web.json_response({"error": str(exc)}, status=502)
except Exception:
return web.json_response({"error": "提示词优化失败,请稍后重试"}, status=500)
# === 视频生成 AI帮写(独立视频预设,仅分析安全的 input 图片描述) ===
@PromptServer.instance.routes.post("/o1key/video/prompt-write")
async def write_o1key_video_prompt(request):
import aiohttp as _aiohttp
try:
data = await request.json()
if not isinstance(data, dict):
raise ValueError("请求体必须是对象")
api_key = get_api_key() or ""
if not api_key:
return web.json_response({"error": "未配置 API Key"}, status=401)
context = {
"generation_mode": data.get("generation_mode", "text"),
"duration": data.get("duration", "auto"),
"aspect_ratio": data.get("aspect_ratio", "auto"),
"generate_audio": data.get("generate_audio", False),
"reference_video_count": data.get("reference_video_count", 0),
"reference_audio_count": data.get("reference_audio_count", 0),
}
timeout = _aiohttp.ClientTimeout(total=PROMPT_OPTIMIZER_TIMEOUT_SECONDS + 15)
async with _aiohttp.ClientSession(timeout=timeout) as session:
written = await write_video_prompt(
session,
NETWORK_ROUTES[get_network_route()],
api_key,
data.get("prompt", ""),
data.get("references", []),
folder_paths.get_input_directory(),
context,
)
return web.json_response(
{
"prompt": written,
"model": "gpt-5.6-sol",
"reasoning_effort": "high",
"preset": "video-default",
},
headers={"Cache-Control": "no-store"},
)
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
except RuntimeError as exc:
return web.json_response({"error": str(exc)}, status=502)
except Exception:
return web.json_response({"error": "视频 AI帮写失败,请稍后重试"}, status=500)
# === AI 聊天代理(流式 SSE 透传) ===
@PromptServer.instance.routes.post("/o1key/chat/completions")
@@ -660,45 +1207,88 @@ try:
if not api_key:
return web.json_response({"error": "未配置 API Key"}, status=401)
route = data.get("route", "CF加速")
base_url = NETWORK_ROUTES.get(route, "https://cf-api.o1key.com")
model = data.get("model", "gpt-5.5")
base_url = NETWORK_ROUTES[get_network_route()]
model = data.get("model", "gpt-6-sol")
messages = data.get("messages", [])
if not isinstance(messages, list) or not messages:
return web.json_response({"error": "缺少对话内容"}, status=400)
try:
messages = expand_xlsx_attachments(messages)
except ValueError as exc:
return web.json_response({"error": str(exc)}, status=400)
url = f"{base_url}/v1/chat/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
reasoning = data.get("reasoning_effort", "high")
if reasoning not in ("low", "medium", "high"):
reasoning = "high"
body = {"model": model, "messages": messages, "stream": True}
if model == "claude-fable-5":
budgets = {"low": 2048, "medium": 8192, "high": 16384}
budget = budgets[reasoning]
body["thinking"] = {"type": "enabled", "budget_tokens": budget}
body["max_tokens"] = budget + 8192
elif model in ("gpt-5.5", "gpt-5.6-sol", "gpt-6-astra", "gpt-6-sol", "gemini-3.1-pro-preview"):
body["reasoning_effort"] = reasoning
resp = web.StreamResponse(
status=200, reason="OK",
headers={
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
)
await resp.prepare(request)
search_trace = None
timeout = _aiohttp.ClientTimeout(total=120)
async with _aiohttp.ClientSession(timeout=timeout) as session:
if data.get("web_search") is True:
raw_query = extract_search_query(messages)
if raw_query:
query = await rewrite_search_query(session, base_url, api_key, raw_query) or raw_query
try:
results = await web_search(session, query)
search_trace = {
"query": query,
"results": [
{"title": item["title"], "url": item["url"]}
for item in results
],
}
messages = list(messages)
messages.insert(max(0, len(messages) - 1), {
"role": "system",
"content": build_search_context(query, results),
})
body["messages"] = messages
except Exception as exc:
search_trace = {"query": query, "results": [], "error": str(exc)}
try:
timeout = _aiohttp.ClientTimeout(total=120)
async with _aiohttp.ClientSession(timeout=timeout) as session:
resp = web.StreamResponse(
status=200, reason="OK",
headers={
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
)
await resp.prepare(request)
try:
if search_trace:
event = _cjson.dumps({"o1key_search": search_trace}, ensure_ascii=False)
await resp.write(f"data: {event}\n\n".encode("utf-8"))
async with session.post(url, headers=headers, json=body) as upstream:
if upstream.status != 200:
err = await upstream.text()
await resp.write(f"data: {_cjson.dumps({'error': err})}\n\n".encode())
event = _cjson.dumps({"error": err}, ensure_ascii=False)
await resp.write(f"data: {event}\n\n".encode("utf-8"))
await resp.write(b"data: [DONE]\n\n")
return resp
async for chunk in upstream.content.iter_any():
await resp.write(chunk)
except Exception as e:
await resp.write(f"data: {_cjson.dumps({'error': str(e)})}\n\n".encode())
await resp.write(b"data: [DONE]\n\n")
except Exception as e:
event = _cjson.dumps({"error": str(e)}, ensure_ascii=False)
await resp.write(f"data: {event}\n\n".encode("utf-8"))
await resp.write(b"data: [DONE]\n\n")
return resp
# === 执行事件 Hook:持久化耗时元数据 ===
import time as _time, json as _json2, os as _os
_execution_tracker = {}