Add Grok and VEO video workflow support

This commit is contained in:
o1key
2026-06-04 16:06:23 +08:00
parent 5d9aff9ca7
commit 30e9603f77
24 changed files with 3752 additions and 528 deletions
+125 -1
View File
@@ -12,6 +12,7 @@ Comfyui_o1key - ComfyUI 自定义节点集合
import ssl import ssl
import logging import logging
import asyncio
# 屏蔽 ComfyUI 资产扫描的终端日志输出 # 屏蔽 ComfyUI 资产扫描的终端日志输出
_seeder_filter = lambda record: not any( _seeder_filter = lambda record: not any(
@@ -20,7 +21,57 @@ _seeder_filter = lambda record: not any(
) )
logging.getLogger().addFilter(_seeder_filter) 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 K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch, SaveImageFormat
from .nodes import O1keySavePSD from .nodes import O1keySavePSD
from .nodes import O1keyRemoveBackground from .nodes import O1keyRemoveBackground
@@ -73,6 +124,7 @@ NODE_CLASS_MAPPINGS = {
"BatchCleanMetadata": BatchCleanMetadata, "BatchCleanMetadata": BatchCleanMetadata,
"VideoPreview": VideoPreview, "VideoPreview": VideoPreview,
"GoogleVeo": GoogleVeo, "GoogleVeo": GoogleVeo,
"Google31Video": Google31Video,
"FluxImageEdit": FluxImageEdit, "FluxImageEdit": FluxImageEdit,
"UniversalLLMChat": UniversalLLMChat, "UniversalLLMChat": UniversalLLMChat,
"KlingVideo": KlingVideo, "KlingVideo": KlingVideo,
@@ -88,6 +140,7 @@ NODE_CLASS_MAPPINGS = {
"O1keyGPTImage": O1keyGPTImage, "O1keyGPTImage": O1keyGPTImage,
"O1keyGPTImageBatch": O1keyGPTImageBatch, "O1keyGPTImageBatch": O1keyGPTImageBatch,
"O1keyGrokImage": O1keyGrokImage, "O1keyGrokImage": O1keyGrokImage,
"O1keyGrokVideo": O1keyGrokVideo,
"KVideoFirstLast": KVideoFirstLast, "KVideoFirstLast": KVideoFirstLast,
"KVideoImage2Video": KVideoImage2Video, "KVideoImage2Video": KVideoImage2Video,
"K3Video": K3Video, "K3Video": K3Video,
@@ -113,6 +166,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"BatchCleanMetadata": "批量任务(防AI识别)", "BatchCleanMetadata": "批量任务(防AI识别)",
"VideoPreview": "预览视频", "VideoPreview": "预览视频",
"GoogleVeo": "Google Veo - ab", "GoogleVeo": "Google Veo - ab",
"Google31Video": "Google 3.1 Video",
"FluxImageEdit": "Flux2 图像编辑", "FluxImageEdit": "Flux2 图像编辑",
"UniversalLLMChat": "全能LLM对话助手", "UniversalLLMChat": "全能LLM对话助手",
"KlingVideo": "文/图生视频 自研模型", "KlingVideo": "文/图生视频 自研模型",
@@ -128,6 +182,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"O1keyGPTImage": "o1key GPT Image", "O1keyGPTImage": "o1key GPT Image",
"O1keyGPTImageBatch": "o1key GPT Image(批量)", "O1keyGPTImageBatch": "o1key GPT Image(批量)",
"O1keyGrokImage": "Grok Image", "O1keyGrokImage": "Grok Image",
"O1keyGrokVideo": "Grok Video",
"KVideoFirstLast": "K26 图生视频(首尾帧)", "KVideoFirstLast": "K26 图生视频(首尾帧)",
"KVideoImage2Video": "K26 图生视频", "KVideoImage2Video": "K26 图生视频",
"K3Video": "K3 图生视频 自研", "K3Video": "K3 图生视频 自研",
@@ -183,6 +238,75 @@ try:
f".o1key_history_{_get_o1key_server_port()}.json", 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") @PromptServer.instance.routes.get("/o1key/input_dir")
async def get_input_dir(request): async def get_input_dir(request):
import os import os
+3 -1
View File
@@ -9,6 +9,8 @@ from .gemini_flash_client import GeminiFlashClient
from .sora_client import SoraClient from .sora_client import SoraClient
from .kling_client import KlingClient from .kling_client import KlingClient
from .veo_client import VeoClient from .veo_client import VeoClient
from .newapi_veo_client import NewAPIVeoClient
from .grok_video_client import GrokVideoClient
from .openai_client import OpenAIAPIClient from .openai_client import OpenAIAPIClient
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient'] __all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient']
+2 -22
View File
@@ -40,7 +40,8 @@ class BaseAPIClient(ABC):
Args: Args:
base_url: API 基础 URL base_url: API 基础 URL
api_key: API 密钥 api_key: API 密钥
max_request_size: 最大请求体大小(字节),默认 100MB max_request_size: 兼容参数;基类不再用它限制 JSON 请求体,
部分子类仍用它作为上传文件大小限制
""" """
self.base_url = base_url self.base_url = base_url
self.api_key = api_key self.api_key = api_key
@@ -116,24 +117,6 @@ class BaseAPIClient(ABC):
"Content-Type": "application/json" "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]: def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
""" """
子类可重写:为指定 HTTP 状态码返回自定义错误文案。 子类可重写:为指定 HTTP 状态码返回自定义错误文案。
@@ -185,9 +168,6 @@ class BaseAPIClient(ABC):
url = f"{self.base_url}{endpoint}" url = f"{self.base_url}{endpoint}"
headers = self.get_headers(use_bearer_token) headers = self.get_headers(use_bearer_token)
# 检查请求大小
self.check_request_size(request_body)
close_session = False close_session = False
if session is None: if session is None:
session = self._make_session() session = self._make_session()
+53 -9
View File
@@ -72,7 +72,9 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
resolution=resolution, resolution=resolution,
enable_grounding=kwargs.get("enable_grounding", False), enable_grounding=kwargs.get("enable_grounding", False),
enable_image_search=kwargs.get("enable_image_search", 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: def extract_task_id(self, response: dict) -> str:
@@ -85,6 +87,23 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
return response.get("status", "UNKNOWN") return response.get("status", "UNKNOWN")
async def parse_result(self, result_data: dict, session) -> List[Image.Image]: 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
image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else "" image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else ""
if image_url: if image_url:
@@ -124,19 +143,44 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
def extract_progress(self, response: dict) -> Optional[float]: def extract_progress(self, response: dict) -> Optional[float]:
"""从轮询响应中提取进度(0.0-1.0""" """从轮询响应中提取进度(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 # 直接字段:progress / percentage
for field in ("progress", "percentage"): for field in ("progress", "percentage", "percent"):
val = response.get(field) progress = _coerce(response.get(field))
if val is not None and isinstance(val, (int, float)): if progress is not None:
return val / 100.0 if val > 1 else float(val) return progress
# 嵌套字段:progressInfo / progress_info # 嵌套字段:progressInfo / progress_info
progress_info = response.get("progressInfo") or response.get("progress_info") progress_info = response.get("progressInfo") or response.get("progress_info")
if isinstance(progress_info, dict): if isinstance(progress_info, dict):
for field in ("progress", "percentage"): for field in ("progress", "percentage", "percent"):
val = progress_info.get(field) progress = _coerce(progress_info.get(field))
if val is not None and isinstance(val, (int, float)): if progress is not None:
return val / 100.0 if val > 1 else float(val) return progress
return None return None
+103 -80
View File
@@ -34,8 +34,7 @@ class GeminiAPIClient(BaseAPIClient):
super().__init__( super().__init__(
base_url=get_api_base_url(), base_url=get_api_base_url(),
api_key=api_key, api_key=api_key
max_request_size=100 * 1024 * 1024
) )
@staticmethod @staticmethod
@@ -188,6 +187,7 @@ class GeminiAPIClient(BaseAPIClient):
enable_image_search: bool = False, enable_image_search: bool = False,
image_compression: str = None, image_compression: str = None,
thinking_level: str = None, thinking_level: str = None,
request_log_enabled: bool = True,
**kwargs **kwargs
) -> Dict[str, Any]: ) -> Dict[str, Any]:
""" """
@@ -204,79 +204,26 @@ class GeminiAPIClient(BaseAPIClient):
Returns: Returns:
请求体字典 请求体字典
""" """
import json
_MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB _MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB
_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.9)
parts = [] parts = []
# 添加文本部分 # 添加文本部分
parts.append({"text": prompt}) parts.append({"text": prompt})
# 添加图像部分(如果有)
if images:
working_images = list(images)
# 编码一次,估算大小,超限则迭代缩放
for _attempt in range(10):
img_parts = []
for img in working_images:
img_base64 = encode_image_to_base64(img)
img_parts.append({
"inline_data": {
"mime_type": "image/png",
"data": img_base64
}
})
# 估算完整 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
}
}
)
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)
# 构建请求体
image_config = {"imageSize": resolution} image_config = {"imageSize": resolution}
if aspect_ratio and aspect_ratio != "智能": if aspect_ratio and aspect_ratio != "智能":
image_config["aspectRatio"] = aspect_ratio image_config["aspectRatio"] = aspect_ratio
request_body = { def _build_request_body(body_parts: List[dict]) -> Dict[str, Any]:
body = {
"contents": [ "contents": [
{ {
"role": "user", "role": "user",
"parts": parts "parts": body_parts
} }
], ],
"generationConfig": { "generationConfig": {
@@ -287,33 +234,110 @@ class GeminiAPIClient(BaseAPIClient):
# 添加思考深度配置 # 添加思考深度配置
if thinking_level: if thinking_level:
request_body["generationConfig"]["thinkingConfig"] = { body["generationConfig"]["thinkingConfig"] = {
"thinkingLevel": thinking_level, "thinkingLevel": thinking_level,
"includeThoughts": True "includeThoughts": True
} }
# 添加图片压缩参数 # 添加图片压缩参数
if image_compression: if image_compression:
request_body["image_compression"] = image_compression body["image_compression"] = image_compression
# 添加 Google Search Grounding 工具(如果启用) # 添加 Google Search Grounding(如果启用)
# 注意:enable_image_search=True 时会自动隐含 enable_grounding # 新异步接口要求直接放在请求体顶层:{"google_search": true}
if enable_grounding or enable_image_search: if enable_grounding or enable_image_search:
if enable_image_search: body["google_search"] = True
# 同时启用网页搜索和图片搜索(仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持)
request_body["tools"] = [ return body
{
"google_search": { def _request_size(body: Dict[str, Any]) -> int:
"searchTypes": { return len(json.dumps(body).encode("utf-8"))
"webSearch": {},
"imageSearch": {} 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"<base64 data, {len(value)} chars>"
else: else:
# 仅启用网页搜索(通用) result[key] = _shorten_base64_for_log(value, max_len)
request_body["tools"] = [{"google_search": {}}] 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:
def _build_image_parts(src_images: List[Image.Image]) -> List[dict]:
img_parts = []
for img in src_images:
img_base64 = encode_image_to_base64(img)
img_parts.append({
"inline_data": {
"mime_type": "image/png",
"data": img_base64
}
})
return img_parts
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"
)
parts.extend(img_parts)
request_body = _build_request_body(parts)
if request_log_enabled and not original_request_logged:
_log_original_request_body(request_body)
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 return request_body
@@ -1106,4 +1130,3 @@ class GeminiAPIClient(BaseAPIClient):
) )
return self.run_async_in_thread(coro) return self.run_async_in_thread(coro)
+636 -10
View File
@@ -1,13 +1,17 @@
""" """
GPT Image API 客户端 GPT Image API 客户端
支持两个接口: 新版节点请求走异步任务接口:
- POST /v1/images/generations/ 文生图 / 图生图(gpt-image-1 / gpt-image-1.5 - POST /async/v1/generateImage
- POST /v1/images/edits/ 图像编辑(带蒙版 inpainting - GET /async/v1/tasks/{task_id}
旧同步接口保留兼容代码,但 GPT Image 节点不再使用:
- POST /v1/images/generations/
- POST /v1/images/edits
设计原则: 设计原则:
- 与 doubao_image_client.py 保持相同的异步 + 同步入口模式 - 对 ComfyUI 节点暴露同步入口,内部提交异步任务并轮询
- generations / edits 接口均使用 multipart/form-data - 图片和蒙版以 data:image/png;base64,... 放入 JSON 请求体
- 响应兼容 SSE 流式、JSON、url 和 b64_json - 响应优先读取 data.images[].url
""" """
import asyncio import asyncio
@@ -16,7 +20,7 @@ import json
import time import time
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from io import BytesIO from io import BytesIO
from typing import List, Optional from typing import Any, Callable, Dict, List, Optional
import aiohttp import aiohttp
import numpy as np import numpy as np
@@ -38,6 +42,8 @@ except ImportError:
# ── 接口端点 ────────────────────────────────────────────────────────────────── # ── 接口端点 ──────────────────────────────────────────────────────────────────
_ENDPOINT_GENERATIONS = "/v1/images/generations/" _ENDPOINT_GENERATIONS = "/v1/images/generations/"
_ENDPOINT_EDITS = "/v1/images/edits" _ENDPOINT_EDITS = "/v1/images/edits"
_ENDPOINT_ASYNC_GENERATE = "/async/v1/generateImage"
_ENDPOINT_ASYNC_TASK = "/async/v1/tasks/{task_id}"
# ── 模型名映射(UI 显示名 → API 实际参数名)───────────────────────────────── # ── 模型名映射(UI 显示名 → API 实际参数名)─────────────────────────────────
_MODEL_NAME_MAP = { _MODEL_NAME_MAP = {
@@ -47,6 +53,49 @@ _MODEL_NAME_MAP = {
# ── 超时 ────────────────────────────────────────────────────────────────────── # ── 超时 ──────────────────────────────────────────────────────────────────────
_REQUEST_TIMEOUT = 900 # 秒 _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: class GptImageClient:
@@ -54,10 +103,10 @@ class GptImageClient:
GPT Image API 客户端 GPT Image API 客户端
接口说明: 接口说明:
generationsmultipart/form-data,支持 quality / size / n / model async generateImageJSON 提交,返回 task_id
editsmultipart/form-data,图片和 mask 使用 PNG 文件上传 tasks/{task_id}:轮询任务状态,成功后读取 data.images[].url
响应支持 JSON 和 SSE 流式格式 旧 generations / edits 同步接口保留为兼容代码
""" """
def __init__(self): def __init__(self):
@@ -132,6 +181,226 @@ class GptImageClient:
) )
return png_bytes 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},<base64 data, {len(data)} chars>"
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}KBscale={best_scale:.3f}"
)
return best_bytes
@staticmethod @staticmethod
def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes: def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes:
""" """
@@ -500,6 +769,316 @@ class GptImageClient:
return request_task.result() 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 接口)─────────────────────────────────── # ── 文生图 / 图生图(generations 接口)───────────────────────────────────
async def _generate_async( async def _generate_async(
@@ -806,6 +1385,53 @@ class GptImageClient:
f"o1key GPT Image 请求超时(>{_REQUEST_TIMEOUT}s),请检查网络或稍后重试" 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: async def _query_balance_async(self) -> dict:
+483
View File
@@ -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"<data-url chars={len(item)}>" 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 Videocontent 下载重试 {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())
+542
View File
@@ -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"<PNG bytes: {len(image_bytes)}>"
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())
+6 -3
View File
@@ -110,10 +110,13 @@ class KVideoFirstLast:
"mode": mode_api, "mode": mode_api,
"duration": 时长, "duration": 时长,
} }
if 生成音频 == "打开": metadata = {}
body["generate_audio"] = True
if 尾帧 is not None: 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_str = json.dumps(body, ensure_ascii=False)
body_size = len(body_str.encode("utf-8")) body_size = len(body_str.encode("utf-8"))
+4 -1
View File
@@ -82,6 +82,9 @@ class KVideoImage2Video:
CATEGORY = "comfyui_o1key/KVideo" CATEGORY = "comfyui_o1key/KVideo"
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", seed=0): async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", seed=0):
if 模式 == "720p" and 生成音频 == "打开":
raise RuntimeError("K26 仅1080p支持音频,请将模式切换为1080p或关闭生成音频。")
api_key = get_api_key_or_raise() api_key = get_api_key_or_raise()
base_url = get_base_url_by_route(网络线路) base_url = get_base_url_by_route(网络线路)
headers = { headers = {
@@ -109,7 +112,7 @@ class KVideoImage2Video:
"duration": 时长, "duration": 时长,
} }
if 生成音频 == "打开": if 生成音频 == "打开":
body["generate_audio"] = True body["metadata"] = {"sound": "on"}
body_str = json.dumps(body, ensure_ascii=False) body_str = json.dumps(body, ensure_ascii=False)
body_size = len(body_str.encode("utf-8")) body_size = len(body_str.encode("utf-8"))
+3 -1
View File
@@ -14,6 +14,7 @@ from .remove_metadata import BatchCleanMetadata
from .video_preview import VideoPreview from .video_preview import VideoPreview
from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset
from .veo_video import GoogleVeo from .veo_video import GoogleVeo
from .newapi_veo_video import Google31Video
from .flux_edit import FluxImageEdit from .flux_edit import FluxImageEdit
from .universal_llm import UniversalLLMChat from .universal_llm import UniversalLLMChat
from .batch_images_o1key import BatchImagesO1key from .batch_images_o1key import BatchImagesO1key
@@ -22,6 +23,7 @@ from .nano_banana_v2 import NanoBananaV2, NanoBananaV2Batch, AsyncImageGenerator
from .doubao_image import DoubaoImage from .doubao_image import DoubaoImage
from .gpt_image import O1keyGPTImage, O1keyGPTImageBatch from .gpt_image import O1keyGPTImage, O1keyGPTImageBatch
from .grok_image import O1keyGrokImage from .grok_image import O1keyGrokImage
from .grok_video import O1keyGrokVideo
from .K_video_firstlast import KVideoFirstLast from .K_video_firstlast import KVideoFirstLast
from .K_video_image2video import KVideoImage2Video from .K_video_image2video import KVideoImage2Video
from .K3_video import K3Video from .K3_video import K3Video
@@ -33,4 +35,4 @@ from .remove_bg import O1keyRemoveBackground
from .color_remove_bg import O1keyColorRemoveBG from .color_remove_bg import O1keyColorRemoveBG
from .grid_splitter import O1keyGridSplitter 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']
+172 -145
View File
@@ -4,23 +4,21 @@ ComfyUI 自定义节点,用于批量处理图像生成任务
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存 支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
""" """
import io as _io
import re
import json
import time import time
import math import math
import base64
import random import random
import asyncio import asyncio
import aiohttp import aiohttp
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Tuple, List from typing import Callable, Optional, Tuple, List
from PIL import Image from PIL import Image
import torch import torch
import numpy as np 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 ( from ..utils.file_utils import (
ImageInfo, ImageInfo,
load_images_from_folder, load_images_from_folder,
@@ -30,7 +28,7 @@ from ..utils.file_utils import (
save_image, save_image,
) )
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route, get_api_key_or_raise 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 ..clients.gemini_client import GeminiAPIClient
from ..models_config import ( from ..models_config import (
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios, get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
@@ -73,64 +71,28 @@ REQUEST_LOG_ENABLED = False
# ============================================================================ # ============================================================================
_NODE = "Nano Banana" _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: def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
return { if pbar is None:
"Authorization": f"Bearer {api_key}", return None
"Content-Type": "application/json",
"X-Accel-Buffering": "no", last_progress = [0.0]
"Cache-Control": "no-cache, no-transform",
} 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( async def _generate_single_async(
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(
session: aiohttp.ClientSession, session: aiohttp.ClientSession,
base_url: str, base_url: str,
api_key: str, api_key: str,
@@ -140,82 +102,25 @@ async def _generate_single_openai(
aspect_ratio: str, aspect_ratio: str,
images: Optional[List[Image.Image]] = None, images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False, enable_grounding: bool = False,
progress_callback: Optional[Callable[[float], None]] = None,
thinking_level: Optional[str] = None,
) -> List[Image.Image]: ) -> List[Image.Image]:
url = f"{base_url}{_ENDPOINT}" result_images, _ = await generate_nano_banana_async(
headers = _get_headers(api_key) session=session,
body = _build_request_body( base_url=base_url,
api_key=api_key,
prompt=prompt, prompt=prompt,
model=model, model=model,
aspect_ratio=aspect_ratio,
resolution=resolution, resolution=resolution,
aspect_ratio=aspect_ratio,
images=images, images=images,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
thinking_level=thinking_level,
node_label="BatchNanoBananaPro",
request_log_enabled=REQUEST_LOG_ENABLED,
progress_callback=progress_callback,
) )
return result_images
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]
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor: 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) return pil_to_tensor(matched)
class BatchNanoBananaPro: class BatchNanoBananaPro(io.ComfyNode):
""" """
批量 Nano Banana 节点 批量 Nano Banana 节点
@@ -291,6 +196,116 @@ class BatchNanoBananaPro:
def __init__(self): def __init__(self):
pass 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( def resize_to_megapixels(
self, self,
image: Image.Image, image: Image.Image,
@@ -352,7 +367,6 @@ class BatchNanoBananaPro:
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, { optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
"default": "不配对" "default": "不配对"
}) })
return { return {
"required": { "required": {
"prompt": ("STRING", { "prompt": ("STRING", {
@@ -532,9 +546,11 @@ class BatchNanoBananaPro:
enable_grounding: bool = False, enable_grounding: bool = False,
base_filename: str = None, base_filename: str = None,
image_format: str = "原始", image_format: str = "原始",
progress_callback: Optional[Callable[[float], None]] = None,
thinking_level: Optional[str] = None,
) -> dict: ) -> dict:
""" """
执行单个生成任务(OpenAI 兼容接口) 执行单个生成任务(异步生图接口)
""" """
result = { result = {
"task_index": task_index, "task_index": task_index,
@@ -550,10 +566,10 @@ class BatchNanoBananaPro:
# 准备输入图片 # 准备输入图片
input_pil_images = [info.image for info in images] input_pil_images = [info.image for info in images]
# 调用 OpenAI 兼容接口生成图片 # 调用异步生图接口生成图片
generated_images = [] generated_images = []
try: try:
gen_images = await _generate_single_openai( gen_images = await _generate_single_async(
session=session, session=session,
base_url=base_url, base_url=base_url,
api_key=api_key, api_key=api_key,
@@ -563,6 +579,8 @@ class BatchNanoBananaPro:
aspect_ratio=aspect_ratio, aspect_ratio=aspect_ratio,
images=input_pil_images if input_pil_images else None, images=input_pil_images if input_pil_images else None,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
progress_callback=progress_callback,
thinking_level=thinking_level,
) )
generated_images.extend(gen_images) generated_images.extend(gen_images)
except Exception as e: except Exception as e:
@@ -661,9 +679,10 @@ class BatchNanoBananaPro:
prompts_per_task: Optional[List[str]] = None, prompts_per_task: Optional[List[str]] = None,
enable_grounding: bool = False, enable_grounding: bool = False,
image_format: str = "原始", image_format: str = "原始",
thinking_level: Optional[str] = None,
) -> List[dict]: ) -> List[dict]:
""" """
异步批量处理所有任务(OpenAI 兼容接口) 异步批量处理所有任务(异步生图接口)
""" """
total_tasks = len(pairs) total_tasks = len(pairs)
@@ -734,6 +753,8 @@ class BatchNanoBananaPro:
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
base_filename=base_filename, base_filename=base_filename,
image_format=image_format, image_format=image_format,
progress_callback=_make_progress_callback(pbar),
thinking_level=thinking_level,
) )
) )
tasks.append(task) tasks.append(task)
@@ -773,13 +794,12 @@ class BatchNanoBananaPro:
else: else:
fail_count += 1 fail_count += 1
# 更新 ComfyUI 原生进度条
if pbar is not None:
pbar.update(1)
# 大任务额外显示百分比里程碑 # 大任务额外显示百分比里程碑
if show_milestone and milestone_index < len(milestones): 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]: if progress >= milestones[milestone_index]:
percentage = int(milestones[milestone_index] * 100) percentage = int(milestones[milestone_index] * 100)
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<") print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
@@ -856,10 +876,15 @@ class BatchNanoBananaPro:
start_time = time.time() start_time = time.time()
# 从 kwargs 提取搜索参数(界面显示为「关闭/打开」,转为 bool 供调用) # 从 kwargs 提取搜索参数(界面显示为「关闭/打开」,转为 bool 供调用)
enable_grounding: bool = False enable_grounding: bool = kwargs.get("谷歌搜索", "关闭") == "打开"
# 拼接实际模型 ID # 拼接实际模型 ID
base_model_id = self.MODEL_ID_MAP.get(模型, "nano-banana-pro") 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 base_model_id == "nano-banana":
if 计费 == "官方": if 计费 == "官方":
raise ValueError(f"模型 \"{模型}\" 仅支持特价计费") raise ValueError(f"模型 \"{模型}\" 仅支持特价计费")
@@ -961,11 +986,12 @@ class BatchNanoBananaPro:
grounding_str = "" grounding_str = ""
if enable_grounding: if enable_grounding:
grounding_str = " | 谷歌搜索接地" grounding_str = " | 谷歌搜索接地"
thinking_str = f" | 思考:{thinking_level}" if thinking_level else ""
if batch_prompts: 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: else:
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}") print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
# 创建 ComfyUI 原生进度条 # 创建 ComfyUI 原生进度条
pbar = None pbar = None
@@ -1025,6 +1051,7 @@ class BatchNanoBananaPro:
prompts_per_task=prompts_per_task, prompts_per_task=prompts_per_task,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
image_format=图片格式, image_format=图片格式,
thinking_level=thinking_level,
) )
) )
except Exception as e: except Exception as e:
@@ -1146,7 +1173,7 @@ class BatchNanoBananaPro:
raise RuntimeError(str(e)) from None raise RuntimeError(str(e)) from None
except Exception as e: except Exception as e:
raise type(e)(str(e)) from None raise RuntimeError(str(e)) from None
finally: finally:
# 查询余额 # 查询余额
+82 -7
View File
@@ -42,6 +42,56 @@ except ImportError:
_FOLDER_PATHS_AVAILABLE = False _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: class O1keyGPTImage:
""" """
o1key GPT Image 节点 o1key GPT Image 节点
@@ -123,6 +173,10 @@ class O1keyGPTImage:
"default": "自动", "default": "自动",
"tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto", "tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto",
}) })
optional_inputs["输出格式"] = (["png", "jpeg", "webp"], {
"default": "jpeg",
"tooltip": "Generated image output format",
})
optional_inputs["seed"] = ("INT", { optional_inputs["seed"] = ("INT", {
"default": 0, "default": 0,
"min": 0, "min": 0,
@@ -160,6 +214,7 @@ class O1keyGPTImage:
网络: str = "全球加速", 网络: str = "全球加速",
分辨率: str = "auto", 分辨率: str = "auto",
质量: str = "自动", 质量: str = "自动",
输出格式: str = "jpeg",
生图数量: int = 1, 生图数量: int = 1,
seed: int = 0, seed: int = 0,
遮罩=None, 遮罩=None,
@@ -190,7 +245,7 @@ class O1keyGPTImage:
raise ValueError("提供了遮罩但未提供图片,请同时提供图片和遮罩") raise ValueError("提供了遮罩但未提供图片,请同时提供图片和遮罩")
# ── 2. 解析分辨率显示值 → API 参数值 ────────────────────────────────── # ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
size = "auto" if 分辨率 == "智能" else 分辨率.split("")[0].strip() size = _resolve_async_size(分辨率)
# ── 2b. 解析模型显示值 → API 参数值 ─────────────────────────────────── # ── 2b. 解析模型显示值 → API 参数值 ───────────────────────────────────
_model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"} _model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"}
@@ -216,6 +271,8 @@ class O1keyGPTImage:
# ── 5. 调用 API ─────────────────────────────────────────────────── # ── 5. 调用 API ───────────────────────────────────────────────────
all_pil_images = [] 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: if batch_prompts:
# 批量模式:逐条提示词调用 # 批量模式:逐条提示词调用
@@ -226,7 +283,7 @@ class O1keyGPTImage:
print("[o1key GPT Image] 用户取消,已中断批量生成") print("[o1key GPT Image] 用户取消,已中断批量生成")
raise InterruptProcessingException() raise InterruptProcessingException()
try: try:
pil_images = client.run_sync( pil_images = client.generate_image_async_sync(
prompt=p, prompt=p,
model=model, model=model,
quality=quality, quality=quality,
@@ -235,6 +292,8 @@ class O1keyGPTImage:
seed=seed, seed=seed,
image_tensor=图片, image_tensor=图片,
mask_tensor=遮罩, mask_tensor=遮罩,
output_format=输出格式,
progress_callback=_make_node_progress_callback(progress_bar, idx, total),
) )
all_pil_images.extend(pil_images) all_pil_images.extend(pil_images)
snippet = p[:30] + ("..." if len(p) >= 30 else "") snippet = p[:30] + ("..." if len(p) >= 30 else "")
@@ -245,12 +304,14 @@ class O1keyGPTImage:
error_msg = str(e).split('\n')[0] error_msg = str(e).split('\n')[0]
snippet = p[:30] + ("..." if len(p) >= 30 else "") snippet = p[:30] + ("..." if len(p) >= 30 else "")
print(f"[o1key GPT Image] [{idx}/{total}] ❌ {snippet}{error_msg}") 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: else:
# 单提示词模式 # 单提示词模式
if not prompt or not prompt.strip(): if not prompt or not prompt.strip():
raise ValueError("提示词不能为空") raise ValueError("提示词不能为空")
try: try:
pil_images = client.run_sync( pil_images = client.generate_image_async_sync(
prompt=prompt, prompt=prompt,
model=model, model=model,
quality=quality, quality=quality,
@@ -259,6 +320,8 @@ class O1keyGPTImage:
seed=seed, seed=seed,
image_tensor=图片, image_tensor=图片,
mask_tensor=遮罩, mask_tensor=遮罩,
output_format=输出格式,
progress_callback=_make_node_progress_callback(progress_bar, 1, 1),
) )
all_pil_images.extend(pil_images) all_pil_images.extend(pil_images)
except InterruptProcessingException: except InterruptProcessingException:
@@ -495,7 +558,7 @@ class O1keyGPTImageBatch:
@staticmethod @staticmethod
def _resolve_size(分辨率: str) -> str: def _resolve_size(分辨率: str) -> str:
return "auto" if 分辨率 == "智能" else 分辨率.split("")[0].strip() return _resolve_async_size(分辨率)
@staticmethod @staticmethod
def _resolve_model(模型: str) -> str: def _resolve_model(模型: str) -> str:
@@ -510,6 +573,15 @@ class O1keyGPTImageBatch:
quality_map = {"": "high", "": "medium", "": "low", "自动": "auto"} quality_map = {"": "high", "": "medium", "": "low", "自动": "auto"}
return quality_map.get(质量, "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 @staticmethod
def _ensure_output_folder(保存路径: str) -> str: def _ensure_output_folder(保存路径: str) -> str:
output_folder = (保存路径 or "").strip() output_folder = (保存路径 or "").strip()
@@ -637,11 +709,12 @@ class O1keyGPTImageBatch:
size = self._resolve_size(分辨率) size = self._resolve_size(分辨率)
model = self._resolve_model(模型) model = self._resolve_model(模型)
quality = self._resolve_quality(质量) quality = self._resolve_quality(质量)
output_format = self._resolve_output_format(图片格式)
client = GptImageClient() client = GptImageClient()
client.base_url = get_base_url_by_route(网络) 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 = [] results = []
all_saved_files = [] all_saved_files = []
@@ -661,7 +734,7 @@ class O1keyGPTImageBatch:
} }
try: try:
pil_images = client.run_sync( pil_images = client.generate_image_async_sync(
prompt=task_prompt, prompt=task_prompt,
model=model, model=model,
quality=quality, quality=quality,
@@ -670,6 +743,8 @@ class O1keyGPTImageBatch:
seed=seed, seed=seed,
image_tensor=self._pair_to_tensors(pair), image_tensor=self._pair_to_tensors(pair),
mask_tensor=遮罩, mask_tensor=遮罩,
output_format=output_format,
progress_callback=_make_node_progress_callback(progress_bar, task_index, total_tasks),
) )
saved_files = self._save_images( saved_files = self._save_images(
images=pil_images, images=pil_images,
@@ -691,7 +766,7 @@ class O1keyGPTImageBatch:
results.append(result) results.append(result)
if progress_bar is not None: 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)) 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) total_generated = sum(result.get("generated_count", 0) for result in results)
+283
View File
@@ -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",
}
+46 -186
View File
@@ -1,20 +1,16 @@
""" """
Nano Banana 节点 (V3) Nano Banana 节点 (V3)
ComfyUI 自定义节点,用于调用生图模型OpenAI 兼容接口) ComfyUI 自定义节点,用于调用异步生图模型
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动 使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
""" """
import io as _io
import re
import json
import time import time
import math import math
import base64
import random import random
import asyncio import asyncio
import aiohttp import aiohttp
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional from typing import Callable, List, Optional
import torch import torch
import numpy as np import numpy as np
@@ -22,21 +18,15 @@ from PIL import Image
from comfy_api.latest import io 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 ( from ..utils.config import (
NETWORK_ROUTE_OPTIONS, NETWORK_ROUTE_OPTIONS,
get_base_url_by_route, get_base_url_by_route,
get_api_key_or_raise, 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 ..clients.gemini_client import GeminiAPIClient
try:
import folder_paths
FOLDER_PATHS_AVAILABLE = True
except ImportError:
FOLDER_PATHS_AVAILABLE = False
try: try:
from comfy.utils import ProgressBar from comfy.utils import ProgressBar
PROGRESS_BAR_AVAILABLE = True PROGRESS_BAR_AVAILABLE = True
@@ -51,17 +41,9 @@ except ImportError:
InterruptProcessingException = RuntimeError InterruptProcessingException = RuntimeError
processing_interrupted = lambda: False 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 REQUEST_LOG_ENABLED = False
_NODE = "Nano Banana" _NODE = "Nano Banana"
_ENDPOINT = "/v1/chat/completions"
_REQUEST_TIMEOUT = 900 _REQUEST_TIMEOUT = 900
_INTERRUPT_CHECK_INTERVAL = 0.2 _INTERRUPT_CHECK_INTERVAL = 0.2
@@ -112,6 +94,25 @@ def _check_interrupt():
raise InterruptProcessingException() 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: def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
if not images: if not images:
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128)) 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" model_id += "-official"
return model_id 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( async def _generate_single(
session: aiohttp.ClientSession, session: aiohttp.ClientSession,
@@ -241,105 +182,25 @@ async def _generate_single(
images: Optional[List[Image.Image]] = None, images: Optional[List[Image.Image]] = None,
enable_grounding: bool = False, enable_grounding: bool = False,
thinking_level: Optional[str] = None, thinking_level: Optional[str] = None,
progress_callback: Optional[Callable[[float], None]] = None,
) -> List[Image.Image]: ) -> List[Image.Image]:
url = f"{base_url}{_ENDPOINT}" result_images, timing = await generate_nano_banana_async(
headers = _get_headers(api_key) session=session,
body = _build_request_body( base_url=base_url,
api_key=api_key,
prompt=prompt, prompt=prompt,
model=model, model=model,
aspect_ratio=aspect_ratio,
resolution=resolution, resolution=resolution,
aspect_ratio=aspect_ratio,
images=images, images=images,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
thinking_level=thinking_level, thinking_level=thinking_level,
node_label="Nano Banana",
request_log_enabled=REQUEST_LOG_ENABLED,
check_interrupt=_check_interrupt,
progress_callback=progress_callback,
) )
return result_images, timing["task_ms"], timing["parse_ms"]
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
async def _generate_single_task( async def _generate_single_task(
@@ -354,6 +215,7 @@ async def _generate_single_task(
global_task_index: int, global_task_index: int,
enable_grounding: bool = False, enable_grounding: bool = False,
thinking_level: Optional[str] = None, thinking_level: Optional[str] = None,
progress_callback: Optional[Callable[[float], None]] = None,
) -> dict: ) -> dict:
result = { result = {
"global_task_index": global_task_index, "global_task_index": global_task_index,
@@ -364,7 +226,7 @@ async def _generate_single_task(
"error": None, "error": None,
} }
try: try:
gen_images, first_token_ms, download_ms = await _generate_single( gen_images, task_ms, parse_ms = await _generate_single(
session=session, session=session,
base_url=base_url, base_url=base_url,
api_key=api_key, api_key=api_key,
@@ -375,7 +237,9 @@ async def _generate_single_task(
images=images if images else None, images=images if images else None,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
thinking_level=thinking_level, thinking_level=thinking_level,
progress_callback=progress_callback,
) )
del task_ms, parse_ms
result["output_images"] = gen_images result["output_images"] = gen_images
result["success"] = True result["success"] = True
result["generated_count"] = len(gen_images) result["generated_count"] = len(gen_images)
@@ -438,6 +302,7 @@ async def _process_batch_async(
global_task_index=i, global_task_index=i,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
thinking_level=thinking_level, thinking_level=thinking_level,
progress_callback=_make_progress_callback(pbar),
) )
) )
tasks.append(task) tasks.append(task)
@@ -473,9 +338,6 @@ async def _process_batch_async(
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误" 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}") 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) all_results.extend(batch_results)
import gc; gc.collect() import gc; gc.collect()
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
@@ -576,7 +438,7 @@ class NanoBanana(io.ComfyNode):
pil_imgs = tensor_to_pil(kwargs[key]) pil_imgs = tensor_to_pil(kwargs[key])
input_images.extend(pil_imgs) 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 张") raise ValueError(f"输入图像数量 {len(input_images)} 超过限制 14 张")
batch_prompts = parse_batch_prompts(prompt) batch_prompts = parse_batch_prompts(prompt)
@@ -667,6 +529,7 @@ class NanoBanana(io.ComfyNode):
images=input_images if input_images else None, images=input_images if input_images else None,
enable_grounding=enable_grounding, enable_grounding=enable_grounding,
thinking_level=thinking_level, thinking_level=thinking_level,
progress_callback=_make_progress_callback(pbar),
) )
return loop.run_until_complete(_run_with_interrupt(_do())) return loop.run_until_complete(_run_with_interrupt(_do()))
finally: finally:
@@ -674,17 +537,14 @@ class NanoBanana(io.ComfyNode):
with ThreadPoolExecutor(max_workers=1) as executor: with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(run_single) future = executor.submit(run_single)
generated_images, first_token_ms, download_ms = future.result(timeout=_REQUEST_TIMEOUT) generated_images, task_ms, parse_ms = future.result(timeout=_REQUEST_TIMEOUT)
if pbar is not None:
pbar.update(1)
output_tensor = _images_to_tensor_safe(generated_images, _NODE) output_tensor = _images_to_tensor_safe(generated_images, _NODE)
elapsed = time.time() - start_time elapsed = time.time() - start_time
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s" time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
ft_str = f"{first_token_ms/1000:.2f}s" task_str = f"{task_ms/1000:.2f}s"
dl_str = f"{download_ms/1000:.2f}s" parse_str = f"{parse_ms/1000:.2f}s"
print(f"完成!总耗时 {time_str} | 首字 {ft_str} | 下载 {dl_str} | 成功 {len(generated_images)}") print(f"完成!总耗时 {time_str} | 异步任务 {task_str} | 解析 {parse_str} | 成功 {len(generated_images)}")
import gc; gc.collect() import gc; gc.collect()
return io.NodeOutput(output_tensor) return io.NodeOutput(output_tensor)
@@ -701,7 +561,7 @@ class NanoBanana(io.ComfyNode):
except RuntimeError as e: except RuntimeError as e:
raise RuntimeError(str(e)) from None raise RuntimeError(str(e)) from None
except Exception as e: except Exception as e:
raise type(e)(str(e)) from None raise RuntimeError(str(e)) from None
finally: finally:
if not was_interrupted: if not was_interrupted:
try: try:
+12 -17
View File
@@ -31,7 +31,7 @@ from PIL import Image
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts 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.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.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 ( from ..models_config import (
get_enabled_async_models, get_enabled_async_models,
get_model_provider, get_model_provider,
@@ -241,6 +241,10 @@ class NanoBananaV2:
@staticmethod @staticmethod
def _friendly_error(error_msg: str) -> str: 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: if "No available channel for model" in error_msg:
return ( return (
"当前分组下模型不可用,请检查分组是否正确。" "当前分组下模型不可用,请检查分组是否正确。"
@@ -343,16 +347,16 @@ class NanoBananaV2:
error_text = await response.text() error_text = await response.text()
if not error_text.strip(): if not error_text.strip():
error_text = "(服务器未返回错误详情)" error_text = "(服务器未返回错误详情)"
raise RuntimeError(f"查询任务失败 ({response.status}): {error_text}") raise RuntimeError(f"查询任务失败: {get_friendly_message(response.status, error_text)}")
result = await response.json() result = await response.json()
status = provider.extract_status(result) status = provider.extract_status(result)
# 提取进度并回调(封顶 1.0 防止异常值导致进度条溢出) # 提取进度并回调;运行中状态不显示 100%,只有 SUCCESS 才补满。
if on_progress and status in ("SUBMITTED", "IN_PROGRESS"): if on_progress and status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
p = provider.extract_progress(result) p = provider.extract_progress(result)
if p is not None: if p is not None:
p = min(p, 1.0) p = min(p, 0.99)
if p > last_progress: if p > last_progress:
on_progress(p - last_progress) on_progress(p - last_progress)
last_progress = p last_progress = p
@@ -374,7 +378,7 @@ class NanoBananaV2:
print(f"{self.NODE_LABEL}: [轮询] FAILURE 但无错误信息,原始响应: {json.dumps(result, ensure_ascii=False)[:500]}") print(f"{self.NODE_LABEL}: [轮询] FAILURE 但无错误信息,原始响应: {json.dumps(result, ensure_ascii=False)[:500]}")
friendly_msg = self._friendly_error(error_msg) friendly_msg = self._friendly_error(error_msg)
raise RuntimeError(f"任务失败: {friendly_msg}") raise RuntimeError(f"任务失败: {friendly_msg}")
elif status in ("SUBMITTED", "IN_PROGRESS"): elif status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
# 分段 sleep,每 0.1 秒检查一次取消信号 # 分段 sleep,每 0.1 秒检查一次取消信号
sleep_iterations = int(_POLL_INTERVAL / _INTERRUPT_CHECK_INTERVAL) sleep_iterations = int(_POLL_INTERVAL / _INTERRUPT_CHECK_INTERVAL)
for _ in range(sleep_iterations): for _ in range(sleep_iterations):
@@ -406,10 +410,7 @@ class NanoBananaV2:
"error": None, "error": None,
} }
contributed = [0.0] # mutable container,追踪本任务已贡献的 pbar 进度
def _track_progress(delta): def _track_progress(delta):
contributed[0] += delta
if on_progress: if on_progress:
on_progress(delta) on_progress(delta)
@@ -435,15 +436,9 @@ class NanoBananaV2:
result["request_time"] = request_time result["request_time"] = request_time
result["download_time"] = download_time result["download_time"] = download_time
except InterruptProcessingException: except InterruptProcessingException:
# 用户取消:补齐进度后向上传播,不吞掉
if contributed[0] < 1.0 and on_progress:
on_progress(1.0 - contributed[0])
raise raise
except Exception as e: except Exception as e:
result["error"] = str(e) or f"{type(e).__name__}(无错误详情)" 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 return result
@@ -704,7 +699,7 @@ class NanoBananaV2:
raise RuntimeError(str(e)) from None raise RuntimeError(str(e)) from None
except Exception as e: except Exception as e:
raise type(e)(str(e)) from None raise RuntimeError(str(e)) from None
finally: finally:
# 查询并打印余额 # 查询并打印余额
@@ -1223,7 +1218,7 @@ class NanoBananaV2Batch(NanoBananaV2):
raise RuntimeError(str(e)) from None raise RuntimeError(str(e)) from None
except Exception as e: except Exception as e:
raise type(e)(str(e)) from None raise RuntimeError(str(e)) from None
finally: finally:
try: try:
+254
View File
@@ -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",
}
+51 -9
View File
@@ -4,7 +4,9 @@ Seedance 视频生成节点
- Seedance: 文生视频 / 图生视频 / 首尾帧生视频(根据图片输入自动切换模式) - Seedance: 文生视频 / 图生视频 / 首尾帧生视频(根据图片输入自动切换模式)
""" """
import base64
import io import io
import json
import os import os
import tempfile import tempfile
@@ -13,7 +15,7 @@ import torch
from ..clients.seedance_client import SeedanceClient from ..clients.seedance_client import SeedanceClient
from ..clients.gemini_client import GeminiAPIClient 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.r2_uploader import upload_video, upload_audio
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
@@ -28,6 +30,9 @@ _MODELS = [
_RESOLUTIONS = ["720p", "1080p", "480p"] _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""" """ComfyUI IMAGE tensor → data:image/png;base64,xxx"""
pil_images = tensor_to_pil(tensor) 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}" 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: async def _url_to_tensor(url: str) -> torch.Tensor:
"""从 URL 下载图片并转为 ComfyUI IMAGE tensor,失败时返回 None""" """从 URL 下载图片并转为 ComfyUI IMAGE tensor,失败时返回 None"""
@@ -201,7 +238,7 @@ class Seedance:
} }
elif mode == "i2v": elif mode == "i2v":
first_url = _tensor_to_base64_url(first_image) first_url = _tensor_to_base64_url(first_image, "首帧图片")
metadata["content"] = [ metadata["content"] = [
{ {
"type": "image_url", "type": "image_url",
@@ -218,8 +255,8 @@ class Seedance:
} }
else: # flipflop else: # flipflop
first_url = _tensor_to_base64_url(first_image) first_url = _tensor_to_base64_url(first_image, "首帧图片")
last_url = _tensor_to_base64_url(last_image) last_url = _tensor_to_base64_url(last_image, "尾帧图片")
metadata["content"] = [ metadata["content"] = [
{ {
"type": "image_url", "type": "image_url",
@@ -240,6 +277,8 @@ class Seedance:
"metadata": metadata, "metadata": metadata,
} }
_validate_request_body_size(body, tag)
# 保存路径(临时文件,避免与下游保存节点重复落盘) # 保存路径(临时文件,避免与下游保存节点重复落盘)
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_") _, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_")
@@ -315,6 +354,7 @@ class SeedanceMultiModal:
web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开" web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开"
return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开" return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开"
seed = _first(kwargs.get("seed"), 0) seed = _first(kwargs.get("seed"), 0)
network_route = _first(kwargs.get("网络线路"), "全球加速")
# 参考图片:INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表,直接保留 # 参考图片:INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表,直接保留
raw_images = kwargs.get("参考图片", None) raw_images = kwargs.get("参考图片", None)
@@ -344,11 +384,11 @@ class SeedanceMultiModal:
imgs = ref_images[:9] imgs = ref_images[:9]
if len(ref_images) > 9: if len(ref_images) > 9:
print(f"[SeedanceMultiModal] 参考图片超过9张,仅取前9张(共{len(ref_images)}张)") 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 维 # 每个 tensor 可能是 [1,H,W,C] 或 [H,W,C],统一确保有 batch 维
if img_tensor.dim() == 3: if img_tensor.dim() == 3:
img_tensor = img_tensor.unsqueeze(0) img_tensor = img_tensor.unsqueeze(0)
url = _tensor_to_base64_url(img_tensor) url = _tensor_to_base64_url(img_tensor, f"参考图片{idx}")
content.append({ content.append({
"type": "image_url", "type": "image_url",
"image_url": {"url": url}, "image_url": {"url": url},
@@ -414,11 +454,13 @@ class SeedanceMultiModal:
if first_image_url: if first_image_url:
body["image"] = first_image_url body["image"] = first_image_url
_validate_request_body_size(body, "Seedance多模态")
# ── 保存路径(临时文件,避免与下游保存节点重复落盘)────────────────── # ── 保存路径(临时文件,避免与下游保存节点重复落盘)──────────────────
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_") _, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_")
client = SeedanceClient() 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() pbar = _make_pbar()
on_stage, on_prog = _make_callbacks("Seedance多模态", pbar) on_stage, on_prog = _make_callbacks("Seedance多模态", pbar)
+2
View File
@@ -148,4 +148,6 @@ def get_async_api_base_url() -> str:
def get_base_url_by_route(route: str) -> str: def get_base_url_by_route(route: str) -> str:
"""根据网络线路选项返回对应域名,未匹配则走 config 垫底""" """根据网络线路选项返回对应域名,未匹配则走 config 垫底"""
if isinstance(route, (list, tuple)):
route = route[0] if route else None
return NETWORK_ROUTES.get(route, get_api_base_url()) return NETWORK_ROUTES.get(route, get_api_base_url())
+54 -2
View File
@@ -9,8 +9,9 @@
""" """
import asyncio import asyncio
import json
import random import random
from typing import Optional from typing import Any, Optional
import aiohttp import aiohttp
@@ -45,15 +46,66 @@ DEFAULT_MAX_DELAY = 30.0 # 最大等待秒数
DEFAULT_BACKOFF_FACTOR = 2.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: def get_friendly_message(status_code: int, raw_message: str = "") -> str:
"""根据状态码/错误内容返回友好文案,未匹配则返回原始信息""" """根据状态码/错误内容返回友好文案,未匹配则返回原始信息"""
if status_code == 524: if status_code == 524:
return "Gateway timed out while waiting for upstream image generation. Please retry, lower resolution/count, or switch network route." return "Gateway timed out while waiting for upstream image generation. Please retry, lower resolution/count, or switch network route."
if raw_message: 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(): for keyword, friendly_msg in ERROR_CONTENT_MESSAGES.items():
if keyword.lower() in raw_message_lower: if keyword.lower() in raw_message_lower:
return friendly_msg return friendly_msg
if structured_message:
return structured_message
if status_code == 500: if status_code == 500:
return "服务器返回 500:上游生成失败或服务端临时异常。请稍后重试;如果多次出现,请降低分辨率/数量,或调整提示词。" return "服务器返回 500:上游生成失败或服务端临时异常。请稍后重试;如果多次出现,请降低分辨率/数量,或调整提示词。"
friendly = HTTP_ERROR_MESSAGES.get(status_code) friendly = HTTP_ERROR_MESSAGES.get(status_code)
+675
View File
@@ -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"<base64 data, {len(item)} chars>"
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"<data image url, {len(value)} chars>"
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),
}
+9 -3
View File
@@ -4,10 +4,16 @@ app.registerExtension({
name: "o1key.hideSidebarItems", name: "o1key.hideSidebarItems",
async setup() { async setup() {
const hide = () => { 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 => { 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))) { if (hiddenLabels.some(k => label.includes(k))) {
btn.style.display = "none"; btn.style.display = "none";
} }
+77 -11
View File
@@ -1,7 +1,9 @@
import { app } from "../../../scripts/app.js"; import { app } from "../../../scripts/app.js";
import { api } from "../../../scripts/api.js";
const STORAGE_KEY = "o1key-notes"; const STORAGE_KEY = "o1key-notes";
const SEEDED_KEY = "o1key-notes-seeded-v2"; const SEEDED_KEY = "o1key-notes-seeded-v2";
const NOTES_API = "/o1key/notes";
const STYLE_ID = "o1key-notes-styles"; const STYLE_ID = "o1key-notes-styles";
let notes = []; let notes = [];
@@ -207,27 +209,91 @@ function injectStyles() {
document.head.appendChild(el); document.head.appendChild(el);
} }
function loadNotes() { function readCachedNotes() {
try { try {
const raw = localStorage.getItem(STORAGE_KEY); 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 { } catch {
notes = []; return { found: false, notes: [] };
}
} }
if (!localStorage.getItem(SEEDED_KEY)) { function writeCachedNotes() {
const existingTitles = new Set(notes.map(n => n.title)); try {
const samples = SAMPLE_NOTES.map(makeNote).filter(n => !existingTitles.has(n.title)); localStorage.setItem(STORAGE_KEY, JSON.stringify(notes));
notes = [...samples, ...notes]; } catch {}
}
function hasSeededNotes() {
try {
return !!localStorage.getItem(SEEDED_KEY);
} catch {
return false;
}
}
function markSeededNotes() {
try {
localStorage.setItem(SEEDED_KEY, "1"); localStorage.setItem(SEEDED_KEY, "1");
saveNotes(); } 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);
}
notes = cached.found ? cached.notes : [];
if (!cached.found && !hasSeededNotes()) {
notes = createSeedNotes();
markSeededNotes();
writeCachedNotes();
} }
} }
function saveNotes() { function saveNotes() {
writeCachedNotes();
void persistNotesToFile();
}
async function persistNotesToFile() {
try { try {
localStorage.setItem(STORAGE_KEY, JSON.stringify(notes)); const resp = await api.fetchApi(NOTES_API, {
} catch {} 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() { function allTags() {
@@ -250,7 +316,7 @@ function filteredNotes() {
app.registerExtension({ app.registerExtension({
name: "o1key.notePanel", name: "o1key.notePanel",
async setup() { async setup() {
loadNotes(); await loadNotes();
app.extensionManager.registerSidebarTab({ app.extensionManager.registerSidebarTab({
id: "o1key-notes", id: "o1key-notes",
title: "笔记", title: "笔记",
+60 -5
View File
@@ -43,6 +43,7 @@ app.registerExtension({
if (!confirm("确定要重启 ComfyUI 吗?")) return; if (!confirm("确定要重启 ComfyUI 吗?")) return;
btn.style.opacity = "0.5"; btn.style.opacity = "0.5";
btn.style.pointerEvents = "none"; btn.style.pointerEvents = "none";
await disableExperimentalAssetApi();
try { await fetch("/o1key/restart", { method: "POST" }); } catch {} try { await fetch("/o1key/restart", { method: "POST" }); } catch {}
pollUntilReady(); pollUntilReady();
}); });
@@ -51,16 +52,70 @@ app.registerExtension({
injected = true; 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() { function pollUntilReady() {
let attempts = 0; let attempts = 0;
const maxAttempts = 40; const maxAttempts = 80;
const minRestartWaitMs = 5000;
const startedAt = Date.now();
let sawUnavailable = false;
const interval = setInterval(async () => { const interval = setInterval(async () => {
attempts++; attempts++;
if (attempts > maxAttempts) { clearInterval(interval); forceReload(); return; } if (attempts > maxAttempts) { clearInterval(interval); forceReload(); return; }
try { const ready = await comfyReady();
const r = await fetch("/api/system_stats", { signal: AbortSignal.timeout(2000) }); if (!ready) {
if (r.ok) { clearInterval(interval); forceReload(); } sawUnavailable = true;
} catch {} return;
}
if (!sawUnavailable && Date.now() - startedAt < minRestartWaitMs) return;
clearInterval(interval);
await disableExperimentalAssetApi();
setTimeout(forceReload, 800);
}, 1500); }, 1500);
} }