""" Comfyui_o1key - ComfyUI 自定义节点集合 通过 api.o1key.cn 调用 AI 模型进行图像生成和文本生成 项目结构: ├── nodes/ # 节点实现 ├── utils/ # 工具模块 ├── clients/ # API 客户端 └── __init__.py # 节点注册入口 """ import ssl import logging import asyncio # 屏蔽 ComfyUI 资产扫描的终端日志输出 _seeder_filter = lambda record: not any( kw in record.getMessage() for kw in ("Seeder start", "Asset scan", "Scan(", "Fast scan") ) logging.getLogger().addFilter(_seeder_filter) def _is_ignored_asyncio_win10054(context): exc = context.get("exception") if not ( isinstance(exc, ConnectionResetError) and getattr(exc, "winerror", None) == 10054 ): return False handle = str(context.get("handle", "")) message = str(context.get("message", "")) marker = "_ProactorBasePipeTransport._call_connection_lost" return marker in handle or marker in message def _install_asyncio_win10054_filter(loop): if getattr(loop, "_o1key_win10054_filter_installed", False): return loop previous_handler = loop.get_exception_handler() def _o1key_asyncio_exception_handler(loop, context): if _is_ignored_asyncio_win10054(context): return if previous_handler is not None: previous_handler(loop, context) else: loop.default_exception_handler(context) loop.set_exception_handler(_o1key_asyncio_exception_handler) setattr(loop, "_o1key_win10054_filter_installed", True) return loop try: _install_asyncio_win10054_filter(asyncio.get_event_loop()) except RuntimeError: pass if not getattr(asyncio, "_o1key_new_event_loop_patched", False): _o1key_original_new_event_loop = asyncio.new_event_loop def _o1key_new_event_loop(*args, **kwargs): return _install_asyncio_win10054_filter( _o1key_original_new_event_loop(*args, **kwargs) ) asyncio.new_event_loop = _o1key_new_event_loop asyncio._o1key_new_event_loop_patched = True 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 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 请求超时，请稍后重试或检查网络。" _MSG_SSL_NETWORK = ( "本地网络不太稳定!解决方案如下:\n" "1. 重启程序再试试看 (优先)\n" "2. 调整一下网络环境,如wifi或宽带等\n" "3. 切换VPN节点,或更换代理模式\n" "4. 关掉杀毒软件或防火墙\n" "5. 关掉浏览器VPN插件,避免冲突" ) def _wrap_generate_for_error_display(cls, attr="generate"): original = getattr(cls, attr, None) if original is None: return def wrapped(self, *args, **kwargs): try: return original(self, *args, **kwargs) except TimeoutError as e: msg = (str(e) or "").strip() if not msg: msg = _MSG_TIMEOUT raise TimeoutError(msg) from None except (ssl.SSLError, OSError) as e: err_str = str(e) if "DECRYPTION_FAILED_OR_BAD_RECORD_MAC" in err_str or "decryption failed or bad record mac" in err_str.lower(): raise RuntimeError(_MSG_SSL_NETWORK) from None raise setattr(cls, attr, wrapped) _wrap_generate_for_error_display(NanoBananaPro) _wrap_generate_for_error_display(BatchNanoBananaPro) # ComfyUI 节点注册 NODE_CLASS_MAPPINGS = { "NanoBanana": NanoBananaPro, "BatchNanoBananaPro": BatchNanoBananaPro, "GoogleGemini": GoogleGemini, "LoadFile": LoadFile, "O1keyLoadImagesFromFolder": LoadImagesFromFolder, "ImageStitchPro": ImageStitchPro, "BatchCleanMetadata": BatchCleanMetadata, "VideoPreview": VideoPreview, "GoogleVeo": GoogleVeo, "Google31Video": Google31Video, "MiniMaxH3Video": MiniMaxH3Video, "FluxImageEdit": FluxImageEdit, "UniversalLLMChat": UniversalLLMChat, "BatchImagesO1key": BatchImagesO1key, "SeedanceMultiModal": SeedanceMultiModal, "StreamPreview": StreamPreview, "DoubaoImage": DoubaoImage, "O1keyGPTImage": O1keyGPTImage, "O1keyGPTImageBatch": O1keyGPTImageBatch, "O1keyGrokImage": O1keyGrokImage, "O1keyGrokVideo": O1keyGrokVideo, "O1keyGrokVideoEdit": O1keyGrokVideoEdit, "K3Video": K3Video, "K3MotionControl": K3MotionControl, "SaveImageFormat": SaveImageFormat, "O1keySavePSD": O1keySavePSD, "O1keyRemoveBackground": O1keyRemoveBackground, "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 批量跑图", "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": "提示词专家", "BatchImagesO1key": "加载图像（批量）", "SeedanceMultiModal": "Seedance 多模态参考生视频", "StreamPreview": "流式文本预览", "DoubaoImage": "豆包生图", "O1keyGPTImage": "gpt image", "O1keyGPTImageBatch": "o1key GPT Image（批量）", "O1keyGrokImage": "Grok Image", "O1keyGrokVideo": "Grok Video", "O1keyGrokVideoEdit": "Grok Video Edit", "K3Video": "K 视频生成", "K3MotionControl": "K 动作模仿", "SaveImageFormat": "保存图像（格式转换）", "O1keySavePSD": "保存 PSD（分层）", "O1keyRemoveBackground": "去背景（rembg）", "O1keyGridSplitter": "合并图智能切割", "O1keyPromptMultiFunction": "提示词（多功能）", "O1keyVideoTrim": "视频裁剪", "SeedanceElementCreate": "Seedance 创建素材", "SeedanceAutoPass": "Seedance 全能生成视频", "SeedanceAutoPassBatch": "Seedance 全能生成视频（批量）", "O1keyAutoRedCast": "自动红偏校正", "O1keyImageGenerator": "o1key 图片生成", "O1keyImageSave": "o1key 保存图像", } WEB_DIRECTORY = "./web" __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] # 注册 /o1key/input_dir 接口，供前端文件上传按钮获取 input 目录绝对路径 try: from aiohttp import web from server import PromptServer import folder_paths 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, check_for_update, update_package _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.get("/o1key/update/check") async def check_o1key_update(request): if not _o1key_update_lock.acquire(blocking=False): return web.json_response( {"code": "update_in_progress", "error": "更新正在进行。"}, status=409, ) try: result = await asyncio.to_thread(check_for_update) 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": "检查更新失败。"}, status=500, ) finally: _o1key_update_lock.release() @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: import comfy.cli_args as _cli_args args = getattr(_cli_args, "args", None) port = getattr(args, "port", None) if args else None port = port or getattr(_cli_args, "server_port", None) or getattr(_cli_args, "port", None) if port is not None: return str(int(port)) except Exception: pass try: import sys as _sys for idx, arg in enumerate(_sys.argv): if arg in ("--port", "--listen-port") and idx + 1 < len(_sys.argv): return str(int(_sys.argv[idx + 1])) for prefix in ("--port=", "--listen-port="): if arg.startswith(prefix): return str(int(arg.split("=", 1)[1])) except Exception: pass return "8188" def _get_o1key_history_meta_file(output_dir): import os as _os_history return _os_history.path.join( output_dir, f".o1key_history_{_get_o1key_server_port()}.json", ) 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) 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): return payload if isinstance(payload, dict) and isinstance(payload.get("notes"), list): 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 import json as _json_notes notes_file = _get_o1key_notes_file() exists = _os_notes.path.isfile(notes_file) notes = [] if exists: try: with open(notes_file, "r", encoding="utf-8") as nf: loaded = _json_notes.load(nf) notes = _extract_o1key_notes(loaded) if notes is None: return web.json_response( {"error": "invalid notes file", "path": notes_file}, status=500, ) except Exception as e: return web.json_response( {"error": str(e), "path": notes_file}, status=500, ) return web.json_response({"notes": notes, "path": notes_file, "exists": exists}) @PromptServer.instance.routes.post("/o1key/notes") async def save_o1key_notes(request): import os as _os_notes import json as _json_notes try: payload = await request.json() notes = _extract_o1key_notes(payload) if notes is None: return web.json_response({"error": "notes must be a list"}, status=400) except Exception as e: return web.json_response({"error": f"invalid notes payload: {str(e)}"}, status=400) notes_file = _get_o1key_notes_file() temp_file = notes_file + ".tmp" try: with open(temp_file, "w", encoding="utf-8") as nf: _json_notes.dump(notes, nf, ensure_ascii=False, indent=2) nf.write("\n") _os_notes.replace(temp_file, notes_file) except Exception as e: return web.json_response({"error": f"save notes failed: {str(e)}"}, status=500) return web.json_response({ "success": True, "path": notes_file, "count": len(notes), }) @PromptServer.instance.routes.get("/o1key/input_dir") async def get_input_dir(request): import os path = os.path.abspath(folder_paths.get_input_directory()) return web.json_response({"path": path}) @PromptServer.instance.routes.get("/o1key/api_key") async def get_api_key_route(request): config = load_config() key = config.get("O1KEY_API_KEY", "") masked = "" if key: if len(key) > 8: masked = key[:3] + "****" + key[-4:] else: 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): data = await request.json() new_key = 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) 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 = str(data.get("api_key") or "").strip() or get_api_key() if not test_key: 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: async with _aiohttp.ClientSession() as session: async with session.get(url, headers=headers, timeout=_aiohttp.ClientTimeout(total=10)) as resp: if resp.status == 200: return web.json_response({"valid": True}) elif resp.status == 401: return web.json_response({"valid": False, "error": "密钥无效或已过期"}) else: text = await resp.text() return web.json_response({"valid": False, "error": f"验证失败 ({resp.status})"}) except Exception as e: return web.json_response({"valid": False, "error": f"网络错误: {str(e)}"}) @PromptServer.instance.routes.delete("/o1key/api_key") async def delete_api_key_route(request): 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 兼容格式""" import os, json as _json limit = int(request.query.get("limit", "200")) offset = int(request.query.get("offset", "0")) output_dir = os.path.abspath(folder_paths.get_output_directory()) meta_file = _get_o1key_history_meta_file(output_dir) meta = {} if os.path.isfile(meta_file): try: with open(meta_file, "r", encoding="utf-8") as mf: meta = _json.load(mf) except Exception: pass supported_ext = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.webm'} # 收集所有文件并按 workflow_id 分组 all_files = [] for fname in meta.keys(): ext = os.path.splitext(fname)[1].lower() if ext not in supported_ext: continue fpath = os.path.join(output_dir, fname) if not os.path.isfile(fpath): continue mtime = os.path.getmtime(fpath) media = "images" if ext in {'.png','.jpg','.jpeg','.webp','.gif'} else "video" all_files.append({"name": fname, "mtime": mtime, "media": media}) # 按 workflow_id 分组（同一次执行合并为一个 job） groups = {} for f in all_files: m = meta.get(f["name"], {}) wid = m.get("workflow_id") if wid: groups.setdefault(wid, []).append((f, m)) # 构建 job 列表 jobs = [] for wid, items in groups.items(): items.sort(key=lambda x: x[0]["mtime"], reverse=True) latest = items[0] f, m = latest start_ms = int(m.get("start_time", f["mtime"]) * 1000) end_ms = int(m.get("end_time", f["mtime"]) * 1000) jobs.append({ "id": wid, "status": "completed", "create_time": start_ms, "execution_start_time": start_ms, "execution_end_time": end_ms, "preview_output": { "filename": f["name"], "subfolder": "", "type": "output", "nodeId": "0", "mediaType": f["media"], }, "outputs_count": len(items), "execution_error": None, "workflow_id": wid, }) # 按时间倒序排列，分页 jobs.sort(key=lambda x: x["create_time"], reverse=True) total = len(jobs) page = jobs[offset:offset+limit] return web.json_response({ "jobs": page, "pagination": {"offset": offset, "limit": limit, "total": total, "has_more": offset + limit < total} }) @PromptServer.instance.routes.get("/o1key/output_workflow") async def get_output_workflow(request): """从 PNG 元数据中读取工作流，供前端恢复使用""" import os, struct, json as _json filename = request.query.get("filename", "") if not filename: return web.json_response({"error": "missing filename"}, status=400) output_dir = os.path.abspath(folder_paths.get_output_directory()) fpath = os.path.join(output_dir, filename) if not os.path.isfile(fpath) or not fpath.lower().endswith(".png"): return web.json_response({"error": "file not found"}, status=404) workflow = None prompt_data = None try: with open(fpath, "rb") as pf: pf.read(8) # PNG signature while True: raw = pf.read(8) if len(raw) < 8: break length = struct.unpack(">I", raw[:4])[0] chunk_type = raw[4:8] data = pf.read(length) pf.read(4) # CRC if chunk_type == b"tEXt": key, val = data.split(b"\x00", 1) k = key.decode("ascii", errors="replace") if k == "workflow": workflow = _json.loads(val) elif k == "prompt": prompt_data = _json.loads(val) elif chunk_type == b"IEND": break except Exception: pass return web.json_response({"workflow": workflow, "prompt": prompt_data}) @PromptServer.instance.routes.get("/o1key/job_detail/{job_id}") async def get_job_detail(request): """根据 job_id 返回当前端口持久化历史中的 job 详情""" import os, struct, json as _json job_id = request.match_info["job_id"] output_dir = os.path.abspath(folder_paths.get_output_directory()) meta_file = _get_o1key_history_meta_file(output_dir) meta = {} if os.path.isfile(meta_file): try: with open(meta_file, "r", encoding="utf-8") as mf: meta = _json.load(mf) except Exception: pass supported_ext = {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.mp4', '.webm'} # 只在当前端口的持久化记录中查找该 job 的文件 matched_files = [] for fname in meta.keys(): ext = os.path.splitext(fname)[1].lower() if ext not in supported_ext: continue m = meta.get(fname, {}) if m.get("workflow_id") == job_id: matched_files.append(fname) if not matched_files: return web.json_response({"error": "not found"}, status=404) # 用最新文件作为代表 matched_files.sort(key=lambda f: os.path.getmtime(os.path.join(output_dir, f)), reverse=True) target_file = matched_files[0] fpath = os.path.join(output_dir, target_file) m = meta.get(target_file, {}) mtime = os.path.getmtime(fpath) start_ms = int(m.get("start_time", mtime) * 1000) end_ms = int(m.get("end_time", mtime) * 1000) ext = os.path.splitext(target_file)[1].lower() media = "images" if ext in {'.png','.jpg','.jpeg','.webp','.gif'} else "video" # 读取 PNG 工作流元数据 workflow = None if ext == ".png": try: with open(fpath, "rb") as pf: pf.read(8) while True: raw = pf.read(8) if len(raw) < 8: break length = struct.unpack(">I", raw[:4])[0] chunk_type = raw[4:8] data = pf.read(length) pf.read(4) if chunk_type == b"tEXt": key, val = data.split(b"\x00", 1) k = key.decode("ascii", errors="replace") if k == "workflow": workflow = _json.loads(val) elif chunk_type == b"IEND": break except Exception: pass # 构建 outputs：包含该执行的所有文件 outputs = {} for i, fname in enumerate(matched_files): e = os.path.splitext(fname)[1].lower() mt = "images" if e in {'.png','.jpg','.jpeg','.webp','.gif'} else "gifs" outputs.setdefault(str(i), {}).setdefault(mt, []).append( {"filename": fname, "subfolder": "", "type": "output"} ) job_detail = { "id": job_id, "status": "completed", "create_time": start_ms, "execution_start_time": start_ms, "execution_end_time": end_ms, "preview_output": { "filename": target_file, "subfolder": "", "type": "output", "nodeId": "0", "mediaType": media, }, "outputs_count": len(matched_files), "execution_error": None, "workflow_id": job_id, "workflow": { "extra_data": { "extra_pnginfo": {"workflow": workflow} } } if workflow else None, "outputs": outputs, } return web.json_response(job_detail) @PromptServer.instance.routes.post("/o1key/delete_history") async def delete_history_item(request): """删除持久化历史记录及对应的输出文件""" import os, json as _json body = await request.json() job_ids = body.get("delete", []) if not job_ids: return web.json_response({"success": False, "error": "missing ids"}, status=400) output_dir = os.path.abspath(folder_paths.get_output_directory()) meta_file = _get_o1key_history_meta_file(output_dir) meta = {} if os.path.isfile(meta_file): try: with open(meta_file, "r", encoding="utf-8") as mf: meta = _json.load(mf) except Exception: pass deleted_files = [] for job_id in job_ids: files_to_remove = [] for fname, m in list(meta.items()): if m.get("workflow_id") == job_id: files_to_remove.append(fname) for fname in files_to_remove: meta.pop(fname, None) fpath = os.path.join(output_dir, fname) if os.path.isfile(fpath): try: os.remove(fpath) deleted_files.append(fname) except Exception: pass try: with open(meta_file, "w", encoding="utf-8") as mf: _json.dump(meta, mf, ensure_ascii=False) except Exception: pass return web.json_response({"success": True, "deleted": deleted_files}) # === 图片生成提示词优化（服务端读取参考图，避免前端接触 API Key） === @PromptServer.instance.routes.post("/o1key/image/prompt-optimize") async def optimize_o1key_image_prompt(request): import aiohttp as _aiohttp 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") async def chat_completions_proxy(request): import aiohttp as _aiohttp import json as _cjson data = await request.json() config = load_config() api_key = config.get("O1KEY_API_KEY", "") if not api_key: return web.json_response({"error": "未配置 API Key"}, status=401) 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 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)} 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() 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: 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 = {} _orig_send_sync = PromptServer.instance.send_sync def _patched_send_sync(event, data, *args, **kwargs): try: if event == "execution_start": pid = data.get("prompt_id", "") if pid: _execution_tracker[pid] = {"start": _time.time(), "outputs": []} elif event == "executed": pid = data.get("prompt_id", "") output = data.get("output") or {} if pid and pid in _execution_tracker: for img in output.get("images", []) + output.get("gifs", []): if img.get("type") == "output" and img.get("filename"): _execution_tracker[pid]["outputs"].append(img["filename"]) elif event == "executing" and data.get("node") is None: pid = data.get("prompt_id", "") tracker = _execution_tracker.pop(pid, None) if tracker and tracker["outputs"]: end_time = _time.time() start_time = tracker["start"] output_dir = _os.path.abspath(folder_paths.get_output_directory()) meta_file = _get_o1key_history_meta_file(output_dir) meta = {} if _os.path.isfile(meta_file): try: with open(meta_file, "r", encoding="utf-8") as mf: meta = _json2.load(mf) except Exception: pass for fname in tracker["outputs"]: meta[fname] = { "start_time": start_time, "end_time": end_time, "outputs_count": len(tracker["outputs"]), "workflow_id": pid, } try: with open(meta_file, "w", encoding="utf-8") as mf: _json2.dump(meta, mf, ensure_ascii=False) except Exception: pass except Exception: pass return _orig_send_sync(event, data, *args, **kwargs) PromptServer.instance.send_sync = _patched_send_sync except Exception: pass