diff --git a/__init__.py b/__init__.py index 947d4c2..1e60d60 100644 --- a/__init__.py +++ b/__init__.py @@ -12,6 +12,7 @@ Comfyui_o1key - ComfyUI 自定义节点集合 import ssl import logging +import asyncio # 屏蔽 ComfyUI 资产扫描的终端日志输出 _seeder_filter = lambda record: not any( @@ -20,7 +21,57 @@ _seeder_filter = lambda record: not any( ) logging.getLogger().addFilter(_seeder_filter) -from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, KVideoFirstLast, KVideoImage2Video + +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, 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 O1keySavePSD from .nodes import O1keyRemoveBackground @@ -73,6 +124,7 @@ NODE_CLASS_MAPPINGS = { "BatchCleanMetadata": BatchCleanMetadata, "VideoPreview": VideoPreview, "GoogleVeo": GoogleVeo, + "Google31Video": Google31Video, "FluxImageEdit": FluxImageEdit, "UniversalLLMChat": UniversalLLMChat, "KlingVideo": KlingVideo, @@ -88,6 +140,7 @@ NODE_CLASS_MAPPINGS = { "O1keyGPTImage": O1keyGPTImage, "O1keyGPTImageBatch": O1keyGPTImageBatch, "O1keyGrokImage": O1keyGrokImage, + "O1keyGrokVideo": O1keyGrokVideo, "KVideoFirstLast": KVideoFirstLast, "KVideoImage2Video": KVideoImage2Video, "K3Video": K3Video, @@ -113,6 +166,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "BatchCleanMetadata": "批量任务（防AI识别）", "VideoPreview": "预览视频", "GoogleVeo": "Google Veo - ab", + "Google31Video": "Google 3.1 Video", "FluxImageEdit": "Flux2 图像编辑", "UniversalLLMChat": "全能LLM对话助手", "KlingVideo": "文/图生视频 自研模型", @@ -128,6 +182,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "O1keyGPTImage": "o1key GPT Image", "O1keyGPTImageBatch": "o1key GPT Image（批量）", "O1keyGrokImage": "Grok Image", + "O1keyGrokVideo": "Grok Video", "KVideoFirstLast": "K26 图生视频（首尾帧）", "KVideoImage2Video": "K26 图生视频", "K3Video": "K3 图生视频 自研", @@ -183,6 +238,75 @@ try: f".o1key_history_{_get_o1key_server_port()}.json", ) + def _get_o1key_notes_file(): + 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") + + 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/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 diff --git a/clients/__init__.py b/clients/__init__.py index e96f850..16fe95d 100644 --- a/clients/__init__.py +++ b/clients/__init__.py @@ -9,6 +9,8 @@ from .gemini_flash_client import GeminiFlashClient from .sora_client import SoraClient from .kling_client import KlingClient from .veo_client import VeoClient +from .newapi_veo_client import NewAPIVeoClient +from .grok_video_client import GrokVideoClient from .openai_client import OpenAIAPIClient -__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient'] +__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient'] diff --git a/clients/base_client.py b/clients/base_client.py index 4d369f7..188c03d 100644 --- a/clients/base_client.py +++ b/clients/base_client.py @@ -40,7 +40,8 @@ class BaseAPIClient(ABC): Args: base_url: API 基础 URL api_key: API 密钥 - max_request_size: 最大请求体大小（字节），默认 100MB + max_request_size: 兼容参数；基类不再用它限制 JSON 请求体， + 部分子类仍用它作为上传文件大小限制 """ self.base_url = base_url self.api_key = api_key @@ -116,24 +117,6 @@ class BaseAPIClient(ABC): "Content-Type": "application/json" } - def check_request_size(self, request_body: Dict[str, Any]) -> None: - """ - 检查请求体大小是否超过限制 - - Args: - request_body: 请求体字典 - - Raises: - ValueError: 如果请求体超过限制 - """ - request_json = json.dumps(request_body) - request_size = len(request_json.encode('utf-8')) - - if request_size > self.max_request_size: - raise ValueError( - "请求体积超过100MB限制，请调整分辨率或减少图片数量" - ) - def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]: """ 子类可重写：为指定 HTTP 状态码返回自定义错误文案。 @@ -185,9 +168,6 @@ class BaseAPIClient(ABC): url = f"{self.base_url}{endpoint}" headers = self.get_headers(use_bearer_token) - # 检查请求大小 - self.check_request_size(request_body) - close_session = False if session is None: session = self._make_session() diff --git a/clients/gemini_async_provider.py b/clients/gemini_async_provider.py index 64cb070..9ce237b 100644 --- a/clients/gemini_async_provider.py +++ b/clients/gemini_async_provider.py @@ -72,7 +72,9 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider): resolution=resolution, enable_grounding=kwargs.get("enable_grounding", False), enable_image_search=kwargs.get("enable_image_search", False), - image_compression=getattr(self, 'image_compression', None), + image_compression=getattr(self, "image_compression", None), + thinking_level=kwargs.get("thinking_level"), + request_log_enabled=False, ) def extract_task_id(self, response: dict) -> str: @@ -85,6 +87,23 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider): return response.get("status", "UNKNOWN") async def parse_result(self, result_data: dict, session) -> List[Image.Image]: + images = result_data.get("images") if isinstance(result_data, dict) else None + if isinstance(images, list) and images: + parsed = [] + for item in images: + if not isinstance(item, dict): + continue + image_url = item.get("url") or item.get("image_url") + if image_url: + async with session.get(image_url) as img_resp: + if img_resp.status == 200: + img_bytes = await img_resp.read() + parsed.append(Image.open(BytesIO(img_bytes)).convert("RGB")) + else: + raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}") + if parsed: + return parsed + # 异步接口可能直接返回 image_url image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else "" if image_url: @@ -124,19 +143,44 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider): def extract_progress(self, response: dict) -> Optional[float]: """从轮询响应中提取进度（0.0-1.0）""" + + def _coerce(val) -> Optional[float]: + if val is None or isinstance(val, bool): + return None + if isinstance(val, (int, float)): + progress = float(val) + elif isinstance(val, str): + text = val.strip() + if not text: + return None + has_percent_suffix = text.endswith("%") + if has_percent_suffix: + text = text[:-1].strip() + try: + progress = float(text) + except ValueError: + return None + if has_percent_suffix: + progress /= 100.0 + else: + return None + if progress > 1.0: + progress /= 100.0 + return max(0.0, min(progress, 1.0)) + # 直接字段：progress / percentage - for field in ("progress", "percentage"): - val = response.get(field) - if val is not None and isinstance(val, (int, float)): - return val / 100.0 if val > 1 else float(val) + for field in ("progress", "percentage", "percent"): + progress = _coerce(response.get(field)) + if progress is not None: + return progress # 嵌套字段：progressInfo / progress_info progress_info = response.get("progressInfo") or response.get("progress_info") if isinstance(progress_info, dict): - for field in ("progress", "percentage"): - val = progress_info.get(field) - if val is not None and isinstance(val, (int, float)): - return val / 100.0 if val > 1 else float(val) + for field in ("progress", "percentage", "percent"): + progress = _coerce(progress_info.get(field)) + if progress is not None: + return progress return None diff --git a/clients/gemini_client.py b/clients/gemini_client.py index a76cc8f..67c0cde 100644 --- a/clients/gemini_client.py +++ b/clients/gemini_client.py @@ -34,8 +34,7 @@ class GeminiAPIClient(BaseAPIClient): super().__init__( base_url=get_api_base_url(), - api_key=api_key, - max_request_size=100 * 1024 * 1024 + api_key=api_key ) @staticmethod @@ -188,6 +187,7 @@ class GeminiAPIClient(BaseAPIClient): enable_image_search: bool = False, image_compression: str = None, thinking_level: str = None, + request_log_enabled: bool = True, **kwargs ) -> Dict[str, Any]: """ @@ -204,21 +204,90 @@ class GeminiAPIClient(BaseAPIClient): Returns: 请求体字典 """ + import json + _MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB + _BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.9) parts = [] # 添加文本部分 parts.append({"text": prompt}) + image_config = {"imageSize": resolution} + if aspect_ratio and aspect_ratio != "智能": + image_config["aspectRatio"] = aspect_ratio + + def _build_request_body(body_parts: List[dict]) -> Dict[str, Any]: + body = { + "contents": [ + { + "role": "user", + "parts": body_parts + } + ], + "generationConfig": { + "responseModalities": ["IMAGE"], + "imageConfig": image_config + } + } + + # 添加思考深度配置 + if thinking_level: + body["generationConfig"]["thinkingConfig"] = { + "thinkingLevel": thinking_level, + "includeThoughts": True + } + + # 添加图片压缩参数 + if image_compression: + body["image_compression"] = image_compression + + # 添加 Google Search Grounding（如果启用） + # 新异步接口要求直接放在请求体顶层：{"google_search": true} + if enable_grounding or enable_image_search: + body["google_search"] = True + + return body + + def _request_size(body: Dict[str, Any]) -> int: + return len(json.dumps(body).encode("utf-8")) + + def _format_size(size: int) -> str: + if size < 1024 * 1024: + return f"{size / 1024:.2f}KB" + return f"{size / 1024 / 1024:.2f}MB" + + def _shorten_base64_for_log(obj, max_len: int = 200): + if isinstance(obj, dict): + result = {} + for key, value in obj.items(): + if key == "data" and isinstance(value, str) and len(value) > max_len: + result[key] = f"" + else: + result[key] = _shorten_base64_for_log(value, max_len) + return result + if isinstance(obj, list): + return [_shorten_base64_for_log(item, max_len) for item in obj] + return obj + + def _log_original_request_body(body: Dict[str, Any]) -> None: + body_size = _request_size(body) + print( + f"\n{'=' * 60}\n" + f"[原始请求体日志] 请求体积: {_format_size(body_size)} " + f"(inline_data.data 已折叠显示 base64 长度)\n" + f"{json.dumps(_shorten_base64_for_log(body), ensure_ascii=False, indent=2)}\n" + f"{'=' * 60}\n" + ) + + original_request_logged = False + # 添加图像部分（如果有） if images: - working_images = list(images) - - # 编码一次，估算大小，超限则迭代缩放 - for _attempt in range(10): + def _build_image_parts(src_images: List[Image.Image]) -> List[dict]: img_parts = [] - for img in working_images: + for img in src_images: img_base64 = encode_image_to_base64(img) img_parts.append({ "inline_data": { @@ -226,94 +295,49 @@ class GeminiAPIClient(BaseAPIClient): "data": img_base64 } }) + return img_parts - # 估算完整 body 大小（不含工具字段，工具字段很小可忽略） - est_image_config = {"imageSize": resolution} - if aspect_ratio and aspect_ratio != "智能": - est_image_config["aspectRatio"] = aspect_ratio - estimated = self._estimate_body_size( - parts + img_parts, - { - "generationConfig": { - "responseModalities": ["IMAGE"], - "imageConfig": est_image_config - } - } + def _estimate_with_images(img_parts: List[dict]) -> int: + return _request_size(_build_request_body(parts + img_parts)) + + working_images = list(images) + img_parts = _build_image_parts(working_images) + original_request_body = _build_request_body(parts + img_parts) + if request_log_enabled: + _log_original_request_body(original_request_body) + original_request_logged = True + estimated = _estimate_with_images(img_parts) + + if estimated > _MAX_BODY_BYTES: + ratio = _BODY_TARGET_BYTES / estimated + scale = ratio ** 0.5 # 面积比 → 线性比 + working_images = self._scale_images_to_fit(working_images, scale) + img_parts = _build_image_parts(working_images) + estimated = _estimate_with_images(img_parts) + + orig_sizes = ", ".join(f"{img.width}×{img.height}" for img in images) + new_sizes = ", ".join(f"{img.width}×{img.height}" for img in working_images) + size_mb = estimated / (1024 * 1024) + target_mb = _BODY_TARGET_BYTES / (1024 * 1024) + print( + f"Nano Banana Pro: 输入图片已按请求体目标大小自动缩放\n" + f" 原始尺寸: {orig_sizes}\n" + f" 缩放后: {new_sizes}\n" + f" 请求体积: {size_mb:.2f}MB（目标 {target_mb:.2f}MB，限制 20MB）" ) - if estimated <= _MAX_BODY_BYTES: - parts.extend(img_parts) - if _attempt > 0: - orig_sizes = ", ".join( - f"{img.width}×{img.height}" for img in images - ) - new_sizes = ", ".join( - f"{img.width}×{img.height}" for img in working_images - ) - size_mb = estimated / (1024 * 1024) - print( - f"Nano Banana Pro: 输入图片已自动缩放以控制请求体积\n" - f" 原始尺寸: {orig_sizes}\n" - f" 缩放后: {new_sizes}\n" - f" 请求体积: {size_mb:.2f}MB（限制 20MB）" - ) - break - else: - # 按像素面积比推算需要的线性缩放系数，留 5% 余量 - ratio = (_MAX_BODY_BYTES * 0.95) / estimated - scale = ratio ** 0.5 # 面积比 → 线性比 - working_images = self._scale_images_to_fit(working_images, scale) - else: - # 10 轮后仍超限，使用最后一次结果（极端情况兜底） - parts.extend(img_parts) + parts.extend(img_parts) - # 构建请求体 - image_config = {"imageSize": resolution} - if aspect_ratio and aspect_ratio != "智能": - image_config["aspectRatio"] = aspect_ratio + request_body = _build_request_body(parts) + if request_log_enabled and not original_request_logged: + _log_original_request_body(request_body) - request_body = { - "contents": [ - { - "role": "user", - "parts": parts - } - ], - "generationConfig": { - "responseModalities": ["IMAGE"], - "imageConfig": image_config - } - } - - # 添加思考深度配置 - if thinking_level: - request_body["generationConfig"]["thinkingConfig"] = { - "thinkingLevel": thinking_level, - "includeThoughts": True - } - - # 添加图片压缩参数 - if image_compression: - request_body["image_compression"] = image_compression - - # 添加 Google Search Grounding 工具（如果启用） - # 注意：enable_image_search=True 时会自动隐含 enable_grounding - if enable_grounding or enable_image_search: - if enable_image_search: - # 同时启用网页搜索和图片搜索（仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持） - request_body["tools"] = [ - { - "google_search": { - "searchTypes": { - "webSearch": {}, - "imageSearch": {} - } - } - } - ] - else: - # 仅启用网页搜索（通用） - request_body["tools"] = [{"google_search": {}}] + request_size = _request_size(request_body) + if request_size > _MAX_BODY_BYTES: + raise ValueError( + f"请求体超过 20MB 限制（当前 {request_size / 1024 / 1024:.2f}MB），" + "已停止提交；请减少参考图数量、降低图片复杂度或缩短提示词" + ) return request_body @@ -924,7 +948,7 @@ class GeminiAPIClient(BaseAPIClient): ) return self.run_async_in_thread(coro) - + async def generate_multi_prompts_async( self, prompts: List[str], @@ -1104,6 +1128,5 @@ class GeminiAPIClient(BaseAPIClient): enable_grounding=enable_grounding, enable_image_search=enable_image_search ) - + return self.run_async_in_thread(coro) - \ No newline at end of file diff --git a/clients/gpt_image_client.py b/clients/gpt_image_client.py index ef129c2..e05439e 100644 --- a/clients/gpt_image_client.py +++ b/clients/gpt_image_client.py @@ -1,13 +1,17 @@ """ GPT Image API 客户端 -支持两个接口： - - POST /v1/images/generations/ 文生图 / 图生图（gpt-image-1 / gpt-image-1.5） - - POST /v1/images/edits/ 图像编辑（带蒙版 inpainting） +新版节点请求走异步任务接口： + - POST /async/v1/generateImage + - GET /async/v1/tasks/{task_id} + +旧同步接口保留兼容代码，但 GPT Image 节点不再使用： + - POST /v1/images/generations/ + - POST /v1/images/edits 设计原则： - - 与 doubao_image_client.py 保持相同的异步 + 同步双入口模式 - - generations / edits 接口均使用 multipart/form-data - - 响应兼容 SSE 流式、JSON、url 和 b64_json + - 对 ComfyUI 节点暴露同步入口，内部提交异步任务并轮询 + - 图片和蒙版以 data:image/png;base64,... 放入 JSON 请求体 + - 响应优先读取 data.images[].url """ import asyncio @@ -16,7 +20,7 @@ import json import time from concurrent.futures import ThreadPoolExecutor from io import BytesIO -from typing import List, Optional +from typing import Any, Callable, Dict, List, Optional import aiohttp import numpy as np @@ -38,6 +42,8 @@ except ImportError: # ── 接口端点 ────────────────────────────────────────────────────────────────── _ENDPOINT_GENERATIONS = "/v1/images/generations/" _ENDPOINT_EDITS = "/v1/images/edits" +_ENDPOINT_ASYNC_GENERATE = "/async/v1/generateImage" +_ENDPOINT_ASYNC_TASK = "/async/v1/tasks/{task_id}" # ── 模型名映射（UI 显示名 → API 实际参数名）───────────────────────────────── _MODEL_NAME_MAP = { @@ -47,6 +53,49 @@ _MODEL_NAME_MAP = { # ── 超时 ────────────────────────────────────────────────────────────────────── _REQUEST_TIMEOUT = 900 # 秒 +_ASYNC_POLL_SCHEDULE = [5.0, 20.0] +_ASYNC_POLL_INTERVAL = 3.0 +_ASYNC_MAX_WAIT = 600.0 +_ASYNC_RETRY_DELAYS = [2.0, 5.0, 10.0] +_ASYNC_RETRYABLE_ERROR_CODES = { + "image_rate_limited", + "image_upstream_busy", + "image_timeout", + "image_storage_failed", + "image_empty_result", + "image_upstream_error", + "image_internal_error", + "image_unknown_error", +} +_ASYNC_RETRYABLE_ERROR_CATEGORIES = { + "rate_limit", + "upstream_busy", + "timeout", + "storage", + "upstream_error", + "internal_error", + "unknown", +} +_ASYNC_NON_RETRYABLE_ERROR_CODES = { + "image_invalid_size", + "image_payload_too_large", + "image_invalid_mask", + "image_invalid_parameter", + "image_safety_blocked", + "image_provider_quota_exceeded", + "image_provider_permission_required", + "image_model_unavailable", + "image_reference_download_failed", +} + +REQUEST_LOG_ENABLED = False +POLL_LOG_ENABLED = False + + +class _AsyncImageTaskFailure(RuntimeError): + def __init__(self, message: str, error_detail: Optional[dict] = None): + super().__init__(message) + self.error_detail = error_detail or {} class GptImageClient: @@ -54,10 +103,10 @@ class GptImageClient: GPT Image API 客户端 接口说明： - generations：multipart/form-data，支持 quality / size / n / model - edits：multipart/form-data，图片和 mask 使用 PNG 文件上传 + async generateImage：JSON 提交，返回 task_id + tasks/{task_id}：轮询任务状态，成功后读取 data.images[].url - 响应支持 JSON 和 SSE 流式格式。 + 旧 generations / edits 同步接口保留为兼容代码。 """ def __init__(self): @@ -132,6 +181,226 @@ class GptImageClient: ) return png_bytes + @staticmethod + def _png_bytes_to_data_url(png_bytes: bytes) -> str: + b64 = base64.b64encode(png_bytes).decode("ascii") + return f"data:image/png;base64,{b64}" + + @staticmethod + def _json_body_size(body: dict) -> int: + return len(json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode("utf-8")) + + @staticmethod + def _format_body_size(size: int) -> str: + if size < 1024 * 1024: + return f"{size / 1024:.2f}KB" + return f"{size / 1024 / 1024:.2f}MB" + + @staticmethod + def _shorten_data_urls_for_log(obj, max_len: int = 200): + if isinstance(obj, dict): + return { + key: GptImageClient._shorten_data_urls_for_log(value, max_len) + for key, value in obj.items() + } + if isinstance(obj, list): + return [GptImageClient._shorten_data_urls_for_log(item, max_len) for item in obj] + if isinstance(obj, str) and obj.startswith("data:image") and len(obj) > max_len: + header, _, data = obj.partition(",") + return f"{header}," + return obj + + def _log_original_request_body(self, label: str, body: dict) -> None: + if not REQUEST_LOG_ENABLED: + return + body_size = self._json_body_size(body) + print( + f"\n{'=' * 60}\n" + f"[o1key GPT Image] 原始请求体日志 | {label} | " + f"请求体积: {self._format_body_size(body_size)} " + f"(data URL 已折叠显示 base64 长度)\n" + f"{json.dumps(self._shorten_data_urls_for_log(body), ensure_ascii=False, indent=2)}\n" + f"{'=' * 60}\n" + ) + + def _log_original_response_body(self, label: str, text: str) -> None: + if not REQUEST_LOG_ENABLED: + return + size = len(text.encode("utf-8")) + print( + f"\n{'=' * 60}\n" + f"[o1key GPT Image] 原始返回响应体日志 | {label} | " + f"响应体积: {self._format_body_size(size)}\n" + f"{text}\n" + f"{'=' * 60}\n" + ) + + @staticmethod + def _extract_error_message(payload_or_text, status_code: int = 0) -> str: + payload = payload_or_text + if isinstance(payload_or_text, str): + try: + payload = json.loads(payload_or_text) + except Exception: + return get_friendly_message(status_code, payload_or_text) + + if isinstance(payload, dict): + error = payload.get("error") + if isinstance(error, str) and error.strip(): + return error.strip() + if isinstance(error, dict): + msg = error.get("message") or error.get("msg") or error.get("error") + if msg: + return str(msg) + return json.dumps(error, ensure_ascii=False) + + msg = payload.get("message") or payload.get("msg") + if msg: + return str(msg) + + return get_friendly_message(status_code, str(payload_or_text)) + + @staticmethod + def _extract_error_detail(payload: dict) -> dict: + if not isinstance(payload, dict): + return {} + detail = payload.get("error_detail") + return detail if isinstance(detail, dict) else {} + + @staticmethod + def _should_retry_async_failure(error_detail: dict, retry_index: int) -> bool: + if not isinstance(error_detail, dict) or not error_detail: + return False + + code = error_detail.get("code") + category = error_detail.get("category") + retryable = error_detail.get("retryable") + + if code in _ASYNC_NON_RETRYABLE_ERROR_CODES: + return False + if code == "image_unknown_error": + return retry_index == 0 + if retryable is True: + return True + if retryable is False: + return False + + return ( + code in _ASYNC_RETRYABLE_ERROR_CODES + or category in _ASYNC_RETRYABLE_ERROR_CATEGORIES + ) + + @staticmethod + def _coerce_progress_percent(value: Any) -> Optional[int]: + if value is None: + return None + if isinstance(value, str): + text = value.strip().rstrip("%") + if not text: + return None + try: + value = float(text) + except ValueError: + return None + elif isinstance(value, (int, float)): + value = float(value) + else: + return None + + if 0 <= value <= 1: + value *= 100 + return max(0, min(100, int(round(value)))) + + @staticmethod + def _emit_progress(progress_callback: Optional[Callable[[int], None]], pct: int) -> None: + if progress_callback is None: + return + try: + progress_callback(pct) + except Exception as error: + print(f"[o1key GPT Image] progress callback failed: {error}") + + @staticmethod + def _resize_png_bytes(source_image: Image.Image, scale: float) -> bytes: + if scale < 0.999: + width, height = source_image.size + new_width = max(1, int(width * scale)) + new_height = max(1, int(height * scale)) + image = source_image.resize((new_width, new_height), Image.LANCZOS) + else: + image = source_image + + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + @staticmethod + def _pil_to_png_bytes(image: Image.Image) -> bytes: + buf = BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + def _fit_png_assets_to_body_limit(self, assets: List[Dict[str, Any]], build_body) -> Dict[str, bytes]: + """ + 根据完整 JSON 请求体大小压缩图片资产，保证最终 body 不超过 20MB。 + 使用同一个缩放比例二分搜索，让压缩结果尽量贴近上限而不是过度压缩。 + """ + asset_bytes = {asset["key"]: asset["bytes"] for asset in assets} + initial_size = self._json_body_size(build_body(asset_bytes)) + if initial_size <= self._MAX_BODY_BYTES: + return asset_bytes + + if not assets: + raise RuntimeError( + f"请求体大小 {initial_size // 1024}KB 超过 20MB，且没有可压缩图片" + ) + + originals = [] + for asset in assets: + image = Image.open(BytesIO(asset["bytes"])) + image.load() + originals.append((asset, image.copy())) + + low = 0.001 + high = 1.0 + best_bytes = None + best_size = 0 + best_scale = 0.0 + + for _ in range(16): + scale = (low + high) / 2 + candidate = {} + for asset, image in originals: + candidate[asset["key"]] = self._resize_png_bytes(image, scale) + + body_size = self._json_body_size(build_body(candidate)) + if body_size <= self._MAX_BODY_BYTES: + best_bytes = candidate + best_size = body_size + best_scale = scale + low = scale + else: + high = scale + + if best_bytes is None: + candidate = {} + for asset, image in originals: + candidate[asset["key"]] = self._resize_png_bytes(image, low) + body_size = self._json_body_size(build_body(candidate)) + if body_size > self._MAX_BODY_BYTES: + raise RuntimeError( + f"图片已压缩到最小比例，但请求体仍超过 20MB：{body_size // 1024}KB" + ) + best_bytes = candidate + best_size = body_size + best_scale = low + + print( + f"[o1key GPT Image] 请求体超过 20MB，已等比压缩图片：" + f"{initial_size // 1024}KB → {best_size // 1024}KB，scale={best_scale:.3f}" + ) + return best_bytes + @staticmethod def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes: """ @@ -500,6 +769,316 @@ class GptImageClient: return request_task.result() + # ── 新版异步 GPT Image 接口 ────────────────────────────────────────────── + + def _build_async_generate_body( + self, + prompt: str, + model: str, + quality: str, + size: str, + n: int, + image_list: Optional[List[torch.Tensor]] = None, + mask_tensor: Optional[torch.Tensor] = None, + output_format: str = "png", + ) -> dict: + api_model = _MODEL_NAME_MAP.get(model, model) + assets: List[Dict[str, Any]] = [] + image_keys: List[str] = [] + mask_key = None + + if image_list: + for idx_img, img_tensor in enumerate(image_list): + pil_images = tensor_to_pil(img_tensor) + if not pil_images: + continue + key = f"image_{idx_img}" + image_keys.append(key) + assets.append({ + "key": key, + "label": f"参考图{idx_img + 1}", + "bytes": self._pil_to_png_bytes(pil_images[0]), + }) + + if mask_tensor is not None: + if not image_list: + raise ValueError("提供了蒙版但未提供图片，请同时提供图片和蒙版") + + first_tensor = image_list[0] + if first_tensor.dim() == 3: + first_tensor = first_tensor.unsqueeze(0) + image_size = (first_tensor.shape[1], first_tensor.shape[2]) + mask_key = "mask" + assets.append({ + "key": mask_key, + "label": "蒙版", + "bytes": self._mask_tensor_to_rgba_png_bytes(mask_tensor, image_size), + }) + + def _make_body(asset_bytes: Dict[str, bytes]) -> dict: + body = { + "model": api_model, + "prompt": prompt, + "images": [ + self._png_bytes_to_data_url(asset_bytes[key]) + for key in image_keys + if key in asset_bytes + ], + "size": size if size else "auto", + "quality": quality, + "n": int(n), + "output_format": output_format or "png", + } + if mask_key and mask_key in asset_bytes: + body["mask"] = { + "image_url": self._png_bytes_to_data_url(asset_bytes[mask_key]) + } + return body + + original_asset_bytes = {asset["key"]: asset["bytes"] for asset in assets} + original_body = _make_body(original_asset_bytes) + self._log_original_request_body("async generateImage", original_body) + + asset_bytes = self._fit_png_assets_to_body_limit(assets, _make_body) + body = _make_body(asset_bytes) + body_size = self._json_body_size(body) + if body_size > self._MAX_BODY_BYTES: + raise RuntimeError(f"请求体超过 20MB：{body_size // 1024}KB") + + return body + + async def _submit_generate_image_task( + self, + session: aiohttp.ClientSession, + payload: dict, + ) -> str: + url = f"{self.base_url}{_ENDPOINT_ASYNC_GENERATE}" + last_status = None + + for attempt in range(DEFAULT_MAX_RETRIES + 1): + t0 = time.time() + async with session.post( + url, + json=payload, + headers=self._json_headers(), + ) as resp: + elapsed = time.time() - t0 + text = await resp.text() + self._log_original_response_body( + f"submit generateImage status={resp.status}", + text, + ) + + if resp.status not in (200, 201, 202): + last_status = resp.status + if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES: + friendly = get_friendly_message(resp.status) + delay = _compute_delay( + attempt, + DEFAULT_BASE_DELAY, + DEFAULT_MAX_DELAY, + DEFAULT_BACKOFF_FACTOR, + ) + print(f"[o1key GPT Image] {friendly} {delay:.1f}s 后重试提交 ({attempt+1}/{DEFAULT_MAX_RETRIES})...") + await asyncio.sleep(delay) + continue + raise RuntimeError(self._extract_error_message(text, resp.status)) + + try: + data = json.loads(text) + except Exception: + raise RuntimeError(f"提交响应 JSON 解析失败，原始内容：{text[:500]}") from None + + task_id = data.get("task_id") + if not task_id: + raise RuntimeError(f"提交响应中未找到 task_id: {data}") + + status = data.get("status", "") + print(f"[o1key GPT Image] 异步任务已提交 | task_id={task_id} | status={status} | 耗时 {elapsed:.1f}s") + return task_id + + if last_status and last_status in HTTP_ERROR_MESSAGES: + raise RuntimeError(HTTP_ERROR_MESSAGES[last_status]) + raise RuntimeError(f"异步任务提交失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败") + + async def _poll_generate_image_task( + self, + session: aiohttp.ClientSession, + task_id: str, + progress_callback: Optional[Callable[[int], None]] = None, + ) -> dict: + url = f"{self.base_url}{_ENDPOINT_ASYNC_TASK.format(task_id=task_id)}" + start_time = time.time() + + poll_count = 0 + last_poll_at = start_time + + while True: + if poll_count < len(_ASYNC_POLL_SCHEDULE): + next_poll_at = start_time + _ASYNC_POLL_SCHEDULE[poll_count] + else: + next_poll_at = last_poll_at + _ASYNC_POLL_INTERVAL + + sleep_time = next_poll_at - time.time() + if sleep_time > 0: + await asyncio.sleep(sleep_time) + + last_poll_at = time.time() + elapsed = last_poll_at - start_time + if elapsed > _ASYNC_MAX_WAIT: + raise RuntimeError(f"任务 {task_id} 超时（>{int(_ASYNC_MAX_WAIT)}秒），请稍后用 task_id 查询结果") + + poll_count += 1 + async with session.get(url, headers=self._auth_headers()) as resp: + text = await resp.text() + self._log_original_response_body( + f"poll task status={resp.status}", + text, + ) + if resp.status != 200: + raise RuntimeError(self._extract_error_message(text, resp.status)) + + try: + task = json.loads(text) + except Exception: + raise RuntimeError(f"任务查询响应 JSON 解析失败，原始内容：{text[:500]}") from None + + status = task.get("status", "UNKNOWN") + progress = task.get("progress") + progress_pct = self._coerce_progress_percent(progress) + if POLL_LOG_ENABLED: + progress_text = f" | progress={progress}" if progress is not None else "" + print(f"[o1key GPT Image] 查询任务 #{poll_count} | task_id={task_id} | status={status}{progress_text}") + + if status == "SUCCESS": + self._emit_progress(progress_callback, 100) + return task + if progress_pct is not None and progress_pct < 100: + self._emit_progress(progress_callback, progress_pct) + if status == "FAILURE": + error_message = self._extract_error_message(task, 500) or "生成失败" + raise _AsyncImageTaskFailure( + error_message, + self._extract_error_detail(task), + ) + if status not in ("SUBMITTED", "IN_PROGRESS"): + raise RuntimeError(f"未知任务状态 {status}: {task}") + + async def _parse_async_task_images( + self, + task: dict, + session: aiohttp.ClientSession, + ) -> List[Image.Image]: + data = task.get("data", {}) + image_items = data.get("images") if isinstance(data, dict) else None + + if not isinstance(image_items, list) or not image_items: + raise RuntimeError(f"任务结果中未找到 data.images: {task}") + + images: List[Image.Image] = [] + for idx, item in enumerate(image_items, 1): + if not isinstance(item, dict): + continue + + url = item.get("url") or item.get("image_url") + b64 = item.get("b64_json", "") + + if url and isinstance(url, str) and url.startswith("data:image"): + try: + _, b64_data = url.split(",", 1) + img = Image.open(BytesIO(base64.b64decode(b64_data))) + images.append(img) + print(f"[o1key GPT Image] 第 {idx} 张 data URL 解码完成 ({img.size[0]}×{img.size[1]})") + except Exception as e: + raise RuntimeError(f"第 {idx} 张 data URL 解码失败: {e}") from None + elif url and isinstance(url, str) and url.startswith("http"): + async with session.get(url, allow_redirects=True) as resp: + if resp.status != 200: + raise RuntimeError(f"图像下载失败 HTTP {resp.status}，URL: {url}") + img_bytes = await resp.read() + img = Image.open(BytesIO(img_bytes)) + images.append(img) + print(f"[o1key GPT Image] 第 {idx} 张 URL 下载完成 ({img.size[0]}×{img.size[1]}) | {url}") + elif b64: + img = self._decode_b64_image(b64, f"第 {idx} 张") + images.append(img) + else: + print(f"[o1key GPT Image] 警告：第 {idx} 条结果既无 url 也无 b64_json，已跳过") + + if not images: + raise RuntimeError("任务成功但没有可用图片结果") + + return images + + async def _generate_image_task_async( + self, + prompt: str, + model: str, + quality: str, + size: str, + n: int, + seed: int, + image_tensor: Optional[List[torch.Tensor]] = None, + mask_tensor: Optional[torch.Tensor] = None, + output_format: str = "png", + progress_callback: Optional[Callable[[int], None]] = None, + ) -> List[Image.Image]: + body = self._build_async_generate_body( + prompt=prompt, + model=model, + quality=quality, + size=size, + n=n, + image_list=image_tensor, + mask_tensor=mask_tensor, + output_format=output_format, + ) + + mode = "图像编辑" if mask_tensor is not None else ("图生图" if image_tensor else "文生图") + body_size = self._json_body_size(body) + print( + f"[o1key GPT Image] {mode} | 新异步接口 | 模型={model} | " + f"quality={quality} | size={size} | n={n} | body={body_size // 1024}KB" + ) + + connector = aiohttp.TCPConnector(ssl=False, force_close=True) + timeout = aiohttp.ClientTimeout(total=_ASYNC_MAX_WAIT + 120) + async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: + last_error = None + for retry_index in range(len(_ASYNC_RETRY_DELAYS) + 1): + try: + task_id = await self._submit_generate_image_task(session, body) + task = await self._poll_generate_image_task( + session, + task_id, + progress_callback=progress_callback, + ) + return await self._parse_async_task_images(task, session) + except _AsyncImageTaskFailure as error: + last_error = error + detail = error.error_detail + if ( + retry_index < len(_ASYNC_RETRY_DELAYS) + and self._should_retry_async_failure(detail, retry_index) + ): + delay = _ASYNC_RETRY_DELAYS[retry_index] + code = detail.get("code", "unknown") + category = detail.get("category", "unknown") + failed_task_id = detail.get("task_id", "") + task_text = f" | failed_task_id={failed_task_id}" if failed_task_id else "" + print( + f"[o1key GPT Image] 任务失败但可重试 | code={code} | " + f"category={category}{task_text} | {delay:.0f}s 后重试 " + f"({retry_index + 1}/{len(_ASYNC_RETRY_DELAYS)})" + ) + await asyncio.sleep(delay) + continue + raise RuntimeError(str(error)) from None + + if last_error is not None: + raise RuntimeError(str(last_error)) from None + raise RuntimeError("生成失败") + # ── 文生图 / 图生图（generations 接口）─────────────────────────────────── async def _generate_async( @@ -806,6 +1385,53 @@ class GptImageClient: f"o1key GPT Image 请求超时（>{_REQUEST_TIMEOUT}s），请检查网络或稍后重试" ) + def generate_image_async_sync( + self, + prompt: str, + model: str, + quality: str, + size: str, + n: int, + seed: int, + image_tensor: Optional[List[torch.Tensor]] = None, + mask_tensor: Optional[torch.Tensor] = None, + output_format: str = "png", + progress_callback: Optional[Callable[[int], None]] = None, + ) -> List[Image.Image]: + """ + 新版异步任务入口，供节点调用。 + 旧 run_sync 保留兼容，但 GPT Image 节点不再使用旧同步接口。 + """ + coro = self._generate_image_task_async( + prompt=prompt, + model=model, + quality=quality, + size=size, + n=n, + seed=seed, + image_tensor=image_tensor, + mask_tensor=mask_tensor, + output_format=output_format, + progress_callback=progress_callback, + ) + + def _run(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete(self._run_with_interrupt(coro)) + finally: + loop.close() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(_run) + try: + return future.result(timeout=_ASYNC_MAX_WAIT + 150) + except TimeoutError: + raise RuntimeError( + f"o1key GPT Image 异步任务超时（>{int(_ASYNC_MAX_WAIT)}秒），请稍后用 task_id 查询结果" + ) + # ── 余额查询 ────────────────────────────────────────────────────────────── async def _query_balance_async(self) -> dict: diff --git a/clients/grok_video_client.py b/clients/grok_video_client.py new file mode 100644 index 0000000..651d127 --- /dev/null +++ b/clients/grok_video_client.py @@ -0,0 +1,483 @@ +""" +Grok Video API client. + +Flow: + 1. POST /v1/videos + 2. GET /v1/videos/{task_id} + 3. GET /v1/videos/{task_id}/content, or download a URL from the status body +""" + +import asyncio +import base64 +import json +import os +import re +import time +from typing import Any, Callable, Dict, List, Optional + +import aiohttp + +from .base_client import BaseAPIClient +from ..utils.config import get_api_base_url, get_api_key_or_raise +from ..utils.http_error import RETRYABLE_STATUS_CODES, get_friendly_message +from ..utils.video_task import ( + check_interrupt, + extract_error_message, + extract_progress, + extract_status, + extract_video_url, + interruptible_sleep, + is_failure_status, + is_success_status, + run_with_interrupt, +) + + +class GrokVideoClient(BaseAPIClient): + CREATE_ENDPOINT = "/v1/videos" + STATUS_ENDPOINT = "/v1/videos/{task_id}" + CONTENT_ENDPOINT = "/v1/videos/{task_id}/content" + + MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"] + ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"] + QUALITY_OPTIONS = ["720p"] + MODEL_SECONDS_OPTIONS = { + "grok-imagine-1.0-video": [6, 10, 12, 16, 20], + } + QUALITY_API_MAP = { + "720p": "high", + "high": "high", + } + + SUCCESS_STATUSES = {"complete", "completed", "succeed", "succeeded", "success", "done", "finished"} + FAILURE_STATUSES = {"fail", "failed", "failure", "error", "expired", "timeout", "cancelled", "canceled"} + + def __init__(self, base_url: Optional[str] = None): + api_key = get_api_key_or_raise("O1KEY_API_KEY") + resolved_base_url = (base_url or "").strip() or get_api_base_url() + super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key) + + def get_endpoint(self, **kwargs) -> str: + return self.CREATE_ENDPOINT + + def build_request_body(self, **kwargs) -> Dict[str, Any]: + return self.build_video_body(**kwargs) + + def parse_response(self, response: Dict[str, Any]) -> Any: + return response + + @classmethod + def build_video_body( + cls, + prompt: str, + model: str, + aspect_ratio: str, + seconds: int, + quality: str = "720p", + images: Optional[List[str]] = None, + ) -> Dict[str, Any]: + prompt = (prompt or "").strip() + if not prompt: + raise ValueError("提示词不能为空。") + + if model not in cls.MODEL_OPTIONS: + raise ValueError(f"模型仅支持: {', '.join(cls.MODEL_OPTIONS)}。") + + if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS: + raise ValueError(f"宽高比仅支持: {', '.join(cls.ASPECT_RATIO_OPTIONS)}。") + + try: + seconds_value = int(seconds) + except (TypeError, ValueError): + raise ValueError("秒数必须是整数。") from None + + allowed_seconds = cls.MODEL_SECONDS_OPTIONS.get(model) + if allowed_seconds is not None: + if seconds_value not in allowed_seconds: + raise ValueError( + f"模型 {model} 仅支持秒数: " + f"{', '.join(str(s) for s in allowed_seconds)}。" + "请修改为正确的秒数后再发起请求。" + ) + elif seconds_value < 5 or seconds_value > 15: + raise ValueError("秒数仅支持 5 到 15。") + + api_quality = cls.QUALITY_API_MAP.get(str(quality), str(quality)) + if api_quality != "high": + raise ValueError("画质仅支持 720p。") + + body: Dict[str, Any] = { + "model": model, + "prompt": prompt, + "aspect_ratio": aspect_ratio, + "seconds": str(seconds_value), + "quality": api_quality, + } + + image_list = [img for img in (images or []) if img] + if image_list: + body["images"] = image_list[:3] + + return body + + @staticmethod + def _safe_task_filename(task_id: str) -> str: + safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._") + return safe or "grok_video" + + @staticmethod + def _mask_body_for_log(body: Dict[str, Any]) -> Dict[str, Any]: + log_body = dict(body) + images = log_body.get("images") + if isinstance(images, list): + log_body["images"] = [f"" for item in images] + return log_body + + @staticmethod + def _extract_task_id(payload: Dict[str, Any]) -> Optional[str]: + sources = [payload] + data = payload.get("data") + if isinstance(data, dict): + sources.append(data) + + for source in sources: + for key in ("id", "task_id", "video_id"): + value = source.get(key) + if value: + return str(value) + return None + + @staticmethod + def _format_http_error(endpoint: str, status: int, error_text: str, task_id: Optional[str] = None) -> str: + message = get_friendly_message(status, error_text) + parts = [ + "Grok Video 请求失败。", + f"endpoint: {endpoint}", + f"http_status: {status}", + ] + if task_id: + parts.append(f"task_id: {task_id}") + if message: + parts.append(f"message: {message}") + return "\n".join(parts) + + @classmethod + def _format_task_failure(cls, task_id: str, payload: Dict[str, Any]) -> str: + return "\n".join( + [ + "Grok Video 任务失败。", + f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}", + f"task_id: {task_id}", + f"message: {extract_error_message(payload)}", + ] + ) + + async def _request_json_with_retry( + self, + method: str, + endpoint: str, + session: aiohttp.ClientSession, + task_id: Optional[str] = None, + json_body: Optional[Dict[str, Any]] = None, + max_retries: int = 3, + timeout_seconds: int = 120, + ) -> Dict[str, Any]: + url = f"{self.base_url}{endpoint}" + headers = self.get_headers(use_bearer_token=True) + timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds) + + last_status = 0 + last_text = "" + + for attempt in range(max_retries + 1): + check_interrupt() + response = None + try: + response = await run_with_interrupt( + session.request(method, url, json=json_body, headers=headers, timeout=timeout) + ) + text = await run_with_interrupt(response.text()) + last_status = response.status + last_text = text + + if 200 <= response.status < 300: + if not text.strip(): + return {} + try: + return json.loads(text) + except Exception: + raise RuntimeError(f"Grok Video 响应 JSON 解析失败，原始内容：{text[:500]}") from None + + if response.status in RETRYABLE_STATUS_CODES and attempt < max_retries: + delay = min(2 ** attempt, 8) + print( + f"Grok Video：{get_friendly_message(response.status)} " + f"{delay}s 后重试 ({attempt + 1}/{max_retries})..." + ) + await interruptible_sleep(delay) + continue + + break + + except (aiohttp.ClientError, asyncio.TimeoutError) as e: + if attempt < max_retries: + delay = min(2 ** attempt, 8) + print(f"Grok Video：网络错误，{delay}s 后重试 ({attempt + 1}/{max_retries})...") + await interruptible_sleep(delay) + continue + raise RuntimeError(f"Grok Video 网络错误: {e}") from None + + finally: + if response is not None: + response.release() + + raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id)) + + async def create_video_async( + self, + body: Dict[str, Any], + session: aiohttp.ClientSession, + ) -> Dict[str, Any]: + print("Grok Video：正在提交任务...") + return await self._request_json_with_retry( + "POST", + self.CREATE_ENDPOINT, + session=session, + json_body=body, + timeout_seconds=180, + ) + + async def poll_video_status_async( + self, + task_id: str, + session: aiohttp.ClientSession, + poll_interval: int = 5, + timeout: int = 900, + progress_callback: Optional[Callable[[int, str, float], None]] = None, + ) -> Dict[str, Any]: + endpoint = self.STATUS_ENDPOINT.format(task_id=task_id) + start = time.time() + interval = max(1, int(poll_interval)) + + await interruptible_sleep(interval) + + while True: + data = await self._request_json_with_retry( + "GET", + endpoint, + session=session, + task_id=task_id, + timeout_seconds=60, + ) + + status = extract_status(data) + progress = extract_progress(data) + elapsed = time.time() - start + + if progress_callback: + progress_callback(progress, status, elapsed) + + if status in self.SUCCESS_STATUSES or is_success_status(status): + return data + + if status in self.FAILURE_STATUSES or is_failure_status(status, data): + raise RuntimeError(self._format_task_failure(task_id, data)) + + if elapsed >= timeout: + raise TimeoutError( + "Grok Video 任务轮询超时；任务未被标记为失败，可用 task_id 继续查询。\n" + f"endpoint: {endpoint}\n" + f"task_id: {task_id}\n" + f"status: {status or 'unknown'}\n" + f"timeout: {timeout}s" + ) + + await interruptible_sleep(min(interval, max(0.0, timeout - elapsed))) + + async def _download_url_to_file( + self, + url: str, + save_path: str, + session: aiohttp.ClientSession, + max_retries: int = 3, + ) -> str: + timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900) + last_status = 0 + last_text = "" + headers = None + resolved_url = url + + if url.startswith("data:"): + if "," not in url: + raise RuntimeError("Grok Video 下载失败：data URL 格式无效。") + _, b64_data = url.split(",", 1) + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + f.write(base64.b64decode(b64_data)) + if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0: + raise RuntimeError("Grok Video 下载失败：保存后的文件为空。") + return save_path + + if url.startswith("/"): + resolved_url = f"{self.base_url}{url}" + headers = self.get_headers(use_bearer_token=True) + + for attempt in range(max_retries + 1): + check_interrupt() + async with session.get( + resolved_url, + headers=headers, + timeout=timeout, + allow_redirects=True, + ) as response: + if 200 <= response.status < 300: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(1024 * 1024): + check_interrupt() + if chunk: + f.write(chunk) + + if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0: + raise RuntimeError("Grok Video 下载失败：保存后的文件为空。") + return save_path + + last_status = response.status + last_text = await response.text() + if response.status not in RETRYABLE_STATUS_CODES or attempt >= max_retries: + break + + delay = min(2 ** attempt, 8) + print(f"Grok Video：下载重试 {attempt + 1}/{max_retries}，{delay}s 后继续...") + await interruptible_sleep(delay) + + raise RuntimeError(self._format_http_error("download_url", last_status, last_text)) + + async def download_video_async( + self, + task_id: str, + save_path: str, + session: aiohttp.ClientSession, + ) -> str: + endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id) + url = f"{self.base_url}{endpoint}" + headers = self.get_headers(use_bearer_token=True) + timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900) + + last_status = 0 + last_text = "" + + for attempt in range(4): + check_interrupt() + async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response: + if 200 <= response.status < 300: + content_type = response.headers.get("Content-Type", "").lower() + if "application/json" in content_type: + data = await response.json(content_type=None) + download_url = extract_video_url(data) + if not download_url: + raise RuntimeError( + "Grok Video 下载失败：content 响应为 JSON，但未包含视频 URL。\n" + f"endpoint: {endpoint}\n" + f"task_id: {task_id}" + ) + return await self._download_url_to_file(download_url, save_path, session) + + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(1024 * 1024): + check_interrupt() + if chunk: + f.write(chunk) + + if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0: + raise RuntimeError( + "Grok Video 下载失败：保存后的文件为空。\n" + f"endpoint: {endpoint}\n" + f"task_id: {task_id}" + ) + return save_path + + last_status = response.status + last_text = await response.text() + if response.status not in RETRYABLE_STATUS_CODES or attempt >= 3: + break + + delay = min(2 ** attempt, 8) + print(f"Grok Video：content 下载重试 {attempt + 1}/3，{delay}s 后继续...") + await interruptible_sleep(delay) + + raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id)) + + def generate_video_sync( + self, + prompt: str, + model: str, + aspect_ratio: str, + seconds: int, + quality: str, + images: Optional[List[str]], + output_dir: Optional[str] = None, + save_path: Optional[str] = None, + poll_interval: int = 5, + timeout: int = 900, + progress_callback: Optional[Callable[[int, str, float], None]] = None, + ) -> Dict[str, Any]: + async def _run(): + async with self._make_session() as session: + body = self.build_video_body( + prompt=prompt, + model=model, + aspect_ratio=aspect_ratio, + seconds=seconds, + quality=quality, + images=images, + ) + + create_response = await self.create_video_async(body, session) + task_id = self._extract_task_id(create_response) or "" + if not task_id: + raise RuntimeError( + "Grok Video 未返回任务 ID。\n" + f"endpoint: {self.CREATE_ENDPOINT}\n" + f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}" + ) + + print(f"Grok Video：任务已提交，任务ID：{task_id}") + print("Grok Video：视频生成中...") + status_response = await self.poll_video_status_async( + task_id=task_id, + session=session, + poll_interval=poll_interval, + timeout=timeout, + progress_callback=progress_callback, + ) + + video_url = extract_video_url(status_response) + print("Grok Video：视频生成完成，正在下载...") + if save_path is None: + resolved_output_dir = output_dir or os.getcwd() + os.makedirs(resolved_output_dir, exist_ok=True) + target_path = os.path.join( + resolved_output_dir, + f"{self._safe_task_filename(task_id)}.mp4", + ) + else: + target_path = save_path + + if video_url: + video_path = await self._download_url_to_file(video_url, target_path, session) + else: + video_path = await self.download_video_async(task_id, target_path, session) + + return { + "task_id": task_id, + "status": extract_status(status_response), + "video_path": video_path, + "raw_json": { + "create": create_response, + "status": status_response, + }, + } + + return self.run_async_in_thread(_run()) diff --git a/clients/newapi_veo_client.py b/clients/newapi_veo_client.py new file mode 100644 index 0000000..5dc4b02 --- /dev/null +++ b/clients/newapi_veo_client.py @@ -0,0 +1,542 @@ +""" +new-api Veo 3.1 video client. + +Implements the OpenAI-compatible /v1/videos task flow: +submit, poll, and stream-download video content. +""" + +import asyncio +import json +import os +import re +import time +from typing import Any, Callable, Dict, Optional + +import aiohttp + +from .base_client import BaseAPIClient +from ..utils.config import get_api_base_url, get_api_key_or_raise + + +class NewAPIVeoClient(BaseAPIClient): + CREATE_ENDPOINT = "/v1/videos" + STATUS_ENDPOINT = "/v1/videos/{task_id}" + CONTENT_ENDPOINT = "/v1/videos/{task_id}/content" + + RETRYABLE_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504} + COMPLETED_STATUSES = {"completed", "succeeded", "success", "done"} + FAILED_STATUSES = {"failed", "error", "cancelled", "canceled"} + + def __init__( + self, + base_url: Optional[str] = None, + ): + api_key = get_api_key_or_raise("O1KEY_API_KEY") + resolved_base_url = (base_url or "").strip() or get_api_base_url() + super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key) + + def get_endpoint(self, **kwargs) -> str: + return self.CREATE_ENDPOINT + + def build_request_body(self, **kwargs) -> Dict[str, Any]: + return self._build_video_body(**kwargs) + + def parse_response(self, response: Dict[str, Any]) -> Any: + return response + + @staticmethod + def _build_video_body( + prompt: str, + model: str, + duration: int, + aspect_ratio: str, + resolution: str, + negative_prompt: str = "", + generate_audio: bool = True, + ) -> Dict[str, Any]: + metadata: Dict[str, Any] = { + "aspectRatio": aspect_ratio, + "resolution": resolution, + "generateAudio": bool(generate_audio), + } + + negative_prompt = (negative_prompt or "").strip() + if negative_prompt: + metadata["negativePrompt"] = negative_prompt + + body: Dict[str, Any] = { + "model": model, + "prompt": prompt, + "duration": int(duration), + "metadata": metadata, + } + + return body + + @staticmethod + def _print_request_body(body: Dict[str, Any], image_bytes: Optional[bytes] = None) -> None: + log_body = dict(body) + if image_bytes is not None: + log_body["input_reference"] = f"" + print( + "NewAPI Veo request body:\n" + f"{json.dumps(log_body, ensure_ascii=False, indent=2)}" + ) + + @staticmethod + def _safe_task_filename(task_id: str) -> str: + safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._") + return safe or "newapi_veo" + + @staticmethod + def _extract_task_id(data: Dict[str, Any]) -> Optional[str]: + for key in ("id", "task_id", "video_id"): + value = data.get(key) + if value: + return str(value) + + nested = data.get("data") + if isinstance(nested, dict): + for key in ("id", "task_id", "video_id"): + value = nested.get(key) + if value: + return str(value) + return None + + @staticmethod + def _extract_status(data: Dict[str, Any]) -> str: + for key in ("status", "state", "task_status"): + value = data.get(key) + if value: + return str(value).lower() + + nested = data.get("data") + if isinstance(nested, dict): + for key in ("status", "state", "task_status"): + value = nested.get(key) + if value: + return str(value).lower() + return "unknown" + + @staticmethod + def _extract_progress(data: Dict[str, Any]) -> int: + progress = data.get("progress") + if progress is None and isinstance(data.get("data"), dict): + progress = data["data"].get("progress") + + if isinstance(progress, str): + progress = progress.rstrip("%").strip() + try: + return int(float(progress)) + except ValueError: + return 0 + if isinstance(progress, (int, float)): + return int(progress) + return 0 + + @classmethod + def _format_http_error( + cls, + endpoint: str, + status: int, + error_text: str, + task_id: Optional[str] = None, + ) -> str: + code = "" + message = error_text + try: + payload = json.loads(error_text) + error = payload.get("error", payload) + if isinstance(error, dict): + code = str(error.get("code") or error.get("type") or "") + message = str(error.get("message") or payload.get("message") or error_text) + elif error is not None: + message = str(error) + except Exception: + pass + + message = (message or "").strip() + if len(message) > 1200: + message = message[:1200] + "...(truncated)" + + if status in (401, 403): + hint = "凭证或分组权限问题，请检查 new-api token、模型分组或渠道权限。" + elif status == 429: + hint = "频率或额度限制，请稍后重试或检查 new-api 额度。" + elif status in (502, 503, 504): + hint = "上游服务暂时不可用或超时，请稍后用 task_id 继续查询。" + elif status == 400: + hint = "请求参数错误，请检查 model、duration、metadata 和图片输入。" + else: + hint = "new-api 视频请求失败。" + + parts = [ + hint, + f"endpoint: {endpoint}", + f"http_status: {status}", + ] + if task_id: + parts.append(f"task_id: {task_id}") + if code: + parts.append(f"error_code: {code}") + if message: + parts.append(f"message: {message}") + return "\n".join(parts) + + @classmethod + def _format_task_failure(cls, task_id: str, data: Dict[str, Any]) -> str: + error = data.get("error") + if error is None and isinstance(data.get("data"), dict): + error = data["data"].get("error") + + if isinstance(error, dict): + code = error.get("code") or error.get("type") or "" + message = error.get("message") or json.dumps(error, ensure_ascii=False) + else: + code = "" + message = str(error or "未知错误") + + return "\n".join( + [ + "Veo 视频任务失败。", + f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}", + f"task_id: {task_id}", + f"error_code: {code}", + f"message: {message}", + ] + ) + + async def create_video_async( + self, + prompt: str, + model: str, + duration: int, + aspect_ratio: str, + resolution: str, + negative_prompt: str = "", + generate_audio: bool = True, + image_bytes: Optional[bytes] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> Dict[str, Any]: + body = self._build_video_body( + prompt=prompt, + model=model, + duration=duration, + aspect_ratio=aspect_ratio, + resolution=resolution, + negative_prompt=negative_prompt, + generate_audio=generate_audio, + ) + + url = f"{self.base_url}{self.CREATE_ENDPOINT}" + close_session = False + if session is None: + session = self._make_session() + close_session = True + + try: + timeout = aiohttp.ClientTimeout(total=120, connect=30, sock_read=120) + headers = {"Authorization": f"Bearer {self.api_key}"} + + if image_bytes is not None: + if len(image_bytes) > self.max_request_size: + raise ValueError( + f"输入图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制" + ) + + self._print_request_body(body, image_bytes=image_bytes) + form = aiohttp.FormData() + form.add_field("model", body["model"]) + form.add_field("prompt", body["prompt"]) + form.add_field("duration", str(body["duration"])) + form.add_field("metadata", json.dumps(body["metadata"], ensure_ascii=False)) + form.add_field( + "input_reference", + image_bytes, + filename="input_reference.png", + content_type="image/png", + ) + request_kwargs = {"data": form, "headers": headers} + print( + "NewAPI Veo: POST /v1/videos multipart " + f"| model={model} | duration={duration}s | {resolution} {aspect_ratio}" + ) + else: + self._print_request_body(body) + headers["Content-Type"] = "application/json" + request_kwargs = {"json": body, "headers": headers} + print( + "NewAPI Veo: POST /v1/videos json " + f"| model={model} | duration={duration}s | {resolution} {aspect_ratio}" + ) + + async with session.post(url, timeout=timeout, **request_kwargs) as response: + if response.status >= 300: + error_text = await response.text() + raise RuntimeError( + self._format_http_error( + self.CREATE_ENDPOINT, + response.status, + error_text, + ) + ) + return await response.json() + finally: + if close_session: + await session.close() + + async def _get_json_with_retry( + self, + endpoint: str, + session: aiohttp.ClientSession, + task_id: Optional[str] = None, + max_retries: int = 3, + ) -> Dict[str, Any]: + url = f"{self.base_url}{endpoint}" + headers = self.get_headers(use_bearer_token=True) + timeout = aiohttp.ClientTimeout(total=60, connect=30, sock_read=60) + + last_error = "" + last_status = 0 + for attempt in range(max_retries + 1): + async with session.get(url, headers=headers, timeout=timeout) as response: + if response.status < 300: + return await response.json() + + last_status = response.status + last_error = await response.text() + if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= max_retries: + break + + await asyncio.sleep(min(2 ** attempt, 8)) + + raise RuntimeError( + self._format_http_error(endpoint, last_status, last_error, task_id=task_id) + ) + + async def poll_video_status_async( + self, + task_id: str, + poll_interval: int = 5, + timeout: int = 900, + progress_callback: Optional[Callable[[int, str, float], None]] = None, + session: Optional[aiohttp.ClientSession] = None, + ) -> Dict[str, Any]: + endpoint = self.STATUS_ENDPOINT.format(task_id=task_id) + close_session = False + if session is None: + session = self._make_session() + close_session = True + + start = time.time() + try: + poll_interval = max(1, int(poll_interval)) + await asyncio.sleep(poll_interval) + + while True: + data = await self._get_json_with_retry(endpoint, session, task_id=task_id) + status = self._extract_status(data) + elapsed = time.time() - start + progress = self._extract_progress(data) + + if status == "unknown": + print( + "NewAPI Veo status response did not include a recognized status field:\n" + f"{json.dumps(data, ensure_ascii=False, indent=2)[:1200]}" + ) + + if progress_callback: + progress_callback(progress, status, elapsed) + + if status in self.COMPLETED_STATUSES: + return data + + if status in self.FAILED_STATUSES: + raise RuntimeError(self._format_task_failure(task_id, data)) + + if elapsed >= timeout: + raise TimeoutError( + "Veo 视频任务轮询超时；任务未被标记为失败，可用 task_id 继续查询。\n" + f"endpoint: {endpoint}\n" + f"task_id: {task_id}\n" + f"status: {status}\n" + f"timeout: {timeout}s" + ) + + remaining = max(0.0, timeout - elapsed) + await asyncio.sleep(min(poll_interval, remaining)) + finally: + if close_session: + await session.close() + + async def _download_url_to_file( + self, + url: str, + save_path: str, + session: aiohttp.ClientSession, + max_retries: int = 3, + ) -> None: + timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900) + last_status = 0 + last_error = "" + + for attempt in range(max_retries + 1): + async with session.get(url, timeout=timeout, allow_redirects=True) as response: + if response.status < 300: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(1024 * 1024): + if chunk: + f.write(chunk) + return + + last_status = response.status + last_error = await response.text() + if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= max_retries: + break + + await asyncio.sleep(min(2 ** attempt, 8)) + + raise RuntimeError( + self._format_http_error("download_url", last_status, last_error) + ) + + async def download_video_async( + self, + task_id: str, + save_path: str, + session: Optional[aiohttp.ClientSession] = None, + ) -> str: + endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id) + url = f"{self.base_url}{endpoint}" + headers = self.get_headers(use_bearer_token=True) + timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900) + + close_session = False + if session is None: + session = self._make_session() + close_session = True + + try: + last_status = 0 + last_error = "" + for attempt in range(4): + async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response: + if response.status < 300: + content_type = response.headers.get("Content-Type", "") + if "application/json" in content_type.lower(): + data = await response.json() + nested = data.get("data") if isinstance(data.get("data"), dict) else {} + download_url = ( + data.get("url") + or data.get("download_url") + or nested.get("url") + or nested.get("download_url") + ) + if not download_url: + raise RuntimeError( + "视频下载失败: content 响应为 JSON，但未包含 url/download_url。\n" + f"endpoint: {endpoint}\n" + f"task_id: {task_id}" + ) + await self._download_url_to_file(download_url, save_path, session) + else: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + with open(save_path, "wb") as f: + async for chunk in response.content.iter_chunked(1024 * 1024): + if chunk: + f.write(chunk) + + if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0: + raise RuntimeError( + "视频下载失败: 保存后的文件为空。\n" + f"endpoint: {endpoint}\n" + f"task_id: {task_id}" + ) + return save_path + + last_status = response.status + last_error = await response.text() + if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= 3: + break + + await asyncio.sleep(min(2 ** attempt, 8)) + + raise RuntimeError( + self._format_http_error(endpoint, last_status, last_error, task_id=task_id) + ) + finally: + if close_session: + await session.close() + + def generate_video_sync( + self, + prompt: str, + model: str, + duration: int, + aspect_ratio: str, + resolution: str, + output_dir: str, + negative_prompt: str = "", + generate_audio: bool = True, + image_bytes: Optional[bytes] = None, + poll_interval: int = 5, + timeout: int = 900, + reuse_task_id: str = "", + progress_callback: Optional[Callable[[int, str, float], None]] = None, + ) -> Dict[str, Any]: + async def _run(): + async with self._make_session() as session: + create_response: Dict[str, Any] = {} + task_id = (reuse_task_id or "").strip() + if task_id: + print(f"NewAPI Veo: reuse task_id={task_id}") + else: + create_response = await self.create_video_async( + prompt=prompt, + model=model, + duration=duration, + aspect_ratio=aspect_ratio, + resolution=resolution, + negative_prompt=negative_prompt, + generate_audio=generate_audio, + image_bytes=image_bytes, + session=session, + ) + task_id = self._extract_task_id(create_response) or "" + if not task_id: + raise RuntimeError( + "new-api 未返回视频任务 ID。\n" + f"endpoint: {self.CREATE_ENDPOINT}\n" + f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}" + ) + + status_response = await self.poll_video_status_async( + task_id=task_id, + poll_interval=poll_interval, + timeout=timeout, + progress_callback=progress_callback, + session=session, + ) + status = self._extract_status(status_response) + + os.makedirs(output_dir, exist_ok=True) + filename = f"{self._safe_task_filename(task_id)}.mp4" + save_path = os.path.join(output_dir, filename) + video_path = await self.download_video_async( + task_id=task_id, + save_path=save_path, + session=session, + ) + + return { + "task_id": task_id, + "status": status, + "video_path": video_path, + "raw_json": { + "create": create_response, + "status": status_response, + }, + } + + return self.run_async_in_thread(_run()) diff --git a/nodes/K_video_firstlast.py b/nodes/K_video_firstlast.py index 5594b99..e07ec08 100644 --- a/nodes/K_video_firstlast.py +++ b/nodes/K_video_firstlast.py @@ -110,10 +110,13 @@ class KVideoFirstLast: "mode": mode_api, "duration": 时长, } - if 生成音频 == "打开": - body["generate_audio"] = True + metadata = {} if 尾帧 is not None: - body["metadata"] = {"image_tail": _image_to_base64(尾帧, scale)} + metadata["image_tail"] = _image_to_base64(尾帧, scale) + if 生成音频 == "打开": + metadata["sound"] = "on" + if metadata: + body["metadata"] = metadata body_str = json.dumps(body, ensure_ascii=False) body_size = len(body_str.encode("utf-8")) diff --git a/nodes/K_video_image2video.py b/nodes/K_video_image2video.py index 6b5f976..23efc66 100644 --- a/nodes/K_video_image2video.py +++ b/nodes/K_video_image2video.py @@ -82,6 +82,9 @@ class KVideoImage2Video: CATEGORY = "comfyui_o1key/KVideo" async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", seed=0): + if 模式 == "720p" and 生成音频 == "打开": + raise RuntimeError("K26 仅1080p支持音频，请将模式切换为1080p或关闭生成音频。") + api_key = get_api_key_or_raise() base_url = get_base_url_by_route(网络线路) headers = { @@ -109,7 +112,7 @@ class KVideoImage2Video: "duration": 时长, } if 生成音频 == "打开": - body["generate_audio"] = True + body["metadata"] = {"sound": "on"} body_str = json.dumps(body, ensure_ascii=False) body_size = len(body_str.encode("utf-8")) diff --git a/nodes/__init__.py b/nodes/__init__.py index 023eeda..2855931 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -14,6 +14,7 @@ from .remove_metadata import BatchCleanMetadata from .video_preview import VideoPreview from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset from .veo_video import GoogleVeo +from .newapi_veo_video import Google31Video from .flux_edit import FluxImageEdit from .universal_llm import UniversalLLMChat from .batch_images_o1key import BatchImagesO1key @@ -22,6 +23,7 @@ from .nano_banana_v2 import NanoBananaV2, NanoBananaV2Batch, AsyncImageGenerator from .doubao_image import DoubaoImage from .gpt_image import O1keyGPTImage, O1keyGPTImageBatch from .grok_image import O1keyGrokImage +from .grok_video import O1keyGrokVideo from .K_video_firstlast import KVideoFirstLast from .K_video_image2video import KVideoImage2Video from .K3_video import K3Video @@ -33,4 +35,4 @@ from .remove_bg import O1keyRemoveBackground from .color_remove_bg import O1keyColorRemoveBG from .grid_splitter import O1keyGridSplitter -__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter'] +__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'Google31Video', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'O1keyGrokVideo', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter'] diff --git a/nodes/batch_nano_banana.py b/nodes/batch_nano_banana.py index 583c5dd..c50639d 100644 --- a/nodes/batch_nano_banana.py +++ b/nodes/batch_nano_banana.py @@ -4,23 +4,21 @@ ComfyUI 自定义节点，用于批量处理图像生成任务 支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存 """ -import io as _io -import re -import json import time import math -import base64 import random import asyncio import aiohttp from concurrent.futures import ThreadPoolExecutor -from typing import Optional, Tuple, List +from typing import Callable, Optional, Tuple, List from PIL import Image import torch import numpy as np -from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_images_for_request_body_limit +from comfy_api.latest import io + +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts from ..utils.file_utils import ( ImageInfo, load_images_from_folder, @@ -30,7 +28,7 @@ from ..utils.file_utils import ( save_image, ) from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route, get_api_key_or_raise -from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR +from ..utils.nano_banana_async import generate_nano_banana_async from ..clients.gemini_client import GeminiAPIClient from ..models_config import ( get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, @@ -73,64 +71,28 @@ REQUEST_LOG_ENABLED = False # ============================================================================ _NODE = "Nano Banana" -_ENDPOINT = "/v1/chat/completions" - -_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)") -def _get_headers(api_key: str) -> dict: - return { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - "X-Accel-Buffering": "no", - "Cache-Control": "no-cache, no-transform", - } +def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]: + if pbar is None: + return None + + last_progress = [0.0] + + def _on_progress(progress: float) -> None: + try: + progress = max(0.0, min(float(progress), 1.0)) + except (TypeError, ValueError): + return + if progress <= last_progress[0]: + return + pbar.update(progress - last_progress[0]) + last_progress[0] = progress + + return _on_progress -def _build_request_body( - prompt: str, - model: str, - aspect_ratio: str, - resolution: str, - images: Optional[List[Image.Image]] = None, - enable_grounding: bool = False, -) -> dict: - google_config = { - "image_config": { - "image_size": resolution, - } - } - if aspect_ratio and aspect_ratio != "智能": - google_config["image_config"]["aspect_ratio"] = aspect_ratio - - def _make_body(encoded_images: Optional[List[tuple]] = None) -> dict: - content_parts = [{"type": "text", "text": prompt}] - if encoded_images: - for mime_type, b64 in encoded_images: - content_parts.append({ - "type": "image_url", - "image_url": {"url": f"data:{mime_type};base64,{b64}"} - }) - - body = { - "model": model, - "stream": True, - "messages": [{"role": "user", "content": content_parts}], - "extra_body": {"google": google_config}, - } - - if enable_grounding: - body["extra_body"]["google_search"] = True - return body - - encoded_images = None - if images: - encoded_images = encode_images_for_request_body_limit(images, _make_body) - - return _make_body(encoded_images) - - -async def _generate_single_openai( +async def _generate_single_async( session: aiohttp.ClientSession, base_url: str, api_key: str, @@ -140,82 +102,25 @@ async def _generate_single_openai( aspect_ratio: str, images: Optional[List[Image.Image]] = None, enable_grounding: bool = False, + progress_callback: Optional[Callable[[float], None]] = None, + thinking_level: Optional[str] = None, ) -> List[Image.Image]: - url = f"{base_url}{_ENDPOINT}" - headers = _get_headers(api_key) - body = _build_request_body( + result_images, _ = await generate_nano_banana_async( + session=session, + base_url=base_url, + api_key=api_key, prompt=prompt, model=model, - aspect_ratio=aspect_ratio, resolution=resolution, + aspect_ratio=aspect_ratio, images=images, enable_grounding=enable_grounding, + thinking_level=thinking_level, + node_label="BatchNanoBananaPro", + request_log_enabled=REQUEST_LOG_ENABLED, + progress_callback=progress_callback, ) - - if REQUEST_LOG_ENABLED: - extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False) - print(f"[请求] POST {url} | model={model} | extra_body={extra}") - - last_status = None - for attempt in range(DEFAULT_MAX_RETRIES + 1): - resp = await session.post(url, headers=headers, json=body) - if resp.status == 200: - break - last_status = resp.status - if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES: - friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})") - delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR) - resp.close() - await asyncio.sleep(delay) - continue - error_text = await resp.text() - resp.close() - if resp.status in HTTP_ERROR_MESSAGES: - raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status]) - try: - err_json = json.loads(error_text) - msg = err_json.get("error", {}).get("message", error_text[:200]) - except Exception: - msg = error_text[:200] - raise RuntimeError(f"API 错误 ({resp.status}): {msg}") - else: - if last_status and last_status in HTTP_ERROR_MESSAGES: - raise RuntimeError(HTTP_ERROR_MESSAGES[last_status]) - raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败") - - full_content = "" - buffer = "" - async for raw_chunk in resp.content.iter_any(): - buffer += raw_chunk.decode("utf-8") - while "\n" in buffer: - line_str, buffer = buffer.split("\n", 1) - line_str = line_str.strip() - if not line_str or not line_str.startswith("data:"): - continue - data_str = line_str[5:].strip() - if data_str == "[DONE]": - break - try: - chunk = json.loads(data_str) - delta = chunk.get("choices", [{}])[0].get("delta", {}) - if "content" in delta: - full_content += delta["content"] - except (json.JSONDecodeError, IndexError): - continue - resp.close() - - if not full_content: - raise RuntimeError("API 未返回有效内容") - - matches = list(_IMAGE_RE.finditer(full_content)) - if not matches: - raise RuntimeError(f"响应中未找到图片: {full_content[:100]}") - - last_match = matches[-1] - img_data = base64.b64decode(last_match.group(2)) - final_image = Image.open(_io.BytesIO(img_data)).convert("RGB") - - return [final_image] + return result_images def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: @@ -244,7 +149,7 @@ def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch. return pil_to_tensor(matched) -class BatchNanoBananaPro: +class BatchNanoBananaPro(io.ComfyNode): """ 批量 Nano Banana 节点 @@ -290,6 +195,116 @@ class BatchNanoBananaPro: def __init__(self): pass + + @classmethod + def define_schema(cls): + normal_aspect_ratios = [ + "智能", "1:1", "2:3", "3:2", "3:4", "4:3", + "4:5", "5:4", "9:16", "16:9", "21:9", + ] + nano_banana_2_aspect_ratios = [ + "智能", "1:1", "1:4", "1:8", "2:3", "3:2", "3:4", + "4:1", "4:3", "4:5", "5:4", "8:1", + "9:16", "16:9", "21:9", + ] + + return io.Schema( + node_id="BatchNanoBananaPro", + display_name="批量 Nano Banana", + category="image/batch", + inputs=[ + io.String.Input( + "prompt", + default="一个中国女子的OOTD", + multiline=True, + ), + io.DynamicCombo.Input("模型", options=[ + io.DynamicCombo.Option("Nano Banana Pro", [ + io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"), + io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"), + io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"), + ]), + io.DynamicCombo.Option("Nano Banana 2", [ + io.Combo.Input("宽高比", options=nano_banana_2_aspect_ratios, default="智能"), + io.Combo.Input("分辨率", options=["512px", "1K", "2K", "4K"], default="2K"), + io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"), + io.Combo.Input("思考深度", options=["高", "低"], default="高"), + ]), + io.DynamicCombo.Option("Nano Banana", [ + io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"), + io.Combo.Input("分辨率", options=["1K"], default="1K"), + io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"), + ]), + ]), + io.Combo.Input("图片格式", options=["原始", "JPEG", "PNG", "WebP"], default="原始"), + io.Combo.Input("计费", options=["特价", "官方"], default="特价"), + io.Combo.Input("网络", options=NETWORK_ROUTE_OPTIONS, default="全球加速"), + io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF), + io.String.Input("文件夹1", default="", multiline=False), + io.String.Input("文件夹2", default="", multiline=False), + io.String.Input("文件夹3", default="", multiline=False), + io.String.Input("文件夹4", default="", multiline=False), + io.String.Input("文件夹5", default="", multiline=False), + io.String.Input("保存路径", default="", multiline=False), + io.Combo.Input("图片配对模式", options=cls.PAIRING_MODES, default="不配对"), + io.Image.Input("参考图1", optional=True), + io.Image.Input("参考图2", optional=True), + io.Image.Input("参考图3", optional=True), + io.Image.Input("参考图4", optional=True), + io.Image.Input("参考图5", optional=True), + ], + outputs=[ + io.Image.Output(display_name="输出图像"), + ], + ) + + @classmethod + def execute( + cls, + prompt, + 模型, + 图片格式, + 计费, + 网络, + seed, + 文件夹1, + 文件夹2, + 文件夹3, + 文件夹4, + 文件夹5, + 保存路径, + 图片配对模式, + **kwargs, + ) -> io.NodeOutput: + model_name = 模型["模型"] + 宽高比 = 模型.get("宽高比", "智能") + 分辨率 = 模型.get("分辨率", "2K") + 思考深度 = 模型.get("思考深度") + 谷歌搜索 = 模型.get("谷歌搜索", "关闭") + if 思考深度: + kwargs["思考深度"] = 思考深度 + kwargs["谷歌搜索"] = 谷歌搜索 + + node = cls() + output_tensor, = node.process_batch( + prompt=prompt, + 文件夹1=文件夹1, + 文件夹2=文件夹2, + 文件夹3=文件夹3, + 文件夹4=文件夹4, + 文件夹5=文件夹5, + seed=seed, + 图片配对模式=图片配对模式, + 模型=model_name, + 计费=计费, + 宽高比=宽高比, + 分辨率=分辨率, + 图片格式=图片格式, + 网络=网络, + 保存路径=保存路径, + **kwargs, + ) + return io.NodeOutput(output_tensor) def resize_to_megapixels( self, @@ -352,7 +367,6 @@ class BatchNanoBananaPro: optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, { "default": "不配对" }) - return { "required": { "prompt": ("STRING", { @@ -532,9 +546,11 @@ class BatchNanoBananaPro: enable_grounding: bool = False, base_filename: str = None, image_format: str = "原始", + progress_callback: Optional[Callable[[float], None]] = None, + thinking_level: Optional[str] = None, ) -> dict: """ - 执行单个生成任务（OpenAI 兼容接口） + 执行单个生成任务（异步生图接口） """ result = { "task_index": task_index, @@ -550,10 +566,10 @@ class BatchNanoBananaPro: # 准备输入图片 input_pil_images = [info.image for info in images] - # 调用 OpenAI 兼容接口生成图片 + # 调用异步生图接口生成图片 generated_images = [] try: - gen_images = await _generate_single_openai( + gen_images = await _generate_single_async( session=session, base_url=base_url, api_key=api_key, @@ -563,6 +579,8 @@ class BatchNanoBananaPro: aspect_ratio=aspect_ratio, images=input_pil_images if input_pil_images else None, enable_grounding=enable_grounding, + progress_callback=progress_callback, + thinking_level=thinking_level, ) generated_images.extend(gen_images) except Exception as e: @@ -661,9 +679,10 @@ class BatchNanoBananaPro: prompts_per_task: Optional[List[str]] = None, enable_grounding: bool = False, image_format: str = "原始", + thinking_level: Optional[str] = None, ) -> List[dict]: """ - 异步批量处理所有任务（OpenAI 兼容接口） + 异步批量处理所有任务（异步生图接口） """ total_tasks = len(pairs) @@ -734,6 +753,8 @@ class BatchNanoBananaPro: enable_grounding=enable_grounding, base_filename=base_filename, image_format=image_format, + progress_callback=_make_progress_callback(pbar), + thinking_level=thinking_level, ) ) tasks.append(task) @@ -773,13 +794,12 @@ class BatchNanoBananaPro: else: fail_count += 1 - # 更新 ComfyUI 原生进度条 - if pbar is not None: - pbar.update(1) - # 大任务额外显示百分比里程碑 if show_milestone and milestone_index < len(milestones): - progress = completed / total_tasks + if pbar is not None and getattr(pbar, "total", 0): + progress = pbar.current / pbar.total + else: + progress = success_count / total_tasks if progress >= milestones[milestone_index]: percentage = int(milestones[milestone_index] * 100) print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<") @@ -856,10 +876,15 @@ class BatchNanoBananaPro: start_time = time.time() # 从 kwargs 提取搜索参数（界面显示为「关闭/打开」，转为 bool 供调用） - enable_grounding: bool = False + enable_grounding: bool = kwargs.get("谷歌搜索", "关闭") == "打开" # 拼接实际模型 ID base_model_id = self.MODEL_ID_MAP.get(模型, "nano-banana-pro") + 思考深度 = kwargs.get("思考深度", "高") + thinking_level = None + if base_model_id == "nano-banana-2": + thinking_level = "High" if 思考深度 == "高" else "Low" + if base_model_id == "nano-banana": if 计费 == "官方": raise ValueError(f"模型 \"{模型}\" 仅支持特价计费") @@ -961,11 +986,12 @@ class BatchNanoBananaPro: grounding_str = "" if enable_grounding: grounding_str = " | 谷歌搜索接地" + thinking_str = f" | 思考:{thinking_level}" if thinking_level else "" if batch_prompts: - print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}") + print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}{thinking_str}") else: - print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}") + print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}{thinking_str}") # 创建 ComfyUI 原生进度条 pbar = None @@ -1025,6 +1051,7 @@ class BatchNanoBananaPro: prompts_per_task=prompts_per_task, enable_grounding=enable_grounding, image_format=图片格式, + thinking_level=thinking_level, ) ) except Exception as e: @@ -1146,7 +1173,7 @@ class BatchNanoBananaPro: raise RuntimeError(str(e)) from None except Exception as e: - raise type(e)(str(e)) from None + raise RuntimeError(str(e)) from None finally: # 查询余额 diff --git a/nodes/gpt_image.py b/nodes/gpt_image.py index 68b7a4e..ddd7b5f 100644 --- a/nodes/gpt_image.py +++ b/nodes/gpt_image.py @@ -42,6 +42,56 @@ except ImportError: _FOLDER_PATHS_AVAILABLE = False +def _make_node_progress_callback(progress_bar, task_index: int, total_tasks: int): + if progress_bar is None: + return None + + total_units = max(1, total_tasks) * 100 + base_units = max(0, task_index - 1) * 100 + last_pct = {"value": -1} + + def _callback(pct: int): + try: + pct_value = int(round(float(pct))) + except (TypeError, ValueError): + return + pct_value = max(0, min(100, pct_value)) + if pct_value < last_pct["value"]: + return + last_pct["value"] = pct_value + progress_bar.update_absolute( + min(total_units, base_units + pct_value), + total_units, + ) + + return _callback + + +def _resolve_async_size(value: str) -> str: + value = (value or "").strip() + if not value or value == "智能" or value.lower() == "auto": + return "auto" + + first_part = value.split("（")[0].strip() + normalized_size = first_part.lower().replace("*", "x").replace("×", "x") + size_parts = [part.strip() for part in normalized_size.split("x")] + if len(size_parts) == 2 and all(part.isdigit() for part in size_parts): + return f"{int(size_parts[0])}x{int(size_parts[1])}" + + allowed = {"auto", "1024x1024", "1K", "2K", "4K"} + if first_part in allowed: + return first_part + + if "4K" in value: + return "4K" + if "2K" in value: + return "2K" + if "1K" in value: + return "1K" + + return "auto" + + class O1keyGPTImage: """ o1key GPT Image 节点 @@ -123,6 +173,10 @@ class O1keyGPTImage: "default": "自动", "tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto", }) + optional_inputs["输出格式"] = (["png", "jpeg", "webp"], { + "default": "jpeg", + "tooltip": "Generated image output format", + }) optional_inputs["seed"] = ("INT", { "default": 0, "min": 0, @@ -160,6 +214,7 @@ class O1keyGPTImage: 网络: str = "全球加速", 分辨率: str = "auto", 质量: str = "自动", + 输出格式: str = "jpeg", 生图数量: int = 1, seed: int = 0, 遮罩=None, @@ -190,7 +245,7 @@ class O1keyGPTImage: raise ValueError("提供了遮罩但未提供图片，请同时提供图片和遮罩") # ── 2. 解析分辨率显示值 → API 参数值 ────────────────────────────────── - size = "auto" if 分辨率 == "智能" else 分辨率.split("（")[0].strip() + size = _resolve_async_size(分辨率) # ── 2b. 解析模型显示值 → API 参数值 ─────────────────────────────────── _model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"} @@ -216,6 +271,8 @@ class O1keyGPTImage: # ── 5. 调用 API ─────────────────────────────────────────────────── all_pil_images = [] + progress_total = len(batch_prompts) if batch_prompts else 1 + progress_bar = ProgressBar(progress_total * 100) if _PROGRESS_BAR_AVAILABLE else None if batch_prompts: # 批量模式：逐条提示词调用 @@ -226,7 +283,7 @@ class O1keyGPTImage: print("[o1key GPT Image] 用户取消，已中断批量生成") raise InterruptProcessingException() try: - pil_images = client.run_sync( + pil_images = client.generate_image_async_sync( prompt=p, model=model, quality=quality, @@ -235,6 +292,8 @@ class O1keyGPTImage: seed=seed, image_tensor=图片, mask_tensor=遮罩, + output_format=输出格式, + progress_callback=_make_node_progress_callback(progress_bar, idx, total), ) all_pil_images.extend(pil_images) snippet = p[:30] + ("..." if len(p) >= 30 else "") @@ -245,12 +304,14 @@ class O1keyGPTImage: error_msg = str(e).split('\n')[0] snippet = p[:30] + ("..." if len(p) >= 30 else "") print(f"[o1key GPT Image] [{idx}/{total}] ❌ {snippet} → {error_msg}") + if progress_bar is not None: + progress_bar.update_absolute(idx * 100, total * 100) else: # 单提示词模式 if not prompt or not prompt.strip(): raise ValueError("提示词不能为空") try: - pil_images = client.run_sync( + pil_images = client.generate_image_async_sync( prompt=prompt, model=model, quality=quality, @@ -259,6 +320,8 @@ class O1keyGPTImage: seed=seed, image_tensor=图片, mask_tensor=遮罩, + output_format=输出格式, + progress_callback=_make_node_progress_callback(progress_bar, 1, 1), ) all_pil_images.extend(pil_images) except InterruptProcessingException: @@ -495,7 +558,7 @@ class O1keyGPTImageBatch: @staticmethod def _resolve_size(分辨率: str) -> str: - return "auto" if 分辨率 == "智能" else 分辨率.split("（")[0].strip() + return _resolve_async_size(分辨率) @staticmethod def _resolve_model(模型: str) -> str: @@ -510,6 +573,15 @@ class O1keyGPTImageBatch: quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"} return quality_map.get(质量, "auto") + @staticmethod + def _resolve_output_format(图片格式: str) -> str: + output_format_map = { + "JPEG": "jpeg", + "PNG": "png", + "WebP": "webp", + } + return output_format_map.get(图片格式, "png") + @staticmethod def _ensure_output_folder(保存路径: str) -> str: output_folder = (保存路径 or "").strip() @@ -637,11 +709,12 @@ class O1keyGPTImageBatch: size = self._resolve_size(分辨率) model = self._resolve_model(模型) quality = self._resolve_quality(质量) + output_format = self._resolve_output_format(图片格式) client = GptImageClient() client.base_url = get_base_url_by_route(网络) - progress_bar = ProgressBar(total_tasks) if _PROGRESS_BAR_AVAILABLE else None + progress_bar = ProgressBar(total_tasks * 100) if _PROGRESS_BAR_AVAILABLE else None results = [] all_saved_files = [] @@ -661,7 +734,7 @@ class O1keyGPTImageBatch: } try: - pil_images = client.run_sync( + pil_images = client.generate_image_async_sync( prompt=task_prompt, model=model, quality=quality, @@ -670,6 +743,8 @@ class O1keyGPTImageBatch: seed=seed, image_tensor=self._pair_to_tensors(pair), mask_tensor=遮罩, + output_format=output_format, + progress_callback=_make_node_progress_callback(progress_bar, task_index, total_tasks), ) saved_files = self._save_images( images=pil_images, @@ -691,7 +766,7 @@ class O1keyGPTImageBatch: results.append(result) if progress_bar is not None: - progress_bar.update(1) + progress_bar.update_absolute(task_index * 100, total_tasks * 100) success_count = sum(1 for result in results if result.get("success", False)) total_generated = sum(result.get("generated_count", 0) for result in results) diff --git a/nodes/grok_video.py b/nodes/grok_video.py new file mode 100644 index 0000000..c91b1c3 --- /dev/null +++ b/nodes/grok_video.py @@ -0,0 +1,283 @@ +""" +Grok Video node. + +Submits a /v1/videos task, polls until completion, downloads the mp4, +and returns ComfyUI's native VIDEO object. +""" + +import json +import os +from typing import List, Optional + +from ..clients.grok_video_client import GrokVideoClient +from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route +from ..utils.image_utils import encode_images_for_request_body_limit, tensor_to_pil + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + ProgressBar = None + PROGRESS_BAR_AVAILABLE = False + +try: + from comfy_api.input_impl import VideoFromFile +except Exception: + try: + from comfy_api.latest import InputImpl + VideoFromFile = InputImpl.VideoFromFile + except Exception: + VideoFromFile = None + + +MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"] +ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"] +QUALITY_OPTIONS = ["720p"] +QUALITY_VALUE_MAP = { + "720p": "high", +} +MODEL_SECONDS_OPTIONS = { + "grok-imagine-1.0-video": [6, 10, 12, 16, 20], +} + +MAX_REFERENCE_IMAGES = 3 +MAX_REQUEST_BODY_BYTES = 20 * 1024 * 1024 + + +def _get_output_dir() -> str: + if FOLDER_PATHS_AVAILABLE: + base = folder_paths.get_output_directory() + else: + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + comfy_root = os.path.dirname(os.path.dirname(plugin_dir)) + base = os.path.join(comfy_root, "output") + + output_dir = os.path.join(base, "grok_video") + os.makedirs(output_dir, exist_ok=True) + return output_dir + + +def _format_mb(size_bytes: int) -> str: + return f"{size_bytes / 1024 / 1024:.2f}MB" + + +def _image_tensor_to_first_pil(image_tensor): + if image_tensor is None: + return None + + pil_images = tensor_to_pil(image_tensor) + if not pil_images: + return None + + image = pil_images[0] + if image.mode not in ("RGB", "L"): + image = image.convert("RGB") + return image + + +def _collect_reference_images(**kwargs) -> List[object]: + images = [] + for i in range(1, MAX_REFERENCE_IMAGES + 1): + image = _image_tensor_to_first_pil(kwargs.get(f"参考图{i}")) + if image is not None: + images.append(image) + return images + + +def _to_data_urls(encoded_images) -> List[str]: + return [f"data:{mime};base64,{b64}" for mime, b64 in encoded_images] + + +def _encode_image_data_urls( + images: List[object], + prompt: str, + model: str, + aspect_ratio: str, + seconds: int, + quality: str, +) -> Optional[List[str]]: + if not images: + return None + + def build_body(encoded_images): + return GrokVideoClient.build_video_body( + prompt=prompt, + model=model, + aspect_ratio=aspect_ratio, + seconds=seconds, + quality=quality, + images=_to_data_urls(encoded_images), + ) + + encoded = encode_images_for_request_body_limit( + images, + build_body=build_body, + max_body_bytes=MAX_REQUEST_BODY_BYTES, + ) + data_urls = _to_data_urls(encoded) + + return data_urls + + +def _validate_request_body_size(body: dict) -> None: + body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8")) + if body_size > MAX_REQUEST_BODY_BYTES: + raise ValueError( + f"Grok Video 请求体大小 {_format_mb(body_size)} 超过 " + f"{_format_mb(MAX_REQUEST_BODY_BYTES)} 限制，请减少参考图片或降低图片尺寸。" + ) + + +class O1keyGrokVideo: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "提示词": ( + "STRING", + { + "default": "", + "multiline": True, + }, + ), + "网络线路": (NETWORK_ROUTE_OPTIONS, {"default": NETWORK_ROUTE_OPTIONS[0]}), + "模型": (MODEL_OPTIONS, {"default": MODEL_OPTIONS[0]}), + "宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}), + "秒数（按模型限制）": ( + "INT", + { + "default": 5, + "min": 5, + "max": 20, + "step": 1, + "display": "number", + }, + ), + "画质": (QUALITY_OPTIONS, {"default": "720p"}), + }, + "optional": { + "参考图1": ("IMAGE",), + "参考图2": ("IMAGE",), + "参考图3": ("IMAGE",), + }, + } + + RETURN_TYPES = ("VIDEO",) + RETURN_NAMES = ("视频",) + FUNCTION = "generate" + CATEGORY = "comfyui_o1key/Video" + + DESCRIPTION = ( + "Grok Video /v1/videos task node. Supports prompt plus up to " + "three image references, multiple aspect ratios, model-specific seconds, 720p output." + ) + + def generate( + self, + **kwargs, + ): + if VideoFromFile is None: + raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。") + + 提示词 = kwargs.get("提示词", "") + 网络线路 = kwargs.get("网络线路", NETWORK_ROUTE_OPTIONS[0]) + 模型 = kwargs.get("模型", MODEL_OPTIONS[0]) + 宽高比 = kwargs.get("宽高比", "16:9") + 秒数 = kwargs.get("秒数（按模型限制）", kwargs.get("秒数（≤15s）", kwargs.get("秒数", 5))) + 画质 = kwargs.get("画质", "720p") + + prompt = (提示词 or "").strip() + if not prompt: + raise ValueError("提示词不能为空。") + if 模型 not in MODEL_OPTIONS: + raise ValueError(f"模型仅支持: {', '.join(MODEL_OPTIONS)}") + if 宽高比 not in ASPECT_RATIO_OPTIONS: + raise ValueError(f"宽高比仅支持: {', '.join(ASPECT_RATIO_OPTIONS)}。") + seconds = int(秒数) + allowed_seconds = MODEL_SECONDS_OPTIONS.get(模型) + if allowed_seconds is not None: + if seconds not in allowed_seconds: + raise ValueError( + f"模型 {模型} 仅支持秒数: " + f"{', '.join(str(s) for s in allowed_seconds)}。" + "请修改为正确的秒数后再发起请求。" + ) + elif seconds < 5 or seconds > 15: + raise ValueError("秒数仅支持 5 到 15。") + if 画质 not in QUALITY_OPTIONS: + raise ValueError("画质仅支持 720p。") + + quality = QUALITY_VALUE_MAP[画质] + reference_images = _collect_reference_images(**kwargs) + image_data_urls = _encode_image_data_urls( + reference_images, + prompt=prompt, + model=模型, + aspect_ratio=宽高比, + seconds=seconds, + quality=quality, + ) + + request_body = GrokVideoClient.build_video_body( + prompt=prompt, + model=模型, + aspect_ratio=宽高比, + seconds=seconds, + quality=quality, + images=image_data_urls, + ) + _validate_request_body_size(request_body) + + pbar = ProgressBar(100) if PROGRESS_BAR_AVAILABLE else None + last_progress = [0] + + def progress_callback(progress: int, status: str, elapsed: float): + progress_value = max(0, min(100, int(progress or 0))) + if pbar is not None and progress_value > last_progress[0]: + pbar.update(progress_value - last_progress[0]) + last_progress[0] = progress_value + + client = GrokVideoClient(base_url=get_base_url_by_route(网络线路)) + + try: + result = client.generate_video_sync( + prompt=prompt, + model=模型, + aspect_ratio=宽高比, + seconds=seconds, + quality=quality, + output_dir=_get_output_dir(), + images=image_data_urls, + poll_interval=5, + timeout=1200, + progress_callback=progress_callback, + ) + + if pbar is not None and last_progress[0] < 100: + pbar.update(100 - last_progress[0]) + + video_path = result["video_path"] + print(f"Grok Video：下载完成：{video_path}") + return (VideoFromFile(video_path),) + finally: + try: + balance_data = client.query_balance_sync() + balance_info = client.format_balance_info(balance_data) + print(f"Grok Video：{balance_info}") + except Exception: + pass + + +NODE_CLASS_MAPPINGS = { + "O1keyGrokVideo": O1keyGrokVideo, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "O1keyGrokVideo": "Grok Video", +} diff --git a/nodes/nano_banana.py b/nodes/nano_banana.py index 3e16679..60782e5 100644 --- a/nodes/nano_banana.py +++ b/nodes/nano_banana.py @@ -1,20 +1,16 @@ """ Nano Banana 节点 (V3) -ComfyUI 自定义节点，用于调用生图模型（OpenAI 兼容接口） +ComfyUI 自定义节点，用于调用异步生图模型 使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动 """ -import io as _io -import re -import json import time import math -import base64 import random import asyncio import aiohttp from concurrent.futures import ThreadPoolExecutor -from typing import List, Optional +from typing import Callable, List, Optional import torch import numpy as np @@ -22,21 +18,15 @@ from PIL import Image from comfy_api.latest import io -from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_images_for_image_size_limit +from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts from ..utils.config import ( NETWORK_ROUTE_OPTIONS, get_base_url_by_route, get_api_key_or_raise, ) -from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR +from ..utils.nano_banana_async import generate_nano_banana_async from ..clients.gemini_client import GeminiAPIClient -try: - import folder_paths - FOLDER_PATHS_AVAILABLE = True -except ImportError: - FOLDER_PATHS_AVAILABLE = False - try: from comfy.utils import ProgressBar PROGRESS_BAR_AVAILABLE = True @@ -51,17 +41,9 @@ except ImportError: InterruptProcessingException = RuntimeError processing_interrupted = lambda: False -try: - import psutil - MEMORY_MONITOR_AVAILABLE = True -except ImportError: - MEMORY_MONITOR_AVAILABLE = False - -DEBUG_LOG_ENABLED = True REQUEST_LOG_ENABLED = False _NODE = "Nano Banana" -_ENDPOINT = "/v1/chat/completions" _REQUEST_TIMEOUT = 900 _INTERRUPT_CHECK_INTERVAL = 0.2 @@ -112,6 +94,25 @@ def _check_interrupt(): raise InterruptProcessingException() +def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]: + if pbar is None: + return None + + last_progress = [0.0] + + def _on_progress(progress: float) -> None: + try: + progress = max(0.0, min(float(progress), 1.0)) + except (TypeError, ValueError): + return + if progress <= last_progress[0]: + return + pbar.update(progress - last_progress[0]) + last_progress[0] = progress + + return _on_progress + + def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: if not images: placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) @@ -169,66 +170,6 @@ def _build_model_id(model_name: str, resolution: str, billing: str) -> str: model_id += "-official" return model_id -_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)") - - -def _get_headers(api_key: str) -> dict: - return { - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - "X-Accel-Buffering": "no", - "Cache-Control": "no-cache, no-transform", - } - - -def _build_request_body( - prompt: str, - model: str, - aspect_ratio: str, - resolution: str, - images: Optional[List[Image.Image]] = None, - enable_grounding: bool = False, - thinking_level: Optional[str] = None, -) -> dict: - google_config = { - "image_config": { - "image_size": resolution, - } - } - if aspect_ratio and aspect_ratio != "智能": - google_config["image_config"]["aspect_ratio"] = aspect_ratio - if thinking_level: - google_config["thinking_config"] = { - "thinking_level": thinking_level.lower(), - "include_thoughts": True, - } - - def _make_body(encoded_images: Optional[List[tuple]] = None) -> dict: - content_parts = [{"type": "text", "text": prompt}] - if encoded_images: - for mime_type, b64 in encoded_images: - content_parts.append({ - "type": "image_url", - "image_url": {"url": f"data:{mime_type};base64,{b64}"} - }) - - body = { - "model": model, - "stream": True, - "messages": [{"role": "user", "content": content_parts}], - "extra_body": {"google": google_config}, - } - - if enable_grounding: - body["extra_body"]["google_search"] = True - return body - - encoded_images = None - if images: - encoded_images = encode_images_for_image_size_limit(images) - - return _make_body(encoded_images) - async def _generate_single( session: aiohttp.ClientSession, @@ -241,105 +182,25 @@ async def _generate_single( images: Optional[List[Image.Image]] = None, enable_grounding: bool = False, thinking_level: Optional[str] = None, + progress_callback: Optional[Callable[[float], None]] = None, ) -> List[Image.Image]: - url = f"{base_url}{_ENDPOINT}" - headers = _get_headers(api_key) - body = _build_request_body( + result_images, timing = await generate_nano_banana_async( + session=session, + base_url=base_url, + api_key=api_key, prompt=prompt, model=model, - aspect_ratio=aspect_ratio, resolution=resolution, + aspect_ratio=aspect_ratio, images=images, enable_grounding=enable_grounding, thinking_level=thinking_level, + node_label="Nano Banana", + request_log_enabled=REQUEST_LOG_ENABLED, + check_interrupt=_check_interrupt, + progress_callback=progress_callback, ) - - if REQUEST_LOG_ENABLED: - extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False) - print(f"[请求] POST {url} | model={model} | extra_body={extra}") - - last_status = None - resp = None - timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT, connect=30, sock_read=_REQUEST_TIMEOUT) - for attempt in range(DEFAULT_MAX_RETRIES + 1): - _check_interrupt() - resp = await session.post(url, headers=headers, json=body, timeout=timeout) - if resp.status == 200: - break - last_status = resp.status - if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES: - friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})") - delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR) - print(f"Nano Banana: {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...") - resp.close() - await asyncio.sleep(delay) - continue - error_text = await resp.text() - resp.close() - if resp.status in HTTP_ERROR_MESSAGES: - raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status]) - try: - err_json = json.loads(error_text) - msg = err_json.get("error", {}).get("message", error_text[:200]) - except Exception: - msg = error_text[:200] - raise RuntimeError(f"API 错误 ({resp.status}): {msg}") - else: - if last_status and last_status in HTTP_ERROR_MESSAGES: - raise RuntimeError(HTTP_ERROR_MESSAGES[last_status]) - raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败") - - full_content = "" - buffer = "" - t_request = time.time() - t_first_token = None - try: - async for raw_chunk in resp.content.iter_any(): - _check_interrupt() - if t_first_token is None: - t_first_token = time.time() - buffer += raw_chunk.decode("utf-8") - while "\n" in buffer: - line_str, buffer = buffer.split("\n", 1) - line_str = line_str.strip() - if not line_str or not line_str.startswith("data:"): - continue - data_str = line_str[5:].strip() - if data_str == "[DONE]": - break - try: - chunk = json.loads(data_str) - delta = chunk.get("choices", [{}])[0].get("delta", {}) - if "content" in delta: - full_content += delta["content"] - except (json.JSONDecodeError, IndexError): - continue - except aiohttp.ClientPayloadError as e: - if full_content and _IMAGE_RE.search(full_content): - print(f"Nano Banana: 响应流提前结束，但已收到完整图片，继续解析 ({e})") - else: - raise RuntimeError(f"响应流下载中断，请重试或检查网络/代理: {e}") from None - finally: - if resp is not None: - resp.close() - t_done = time.time() - - if not full_content: - raise RuntimeError("API 未返回有效内容") - - # 思考模型可能输出多张临时图片，最终图片始终是最后一张 - matches = list(_IMAGE_RE.finditer(full_content)) - if not matches: - raise RuntimeError(f"响应中未找到图片: {full_content[:100]}") - - last_match = matches[-1] - img_data = base64.b64decode(last_match.group(2)) - final_image = Image.open(_io.BytesIO(img_data)).convert("RGB") - - first_token_ms = (t_first_token - t_request) * 1000 if t_first_token else 0 - download_ms = (t_done - t_first_token) * 1000 if t_first_token else 0 - - return [final_image], first_token_ms, download_ms + return result_images, timing["task_ms"], timing["parse_ms"] async def _generate_single_task( @@ -354,6 +215,7 @@ async def _generate_single_task( global_task_index: int, enable_grounding: bool = False, thinking_level: Optional[str] = None, + progress_callback: Optional[Callable[[float], None]] = None, ) -> dict: result = { "global_task_index": global_task_index, @@ -364,7 +226,7 @@ async def _generate_single_task( "error": None, } try: - gen_images, first_token_ms, download_ms = await _generate_single( + gen_images, task_ms, parse_ms = await _generate_single( session=session, base_url=base_url, api_key=api_key, @@ -375,7 +237,9 @@ async def _generate_single_task( images=images if images else None, enable_grounding=enable_grounding, thinking_level=thinking_level, + progress_callback=progress_callback, ) + del task_ms, parse_ms result["output_images"] = gen_images result["success"] = True result["generated_count"] = len(gen_images) @@ -438,6 +302,7 @@ async def _process_batch_async( global_task_index=i, enable_grounding=enable_grounding, thinking_level=thinking_level, + progress_callback=_make_progress_callback(pbar), ) ) tasks.append(task) @@ -473,9 +338,6 @@ async def _process_batch_async( error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" print(f"Nano Banana: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}") - if pbar is not None: - pbar.update(1) - all_results.extend(batch_results) import gc; gc.collect() await asyncio.sleep(0.1) @@ -576,7 +438,7 @@ class NanoBanana(io.ComfyNode): pil_imgs = tensor_to_pil(kwargs[key]) input_images.extend(pil_imgs) - if input_images and len(input_images) > 14: + if len(input_images) > 14: raise ValueError(f"输入图像数量 {len(input_images)} 超过限制 14 张") batch_prompts = parse_batch_prompts(prompt) @@ -667,6 +529,7 @@ class NanoBanana(io.ComfyNode): images=input_images if input_images else None, enable_grounding=enable_grounding, thinking_level=thinking_level, + progress_callback=_make_progress_callback(pbar), ) return loop.run_until_complete(_run_with_interrupt(_do())) finally: @@ -674,17 +537,14 @@ class NanoBanana(io.ComfyNode): with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(run_single) - generated_images, first_token_ms, download_ms = future.result(timeout=_REQUEST_TIMEOUT) - - if pbar is not None: - pbar.update(1) + generated_images, task_ms, parse_ms = future.result(timeout=_REQUEST_TIMEOUT) output_tensor = _images_to_tensor_safe(generated_images, _NODE) elapsed = time.time() - start_time time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" - ft_str = f"{first_token_ms/1000:.2f}s" - dl_str = f"{download_ms/1000:.2f}s" - print(f"完成！总耗时 {time_str} | 首字 {ft_str} | 下载 {dl_str} | 成功 {len(generated_images)}张") + task_str = f"{task_ms/1000:.2f}s" + parse_str = f"{parse_ms/1000:.2f}s" + print(f"完成！总耗时 {time_str} | 异步任务 {task_str} | 解析 {parse_str} | 成功 {len(generated_images)}张") import gc; gc.collect() return io.NodeOutput(output_tensor) @@ -701,7 +561,7 @@ class NanoBanana(io.ComfyNode): except RuntimeError as e: raise RuntimeError(str(e)) from None except Exception as e: - raise type(e)(str(e)) from None + raise RuntimeError(str(e)) from None finally: if not was_interrupted: try: diff --git a/nodes/nano_banana_v2.py b/nodes/nano_banana_v2.py index 91a6027..d263c19 100644 --- a/nodes/nano_banana_v2.py +++ b/nodes/nano_banana_v2.py @@ -31,7 +31,7 @@ from PIL import Image from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts from ..utils.file_utils import load_images_from_folder, pair_images_by_name, pair_images_cartesian from ..utils.config import get_api_key_or_raise, NETWORK_ROUTE_OPTIONS, get_base_url_by_route -from ..utils.http_error import async_request_with_retry +from ..utils.http_error import async_request_with_retry, extract_structured_error_message, get_friendly_message from ..models_config import ( get_enabled_async_models, get_model_provider, @@ -241,6 +241,10 @@ class NanoBananaV2: @staticmethod def _friendly_error(error_msg: str) -> str: """将上游错误转化为用户友好的提示""" + structured_message = extract_structured_error_message(error_msg) + if structured_message and structured_message != error_msg: + error_msg = structured_message + if "No available channel for model" in error_msg: return ( "当前分组下模型不可用，请检查分组是否正确。" @@ -343,16 +347,16 @@ class NanoBananaV2: error_text = await response.text() if not error_text.strip(): error_text = "(服务器未返回错误详情)" - raise RuntimeError(f"查询任务失败 ({response.status}): {error_text}") + raise RuntimeError(f"查询任务失败: {get_friendly_message(response.status, error_text)}") result = await response.json() status = provider.extract_status(result) - # 提取进度并回调（封顶 1.0 防止异常值导致进度条溢出） - if on_progress and status in ("SUBMITTED", "IN_PROGRESS"): + # 提取进度并回调；运行中状态不显示 100%，只有 SUCCESS 才补满。 + if on_progress and status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"): p = provider.extract_progress(result) if p is not None: - p = min(p, 1.0) + p = min(p, 0.99) if p > last_progress: on_progress(p - last_progress) last_progress = p @@ -374,7 +378,7 @@ class NanoBananaV2: print(f"{self.NODE_LABEL}: [轮询] FAILURE 但无错误信息，原始响应: {json.dumps(result, ensure_ascii=False)[:500]}") friendly_msg = self._friendly_error(error_msg) raise RuntimeError(f"任务失败: {friendly_msg}") - elif status in ("SUBMITTED", "IN_PROGRESS"): + elif status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"): # 分段 sleep，每 0.1 秒检查一次取消信号 sleep_iterations = int(_POLL_INTERVAL / _INTERRUPT_CHECK_INTERVAL) for _ in range(sleep_iterations): @@ -406,10 +410,7 @@ class NanoBananaV2: "error": None, } - contributed = [0.0] # mutable container，追踪本任务已贡献的 pbar 进度 - def _track_progress(delta): - contributed[0] += delta if on_progress: on_progress(delta) @@ -435,15 +436,9 @@ class NanoBananaV2: result["request_time"] = request_time result["download_time"] = download_time except InterruptProcessingException: - # 用户取消：补齐进度后向上传播，不吞掉 - if contributed[0] < 1.0 and on_progress: - on_progress(1.0 - contributed[0]) raise except Exception as e: result["error"] = str(e) or f"{type(e).__name__}(无错误详情)" - # 失败也补齐 1.0 进度，保证进度条总数正确 - if contributed[0] < 1.0 and on_progress: - on_progress(1.0 - contributed[0]) return result @@ -704,7 +699,7 @@ class NanoBananaV2: raise RuntimeError(str(e)) from None except Exception as e: - raise type(e)(str(e)) from None + raise RuntimeError(str(e)) from None finally: # 查询并打印余额 @@ -1223,7 +1218,7 @@ class NanoBananaV2Batch(NanoBananaV2): raise RuntimeError(str(e)) from None except Exception as e: - raise type(e)(str(e)) from None + raise RuntimeError(str(e)) from None finally: try: diff --git a/nodes/newapi_veo_video.py b/nodes/newapi_veo_video.py new file mode 100644 index 0000000..2f870b4 --- /dev/null +++ b/nodes/newapi_veo_video.py @@ -0,0 +1,254 @@ +""" +Single-node new-api Veo 3.1 generator. + +The node submits a /v1/videos task, waits for completion, downloads the mp4, +and returns ComfyUI's native VIDEO object for the built-in Save Video node. +""" + +import os +from io import BytesIO +from typing import Optional, Tuple + +from ..clients.newapi_veo_client import NewAPIVeoClient +from ..utils.image_utils import tensor_to_pil +from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route + +try: + import folder_paths + FOLDER_PATHS_AVAILABLE = True +except ImportError: + FOLDER_PATHS_AVAILABLE = False + +try: + from comfy.utils import ProgressBar + PROGRESS_BAR_AVAILABLE = True +except ImportError: + ProgressBar = None + PROGRESS_BAR_AVAILABLE = False + +try: + from comfy_api.input_impl import VideoFromFile +except Exception: + VideoFromFile = None + + +MODEL_OPTIONS = [ + "veo-3.1", +] + +DURATION_OPTIONS = ["4", "6", "8"] +ASPECT_RATIO_OPTIONS = ["16:9", "9:16"] +RESOLUTION_OPTIONS = ["720p", "1080p"] + +TARGET_SIZE_MAP = { + ("720p", "16:9"): (1280, 720), + ("720p", "9:16"): (720, 1280), + ("1080p", "16:9"): (1920, 1080), + ("1080p", "9:16"): (1080, 1920), +} + + +def _get_output_dir() -> str: + if FOLDER_PATHS_AVAILABLE: + return folder_paths.get_output_directory() + + plugin_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + comfy_root = os.path.dirname(os.path.dirname(plugin_dir)) + return os.path.join(comfy_root, "output") + + +def _get_download_dir() -> str: + output_dir = _get_output_dir() + video_dir = os.path.join(output_dir, "newapi_veo") + os.makedirs(video_dir, exist_ok=True) + return video_dir + + +def _fit_image_to_target(image, target_size: Tuple[int, int]): + from PIL import Image as PILImage + + target_w, target_h = target_size + src_w, src_h = image.size + src_ratio = src_w / src_h + target_ratio = target_w / target_h + + if src_w == target_w and src_h == target_h: + return image + + resample = PILImage.Resampling.LANCZOS if hasattr(PILImage, "Resampling") else PILImage.LANCZOS + + if src_ratio > target_ratio: + scale = target_h / src_h + new_w = round(src_w * scale) + image = image.resize((new_w, target_h), resample=resample) + left = max(0, (new_w - target_w) // 2) + image = image.crop((left, 0, left + target_w, target_h)) + else: + scale = target_w / src_w + new_h = round(src_h * scale) + image = image.resize((target_w, new_h), resample=resample) + top = max(0, (new_h - target_h) // 2) + image = image.crop((0, top, target_w, top + target_h)) + + return image + + +def _image_to_png_bytes(image_tensor, resolution: str, aspect_ratio: str) -> Optional[bytes]: + if image_tensor is None: + return None + + pil_images = tensor_to_pil(image_tensor) + if not pil_images: + return None + + image = pil_images[0] + if image.mode != "RGB": + image = image.convert("RGB") + + target_size = TARGET_SIZE_MAP.get((resolution, aspect_ratio)) + if target_size is not None: + original_size = image.size + image = _fit_image_to_target(image, target_size) + if image.size != original_size: + print( + "NewAPI Veo: input image fitted " + f"{original_size[0]}x{original_size[1]} -> {image.size[0]}x{image.size[1]}" + ) + + buffer = BytesIO() + image.save(buffer, format="PNG") + image_bytes = buffer.getvalue() + print( + "NewAPI Veo: input_reference PNG " + f"{len(image_bytes) / 1024:.0f} KB ({image.size[0]}x{image.size[1]})" + ) + return image_bytes + + +class Google31Video: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "提示词": ( + "STRING", + { + "default": "A cinematic shot of a small robot walking through a rainy neon street.", + "multiline": True, + }, + ), + "负向提示词": ("STRING", {"default": "", "multiline": True}), + "网络线路": (NETWORK_ROUTE_OPTIONS, {"default": "全球加速"}), + "模型": (MODEL_OPTIONS, {"default": MODEL_OPTIONS[0]}), + "时长": (DURATION_OPTIONS, {"default": "8"}), + "宽高比": (ASPECT_RATIO_OPTIONS, {"default": "16:9"}), + "分辨率": (RESOLUTION_OPTIONS, {"default": "1080p"}), + "生成音频": (["打开", "关闭"], {"default": "打开"}), + "seed": ( + "INT", + { + "default": -1, + "min": -1, + "max": 0xFFFFFFFFFFFFFFFF, + "step": 1, + }, + ), + }, + "optional": { + "参考图像": ("IMAGE",), + }, + } + + RETURN_TYPES = ("VIDEO",) + RETURN_NAMES = ("视频",) + FUNCTION = "generate" + CATEGORY = "comfyui_o1key/Video" + + DESCRIPTION = ( + "Submit a new-api /v1/videos Veo 3.1 task, poll until complete, " + "download the mp4, and output native VIDEO for ComfyUI Save Video." + ) + + def generate( + self, + 提示词: str, + 负向提示词: str, + 网络线路: str, + 模型: str, + 时长: str, + 宽高比: str, + 分辨率: str, + 生成音频: str, + seed: int, + 参考图像=None, + ): + if VideoFromFile is None: + raise RuntimeError("当前 ComfyUI 版本不支持原生 VIDEO 输入实现 VideoFromFile。") + + prompt = (提示词 or "").strip() + if not prompt: + raise ValueError("提示词不能为空。") + + duration_value = int(时长) + if duration_value not in (4, 6, 8): + raise ValueError("时长仅支持 4、6、8。") + if 宽高比 not in ASPECT_RATIO_OPTIONS: + raise ValueError("宽高比仅支持 16:9 或 9:16。") + if 分辨率 not in RESOLUTION_OPTIONS: + raise ValueError("分辨率仅支持 720p 或 1080p。") + + output_dir = _get_download_dir() + image_bytes = _image_to_png_bytes(参考图像, 分辨率, 宽高比) + + pbar = ProgressBar(100) if PROGRESS_BAR_AVAILABLE else None + last_progress = [0] + last_status = [""] + + def progress_callback(progress: int, status: str, elapsed: float): + if status != last_status[0]: + print( + "NewAPI Veo: polling " + f"status={status} | elapsed={elapsed:.0f}s" + ) + last_status[0] = status + + progress = max(0, min(100, int(progress or 0))) + if pbar is not None and progress > last_progress[0]: + pbar.update(progress - last_progress[0]) + last_progress[0] = progress + + client = NewAPIVeoClient(base_url=get_base_url_by_route(网络线路)) + + result = client.generate_video_sync( + prompt=prompt, + model=模型, + duration=duration_value, + aspect_ratio=宽高比, + resolution=分辨率, + output_dir=output_dir, + negative_prompt=负向提示词, + generate_audio=(生成音频 == "打开"), + image_bytes=image_bytes, + poll_interval=10, + timeout=900, + progress_callback=progress_callback, + ) + + video_path = result["video_path"] + video = VideoFromFile(video_path) + + print( + "NewAPI Veo: completed " + f"| task_id={result['task_id']} | video={video_path}" + ) + + return (video,) + + +NODE_CLASS_MAPPINGS = { + "Google31Video": Google31Video, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Google31Video": "Google 3.1 Video", +} diff --git a/nodes/seedance_video.py b/nodes/seedance_video.py index 420987e..7b09ec9 100644 --- a/nodes/seedance_video.py +++ b/nodes/seedance_video.py @@ -4,7 +4,9 @@ Seedance 视频生成节点 - Seedance: 文生视频 / 图生视频 / 首尾帧生视频（根据图片输入自动切换模式） """ +import base64 import io +import json import os import tempfile @@ -13,7 +15,7 @@ import torch from ..clients.seedance_client import SeedanceClient from ..clients.gemini_client import GeminiAPIClient -from ..utils.image_utils import tensor_to_pil, encode_image_to_base64, pil_to_tensor +from ..utils.image_utils import tensor_to_pil, pil_to_tensor from ..utils.r2_uploader import upload_video, upload_audio from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route @@ -28,6 +30,9 @@ _MODELS = [ _RESOLUTIONS = ["720p", "1080p", "480p"] +_MAX_IMAGE_BYTES = 30 * 1024 * 1024 +_MAX_REQUEST_BODY_BYTES = 64 * 1024 * 1024 + # ── 模型能力判断 ────────────────────────────────────────────────────────────── @@ -38,13 +43,45 @@ def _supports_camera_fixed(model: str) -> bool: # ── 工具函数 ────────────────────────────────────────────────────────────────── -def _tensor_to_base64_url(tensor) -> str: +def _format_mb(size_bytes: int) -> str: + return f"{size_bytes / 1024 / 1024:.2f}MB" + + +def _tensor_to_base64_url(tensor, label: str = "图片") -> str: """ComfyUI IMAGE tensor → data:image/png;base64,xxx""" pil_images = tensor_to_pil(tensor) - b64 = encode_image_to_base64(pil_images[0], format="PNG") + image = pil_images[0] + if image.mode == "RGBA": + image = image.convert("RGB") + + buffered = io.BytesIO() + image.save(buffered, format="PNG") + image_bytes = buffered.getvalue() + image_size = len(image_bytes) + + if image_size > _MAX_IMAGE_BYTES: + raise ValueError( + f"Seedance {label}大小 {_format_mb(image_size)} 超过单张图片 " + f"{_format_mb(_MAX_IMAGE_BYTES)} 限制，请先压缩或缩小图片。" + ) + + b64 = base64.b64encode(image_bytes).decode("utf-8") return f"data:image/png;base64,{b64}" +def _validate_request_body_size(body: dict, tag: str): + body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8")) + if body_size > _MAX_REQUEST_BODY_BYTES: + raise ValueError( + f"{tag} 请求体大小 {_format_mb(body_size)} 超过 " + f"{_format_mb(_MAX_REQUEST_BODY_BYTES)} 限制，请减少参考图片数量或降低图片尺寸。" + ) + print( + f"[{tag}] 请求体大小: {_format_mb(body_size)} " + f"(限制 {_format_mb(_MAX_REQUEST_BODY_BYTES)})" + ) + + async def _url_to_tensor(url: str) -> torch.Tensor: """从 URL 下载图片并转为 ComfyUI IMAGE tensor，失败时返回 None""" @@ -201,7 +238,7 @@ class Seedance: } elif mode == "i2v": - first_url = _tensor_to_base64_url(first_image) + first_url = _tensor_to_base64_url(first_image, "首帧图片") metadata["content"] = [ { "type": "image_url", @@ -218,8 +255,8 @@ class Seedance: } else: # flipflop - first_url = _tensor_to_base64_url(first_image) - last_url = _tensor_to_base64_url(last_image) + first_url = _tensor_to_base64_url(first_image, "首帧图片") + last_url = _tensor_to_base64_url(last_image, "尾帧图片") metadata["content"] = [ { "type": "image_url", @@ -240,6 +277,8 @@ class Seedance: "metadata": metadata, } + _validate_request_body_size(body, tag) + # 保存路径（临时文件，避免与下游保存节点重复落盘） _, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_") @@ -315,6 +354,7 @@ class SeedanceMultiModal: web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开" return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开" seed = _first(kwargs.get("seed"), 0) + network_route = _first(kwargs.get("网络线路"), "全球加速") # 参考图片：INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表，直接保留 raw_images = kwargs.get("参考图片", None) @@ -344,11 +384,11 @@ class SeedanceMultiModal: imgs = ref_images[:9] if len(ref_images) > 9: print(f"[SeedanceMultiModal] 参考图片超过9张，仅取前9张（共{len(ref_images)}张）") - for img_tensor in imgs: + for idx, img_tensor in enumerate(imgs, start=1): # 每个 tensor 可能是 [1,H,W,C] 或 [H,W,C]，统一确保有 batch 维 if img_tensor.dim() == 3: img_tensor = img_tensor.unsqueeze(0) - url = _tensor_to_base64_url(img_tensor) + url = _tensor_to_base64_url(img_tensor, f"参考图片{idx}") content.append({ "type": "image_url", "image_url": {"url": url}, @@ -414,11 +454,13 @@ class SeedanceMultiModal: if first_image_url: body["image"] = first_image_url + _validate_request_body_size(body, "Seedance多模态") + # ── 保存路径（临时文件，避免与下游保存节点重复落盘）────────────────── _, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_") client = SeedanceClient() - client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速")) + client.base_url = get_base_url_by_route(network_route) pbar = _make_pbar() on_stage, on_prog = _make_callbacks("Seedance多模态", pbar) diff --git a/utils/config.py b/utils/config.py index f370efa..f72dab2 100644 --- a/utils/config.py +++ b/utils/config.py @@ -148,4 +148,6 @@ def get_async_api_base_url() -> str: def get_base_url_by_route(route: str) -> str: """根据网络线路选项返回对应域名，未匹配则走 config 垫底""" + if isinstance(route, (list, tuple)): + route = route[0] if route else None return NETWORK_ROUTES.get(route, get_api_base_url()) diff --git a/utils/http_error.py b/utils/http_error.py index dfcb11b..237d7c9 100644 --- a/utils/http_error.py +++ b/utils/http_error.py @@ -9,8 +9,9 @@ """ import asyncio +import json import random -from typing import Optional +from typing import Any, Optional import aiohttp @@ -45,15 +46,66 @@ DEFAULT_MAX_DELAY = 30.0 # 最大等待秒数 DEFAULT_BACKOFF_FACTOR = 2.0 # 指数退避因子 +def _extract_message_from_payload(payload: Any) -> str: + if isinstance(payload, str): + text = payload.strip() + if not text: + return "" + if text.startswith("{") or text.startswith("["): + try: + return _extract_message_from_payload(json.loads(text)) + except Exception: + return text + return text + + if not isinstance(payload, dict): + return "" + + error = payload.get("error") + if isinstance(error, dict): + for key in ("message", "msg", "detail", "reason"): + value = error.get(key) + if value: + return _extract_message_from_payload(value) + elif error: + return _extract_message_from_payload(error) + + for key in ("message", "msg", "detail", "reason", "error_message"): + value = payload.get(key) + if value: + return _extract_message_from_payload(value) + + for key in ("data", "result", "response", "output"): + value = payload.get(key) + nested = _extract_message_from_payload(value) + if nested: + return nested + + return "" + + +def extract_structured_error_message(raw_message: str) -> str: + if not isinstance(raw_message, str): + return "" + text = raw_message.strip() + if not (text.startswith("{") or text.startswith("[")): + return "" + return _extract_message_from_payload(text) + + def get_friendly_message(status_code: int, raw_message: str = "") -> str: """根据状态码/错误内容返回友好文案，未匹配则返回原始信息""" if status_code == 524: return "Gateway timed out while waiting for upstream image generation. Please retry, lower resolution/count, or switch network route." if raw_message: - raw_message_lower = raw_message.lower() + structured_message = extract_structured_error_message(raw_message) + message_for_matching = structured_message or raw_message + raw_message_lower = message_for_matching.lower() for keyword, friendly_msg in ERROR_CONTENT_MESSAGES.items(): if keyword.lower() in raw_message_lower: return friendly_msg + if structured_message: + return structured_message if status_code == 500: return "服务器返回 500：上游生成失败或服务端临时异常。请稍后重试；如果多次出现，请降低分辨率/数量，或调整提示词。" friendly = HTTP_ERROR_MESSAGES.get(status_code) diff --git a/utils/nano_banana_async.py b/utils/nano_banana_async.py new file mode 100644 index 0000000..6bf45a9 --- /dev/null +++ b/utils/nano_banana_async.py @@ -0,0 +1,675 @@ +import asyncio +import base64 +import json +import time +from io import BytesIO +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple + +import aiohttp +from PIL import Image + +from .http_error import ( + DEFAULT_BACKOFF_FACTOR, + DEFAULT_BASE_DELAY, + DEFAULT_MAX_DELAY, + DEFAULT_MAX_RETRIES, + RETRYABLE_STATUS_CODES, + _compute_delay, + extract_structured_error_message, + get_friendly_message, +) +from ..clients.gemini_client import GeminiAPIClient + + +_MAX_BODY_BYTES = 20_000_000 +_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.8) +_SUBMIT_ENDPOINT = "/async/v1/generateImage" +_TASK_ENDPOINT = "/async/v1/tasks/{task_id}" +_POLL_SCHEDULE = [5.0, 20.0] +_POLL_INTERVAL = 3.0 +_MAX_WAIT_SECONDS = 900.0 +_INTERRUPT_STEP = 0.2 +_RUNNING_PROGRESS_MAX = 0.99 +_POLL_LOG_ENABLED = False + +_SUCCESS_STATUSES = {"success", "succeed", "succeeded", "completed", "done", "finished"} +_FAILURE_STATUSES = { + "failure", + "fail", + "failed", + "error", + "expired", + "timeout", + "timed_out", + "cancel", + "canceled", + "cancelled", + "rejected", +} +_RUNNING_STATUSES = { + "submitted", + "queued", + "pending", + "running", + "processing", + "in_progress", + "in-progress", + "created", +} + + +def _headers(api_key: str) -> Dict[str, str]: + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + + +def _json_dumps(body: Dict[str, Any]) -> str: + return json.dumps(body, ensure_ascii=False, separators=(",", ":")) + + +def _json_size(body: Dict[str, Any]) -> int: + return len(_json_dumps(body).encode("utf-8")) + + +def _scale_images(images: List[Image.Image], scale: float) -> List[Image.Image]: + if scale >= 1.0: + return images + scaled = [] + for img in images: + new_w = max(1, int(img.width * scale)) + new_h = max(1, int(img.height * scale)) + scaled.append(img.resize((new_w, new_h), Image.Resampling.LANCZOS)) + return scaled + + +def _encode_image_data_url( + image: Image.Image, + image_format: str, + quality: Optional[int] = None, +) -> str: + buffered = BytesIO() + working = image + fmt = image_format.upper() + save_kwargs = {"format": fmt} + + if fmt == "JPEG": + if working.mode != "RGB": + working = working.convert("RGB") + save_kwargs.update({"quality": quality or 90, "optimize": True, "subsampling": 2}) + mime_type = "image/jpeg" + else: + if working.mode == "RGBA": + working = working.convert("RGB") + mime_type = "image/png" + + working.save(buffered, **save_kwargs) + encoded = base64.b64encode(buffered.getvalue()).decode("ascii") + return f"data:{mime_type};base64,{encoded}" + + +def _encode_image_data_urls( + images: Sequence[Image.Image], + image_format: str, + quality: Optional[int] = None, +) -> List[str]: + return [_encode_image_data_url(img, image_format, quality) for img in images] + + +def _fit_image_data_urls_to_body_limit( + images: Sequence[Image.Image], + build_body: Callable[[List[str]], Dict[str, Any]], +) -> Tuple[List[str], List[Image.Image], str, int]: + working_images = list(images) + image_urls = _encode_image_data_urls(working_images, "PNG") + body_size = _json_size(build_body(image_urls)) + if body_size <= _BODY_TARGET_BYTES: + return image_urls, working_images, "PNG", body_size + + for _ in range(10): + if body_size <= _BODY_TARGET_BYTES: + break + ratio = _BODY_TARGET_BYTES / max(body_size, 1) + scale = min(0.98, ratio ** 0.5) + working_images = _scale_images(working_images, scale) + image_urls = _encode_image_data_urls(working_images, "PNG") + body_size = _json_size(build_body(image_urls)) + + return image_urls, working_images, "PNG", body_size + + +def _shorten_base64_for_log(value: Any, max_len: int = 160) -> Any: + if isinstance(value, dict): + result = {} + for key, item in value.items(): + if key in ("data", "b64_json", "base64", "image_base64") and isinstance(item, str) and len(item) > max_len: + result[key] = f"" + else: + result[key] = _shorten_base64_for_log(item, max_len) + return result + if isinstance(value, list): + return [_shorten_base64_for_log(item, max_len) for item in value] + if isinstance(value, str) and value.startswith("data:image") and len(value) > max_len: + return f"" + return value + + +def _log_body(label: str, text_or_body: Any) -> None: + if isinstance(text_or_body, str): + try: + text_or_body = json.loads(text_or_body) + except Exception: + print(f"{label}\n{text_or_body}") + return + print( + f"{label}\n" + f"{json.dumps(_shorten_base64_for_log(text_or_body), ensure_ascii=False, indent=2)}" + ) + + +def build_nano_banana_submit_body( + model: str, + prompt: str, + resolution: str, + aspect_ratio: str, + images: Optional[List[Image.Image]] = None, + enable_grounding: bool = False, + thinking_level: Optional[str] = None, + request_log_enabled: bool = False, + node_label: str = "Nano Banana", +) -> Dict[str, Any]: + def _make_body(image_urls: List[str]) -> Dict[str, Any]: + body: Dict[str, Any] = { + "model": model, + "prompt": prompt, + "size": resolution, + } + if aspect_ratio and aspect_ratio != "智能": + body["aspect_ratio"] = aspect_ratio + if image_urls: + body["images"] = image_urls + if enable_grounding: + body["google_search"] = True + if thinking_level: + body["thinking_level"] = thinking_level + return body + + working_images = list(images or []) + image_urls: List[str] = [] + + if working_images: + image_urls, working_images, _, _ = _fit_image_data_urls_to_body_limit( + working_images, + _make_body, + ) + + body = _make_body(image_urls) + body_size = _json_size(body) + + if working_images and body_size > _MAX_BODY_BYTES: + raise ValueError( + f"Request body exceeds the 20MB limit after compression " + f"({body_size / 1_000_000:.2f}MB). Reduce reference image count, " + "image complexity, or prompt length." + ) + if not working_images and body_size > _MAX_BODY_BYTES: + raise ValueError( + f"Request body exceeds the 20MB limit ({body_size / 1_000_000:.2f}MB). " + "Shorten the prompt or system instructions." + ) + + if request_log_enabled: + print( + f"[{node_label} 异步请求体] {body_size / 1024:.1f}KB\n" + f"{json.dumps(_shorten_base64_for_log(body), ensure_ascii=False, indent=2)}" + ) + + return body + + +async def _interruptible_sleep( + seconds: float, + check_interrupt: Optional[Callable[[], None]] = None, +) -> None: + elapsed = 0.0 + while elapsed < seconds: + if check_interrupt: + check_interrupt() + delay = min(_INTERRUPT_STEP, seconds - elapsed) + await asyncio.sleep(delay) + elapsed += delay + if check_interrupt: + check_interrupt() + + +def _payload_sources(payload: Dict[str, Any]) -> Iterable[Dict[str, Any]]: + queue = [payload] + seen = set() + while queue: + current = queue.pop(0) + if not isinstance(current, dict): + continue + obj_id = id(current) + if obj_id in seen: + continue + seen.add(obj_id) + yield current + for key in ("data", "result", "response", "output", "task_result", "content"): + value = current.get(key) + if isinstance(value, dict): + queue.append(value) + + +def _extract_task_id(payload: Dict[str, Any]) -> str: + for source in _payload_sources(payload): + for key in ("task_id", "taskId", "id"): + value = source.get(key) + if value: + return str(value) + raise RuntimeError(f"提交响应中未找到 task_id: {payload}") + + +def _extract_status(payload: Dict[str, Any]) -> str: + statuses = [] + for source in _payload_sources(payload): + for key in ("status", "task_status", "state", "task_state"): + value = source.get(key) + if value is not None and str(value).strip(): + statuses.append(str(value).strip()) + + for status in statuses: + normalized = status.lower() + if normalized in _FAILURE_STATUSES or any( + token in normalized for token in ("fail", "error", "reject", "timeout", "cancel") + ): + return status + for status in statuses: + if status.lower() in _RUNNING_STATUSES: + return status + for status in statuses: + if status.lower() in _SUCCESS_STATUSES: + return status + return statuses[0] if statuses else "" + + +def _coerce_progress_fraction(value: Any) -> Optional[float]: + if value is None or isinstance(value, bool): + return None + + if isinstance(value, (int, float)): + progress = float(value) + elif isinstance(value, str): + text = value.strip() + if not text: + return None + has_percent_suffix = text.endswith("%") + if has_percent_suffix: + text = text[:-1].strip() + try: + progress = float(text) + except ValueError: + return None + if has_percent_suffix: + progress /= 100.0 + else: + return None + + if progress > 1.0: + progress /= 100.0 + return max(0.0, min(progress, 1.0)) + + +def _extract_progress(payload: Dict[str, Any]) -> Optional[float]: + for source in _payload_sources(payload): + for key in ("progress", "percentage", "percent"): + progress = _coerce_progress_fraction(source.get(key)) + if progress is not None: + return progress + + for key in ("progressInfo", "progress_info"): + info = source.get(key) + if not isinstance(info, dict): + continue + for field in ("progress", "percentage", "percent"): + progress = _coerce_progress_fraction(info.get(field)) + if progress is not None: + return progress + + return None + + +def _is_failure_status(normalized_status: str) -> bool: + return normalized_status in _FAILURE_STATUSES or any( + token in normalized_status for token in ("fail", "error", "reject", "timeout", "cancel") + ) + + +def _extract_error_message(payload: Dict[str, Any]) -> str: + for source in _payload_sources(payload): + error = source.get("error") + if isinstance(error, dict): + for key in ("message", "msg", "detail", "reason", "code"): + value = error.get(key) + if value: + return str(value) + elif error: + message = extract_structured_error_message(str(error)) + return message or str(error) + + for key in ( + "fail_reason", + "failure_reason", + "task_status_msg", + "status_msg", + "error_message", + "message", + "msg", + "reason", + "detail", + ): + value = source.get(key) + if value: + message = extract_structured_error_message(str(value)) + return message or str(value) + return "未知错误" + + +async def _submit_task( + session: aiohttp.ClientSession, + base_url: str, + api_key: str, + body: Dict[str, Any], + node_label: str, + log_body_enabled: bool = False, +) -> str: + url = f"{base_url}{_SUBMIT_ENDPOINT}" + timeout = aiohttp.ClientTimeout(total=120, connect=30, sock_read=120) + last_status = None + last_text = "" + + for attempt in range(DEFAULT_MAX_RETRIES + 1): + try: + async with session.post( + url, + headers=_headers(api_key), + data=_json_dumps(body).encode("utf-8"), + timeout=timeout, + ) as resp: + text = await resp.text() + if log_body_enabled: + _log_body(f"[{node_label} 异步提交响应] HTTP {resp.status}", text) + if resp.status not in (200, 201, 202): + last_status = resp.status + last_text = text + if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES: + friendly = get_friendly_message(resp.status) + delay = _compute_delay( + attempt, + DEFAULT_BASE_DELAY, + DEFAULT_MAX_DELAY, + DEFAULT_BACKOFF_FACTOR, + ) + print(f"{node_label}: {friendly} {delay:.1f}s 后重试提交 ({attempt + 1}/{DEFAULT_MAX_RETRIES})...") + await asyncio.sleep(delay) + continue + raise RuntimeError(get_friendly_message(resp.status, text)) + + try: + data = json.loads(text) + except Exception: + raise RuntimeError(f"提交响应 JSON 解析失败: {text[:500]}") from None + task_id = _extract_task_id(data) + status = _extract_status(data) or "SUBMITTED" + print(f"{node_label}: 异步任务已提交 | task_id={task_id} | status={status}") + return task_id + except ( + aiohttp.ClientConnectorError, + aiohttp.ClientOSError, + aiohttp.ServerDisconnectedError, + asyncio.TimeoutError, + ) as e: + last_text = str(e) + if attempt < DEFAULT_MAX_RETRIES: + delay = _compute_delay( + attempt, + DEFAULT_BASE_DELAY, + DEFAULT_MAX_DELAY, + DEFAULT_BACKOFF_FACTOR, + ) + print(f"{node_label}: 网络连接失败，{delay:.1f}s 后重试提交 ({attempt + 1}/{DEFAULT_MAX_RETRIES})...") + await asyncio.sleep(delay) + continue + raise RuntimeError( + f"网络连接失败，无法连接 {url}: {str(e)}。请切换节点里的网络线路，或检查 VPN/代理/防火墙。" + ) from None + + raise RuntimeError(get_friendly_message(last_status or 0, last_text)) + + +async def _poll_task( + session: aiohttp.ClientSession, + base_url: str, + api_key: str, + task_id: str, + node_label: str, + check_interrupt: Optional[Callable[[], None]] = None, + log_body_enabled: bool = False, + progress_callback: Optional[Callable[[float], None]] = None, +) -> Dict[str, Any]: + url = f"{base_url}{_TASK_ENDPOINT.format(task_id=task_id)}" + start_time = time.time() + last_poll_at = start_time + poll_count = 0 + + while True: + if check_interrupt: + check_interrupt() + + if poll_count < len(_POLL_SCHEDULE): + next_poll_at = start_time + _POLL_SCHEDULE[poll_count] + else: + next_poll_at = last_poll_at + _POLL_INTERVAL + + sleep_time = next_poll_at - time.time() + if sleep_time > 0: + await _interruptible_sleep(sleep_time, check_interrupt=check_interrupt) + + last_poll_at = time.time() + elapsed = last_poll_at - start_time + if elapsed > _MAX_WAIT_SECONDS: + raise RuntimeError(f"任务 {task_id} 超时（>{int(_MAX_WAIT_SECONDS)}秒），请稍后用 task_id 查询结果") + + poll_count += 1 + async with session.get(url, headers=_headers(api_key)) as resp: + text = await resp.text() + if log_body_enabled: + _log_body(f"[{node_label} 任务查询响应 #{poll_count}] HTTP {resp.status}", text) + if resp.status != 200: + raise RuntimeError(get_friendly_message(resp.status, text)) + try: + payload = json.loads(text) + except Exception: + raise RuntimeError(f"任务查询响应 JSON 解析失败: {text[:500]}") from None + + status = _extract_status(payload) or "UNKNOWN" + normalized = status.lower() + progress = _extract_progress(payload) + is_failure = _is_failure_status(normalized) + if _POLL_LOG_ENABLED: + progress_text = "" + if progress is not None and not is_failure: + displayed_progress = 1.0 if normalized in _SUCCESS_STATUSES else min(progress, _RUNNING_PROGRESS_MAX) + progress_text = f" | progress={displayed_progress * 100:.0f}%" + print(f"{node_label}: 查询任务 #{poll_count} | task_id={task_id} | status={status}{progress_text}") + + if normalized in _SUCCESS_STATUSES: + if progress_callback: + progress_callback(1.0) + return payload + if is_failure: + raise RuntimeError(f"任务失败: {_extract_error_message(payload)}") + if normalized not in _RUNNING_STATUSES: + raise RuntimeError(f"未知任务状态 {status}: {payload}") + if progress_callback and progress is not None: + progress_callback(min(progress, _RUNNING_PROGRESS_MAX)) + + +async def _image_from_url_or_data( + value: str, + session: aiohttp.ClientSession, +) -> Optional[Image.Image]: + if not value: + return None + if value.startswith("data:image"): + try: + _, b64_data = value.split(",", 1) + return Image.open(BytesIO(base64.b64decode(b64_data))).convert("RGB") + except Exception as e: + raise RuntimeError(f"data URL 图片解码失败: {e}") from None + if value.startswith("http"): + async with session.get(value, allow_redirects=True) as resp: + if resp.status != 200: + raise RuntimeError(f"图片下载失败 ({resp.status}): {value}") + img_bytes = await resp.read() + return Image.open(BytesIO(img_bytes)).convert("RGB") + return None + + +async def _parse_direct_images( + payload: Dict[str, Any], + session: aiohttp.ClientSession, +) -> List[Image.Image]: + images: List[Image.Image] = [] + + async def _try_item(item: Any) -> None: + if isinstance(item, str): + img = await _image_from_url_or_data(item, session) + if img: + images.append(img) + return + if not isinstance(item, dict): + return + + for key in ("url", "image_url", "result_url", "download_url"): + img = await _image_from_url_or_data(str(item.get(key) or ""), session) + if img: + images.append(img) + return + + b64_data = item.get("b64_json") or item.get("base64") or item.get("image_base64") + if b64_data: + images.append(Image.open(BytesIO(base64.b64decode(str(b64_data)))).convert("RGB")) + return + + for inline_key in ("inline_data", "inlineData"): + inline = item.get(inline_key) + if isinstance(inline, dict) and inline.get("data"): + images.append(Image.open(BytesIO(base64.b64decode(str(inline["data"])))).convert("RGB")) + return + + for source in _payload_sources(payload): + for key in ("image_url", "result_url", "url", "download_url"): + img = await _image_from_url_or_data(str(source.get(key) or ""), session) + if img: + images.append(img) + + for key in ("images", "output_images", "outputs"): + value = source.get(key) + if isinstance(value, list): + for item in value: + await _try_item(item) + elif value: + await _try_item(value) + + return images + + +async def _parse_task_images( + task_payload: Dict[str, Any], + session: aiohttp.ClientSession, + api_key: str, +) -> List[Image.Image]: + direct_images = await _parse_direct_images(task_payload, session) + if direct_images: + return direct_images + + client = GeminiAPIClient(api_key=api_key) + last_error = None + for source in _payload_sources(task_payload): + if "candidates" not in source: + continue + try: + images, _ = await client.parse_response_async(source, session=session) + if images: + return [img.convert("RGB") for img in images] + except Exception as e: + last_error = e + + if last_error is not None: + raise RuntimeError(str(last_error)) from None + raise RuntimeError(f"任务成功但未找到图片结果: {task_payload}") + + +async def generate_nano_banana_async( + session: aiohttp.ClientSession, + base_url: str, + api_key: str, + prompt: str, + model: str, + resolution: str, + aspect_ratio: str, + images: Optional[List[Image.Image]] = None, + enable_grounding: bool = False, + thinking_level: Optional[str] = None, + node_label: str = "Nano Banana", + request_log_enabled: bool = False, + check_interrupt: Optional[Callable[[], None]] = None, + progress_callback: Optional[Callable[[float], None]] = None, +) -> Tuple[List[Image.Image], Dict[str, Any]]: + if check_interrupt: + check_interrupt() + + body = build_nano_banana_submit_body( + model=model, + prompt=prompt, + resolution=resolution, + aspect_ratio=aspect_ratio, + images=images, + enable_grounding=enable_grounding, + thinking_level=thinking_level, + request_log_enabled=request_log_enabled, + node_label=node_label, + ) + + task_start = time.time() + task_id = await _submit_task( + session, + base_url, + api_key, + body, + node_label, + log_body_enabled=request_log_enabled, + ) + task_payload = await _poll_task( + session, + base_url, + api_key, + task_id, + node_label, + check_interrupt=check_interrupt, + log_body_enabled=request_log_enabled, + progress_callback=progress_callback, + ) + task_done = time.time() + + parse_start = time.time() + images_list = await _parse_task_images(task_payload, session, api_key) + parse_done = time.time() + + return images_list, { + "task_id": task_id, + "task_ms": (task_done - task_start) * 1000, + "parse_ms": (parse_done - parse_start) * 1000, + "request_bytes": _json_size(body), + } diff --git a/web/js/hideSidebar.js b/web/js/hideSidebar.js index 8049f44..fa8cb7d 100644 --- a/web/js/hideSidebar.js +++ b/web/js/hideSidebar.js @@ -4,10 +4,16 @@ app.registerExtension({ name: "o1key.hideSidebarItems", async setup() { const hide = () => { - // 隐藏侧边栏的"说明"、"应用"、"模型"、"节点"按钮 - const hiddenLabels = ["说明", "帮助", "Help", "应用", "Apps", "模型", "Models", "节点", "Nodes"]; + // 隐藏侧边栏的"说明"、"应用"、"模型"、"节点"、"模板"按钮 + const hiddenLabels = ["说明", "帮助", "help", "应用", "apps", "模型", "models", "节点", "nodes", "模板", "templates", "template"]; document.querySelectorAll(".side-bar-button, [class*='sidebar'] button, .p-togglebutton").forEach(btn => { - const label = btn.getAttribute("aria-label") || btn.textContent || ""; + const label = [ + btn.getAttribute("aria-label"), + btn.getAttribute("title"), + btn.getAttribute("data-title"), + btn.getAttribute("data-label"), + btn.textContent, + ].filter(Boolean).join(" ").toLowerCase(); if (hiddenLabels.some(k => label.includes(k))) { btn.style.display = "none"; } @@ -212,4 +218,4 @@ app.registerExtension({ setTimeout(hide, 1000); setTimeout(hide, 3000); }, -}); \ No newline at end of file +}); diff --git a/web/js/notePanel.js b/web/js/notePanel.js index 6a0e58a..1ed52cd 100644 --- a/web/js/notePanel.js +++ b/web/js/notePanel.js @@ -1,7 +1,9 @@ import { app } from "../../../scripts/app.js"; +import { api } from "../../../scripts/api.js"; const STORAGE_KEY = "o1key-notes"; const SEEDED_KEY = "o1key-notes-seeded-v2"; +const NOTES_API = "/o1key/notes"; const STYLE_ID = "o1key-notes-styles"; let notes = []; @@ -207,27 +209,91 @@ function injectStyles() { document.head.appendChild(el); } -function loadNotes() { +function readCachedNotes() { try { const raw = localStorage.getItem(STORAGE_KEY); - notes = raw ? JSON.parse(raw).map(makeNote) : []; + if (raw === null) return { found: false, notes: [] }; + const parsed = JSON.parse(raw); + return { found: true, notes: Array.isArray(parsed) ? parsed.map(makeNote) : [] }; } catch { - notes = []; + return { found: false, notes: [] }; + } +} + +function writeCachedNotes() { + try { + localStorage.setItem(STORAGE_KEY, JSON.stringify(notes)); + } catch {} +} + +function hasSeededNotes() { + try { + return !!localStorage.getItem(SEEDED_KEY); + } catch { + return false; + } +} + +function markSeededNotes() { + try { + localStorage.setItem(SEEDED_KEY, "1"); + } catch {} +} + +function createSeedNotes() { + return SAMPLE_NOTES.map(makeNote); +} + +async function loadNotes() { + const cached = readCachedNotes(); + + try { + const resp = await api.fetchApi(NOTES_API); + if (!resp.ok) throw new Error(`HTTP ${resp.status}`); + const data = await resp.json(); + + if (data.exists) { + notes = Array.isArray(data.notes) ? data.notes.map(makeNote) : []; + writeCachedNotes(); + return; + } + + notes = cached.found ? cached.notes : createSeedNotes(); + markSeededNotes(); + writeCachedNotes(); + await persistNotesToFile(); + return; + } catch (e) { + console.warn("[o1key notes] file storage unavailable, using localStorage", e); } - if (!localStorage.getItem(SEEDED_KEY)) { - const existingTitles = new Set(notes.map(n => n.title)); - const samples = SAMPLE_NOTES.map(makeNote).filter(n => !existingTitles.has(n.title)); - notes = [...samples, ...notes]; - localStorage.setItem(SEEDED_KEY, "1"); - saveNotes(); + notes = cached.found ? cached.notes : []; + if (!cached.found && !hasSeededNotes()) { + notes = createSeedNotes(); + markSeededNotes(); + writeCachedNotes(); } } function saveNotes() { + writeCachedNotes(); + void persistNotesToFile(); +} + +async function persistNotesToFile() { try { - localStorage.setItem(STORAGE_KEY, JSON.stringify(notes)); - } catch {} + const resp = await api.fetchApi(NOTES_API, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ notes }), + }); + if (!resp.ok) throw new Error(`HTTP ${resp.status}`); + return true; + } catch (e) { + console.warn("[o1key notes] failed to save notes file", e); + setPanelStatus("笔记文件保存失败，已保存在浏览器缓存"); + return false; + } } function allTags() { @@ -250,7 +316,7 @@ function filteredNotes() { app.registerExtension({ name: "o1key.notePanel", async setup() { - loadNotes(); + await loadNotes(); app.extensionManager.registerSidebarTab({ id: "o1key-notes", title: "笔记", diff --git a/web/js/restartButton.js b/web/js/restartButton.js index c2b48b8..cf1f5cd 100644 --- a/web/js/restartButton.js +++ b/web/js/restartButton.js @@ -43,6 +43,7 @@ app.registerExtension({ if (!confirm("确定要重启 ComfyUI 吗？")) return; btn.style.opacity = "0.5"; btn.style.pointerEvents = "none"; + await disableExperimentalAssetApi(); try { await fetch("/o1key/restart", { method: "POST" }); } catch {} pollUntilReady(); }); @@ -51,16 +52,70 @@ app.registerExtension({ injected = true; } + async function disableExperimentalAssetApi() { + if (!(await shouldDisableExperimentalAssetApi())) return; + try { + await fetch("/api/settings/Comfy.Assets.UseAssetAPI", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(false), + signal: AbortSignal.timeout(2000), + }); + } catch {} + } + + async function shouldDisableExperimentalAssetApi() { + try { + const r = await fetch("/api/settings/Comfy.Assets.UseAssetAPI", { + cache: "no-store", + signal: AbortSignal.timeout(2000), + }); + if (!r.ok || !(await r.json())) return false; + } catch { + return false; + } + return !(await fetchOk("/api/assets/seed/status", 2000)); + } + + async function fetchOk(url, timeout = 2500) { + try { + const r = await fetch(url, { + cache: "no-store", + signal: AbortSignal.timeout(timeout), + }); + return r.ok; + } catch { + return false; + } + } + + async function comfyReady() { + const [statsOk, modelFoldersOk] = await Promise.all([ + fetchOk("/api/system_stats"), + fetchOk("/api/experiment/models"), + ]); + return statsOk && modelFoldersOk; + } + function pollUntilReady() { let attempts = 0; - const maxAttempts = 40; + const maxAttempts = 80; + const minRestartWaitMs = 5000; + const startedAt = Date.now(); + let sawUnavailable = false; const interval = setInterval(async () => { attempts++; if (attempts > maxAttempts) { clearInterval(interval); forceReload(); return; } - try { - const r = await fetch("/api/system_stats", { signal: AbortSignal.timeout(2000) }); - if (r.ok) { clearInterval(interval); forceReload(); } - } catch {} + const ready = await comfyReady(); + if (!ready) { + sawUnavailable = true; + return; + } + if (!sawUnavailable && Date.now() - startedAt < minRestartWaitMs) return; + + clearInterval(interval); + await disableExperimentalAssetApi(); + setTimeout(forceReload, 800); }, 1500); }