Add Grok and VEO video workflow support
This commit is contained in:
+125
-1
@@ -12,6 +12,7 @@ Comfyui_o1key - ComfyUI 自定义节点集合
|
||||
|
||||
import ssl
|
||||
import logging
|
||||
import asyncio
|
||||
|
||||
# 屏蔽 ComfyUI 资产扫描的终端日志输出
|
||||
_seeder_filter = lambda record: not any(
|
||||
@@ -20,7 +21,57 @@ _seeder_filter = lambda record: not any(
|
||||
)
|
||||
logging.getLogger().addFilter(_seeder_filter)
|
||||
|
||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, KVideoFirstLast, KVideoImage2Video
|
||||
|
||||
def _is_ignored_asyncio_win10054(context):
|
||||
exc = context.get("exception")
|
||||
if not (
|
||||
isinstance(exc, ConnectionResetError)
|
||||
and getattr(exc, "winerror", None) == 10054
|
||||
):
|
||||
return False
|
||||
|
||||
handle = str(context.get("handle", ""))
|
||||
message = str(context.get("message", ""))
|
||||
marker = "_ProactorBasePipeTransport._call_connection_lost"
|
||||
return marker in handle or marker in message
|
||||
|
||||
|
||||
def _install_asyncio_win10054_filter(loop):
|
||||
if getattr(loop, "_o1key_win10054_filter_installed", False):
|
||||
return loop
|
||||
|
||||
previous_handler = loop.get_exception_handler()
|
||||
|
||||
def _o1key_asyncio_exception_handler(loop, context):
|
||||
if _is_ignored_asyncio_win10054(context):
|
||||
return
|
||||
if previous_handler is not None:
|
||||
previous_handler(loop, context)
|
||||
else:
|
||||
loop.default_exception_handler(context)
|
||||
|
||||
loop.set_exception_handler(_o1key_asyncio_exception_handler)
|
||||
setattr(loop, "_o1key_win10054_filter_installed", True)
|
||||
return loop
|
||||
|
||||
|
||||
try:
|
||||
_install_asyncio_win10054_filter(asyncio.get_event_loop())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
if not getattr(asyncio, "_o1key_new_event_loop_patched", False):
|
||||
_o1key_original_new_event_loop = asyncio.new_event_loop
|
||||
|
||||
def _o1key_new_event_loop(*args, **kwargs):
|
||||
return _install_asyncio_win10054_filter(
|
||||
_o1key_original_new_event_loop(*args, **kwargs)
|
||||
)
|
||||
|
||||
asyncio.new_event_loop = _o1key_new_event_loop
|
||||
asyncio._o1key_new_event_loop_patched = True
|
||||
|
||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, Google31Video, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGPTImageBatch, O1keyGrokImage, O1keyGrokVideo, KVideoFirstLast, KVideoImage2Video
|
||||
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch, SaveImageFormat
|
||||
from .nodes import O1keySavePSD
|
||||
from .nodes import O1keyRemoveBackground
|
||||
@@ -73,6 +124,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"BatchCleanMetadata": BatchCleanMetadata,
|
||||
"VideoPreview": VideoPreview,
|
||||
"GoogleVeo": GoogleVeo,
|
||||
"Google31Video": Google31Video,
|
||||
"FluxImageEdit": FluxImageEdit,
|
||||
"UniversalLLMChat": UniversalLLMChat,
|
||||
"KlingVideo": KlingVideo,
|
||||
@@ -88,6 +140,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"O1keyGPTImage": O1keyGPTImage,
|
||||
"O1keyGPTImageBatch": O1keyGPTImageBatch,
|
||||
"O1keyGrokImage": O1keyGrokImage,
|
||||
"O1keyGrokVideo": O1keyGrokVideo,
|
||||
"KVideoFirstLast": KVideoFirstLast,
|
||||
"KVideoImage2Video": KVideoImage2Video,
|
||||
"K3Video": K3Video,
|
||||
@@ -113,6 +166,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BatchCleanMetadata": "批量任务(防AI识别)",
|
||||
"VideoPreview": "预览视频",
|
||||
"GoogleVeo": "Google Veo - ab",
|
||||
"Google31Video": "Google 3.1 Video",
|
||||
"FluxImageEdit": "Flux2 图像编辑",
|
||||
"UniversalLLMChat": "全能LLM对话助手",
|
||||
"KlingVideo": "文/图生视频 自研模型",
|
||||
@@ -128,6 +182,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"O1keyGPTImage": "o1key GPT Image",
|
||||
"O1keyGPTImageBatch": "o1key GPT Image(批量)",
|
||||
"O1keyGrokImage": "Grok Image",
|
||||
"O1keyGrokVideo": "Grok Video",
|
||||
"KVideoFirstLast": "K26 图生视频(首尾帧)",
|
||||
"KVideoImage2Video": "K26 图生视频",
|
||||
"K3Video": "K3 图生视频 自研",
|
||||
@@ -183,6 +238,75 @@ try:
|
||||
f".o1key_history_{_get_o1key_server_port()}.json",
|
||||
)
|
||||
|
||||
def _get_o1key_notes_file():
|
||||
import os as _os_notes
|
||||
input_dir = _os_notes.path.abspath(folder_paths.get_input_directory())
|
||||
_os_notes.makedirs(input_dir, exist_ok=True)
|
||||
return _os_notes.path.join(input_dir, "o1key-notes.json")
|
||||
|
||||
def _extract_o1key_notes(payload):
|
||||
if isinstance(payload, list):
|
||||
return payload
|
||||
if isinstance(payload, dict) and isinstance(payload.get("notes"), list):
|
||||
return payload["notes"]
|
||||
return None
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/notes")
|
||||
async def get_o1key_notes(request):
|
||||
import os as _os_notes
|
||||
import json as _json_notes
|
||||
|
||||
notes_file = _get_o1key_notes_file()
|
||||
exists = _os_notes.path.isfile(notes_file)
|
||||
notes = []
|
||||
|
||||
if exists:
|
||||
try:
|
||||
with open(notes_file, "r", encoding="utf-8") as nf:
|
||||
loaded = _json_notes.load(nf)
|
||||
notes = _extract_o1key_notes(loaded)
|
||||
if notes is None:
|
||||
return web.json_response(
|
||||
{"error": "invalid notes file", "path": notes_file},
|
||||
status=500,
|
||||
)
|
||||
except Exception as e:
|
||||
return web.json_response(
|
||||
{"error": str(e), "path": notes_file},
|
||||
status=500,
|
||||
)
|
||||
|
||||
return web.json_response({"notes": notes, "path": notes_file, "exists": exists})
|
||||
|
||||
@PromptServer.instance.routes.post("/o1key/notes")
|
||||
async def save_o1key_notes(request):
|
||||
import os as _os_notes
|
||||
import json as _json_notes
|
||||
|
||||
try:
|
||||
payload = await request.json()
|
||||
notes = _extract_o1key_notes(payload)
|
||||
if notes is None:
|
||||
return web.json_response({"error": "notes must be a list"}, status=400)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": f"invalid notes payload: {str(e)}"}, status=400)
|
||||
|
||||
notes_file = _get_o1key_notes_file()
|
||||
temp_file = notes_file + ".tmp"
|
||||
try:
|
||||
with open(temp_file, "w", encoding="utf-8") as nf:
|
||||
_json_notes.dump(notes, nf, ensure_ascii=False, indent=2)
|
||||
nf.write("\n")
|
||||
_os_notes.replace(temp_file, notes_file)
|
||||
except Exception as e:
|
||||
return web.json_response({"error": f"save notes failed: {str(e)}"}, status=500)
|
||||
|
||||
return web.json_response({
|
||||
"success": True,
|
||||
"path": notes_file,
|
||||
"count": len(notes),
|
||||
})
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/input_dir")
|
||||
async def get_input_dir(request):
|
||||
import os
|
||||
|
||||
+3
-1
@@ -9,6 +9,8 @@ from .gemini_flash_client import GeminiFlashClient
|
||||
from .sora_client import SoraClient
|
||||
from .kling_client import KlingClient
|
||||
from .veo_client import VeoClient
|
||||
from .newapi_veo_client import NewAPIVeoClient
|
||||
from .grok_video_client import GrokVideoClient
|
||||
from .openai_client import OpenAIAPIClient
|
||||
|
||||
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient']
|
||||
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient']
|
||||
|
||||
+2
-22
@@ -40,7 +40,8 @@ class BaseAPIClient(ABC):
|
||||
Args:
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
max_request_size: 最大请求体大小(字节),默认 100MB
|
||||
max_request_size: 兼容参数;基类不再用它限制 JSON 请求体,
|
||||
部分子类仍用它作为上传文件大小限制
|
||||
"""
|
||||
self.base_url = base_url
|
||||
self.api_key = api_key
|
||||
@@ -116,24 +117,6 @@ class BaseAPIClient(ABC):
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def check_request_size(self, request_body: Dict[str, Any]) -> None:
|
||||
"""
|
||||
检查请求体大小是否超过限制
|
||||
|
||||
Args:
|
||||
request_body: 请求体字典
|
||||
|
||||
Raises:
|
||||
ValueError: 如果请求体超过限制
|
||||
"""
|
||||
request_json = json.dumps(request_body)
|
||||
request_size = len(request_json.encode('utf-8'))
|
||||
|
||||
if request_size > self.max_request_size:
|
||||
raise ValueError(
|
||||
"请求体积超过100MB限制,请调整分辨率或减少图片数量"
|
||||
)
|
||||
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""
|
||||
子类可重写:为指定 HTTP 状态码返回自定义错误文案。
|
||||
@@ -185,9 +168,6 @@ class BaseAPIClient(ABC):
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token)
|
||||
|
||||
# 检查请求大小
|
||||
self.check_request_size(request_body)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = self._make_session()
|
||||
|
||||
@@ -72,7 +72,9 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
resolution=resolution,
|
||||
enable_grounding=kwargs.get("enable_grounding", False),
|
||||
enable_image_search=kwargs.get("enable_image_search", False),
|
||||
image_compression=getattr(self, 'image_compression', None),
|
||||
image_compression=getattr(self, "image_compression", None),
|
||||
thinking_level=kwargs.get("thinking_level"),
|
||||
request_log_enabled=False,
|
||||
)
|
||||
|
||||
def extract_task_id(self, response: dict) -> str:
|
||||
@@ -85,6 +87,23 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
return response.get("status", "UNKNOWN")
|
||||
|
||||
async def parse_result(self, result_data: dict, session) -> List[Image.Image]:
|
||||
images = result_data.get("images") if isinstance(result_data, dict) else None
|
||||
if isinstance(images, list) and images:
|
||||
parsed = []
|
||||
for item in images:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
image_url = item.get("url") or item.get("image_url")
|
||||
if image_url:
|
||||
async with session.get(image_url) as img_resp:
|
||||
if img_resp.status == 200:
|
||||
img_bytes = await img_resp.read()
|
||||
parsed.append(Image.open(BytesIO(img_bytes)).convert("RGB"))
|
||||
else:
|
||||
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
|
||||
if parsed:
|
||||
return parsed
|
||||
|
||||
# 异步接口可能直接返回 image_url
|
||||
image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else ""
|
||||
if image_url:
|
||||
@@ -124,19 +143,44 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
|
||||
def extract_progress(self, response: dict) -> Optional[float]:
|
||||
"""从轮询响应中提取进度(0.0-1.0)"""
|
||||
|
||||
def _coerce(val) -> Optional[float]:
|
||||
if val is None or isinstance(val, bool):
|
||||
return None
|
||||
if isinstance(val, (int, float)):
|
||||
progress = float(val)
|
||||
elif isinstance(val, str):
|
||||
text = val.strip()
|
||||
if not text:
|
||||
return None
|
||||
has_percent_suffix = text.endswith("%")
|
||||
if has_percent_suffix:
|
||||
text = text[:-1].strip()
|
||||
try:
|
||||
progress = float(text)
|
||||
except ValueError:
|
||||
return None
|
||||
if has_percent_suffix:
|
||||
progress /= 100.0
|
||||
else:
|
||||
return None
|
||||
if progress > 1.0:
|
||||
progress /= 100.0
|
||||
return max(0.0, min(progress, 1.0))
|
||||
|
||||
# 直接字段:progress / percentage
|
||||
for field in ("progress", "percentage"):
|
||||
val = response.get(field)
|
||||
if val is not None and isinstance(val, (int, float)):
|
||||
return val / 100.0 if val > 1 else float(val)
|
||||
for field in ("progress", "percentage", "percent"):
|
||||
progress = _coerce(response.get(field))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
# 嵌套字段:progressInfo / progress_info
|
||||
progress_info = response.get("progressInfo") or response.get("progress_info")
|
||||
if isinstance(progress_info, dict):
|
||||
for field in ("progress", "percentage"):
|
||||
val = progress_info.get(field)
|
||||
if val is not None and isinstance(val, (int, float)):
|
||||
return val / 100.0 if val > 1 else float(val)
|
||||
for field in ("progress", "percentage", "percent"):
|
||||
progress = _coerce(progress_info.get(field))
|
||||
if progress is not None:
|
||||
return progress
|
||||
|
||||
return None
|
||||
|
||||
|
||||
+103
-80
@@ -34,8 +34,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
|
||||
super().__init__(
|
||||
base_url=get_api_base_url(),
|
||||
api_key=api_key,
|
||||
max_request_size=100 * 1024 * 1024
|
||||
api_key=api_key
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@@ -188,6 +187,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_image_search: bool = False,
|
||||
image_compression: str = None,
|
||||
thinking_level: str = None,
|
||||
request_log_enabled: bool = True,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -204,79 +204,26 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
Returns:
|
||||
请求体字典
|
||||
"""
|
||||
import json
|
||||
|
||||
_MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB
|
||||
_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.9)
|
||||
|
||||
parts = []
|
||||
|
||||
# 添加文本部分
|
||||
parts.append({"text": prompt})
|
||||
|
||||
# 添加图像部分(如果有)
|
||||
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}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
image_config["aspectRatio"] = aspect_ratio
|
||||
|
||||
request_body = {
|
||||
def _build_request_body(body_parts: List[dict]) -> Dict[str, Any]:
|
||||
body = {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": parts
|
||||
"parts": body_parts
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
@@ -287,33 +234,110 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
|
||||
# 添加思考深度配置
|
||||
if thinking_level:
|
||||
request_body["generationConfig"]["thinkingConfig"] = {
|
||||
body["generationConfig"]["thinkingConfig"] = {
|
||||
"thinkingLevel": thinking_level,
|
||||
"includeThoughts": True
|
||||
}
|
||||
|
||||
# 添加图片压缩参数
|
||||
if image_compression:
|
||||
request_body["image_compression"] = image_compression
|
||||
body["image_compression"] = image_compression
|
||||
|
||||
# 添加 Google Search Grounding 工具(如果启用)
|
||||
# 注意:enable_image_search=True 时会自动隐含 enable_grounding
|
||||
# 添加 Google Search Grounding(如果启用)
|
||||
# 新异步接口要求直接放在请求体顶层:{"google_search": true}
|
||||
if enable_grounding or enable_image_search:
|
||||
if enable_image_search:
|
||||
# 同时启用网页搜索和图片搜索(仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持)
|
||||
request_body["tools"] = [
|
||||
{
|
||||
"google_search": {
|
||||
"searchTypes": {
|
||||
"webSearch": {},
|
||||
"imageSearch": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
body["google_search"] = True
|
||||
|
||||
return body
|
||||
|
||||
def _request_size(body: Dict[str, Any]) -> int:
|
||||
return len(json.dumps(body).encode("utf-8"))
|
||||
|
||||
def _format_size(size: int) -> str:
|
||||
if size < 1024 * 1024:
|
||||
return f"{size / 1024:.2f}KB"
|
||||
return f"{size / 1024 / 1024:.2f}MB"
|
||||
|
||||
def _shorten_base64_for_log(obj, max_len: int = 200):
|
||||
if isinstance(obj, dict):
|
||||
result = {}
|
||||
for key, value in obj.items():
|
||||
if key == "data" and isinstance(value, str) and len(value) > max_len:
|
||||
result[key] = f"<base64 data, {len(value)} chars>"
|
||||
else:
|
||||
# 仅启用网页搜索(通用)
|
||||
request_body["tools"] = [{"google_search": {}}]
|
||||
result[key] = _shorten_base64_for_log(value, max_len)
|
||||
return result
|
||||
if isinstance(obj, list):
|
||||
return [_shorten_base64_for_log(item, max_len) for item in obj]
|
||||
return obj
|
||||
|
||||
def _log_original_request_body(body: Dict[str, Any]) -> None:
|
||||
body_size = _request_size(body)
|
||||
print(
|
||||
f"\n{'=' * 60}\n"
|
||||
f"[原始请求体日志] 请求体积: {_format_size(body_size)} "
|
||||
f"(inline_data.data 已折叠显示 base64 长度)\n"
|
||||
f"{json.dumps(_shorten_base64_for_log(body), ensure_ascii=False, indent=2)}\n"
|
||||
f"{'=' * 60}\n"
|
||||
)
|
||||
|
||||
original_request_logged = False
|
||||
|
||||
# 添加图像部分(如果有)
|
||||
if images:
|
||||
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
|
||||
|
||||
@@ -1106,4 +1130,3 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
+636
-10
@@ -1,13 +1,17 @@
|
||||
"""
|
||||
GPT Image API 客户端
|
||||
支持两个接口:
|
||||
- POST /v1/images/generations/ 文生图 / 图生图(gpt-image-1 / gpt-image-1.5)
|
||||
- POST /v1/images/edits/ 图像编辑(带蒙版 inpainting)
|
||||
新版节点请求走异步任务接口:
|
||||
- POST /async/v1/generateImage
|
||||
- GET /async/v1/tasks/{task_id}
|
||||
|
||||
旧同步接口保留兼容代码,但 GPT Image 节点不再使用:
|
||||
- POST /v1/images/generations/
|
||||
- POST /v1/images/edits
|
||||
|
||||
设计原则:
|
||||
- 与 doubao_image_client.py 保持相同的异步 + 同步双入口模式
|
||||
- generations / edits 接口均使用 multipart/form-data
|
||||
- 响应兼容 SSE 流式、JSON、url 和 b64_json
|
||||
- 对 ComfyUI 节点暴露同步入口,内部提交异步任务并轮询
|
||||
- 图片和蒙版以 data:image/png;base64,... 放入 JSON 请求体
|
||||
- 响应优先读取 data.images[].url
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -16,7 +20,7 @@ import json
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from io import BytesIO
|
||||
from typing import List, Optional
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
@@ -38,6 +42,8 @@ except ImportError:
|
||||
# ── 接口端点 ──────────────────────────────────────────────────────────────────
|
||||
_ENDPOINT_GENERATIONS = "/v1/images/generations/"
|
||||
_ENDPOINT_EDITS = "/v1/images/edits"
|
||||
_ENDPOINT_ASYNC_GENERATE = "/async/v1/generateImage"
|
||||
_ENDPOINT_ASYNC_TASK = "/async/v1/tasks/{task_id}"
|
||||
|
||||
# ── 模型名映射(UI 显示名 → API 实际参数名)─────────────────────────────────
|
||||
_MODEL_NAME_MAP = {
|
||||
@@ -47,6 +53,49 @@ _MODEL_NAME_MAP = {
|
||||
|
||||
# ── 超时 ──────────────────────────────────────────────────────────────────────
|
||||
_REQUEST_TIMEOUT = 900 # 秒
|
||||
_ASYNC_POLL_SCHEDULE = [5.0, 20.0]
|
||||
_ASYNC_POLL_INTERVAL = 3.0
|
||||
_ASYNC_MAX_WAIT = 600.0
|
||||
_ASYNC_RETRY_DELAYS = [2.0, 5.0, 10.0]
|
||||
_ASYNC_RETRYABLE_ERROR_CODES = {
|
||||
"image_rate_limited",
|
||||
"image_upstream_busy",
|
||||
"image_timeout",
|
||||
"image_storage_failed",
|
||||
"image_empty_result",
|
||||
"image_upstream_error",
|
||||
"image_internal_error",
|
||||
"image_unknown_error",
|
||||
}
|
||||
_ASYNC_RETRYABLE_ERROR_CATEGORIES = {
|
||||
"rate_limit",
|
||||
"upstream_busy",
|
||||
"timeout",
|
||||
"storage",
|
||||
"upstream_error",
|
||||
"internal_error",
|
||||
"unknown",
|
||||
}
|
||||
_ASYNC_NON_RETRYABLE_ERROR_CODES = {
|
||||
"image_invalid_size",
|
||||
"image_payload_too_large",
|
||||
"image_invalid_mask",
|
||||
"image_invalid_parameter",
|
||||
"image_safety_blocked",
|
||||
"image_provider_quota_exceeded",
|
||||
"image_provider_permission_required",
|
||||
"image_model_unavailable",
|
||||
"image_reference_download_failed",
|
||||
}
|
||||
|
||||
REQUEST_LOG_ENABLED = False
|
||||
POLL_LOG_ENABLED = False
|
||||
|
||||
|
||||
class _AsyncImageTaskFailure(RuntimeError):
|
||||
def __init__(self, message: str, error_detail: Optional[dict] = None):
|
||||
super().__init__(message)
|
||||
self.error_detail = error_detail or {}
|
||||
|
||||
|
||||
class GptImageClient:
|
||||
@@ -54,10 +103,10 @@ class GptImageClient:
|
||||
GPT Image API 客户端
|
||||
|
||||
接口说明:
|
||||
generations:multipart/form-data,支持 quality / size / n / model
|
||||
edits:multipart/form-data,图片和 mask 使用 PNG 文件上传
|
||||
async generateImage:JSON 提交,返回 task_id
|
||||
tasks/{task_id}:轮询任务状态,成功后读取 data.images[].url
|
||||
|
||||
响应支持 JSON 和 SSE 流式格式。
|
||||
旧 generations / edits 同步接口保留为兼容代码。
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
@@ -132,6 +181,226 @@ class GptImageClient:
|
||||
)
|
||||
return png_bytes
|
||||
|
||||
@staticmethod
|
||||
def _png_bytes_to_data_url(png_bytes: bytes) -> str:
|
||||
b64 = base64.b64encode(png_bytes).decode("ascii")
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
@staticmethod
|
||||
def _json_body_size(body: dict) -> int:
|
||||
return len(json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
|
||||
|
||||
@staticmethod
|
||||
def _format_body_size(size: int) -> str:
|
||||
if size < 1024 * 1024:
|
||||
return f"{size / 1024:.2f}KB"
|
||||
return f"{size / 1024 / 1024:.2f}MB"
|
||||
|
||||
@staticmethod
|
||||
def _shorten_data_urls_for_log(obj, max_len: int = 200):
|
||||
if isinstance(obj, dict):
|
||||
return {
|
||||
key: GptImageClient._shorten_data_urls_for_log(value, max_len)
|
||||
for key, value in obj.items()
|
||||
}
|
||||
if isinstance(obj, list):
|
||||
return [GptImageClient._shorten_data_urls_for_log(item, max_len) for item in obj]
|
||||
if isinstance(obj, str) and obj.startswith("data:image") and len(obj) > max_len:
|
||||
header, _, data = obj.partition(",")
|
||||
return f"{header},<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}KB,scale={best_scale:.3f}"
|
||||
)
|
||||
return best_bytes
|
||||
|
||||
@staticmethod
|
||||
def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes:
|
||||
"""
|
||||
@@ -500,6 +769,316 @@ class GptImageClient:
|
||||
|
||||
return request_task.result()
|
||||
|
||||
# ── 新版异步 GPT Image 接口 ──────────────────────────────────────────────
|
||||
|
||||
def _build_async_generate_body(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
quality: str,
|
||||
size: str,
|
||||
n: int,
|
||||
image_list: Optional[List[torch.Tensor]] = None,
|
||||
mask_tensor: Optional[torch.Tensor] = None,
|
||||
output_format: str = "png",
|
||||
) -> dict:
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
assets: List[Dict[str, Any]] = []
|
||||
image_keys: List[str] = []
|
||||
mask_key = None
|
||||
|
||||
if image_list:
|
||||
for idx_img, img_tensor in enumerate(image_list):
|
||||
pil_images = tensor_to_pil(img_tensor)
|
||||
if not pil_images:
|
||||
continue
|
||||
key = f"image_{idx_img}"
|
||||
image_keys.append(key)
|
||||
assets.append({
|
||||
"key": key,
|
||||
"label": f"参考图{idx_img + 1}",
|
||||
"bytes": self._pil_to_png_bytes(pil_images[0]),
|
||||
})
|
||||
|
||||
if mask_tensor is not None:
|
||||
if not image_list:
|
||||
raise ValueError("提供了蒙版但未提供图片,请同时提供图片和蒙版")
|
||||
|
||||
first_tensor = image_list[0]
|
||||
if first_tensor.dim() == 3:
|
||||
first_tensor = first_tensor.unsqueeze(0)
|
||||
image_size = (first_tensor.shape[1], first_tensor.shape[2])
|
||||
mask_key = "mask"
|
||||
assets.append({
|
||||
"key": mask_key,
|
||||
"label": "蒙版",
|
||||
"bytes": self._mask_tensor_to_rgba_png_bytes(mask_tensor, image_size),
|
||||
})
|
||||
|
||||
def _make_body(asset_bytes: Dict[str, bytes]) -> dict:
|
||||
body = {
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"images": [
|
||||
self._png_bytes_to_data_url(asset_bytes[key])
|
||||
for key in image_keys
|
||||
if key in asset_bytes
|
||||
],
|
||||
"size": size if size else "auto",
|
||||
"quality": quality,
|
||||
"n": int(n),
|
||||
"output_format": output_format or "png",
|
||||
}
|
||||
if mask_key and mask_key in asset_bytes:
|
||||
body["mask"] = {
|
||||
"image_url": self._png_bytes_to_data_url(asset_bytes[mask_key])
|
||||
}
|
||||
return body
|
||||
|
||||
original_asset_bytes = {asset["key"]: asset["bytes"] for asset in assets}
|
||||
original_body = _make_body(original_asset_bytes)
|
||||
self._log_original_request_body("async generateImage", original_body)
|
||||
|
||||
asset_bytes = self._fit_png_assets_to_body_limit(assets, _make_body)
|
||||
body = _make_body(asset_bytes)
|
||||
body_size = self._json_body_size(body)
|
||||
if body_size > self._MAX_BODY_BYTES:
|
||||
raise RuntimeError(f"请求体超过 20MB:{body_size // 1024}KB")
|
||||
|
||||
return body
|
||||
|
||||
async def _submit_generate_image_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
payload: dict,
|
||||
) -> str:
|
||||
url = f"{self.base_url}{_ENDPOINT_ASYNC_GENERATE}"
|
||||
last_status = None
|
||||
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
t0 = time.time()
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers=self._json_headers(),
|
||||
) as resp:
|
||||
elapsed = time.time() - t0
|
||||
text = await resp.text()
|
||||
self._log_original_response_body(
|
||||
f"submit generateImage status={resp.status}",
|
||||
text,
|
||||
)
|
||||
|
||||
if resp.status not in (200, 201, 202):
|
||||
last_status = resp.status
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = get_friendly_message(resp.status)
|
||||
delay = _compute_delay(
|
||||
attempt,
|
||||
DEFAULT_BASE_DELAY,
|
||||
DEFAULT_MAX_DELAY,
|
||||
DEFAULT_BACKOFF_FACTOR,
|
||||
)
|
||||
print(f"[o1key GPT Image] {friendly} {delay:.1f}s 后重试提交 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(self._extract_error_message(text, resp.status))
|
||||
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"提交响应 JSON 解析失败,原始内容:{text[:500]}") from None
|
||||
|
||||
task_id = data.get("task_id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"提交响应中未找到 task_id: {data}")
|
||||
|
||||
status = data.get("status", "")
|
||||
print(f"[o1key GPT Image] 异步任务已提交 | task_id={task_id} | status={status} | 耗时 {elapsed:.1f}s")
|
||||
return task_id
|
||||
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
|
||||
raise RuntimeError(f"异步任务提交失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
|
||||
|
||||
async def _poll_generate_image_task(
|
||||
self,
|
||||
session: aiohttp.ClientSession,
|
||||
task_id: str,
|
||||
progress_callback: Optional[Callable[[int], None]] = None,
|
||||
) -> dict:
|
||||
url = f"{self.base_url}{_ENDPOINT_ASYNC_TASK.format(task_id=task_id)}"
|
||||
start_time = time.time()
|
||||
|
||||
poll_count = 0
|
||||
last_poll_at = start_time
|
||||
|
||||
while True:
|
||||
if poll_count < len(_ASYNC_POLL_SCHEDULE):
|
||||
next_poll_at = start_time + _ASYNC_POLL_SCHEDULE[poll_count]
|
||||
else:
|
||||
next_poll_at = last_poll_at + _ASYNC_POLL_INTERVAL
|
||||
|
||||
sleep_time = next_poll_at - time.time()
|
||||
if sleep_time > 0:
|
||||
await asyncio.sleep(sleep_time)
|
||||
|
||||
last_poll_at = time.time()
|
||||
elapsed = last_poll_at - start_time
|
||||
if elapsed > _ASYNC_MAX_WAIT:
|
||||
raise RuntimeError(f"任务 {task_id} 超时(>{int(_ASYNC_MAX_WAIT)}秒),请稍后用 task_id 查询结果")
|
||||
|
||||
poll_count += 1
|
||||
async with session.get(url, headers=self._auth_headers()) as resp:
|
||||
text = await resp.text()
|
||||
self._log_original_response_body(
|
||||
f"poll task status={resp.status}",
|
||||
text,
|
||||
)
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(self._extract_error_message(text, resp.status))
|
||||
|
||||
try:
|
||||
task = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"任务查询响应 JSON 解析失败,原始内容:{text[:500]}") from None
|
||||
|
||||
status = task.get("status", "UNKNOWN")
|
||||
progress = task.get("progress")
|
||||
progress_pct = self._coerce_progress_percent(progress)
|
||||
if POLL_LOG_ENABLED:
|
||||
progress_text = f" | progress={progress}" if progress is not None else ""
|
||||
print(f"[o1key GPT Image] 查询任务 #{poll_count} | task_id={task_id} | status={status}{progress_text}")
|
||||
|
||||
if status == "SUCCESS":
|
||||
self._emit_progress(progress_callback, 100)
|
||||
return task
|
||||
if progress_pct is not None and progress_pct < 100:
|
||||
self._emit_progress(progress_callback, progress_pct)
|
||||
if status == "FAILURE":
|
||||
error_message = self._extract_error_message(task, 500) or "生成失败"
|
||||
raise _AsyncImageTaskFailure(
|
||||
error_message,
|
||||
self._extract_error_detail(task),
|
||||
)
|
||||
if status not in ("SUBMITTED", "IN_PROGRESS"):
|
||||
raise RuntimeError(f"未知任务状态 {status}: {task}")
|
||||
|
||||
async def _parse_async_task_images(
|
||||
self,
|
||||
task: dict,
|
||||
session: aiohttp.ClientSession,
|
||||
) -> List[Image.Image]:
|
||||
data = task.get("data", {})
|
||||
image_items = data.get("images") if isinstance(data, dict) else None
|
||||
|
||||
if not isinstance(image_items, list) or not image_items:
|
||||
raise RuntimeError(f"任务结果中未找到 data.images: {task}")
|
||||
|
||||
images: List[Image.Image] = []
|
||||
for idx, item in enumerate(image_items, 1):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
|
||||
url = item.get("url") or item.get("image_url")
|
||||
b64 = item.get("b64_json", "")
|
||||
|
||||
if url and isinstance(url, str) and url.startswith("data:image"):
|
||||
try:
|
||||
_, b64_data = url.split(",", 1)
|
||||
img = Image.open(BytesIO(base64.b64decode(b64_data)))
|
||||
images.append(img)
|
||||
print(f"[o1key GPT Image] 第 {idx} 张 data URL 解码完成 ({img.size[0]}×{img.size[1]})")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"第 {idx} 张 data URL 解码失败: {e}") from None
|
||||
elif url and isinstance(url, str) and url.startswith("http"):
|
||||
async with session.get(url, allow_redirects=True) as resp:
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(f"图像下载失败 HTTP {resp.status},URL: {url}")
|
||||
img_bytes = await resp.read()
|
||||
img = Image.open(BytesIO(img_bytes))
|
||||
images.append(img)
|
||||
print(f"[o1key GPT Image] 第 {idx} 张 URL 下载完成 ({img.size[0]}×{img.size[1]}) | {url}")
|
||||
elif b64:
|
||||
img = self._decode_b64_image(b64, f"第 {idx} 张")
|
||||
images.append(img)
|
||||
else:
|
||||
print(f"[o1key GPT Image] 警告:第 {idx} 条结果既无 url 也无 b64_json,已跳过")
|
||||
|
||||
if not images:
|
||||
raise RuntimeError("任务成功但没有可用图片结果")
|
||||
|
||||
return images
|
||||
|
||||
async def _generate_image_task_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
quality: str,
|
||||
size: str,
|
||||
n: int,
|
||||
seed: int,
|
||||
image_tensor: Optional[List[torch.Tensor]] = None,
|
||||
mask_tensor: Optional[torch.Tensor] = None,
|
||||
output_format: str = "png",
|
||||
progress_callback: Optional[Callable[[int], None]] = None,
|
||||
) -> List[Image.Image]:
|
||||
body = self._build_async_generate_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
quality=quality,
|
||||
size=size,
|
||||
n=n,
|
||||
image_list=image_tensor,
|
||||
mask_tensor=mask_tensor,
|
||||
output_format=output_format,
|
||||
)
|
||||
|
||||
mode = "图像编辑" if mask_tensor is not None else ("图生图" if image_tensor else "文生图")
|
||||
body_size = self._json_body_size(body)
|
||||
print(
|
||||
f"[o1key GPT Image] {mode} | 新异步接口 | 模型={model} | "
|
||||
f"quality={quality} | size={size} | n={n} | body={body_size // 1024}KB"
|
||||
)
|
||||
|
||||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||
timeout = aiohttp.ClientTimeout(total=_ASYNC_MAX_WAIT + 120)
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
last_error = None
|
||||
for retry_index in range(len(_ASYNC_RETRY_DELAYS) + 1):
|
||||
try:
|
||||
task_id = await self._submit_generate_image_task(session, body)
|
||||
task = await self._poll_generate_image_task(
|
||||
session,
|
||||
task_id,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
return await self._parse_async_task_images(task, session)
|
||||
except _AsyncImageTaskFailure as error:
|
||||
last_error = error
|
||||
detail = error.error_detail
|
||||
if (
|
||||
retry_index < len(_ASYNC_RETRY_DELAYS)
|
||||
and self._should_retry_async_failure(detail, retry_index)
|
||||
):
|
||||
delay = _ASYNC_RETRY_DELAYS[retry_index]
|
||||
code = detail.get("code", "unknown")
|
||||
category = detail.get("category", "unknown")
|
||||
failed_task_id = detail.get("task_id", "")
|
||||
task_text = f" | failed_task_id={failed_task_id}" if failed_task_id else ""
|
||||
print(
|
||||
f"[o1key GPT Image] 任务失败但可重试 | code={code} | "
|
||||
f"category={category}{task_text} | {delay:.0f}s 后重试 "
|
||||
f"({retry_index + 1}/{len(_ASYNC_RETRY_DELAYS)})"
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
raise RuntimeError(str(error)) from None
|
||||
|
||||
if last_error is not None:
|
||||
raise RuntimeError(str(last_error)) from None
|
||||
raise RuntimeError("生成失败")
|
||||
|
||||
# ── 文生图 / 图生图(generations 接口)───────────────────────────────────
|
||||
|
||||
async def _generate_async(
|
||||
@@ -806,6 +1385,53 @@ class GptImageClient:
|
||||
f"o1key GPT Image 请求超时(>{_REQUEST_TIMEOUT}s),请检查网络或稍后重试"
|
||||
)
|
||||
|
||||
def generate_image_async_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
quality: str,
|
||||
size: str,
|
||||
n: int,
|
||||
seed: int,
|
||||
image_tensor: Optional[List[torch.Tensor]] = None,
|
||||
mask_tensor: Optional[torch.Tensor] = None,
|
||||
output_format: str = "png",
|
||||
progress_callback: Optional[Callable[[int], None]] = None,
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
新版异步任务入口,供节点调用。
|
||||
旧 run_sync 保留兼容,但 GPT Image 节点不再使用旧同步接口。
|
||||
"""
|
||||
coro = self._generate_image_task_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
quality=quality,
|
||||
size=size,
|
||||
n=n,
|
||||
seed=seed,
|
||||
image_tensor=image_tensor,
|
||||
mask_tensor=mask_tensor,
|
||||
output_format=output_format,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
def _run():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(self._run_with_interrupt(coro))
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(_run)
|
||||
try:
|
||||
return future.result(timeout=_ASYNC_MAX_WAIT + 150)
|
||||
except TimeoutError:
|
||||
raise RuntimeError(
|
||||
f"o1key GPT Image 异步任务超时(>{int(_ASYNC_MAX_WAIT)}秒),请稍后用 task_id 查询结果"
|
||||
)
|
||||
|
||||
# ── 余额查询 ──────────────────────────────────────────────────────────────
|
||||
|
||||
async def _query_balance_async(self) -> dict:
|
||||
|
||||
@@ -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 Video:content 下载重试 {attempt + 1}/3,{delay}s 后继续...")
|
||||
await interruptible_sleep(delay)
|
||||
|
||||
raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
|
||||
|
||||
def generate_video_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
seconds: int,
|
||||
quality: str,
|
||||
images: Optional[List[str]],
|
||||
output_dir: Optional[str] = None,
|
||||
save_path: Optional[str] = None,
|
||||
poll_interval: int = 5,
|
||||
timeout: int = 900,
|
||||
progress_callback: Optional[Callable[[int, str, float], None]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
async def _run():
|
||||
async with self._make_session() as session:
|
||||
body = self.build_video_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
quality=quality,
|
||||
images=images,
|
||||
)
|
||||
|
||||
create_response = await self.create_video_async(body, session)
|
||||
task_id = self._extract_task_id(create_response) or ""
|
||||
if not task_id:
|
||||
raise RuntimeError(
|
||||
"Grok Video 未返回任务 ID。\n"
|
||||
f"endpoint: {self.CREATE_ENDPOINT}\n"
|
||||
f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}"
|
||||
)
|
||||
|
||||
print(f"Grok Video:任务已提交,任务ID:{task_id}")
|
||||
print("Grok Video:视频生成中...")
|
||||
status_response = await self.poll_video_status_async(
|
||||
task_id=task_id,
|
||||
session=session,
|
||||
poll_interval=poll_interval,
|
||||
timeout=timeout,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
video_url = extract_video_url(status_response)
|
||||
print("Grok Video:视频生成完成,正在下载...")
|
||||
if save_path is None:
|
||||
resolved_output_dir = output_dir or os.getcwd()
|
||||
os.makedirs(resolved_output_dir, exist_ok=True)
|
||||
target_path = os.path.join(
|
||||
resolved_output_dir,
|
||||
f"{self._safe_task_filename(task_id)}.mp4",
|
||||
)
|
||||
else:
|
||||
target_path = save_path
|
||||
|
||||
if video_url:
|
||||
video_path = await self._download_url_to_file(video_url, target_path, session)
|
||||
else:
|
||||
video_path = await self.download_video_async(task_id, target_path, session)
|
||||
|
||||
return {
|
||||
"task_id": task_id,
|
||||
"status": extract_status(status_response),
|
||||
"video_path": video_path,
|
||||
"raw_json": {
|
||||
"create": create_response,
|
||||
"status": status_response,
|
||||
},
|
||||
}
|
||||
|
||||
return self.run_async_in_thread(_run())
|
||||
@@ -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())
|
||||
@@ -110,10 +110,13 @@ class KVideoFirstLast:
|
||||
"mode": mode_api,
|
||||
"duration": 时长,
|
||||
}
|
||||
if 生成音频 == "打开":
|
||||
body["generate_audio"] = True
|
||||
metadata = {}
|
||||
if 尾帧 is not None:
|
||||
body["metadata"] = {"image_tail": _image_to_base64(尾帧, scale)}
|
||||
metadata["image_tail"] = _image_to_base64(尾帧, scale)
|
||||
if 生成音频 == "打开":
|
||||
metadata["sound"] = "on"
|
||||
if metadata:
|
||||
body["metadata"] = metadata
|
||||
|
||||
body_str = json.dumps(body, ensure_ascii=False)
|
||||
body_size = len(body_str.encode("utf-8"))
|
||||
|
||||
@@ -82,6 +82,9 @@ class KVideoImage2Video:
|
||||
CATEGORY = "comfyui_o1key/KVideo"
|
||||
|
||||
async def generate(self, 起始帧, 提示词, 模式, 时长, 生成音频="关闭", 网络线路="全球加速", seed=0):
|
||||
if 模式 == "720p" and 生成音频 == "打开":
|
||||
raise RuntimeError("K26 仅1080p支持音频,请将模式切换为1080p或关闭生成音频。")
|
||||
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_base_url_by_route(网络线路)
|
||||
headers = {
|
||||
@@ -109,7 +112,7 @@ class KVideoImage2Video:
|
||||
"duration": 时长,
|
||||
}
|
||||
if 生成音频 == "打开":
|
||||
body["generate_audio"] = True
|
||||
body["metadata"] = {"sound": "on"}
|
||||
|
||||
body_str = json.dumps(body, ensure_ascii=False)
|
||||
body_size = len(body_str.encode("utf-8"))
|
||||
|
||||
+3
-1
@@ -14,6 +14,7 @@ from .remove_metadata import BatchCleanMetadata
|
||||
from .video_preview import VideoPreview
|
||||
from .kling_video import KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset
|
||||
from .veo_video import GoogleVeo
|
||||
from .newapi_veo_video import Google31Video
|
||||
from .flux_edit import FluxImageEdit
|
||||
from .universal_llm import UniversalLLMChat
|
||||
from .batch_images_o1key import BatchImagesO1key
|
||||
@@ -22,6 +23,7 @@ from .nano_banana_v2 import NanoBananaV2, NanoBananaV2Batch, AsyncImageGenerator
|
||||
from .doubao_image import DoubaoImage
|
||||
from .gpt_image import O1keyGPTImage, O1keyGPTImageBatch
|
||||
from .grok_image import O1keyGrokImage
|
||||
from .grok_video import O1keyGrokVideo
|
||||
from .K_video_firstlast import KVideoFirstLast
|
||||
from .K_video_image2video import KVideoImage2Video
|
||||
from .K3_video import K3Video
|
||||
@@ -33,4 +35,4 @@ from .remove_bg import O1keyRemoveBackground
|
||||
from .color_remove_bg import O1keyColorRemoveBG
|
||||
from .grid_splitter import O1keyGridSplitter
|
||||
|
||||
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter']
|
||||
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'Google31Video', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGPTImageBatch', 'O1keyGrokImage', 'O1keyGrokVideo', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG', 'O1keyGridSplitter']
|
||||
|
||||
+172
-145
@@ -4,23 +4,21 @@ ComfyUI 自定义节点,用于批量处理图像生成任务
|
||||
支持多文件夹加载、1:1/笛卡尔积配对、智能命名保存
|
||||
"""
|
||||
|
||||
import io as _io
|
||||
import re
|
||||
import json
|
||||
import time
|
||||
import math
|
||||
import base64
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Optional, Tuple, List
|
||||
from typing import Callable, Optional, Tuple, List
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_images_for_request_body_limit
|
||||
from comfy_api.latest import io
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import (
|
||||
ImageInfo,
|
||||
load_images_from_folder,
|
||||
@@ -30,7 +28,7 @@ from ..utils.file_utils import (
|
||||
save_image,
|
||||
)
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route, get_api_key_or_raise
|
||||
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
|
||||
from ..utils.nano_banana_async import generate_nano_banana_async
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..models_config import (
|
||||
get_model_supported_aspect_ratios, get_all_supported_aspect_ratios,
|
||||
@@ -73,64 +71,28 @@ REQUEST_LOG_ENABLED = False
|
||||
# ============================================================================
|
||||
|
||||
_NODE = "Nano Banana"
|
||||
_ENDPOINT = "/v1/chat/completions"
|
||||
|
||||
_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")
|
||||
|
||||
|
||||
def _get_headers(api_key: str) -> dict:
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"X-Accel-Buffering": "no",
|
||||
"Cache-Control": "no-cache, no-transform",
|
||||
}
|
||||
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
|
||||
if pbar is None:
|
||||
return None
|
||||
|
||||
last_progress = [0.0]
|
||||
|
||||
def _on_progress(progress: float) -> None:
|
||||
try:
|
||||
progress = max(0.0, min(float(progress), 1.0))
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
if progress <= last_progress[0]:
|
||||
return
|
||||
pbar.update(progress - last_progress[0])
|
||||
last_progress[0] = progress
|
||||
|
||||
return _on_progress
|
||||
|
||||
|
||||
def _build_request_body(
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
) -> dict:
|
||||
google_config = {
|
||||
"image_config": {
|
||||
"image_size": resolution,
|
||||
}
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
google_config["image_config"]["aspect_ratio"] = aspect_ratio
|
||||
|
||||
def _make_body(encoded_images: Optional[List[tuple]] = None) -> dict:
|
||||
content_parts = [{"type": "text", "text": prompt}]
|
||||
if encoded_images:
|
||||
for mime_type, b64 in encoded_images:
|
||||
content_parts.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:{mime_type};base64,{b64}"}
|
||||
})
|
||||
|
||||
body = {
|
||||
"model": model,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": content_parts}],
|
||||
"extra_body": {"google": google_config},
|
||||
}
|
||||
|
||||
if enable_grounding:
|
||||
body["extra_body"]["google_search"] = True
|
||||
return body
|
||||
|
||||
encoded_images = None
|
||||
if images:
|
||||
encoded_images = encode_images_for_request_body_limit(images, _make_body)
|
||||
|
||||
return _make_body(encoded_images)
|
||||
|
||||
|
||||
async def _generate_single_openai(
|
||||
async def _generate_single_async(
|
||||
session: aiohttp.ClientSession,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
@@ -140,82 +102,25 @@ async def _generate_single_openai(
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> List[Image.Image]:
|
||||
url = f"{base_url}{_ENDPOINT}"
|
||||
headers = _get_headers(api_key)
|
||||
body = _build_request_body(
|
||||
result_images, _ = await generate_nano_banana_async(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
node_label="BatchNanoBananaPro",
|
||||
request_log_enabled=REQUEST_LOG_ENABLED,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
if REQUEST_LOG_ENABLED:
|
||||
extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
|
||||
print(f"[请求] POST {url} | model={model} | extra_body={extra}")
|
||||
|
||||
last_status = None
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
resp = await session.post(url, headers=headers, json=body)
|
||||
if resp.status == 200:
|
||||
break
|
||||
last_status = resp.status
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
resp.close()
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
error_text = await resp.text()
|
||||
resp.close()
|
||||
if resp.status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
|
||||
try:
|
||||
err_json = json.loads(error_text)
|
||||
msg = err_json.get("error", {}).get("message", error_text[:200])
|
||||
except Exception:
|
||||
msg = error_text[:200]
|
||||
raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
|
||||
else:
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
|
||||
raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
|
||||
|
||||
full_content = ""
|
||||
buffer = ""
|
||||
async for raw_chunk in resp.content.iter_any():
|
||||
buffer += raw_chunk.decode("utf-8")
|
||||
while "\n" in buffer:
|
||||
line_str, buffer = buffer.split("\n", 1)
|
||||
line_str = line_str.strip()
|
||||
if not line_str or not line_str.startswith("data:"):
|
||||
continue
|
||||
data_str = line_str[5:].strip()
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
try:
|
||||
chunk = json.loads(data_str)
|
||||
delta = chunk.get("choices", [{}])[0].get("delta", {})
|
||||
if "content" in delta:
|
||||
full_content += delta["content"]
|
||||
except (json.JSONDecodeError, IndexError):
|
||||
continue
|
||||
resp.close()
|
||||
|
||||
if not full_content:
|
||||
raise RuntimeError("API 未返回有效内容")
|
||||
|
||||
matches = list(_IMAGE_RE.finditer(full_content))
|
||||
if not matches:
|
||||
raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")
|
||||
|
||||
last_match = matches[-1]
|
||||
img_data = base64.b64decode(last_match.group(2))
|
||||
final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
return [final_image]
|
||||
return result_images
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
@@ -244,7 +149,7 @@ def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.
|
||||
return pil_to_tensor(matched)
|
||||
|
||||
|
||||
class BatchNanoBananaPro:
|
||||
class BatchNanoBananaPro(io.ComfyNode):
|
||||
"""
|
||||
批量 Nano Banana 节点
|
||||
|
||||
@@ -291,6 +196,116 @@ class BatchNanoBananaPro:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
normal_aspect_ratios = [
|
||||
"智能", "1:1", "2:3", "3:2", "3:4", "4:3",
|
||||
"4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
]
|
||||
nano_banana_2_aspect_ratios = [
|
||||
"智能", "1:1", "1:4", "1:8", "2:3", "3:2", "3:4",
|
||||
"4:1", "4:3", "4:5", "5:4", "8:1",
|
||||
"9:16", "16:9", "21:9",
|
||||
]
|
||||
|
||||
return io.Schema(
|
||||
node_id="BatchNanoBananaPro",
|
||||
display_name="批量 Nano Banana",
|
||||
category="image/batch",
|
||||
inputs=[
|
||||
io.String.Input(
|
||||
"prompt",
|
||||
default="一个中国女子的OOTD",
|
||||
multiline=True,
|
||||
),
|
||||
io.DynamicCombo.Input("模型", options=[
|
||||
io.DynamicCombo.Option("Nano Banana Pro", [
|
||||
io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"),
|
||||
io.Combo.Input("分辨率", options=["1K", "2K", "4K"], default="2K"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
]),
|
||||
io.DynamicCombo.Option("Nano Banana 2", [
|
||||
io.Combo.Input("宽高比", options=nano_banana_2_aspect_ratios, default="智能"),
|
||||
io.Combo.Input("分辨率", options=["512px", "1K", "2K", "4K"], default="2K"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
io.Combo.Input("思考深度", options=["高", "低"], default="高"),
|
||||
]),
|
||||
io.DynamicCombo.Option("Nano Banana", [
|
||||
io.Combo.Input("宽高比", options=normal_aspect_ratios, default="智能"),
|
||||
io.Combo.Input("分辨率", options=["1K"], default="1K"),
|
||||
io.Combo.Input("谷歌搜索", options=["关闭", "打开"], default="关闭"),
|
||||
]),
|
||||
]),
|
||||
io.Combo.Input("图片格式", options=["原始", "JPEG", "PNG", "WebP"], default="原始"),
|
||||
io.Combo.Input("计费", options=["特价", "官方"], default="特价"),
|
||||
io.Combo.Input("网络", options=NETWORK_ROUTE_OPTIONS, default="全球加速"),
|
||||
io.Int.Input("seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF),
|
||||
io.String.Input("文件夹1", default="", multiline=False),
|
||||
io.String.Input("文件夹2", default="", multiline=False),
|
||||
io.String.Input("文件夹3", default="", multiline=False),
|
||||
io.String.Input("文件夹4", default="", multiline=False),
|
||||
io.String.Input("文件夹5", default="", multiline=False),
|
||||
io.String.Input("保存路径", default="", multiline=False),
|
||||
io.Combo.Input("图片配对模式", options=cls.PAIRING_MODES, default="不配对"),
|
||||
io.Image.Input("参考图1", optional=True),
|
||||
io.Image.Input("参考图2", optional=True),
|
||||
io.Image.Input("参考图3", optional=True),
|
||||
io.Image.Input("参考图4", optional=True),
|
||||
io.Image.Input("参考图5", optional=True),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="输出图像"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
prompt,
|
||||
模型,
|
||||
图片格式,
|
||||
计费,
|
||||
网络,
|
||||
seed,
|
||||
文件夹1,
|
||||
文件夹2,
|
||||
文件夹3,
|
||||
文件夹4,
|
||||
文件夹5,
|
||||
保存路径,
|
||||
图片配对模式,
|
||||
**kwargs,
|
||||
) -> io.NodeOutput:
|
||||
model_name = 模型["模型"]
|
||||
宽高比 = 模型.get("宽高比", "智能")
|
||||
分辨率 = 模型.get("分辨率", "2K")
|
||||
思考深度 = 模型.get("思考深度")
|
||||
谷歌搜索 = 模型.get("谷歌搜索", "关闭")
|
||||
if 思考深度:
|
||||
kwargs["思考深度"] = 思考深度
|
||||
kwargs["谷歌搜索"] = 谷歌搜索
|
||||
|
||||
node = cls()
|
||||
output_tensor, = node.process_batch(
|
||||
prompt=prompt,
|
||||
文件夹1=文件夹1,
|
||||
文件夹2=文件夹2,
|
||||
文件夹3=文件夹3,
|
||||
文件夹4=文件夹4,
|
||||
文件夹5=文件夹5,
|
||||
seed=seed,
|
||||
图片配对模式=图片配对模式,
|
||||
模型=model_name,
|
||||
计费=计费,
|
||||
宽高比=宽高比,
|
||||
分辨率=分辨率,
|
||||
图片格式=图片格式,
|
||||
网络=网络,
|
||||
保存路径=保存路径,
|
||||
**kwargs,
|
||||
)
|
||||
return io.NodeOutput(output_tensor)
|
||||
|
||||
def resize_to_megapixels(
|
||||
self,
|
||||
image: Image.Image,
|
||||
@@ -352,7 +367,6 @@ class BatchNanoBananaPro:
|
||||
optional_inputs["图片配对模式"] = (cls.PAIRING_MODES, {
|
||||
"default": "不配对"
|
||||
})
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"prompt": ("STRING", {
|
||||
@@ -532,9 +546,11 @@ class BatchNanoBananaPro:
|
||||
enable_grounding: bool = False,
|
||||
base_filename: str = None,
|
||||
image_format: str = "原始",
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
执行单个生成任务(OpenAI 兼容接口)
|
||||
执行单个生成任务(异步生图接口)
|
||||
"""
|
||||
result = {
|
||||
"task_index": task_index,
|
||||
@@ -550,10 +566,10 @@ class BatchNanoBananaPro:
|
||||
# 准备输入图片
|
||||
input_pil_images = [info.image for info in images]
|
||||
|
||||
# 调用 OpenAI 兼容接口生成图片
|
||||
# 调用异步生图接口生成图片
|
||||
generated_images = []
|
||||
try:
|
||||
gen_images = await _generate_single_openai(
|
||||
gen_images = await _generate_single_async(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
@@ -563,6 +579,8 @@ class BatchNanoBananaPro:
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=input_pil_images if input_pil_images else None,
|
||||
enable_grounding=enable_grounding,
|
||||
progress_callback=progress_callback,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
generated_images.extend(gen_images)
|
||||
except Exception as e:
|
||||
@@ -661,9 +679,10 @@ class BatchNanoBananaPro:
|
||||
prompts_per_task: Optional[List[str]] = None,
|
||||
enable_grounding: bool = False,
|
||||
image_format: str = "原始",
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> List[dict]:
|
||||
"""
|
||||
异步批量处理所有任务(OpenAI 兼容接口)
|
||||
异步批量处理所有任务(异步生图接口)
|
||||
"""
|
||||
total_tasks = len(pairs)
|
||||
|
||||
@@ -734,6 +753,8 @@ class BatchNanoBananaPro:
|
||||
enable_grounding=enable_grounding,
|
||||
base_filename=base_filename,
|
||||
image_format=image_format,
|
||||
progress_callback=_make_progress_callback(pbar),
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
@@ -773,13 +794,12 @@ class BatchNanoBananaPro:
|
||||
else:
|
||||
fail_count += 1
|
||||
|
||||
# 更新 ComfyUI 原生进度条
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
# 大任务额外显示百分比里程碑
|
||||
if show_milestone and milestone_index < len(milestones):
|
||||
progress = completed / total_tasks
|
||||
if pbar is not None and getattr(pbar, "total", 0):
|
||||
progress = pbar.current / pbar.total
|
||||
else:
|
||||
progress = success_count / total_tasks
|
||||
if progress >= milestones[milestone_index]:
|
||||
percentage = int(milestones[milestone_index] * 100)
|
||||
print(f"BatchNanoBananaPro: >>> 进度 {percentage}% <<<")
|
||||
@@ -856,10 +876,15 @@ class BatchNanoBananaPro:
|
||||
start_time = time.time()
|
||||
|
||||
# 从 kwargs 提取搜索参数(界面显示为「关闭/打开」,转为 bool 供调用)
|
||||
enable_grounding: bool = False
|
||||
enable_grounding: bool = kwargs.get("谷歌搜索", "关闭") == "打开"
|
||||
|
||||
# 拼接实际模型 ID
|
||||
base_model_id = self.MODEL_ID_MAP.get(模型, "nano-banana-pro")
|
||||
思考深度 = kwargs.get("思考深度", "高")
|
||||
thinking_level = None
|
||||
if base_model_id == "nano-banana-2":
|
||||
thinking_level = "High" if 思考深度 == "高" else "Low"
|
||||
|
||||
if base_model_id == "nano-banana":
|
||||
if 计费 == "官方":
|
||||
raise ValueError(f"模型 \"{模型}\" 仅支持特价计费")
|
||||
@@ -961,11 +986,12 @@ class BatchNanoBananaPro:
|
||||
grounding_str = ""
|
||||
if enable_grounding:
|
||||
grounding_str = " | 谷歌搜索接地"
|
||||
thinking_str = f" | 思考:{thinking_level}" if thinking_level else ""
|
||||
|
||||
if batch_prompts:
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}")
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 × {len(batch_prompts)}个提示词 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
|
||||
else:
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}")
|
||||
print(f"BatchNanoBananaPro: 批量任务 | {图片配对模式} 配对模式 | 共 {total_tasks} 任务{grounding_str}{thinking_str}")
|
||||
|
||||
# 创建 ComfyUI 原生进度条
|
||||
pbar = None
|
||||
@@ -1025,6 +1051,7 @@ class BatchNanoBananaPro:
|
||||
prompts_per_task=prompts_per_task,
|
||||
enable_grounding=enable_grounding,
|
||||
image_format=图片格式,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -1146,7 +1173,7 @@ class BatchNanoBananaPro:
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
finally:
|
||||
# 查询余额
|
||||
|
||||
+82
-7
@@ -42,6 +42,56 @@ except ImportError:
|
||||
_FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
|
||||
def _make_node_progress_callback(progress_bar, task_index: int, total_tasks: int):
|
||||
if progress_bar is None:
|
||||
return None
|
||||
|
||||
total_units = max(1, total_tasks) * 100
|
||||
base_units = max(0, task_index - 1) * 100
|
||||
last_pct = {"value": -1}
|
||||
|
||||
def _callback(pct: int):
|
||||
try:
|
||||
pct_value = int(round(float(pct)))
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
pct_value = max(0, min(100, pct_value))
|
||||
if pct_value < last_pct["value"]:
|
||||
return
|
||||
last_pct["value"] = pct_value
|
||||
progress_bar.update_absolute(
|
||||
min(total_units, base_units + pct_value),
|
||||
total_units,
|
||||
)
|
||||
|
||||
return _callback
|
||||
|
||||
|
||||
def _resolve_async_size(value: str) -> str:
|
||||
value = (value or "").strip()
|
||||
if not value or value == "智能" or value.lower() == "auto":
|
||||
return "auto"
|
||||
|
||||
first_part = value.split("(")[0].strip()
|
||||
normalized_size = first_part.lower().replace("*", "x").replace("×", "x")
|
||||
size_parts = [part.strip() for part in normalized_size.split("x")]
|
||||
if len(size_parts) == 2 and all(part.isdigit() for part in size_parts):
|
||||
return f"{int(size_parts[0])}x{int(size_parts[1])}"
|
||||
|
||||
allowed = {"auto", "1024x1024", "1K", "2K", "4K"}
|
||||
if first_part in allowed:
|
||||
return first_part
|
||||
|
||||
if "4K" in value:
|
||||
return "4K"
|
||||
if "2K" in value:
|
||||
return "2K"
|
||||
if "1K" in value:
|
||||
return "1K"
|
||||
|
||||
return "auto"
|
||||
|
||||
|
||||
class O1keyGPTImage:
|
||||
"""
|
||||
o1key GPT Image 节点
|
||||
@@ -123,6 +173,10 @@ class O1keyGPTImage:
|
||||
"default": "自动",
|
||||
"tooltip": "Image quality: 高=high, 中=medium, 低=low, 自动=auto",
|
||||
})
|
||||
optional_inputs["输出格式"] = (["png", "jpeg", "webp"], {
|
||||
"default": "jpeg",
|
||||
"tooltip": "Generated image output format",
|
||||
})
|
||||
optional_inputs["seed"] = ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
@@ -160,6 +214,7 @@ class O1keyGPTImage:
|
||||
网络: str = "全球加速",
|
||||
分辨率: str = "auto",
|
||||
质量: str = "自动",
|
||||
输出格式: str = "jpeg",
|
||||
生图数量: int = 1,
|
||||
seed: int = 0,
|
||||
遮罩=None,
|
||||
@@ -190,7 +245,7 @@ class O1keyGPTImage:
|
||||
raise ValueError("提供了遮罩但未提供图片,请同时提供图片和遮罩")
|
||||
|
||||
# ── 2. 解析分辨率显示值 → API 参数值 ──────────────────────────────────
|
||||
size = "auto" if 分辨率 == "智能" else 分辨率.split("(")[0].strip()
|
||||
size = _resolve_async_size(分辨率)
|
||||
|
||||
# ── 2b. 解析模型显示值 → API 参数值 ───────────────────────────────────
|
||||
_model_map = {"gpt-image-2-次卡": "gpt-image-2-c", "gpt-image-2-按量": "gpt-image-2"}
|
||||
@@ -216,6 +271,8 @@ class O1keyGPTImage:
|
||||
|
||||
# ── 5. 调用 API ───────────────────────────────────────────────────
|
||||
all_pil_images = []
|
||||
progress_total = len(batch_prompts) if batch_prompts else 1
|
||||
progress_bar = ProgressBar(progress_total * 100) if _PROGRESS_BAR_AVAILABLE else None
|
||||
|
||||
if batch_prompts:
|
||||
# 批量模式:逐条提示词调用
|
||||
@@ -226,7 +283,7 @@ class O1keyGPTImage:
|
||||
print("[o1key GPT Image] 用户取消,已中断批量生成")
|
||||
raise InterruptProcessingException()
|
||||
try:
|
||||
pil_images = client.run_sync(
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=p,
|
||||
model=model,
|
||||
quality=quality,
|
||||
@@ -235,6 +292,8 @@ class O1keyGPTImage:
|
||||
seed=seed,
|
||||
image_tensor=图片,
|
||||
mask_tensor=遮罩,
|
||||
output_format=输出格式,
|
||||
progress_callback=_make_node_progress_callback(progress_bar, idx, total),
|
||||
)
|
||||
all_pil_images.extend(pil_images)
|
||||
snippet = p[:30] + ("..." if len(p) >= 30 else "")
|
||||
@@ -245,12 +304,14 @@ class O1keyGPTImage:
|
||||
error_msg = str(e).split('\n')[0]
|
||||
snippet = p[:30] + ("..." if len(p) >= 30 else "")
|
||||
print(f"[o1key GPT Image] [{idx}/{total}] ❌ {snippet} → {error_msg}")
|
||||
if progress_bar is not None:
|
||||
progress_bar.update_absolute(idx * 100, total * 100)
|
||||
else:
|
||||
# 单提示词模式
|
||||
if not prompt or not prompt.strip():
|
||||
raise ValueError("提示词不能为空")
|
||||
try:
|
||||
pil_images = client.run_sync(
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
quality=quality,
|
||||
@@ -259,6 +320,8 @@ class O1keyGPTImage:
|
||||
seed=seed,
|
||||
image_tensor=图片,
|
||||
mask_tensor=遮罩,
|
||||
output_format=输出格式,
|
||||
progress_callback=_make_node_progress_callback(progress_bar, 1, 1),
|
||||
)
|
||||
all_pil_images.extend(pil_images)
|
||||
except InterruptProcessingException:
|
||||
@@ -495,7 +558,7 @@ class O1keyGPTImageBatch:
|
||||
|
||||
@staticmethod
|
||||
def _resolve_size(分辨率: str) -> str:
|
||||
return "auto" if 分辨率 == "智能" else 分辨率.split("(")[0].strip()
|
||||
return _resolve_async_size(分辨率)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_model(模型: str) -> str:
|
||||
@@ -510,6 +573,15 @@ class O1keyGPTImageBatch:
|
||||
quality_map = {"高": "high", "中": "medium", "低": "low", "自动": "auto"}
|
||||
return quality_map.get(质量, "auto")
|
||||
|
||||
@staticmethod
|
||||
def _resolve_output_format(图片格式: str) -> str:
|
||||
output_format_map = {
|
||||
"JPEG": "jpeg",
|
||||
"PNG": "png",
|
||||
"WebP": "webp",
|
||||
}
|
||||
return output_format_map.get(图片格式, "png")
|
||||
|
||||
@staticmethod
|
||||
def _ensure_output_folder(保存路径: str) -> str:
|
||||
output_folder = (保存路径 or "").strip()
|
||||
@@ -637,11 +709,12 @@ class O1keyGPTImageBatch:
|
||||
size = self._resolve_size(分辨率)
|
||||
model = self._resolve_model(模型)
|
||||
quality = self._resolve_quality(质量)
|
||||
output_format = self._resolve_output_format(图片格式)
|
||||
|
||||
client = GptImageClient()
|
||||
client.base_url = get_base_url_by_route(网络)
|
||||
|
||||
progress_bar = ProgressBar(total_tasks) if _PROGRESS_BAR_AVAILABLE else None
|
||||
progress_bar = ProgressBar(total_tasks * 100) if _PROGRESS_BAR_AVAILABLE else None
|
||||
results = []
|
||||
all_saved_files = []
|
||||
|
||||
@@ -661,7 +734,7 @@ class O1keyGPTImageBatch:
|
||||
}
|
||||
|
||||
try:
|
||||
pil_images = client.run_sync(
|
||||
pil_images = client.generate_image_async_sync(
|
||||
prompt=task_prompt,
|
||||
model=model,
|
||||
quality=quality,
|
||||
@@ -670,6 +743,8 @@ class O1keyGPTImageBatch:
|
||||
seed=seed,
|
||||
image_tensor=self._pair_to_tensors(pair),
|
||||
mask_tensor=遮罩,
|
||||
output_format=output_format,
|
||||
progress_callback=_make_node_progress_callback(progress_bar, task_index, total_tasks),
|
||||
)
|
||||
saved_files = self._save_images(
|
||||
images=pil_images,
|
||||
@@ -691,7 +766,7 @@ class O1keyGPTImageBatch:
|
||||
|
||||
results.append(result)
|
||||
if progress_bar is not None:
|
||||
progress_bar.update(1)
|
||||
progress_bar.update_absolute(task_index * 100, total_tasks * 100)
|
||||
|
||||
success_count = sum(1 for result in results if result.get("success", False))
|
||||
total_generated = sum(result.get("generated_count", 0) for result in results)
|
||||
|
||||
@@ -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
@@ -1,20 +1,16 @@
|
||||
"""
|
||||
Nano Banana 节点 (V3)
|
||||
ComfyUI 自定义节点,用于调用生图模型(OpenAI 兼容接口)
|
||||
ComfyUI 自定义节点,用于调用异步生图模型
|
||||
使用 V3 DynamicCombo 实现模型-宽高比-分辨率动态联动
|
||||
"""
|
||||
|
||||
import io as _io
|
||||
import re
|
||||
import json
|
||||
import time
|
||||
import math
|
||||
import base64
|
||||
import random
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
@@ -22,21 +18,15 @@ from PIL import Image
|
||||
|
||||
from comfy_api.latest import io
|
||||
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts, encode_images_for_image_size_limit
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.config import (
|
||||
NETWORK_ROUTE_OPTIONS,
|
||||
get_base_url_by_route,
|
||||
get_api_key_or_raise,
|
||||
)
|
||||
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
|
||||
from ..utils.nano_banana_async import generate_nano_banana_async
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
|
||||
try:
|
||||
import folder_paths
|
||||
FOLDER_PATHS_AVAILABLE = True
|
||||
except ImportError:
|
||||
FOLDER_PATHS_AVAILABLE = False
|
||||
|
||||
try:
|
||||
from comfy.utils import ProgressBar
|
||||
PROGRESS_BAR_AVAILABLE = True
|
||||
@@ -51,17 +41,9 @@ except ImportError:
|
||||
InterruptProcessingException = RuntimeError
|
||||
processing_interrupted = lambda: False
|
||||
|
||||
try:
|
||||
import psutil
|
||||
MEMORY_MONITOR_AVAILABLE = True
|
||||
except ImportError:
|
||||
MEMORY_MONITOR_AVAILABLE = False
|
||||
|
||||
DEBUG_LOG_ENABLED = True
|
||||
REQUEST_LOG_ENABLED = False
|
||||
|
||||
_NODE = "Nano Banana"
|
||||
_ENDPOINT = "/v1/chat/completions"
|
||||
_REQUEST_TIMEOUT = 900
|
||||
_INTERRUPT_CHECK_INTERVAL = 0.2
|
||||
|
||||
@@ -112,6 +94,25 @@ def _check_interrupt():
|
||||
raise InterruptProcessingException()
|
||||
|
||||
|
||||
def _make_progress_callback(pbar) -> Optional[Callable[[float], None]]:
|
||||
if pbar is None:
|
||||
return None
|
||||
|
||||
last_progress = [0.0]
|
||||
|
||||
def _on_progress(progress: float) -> None:
|
||||
try:
|
||||
progress = max(0.0, min(float(progress), 1.0))
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
if progress <= last_progress[0]:
|
||||
return
|
||||
pbar.update(progress - last_progress[0])
|
||||
last_progress[0] = progress
|
||||
|
||||
return _on_progress
|
||||
|
||||
|
||||
def _images_to_tensor_safe(images: List[Image.Image], node_label: str) -> torch.Tensor:
|
||||
if not images:
|
||||
placeholder = Image.new('RGB', (512, 512), color=(128, 128, 128))
|
||||
@@ -169,66 +170,6 @@ def _build_model_id(model_name: str, resolution: str, billing: str) -> str:
|
||||
model_id += "-official"
|
||||
return model_id
|
||||
|
||||
_IMAGE_RE = re.compile(r"!\[.*?\]\(data:image/(\w+);base64,([A-Za-z0-9+/=]+)\)")
|
||||
|
||||
|
||||
def _get_headers(api_key: str) -> dict:
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"X-Accel-Buffering": "no",
|
||||
"Cache-Control": "no-cache, no-transform",
|
||||
}
|
||||
|
||||
|
||||
def _build_request_body(
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
) -> dict:
|
||||
google_config = {
|
||||
"image_config": {
|
||||
"image_size": resolution,
|
||||
}
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
google_config["image_config"]["aspect_ratio"] = aspect_ratio
|
||||
if thinking_level:
|
||||
google_config["thinking_config"] = {
|
||||
"thinking_level": thinking_level.lower(),
|
||||
"include_thoughts": True,
|
||||
}
|
||||
|
||||
def _make_body(encoded_images: Optional[List[tuple]] = None) -> dict:
|
||||
content_parts = [{"type": "text", "text": prompt}]
|
||||
if encoded_images:
|
||||
for mime_type, b64 in encoded_images:
|
||||
content_parts.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:{mime_type};base64,{b64}"}
|
||||
})
|
||||
|
||||
body = {
|
||||
"model": model,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": content_parts}],
|
||||
"extra_body": {"google": google_config},
|
||||
}
|
||||
|
||||
if enable_grounding:
|
||||
body["extra_body"]["google_search"] = True
|
||||
return body
|
||||
|
||||
encoded_images = None
|
||||
if images:
|
||||
encoded_images = encode_images_for_image_size_limit(images)
|
||||
|
||||
return _make_body(encoded_images)
|
||||
|
||||
|
||||
async def _generate_single(
|
||||
session: aiohttp.ClientSession,
|
||||
@@ -241,105 +182,25 @@ async def _generate_single(
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
) -> List[Image.Image]:
|
||||
url = f"{base_url}{_ENDPOINT}"
|
||||
headers = _get_headers(api_key)
|
||||
body = _build_request_body(
|
||||
result_images, timing = await generate_nano_banana_async(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
node_label="Nano Banana",
|
||||
request_log_enabled=REQUEST_LOG_ENABLED,
|
||||
check_interrupt=_check_interrupt,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
if REQUEST_LOG_ENABLED:
|
||||
extra = json.dumps(body.get("extra_body", {}), ensure_ascii=False)
|
||||
print(f"[请求] POST {url} | model={model} | extra_body={extra}")
|
||||
|
||||
last_status = None
|
||||
resp = None
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT, connect=30, sock_read=_REQUEST_TIMEOUT)
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
_check_interrupt()
|
||||
resp = await session.post(url, headers=headers, json=body, timeout=timeout)
|
||||
if resp.status == 200:
|
||||
break
|
||||
last_status = resp.status
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(resp.status, f"请求失败 ({resp.status})")
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
print(f"Nano Banana: {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
resp.close()
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
error_text = await resp.text()
|
||||
resp.close()
|
||||
if resp.status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
|
||||
try:
|
||||
err_json = json.loads(error_text)
|
||||
msg = err_json.get("error", {}).get("message", error_text[:200])
|
||||
except Exception:
|
||||
msg = error_text[:200]
|
||||
raise RuntimeError(f"API 错误 ({resp.status}): {msg}")
|
||||
else:
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
|
||||
raise RuntimeError(f"API 错误: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
|
||||
|
||||
full_content = ""
|
||||
buffer = ""
|
||||
t_request = time.time()
|
||||
t_first_token = None
|
||||
try:
|
||||
async for raw_chunk in resp.content.iter_any():
|
||||
_check_interrupt()
|
||||
if t_first_token is None:
|
||||
t_first_token = time.time()
|
||||
buffer += raw_chunk.decode("utf-8")
|
||||
while "\n" in buffer:
|
||||
line_str, buffer = buffer.split("\n", 1)
|
||||
line_str = line_str.strip()
|
||||
if not line_str or not line_str.startswith("data:"):
|
||||
continue
|
||||
data_str = line_str[5:].strip()
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
try:
|
||||
chunk = json.loads(data_str)
|
||||
delta = chunk.get("choices", [{}])[0].get("delta", {})
|
||||
if "content" in delta:
|
||||
full_content += delta["content"]
|
||||
except (json.JSONDecodeError, IndexError):
|
||||
continue
|
||||
except aiohttp.ClientPayloadError as e:
|
||||
if full_content and _IMAGE_RE.search(full_content):
|
||||
print(f"Nano Banana: 响应流提前结束,但已收到完整图片,继续解析 ({e})")
|
||||
else:
|
||||
raise RuntimeError(f"响应流下载中断,请重试或检查网络/代理: {e}") from None
|
||||
finally:
|
||||
if resp is not None:
|
||||
resp.close()
|
||||
t_done = time.time()
|
||||
|
||||
if not full_content:
|
||||
raise RuntimeError("API 未返回有效内容")
|
||||
|
||||
# 思考模型可能输出多张临时图片,最终图片始终是最后一张
|
||||
matches = list(_IMAGE_RE.finditer(full_content))
|
||||
if not matches:
|
||||
raise RuntimeError(f"响应中未找到图片: {full_content[:100]}")
|
||||
|
||||
last_match = matches[-1]
|
||||
img_data = base64.b64decode(last_match.group(2))
|
||||
final_image = Image.open(_io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
first_token_ms = (t_first_token - t_request) * 1000 if t_first_token else 0
|
||||
download_ms = (t_done - t_first_token) * 1000 if t_first_token else 0
|
||||
|
||||
return [final_image], first_token_ms, download_ms
|
||||
return result_images, timing["task_ms"], timing["parse_ms"]
|
||||
|
||||
|
||||
async def _generate_single_task(
|
||||
@@ -354,6 +215,7 @@ async def _generate_single_task(
|
||||
global_task_index: int,
|
||||
enable_grounding: bool = False,
|
||||
thinking_level: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[float], None]] = None,
|
||||
) -> dict:
|
||||
result = {
|
||||
"global_task_index": global_task_index,
|
||||
@@ -364,7 +226,7 @@ async def _generate_single_task(
|
||||
"error": None,
|
||||
}
|
||||
try:
|
||||
gen_images, first_token_ms, download_ms = await _generate_single(
|
||||
gen_images, task_ms, parse_ms = await _generate_single(
|
||||
session=session,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
@@ -375,7 +237,9 @@ async def _generate_single_task(
|
||||
images=images if images else None,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
del task_ms, parse_ms
|
||||
result["output_images"] = gen_images
|
||||
result["success"] = True
|
||||
result["generated_count"] = len(gen_images)
|
||||
@@ -438,6 +302,7 @@ async def _process_batch_async(
|
||||
global_task_index=i,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
progress_callback=_make_progress_callback(pbar),
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
@@ -473,9 +338,6 @@ async def _process_batch_async(
|
||||
error_msg = result_data.get("error", "未知错误") if result_data else "未知错误"
|
||||
print(f"Nano Banana: [{completed}/{total_tasks}] {prompt_snippet}{'...' if len(prompt_snippet) >= 30 else ''} → ✗失败: {error_msg}")
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
|
||||
all_results.extend(batch_results)
|
||||
import gc; gc.collect()
|
||||
await asyncio.sleep(0.1)
|
||||
@@ -576,7 +438,7 @@ class NanoBanana(io.ComfyNode):
|
||||
pil_imgs = tensor_to_pil(kwargs[key])
|
||||
input_images.extend(pil_imgs)
|
||||
|
||||
if input_images and len(input_images) > 14:
|
||||
if len(input_images) > 14:
|
||||
raise ValueError(f"输入图像数量 {len(input_images)} 超过限制 14 张")
|
||||
|
||||
batch_prompts = parse_batch_prompts(prompt)
|
||||
@@ -667,6 +529,7 @@ class NanoBanana(io.ComfyNode):
|
||||
images=input_images if input_images else None,
|
||||
enable_grounding=enable_grounding,
|
||||
thinking_level=thinking_level,
|
||||
progress_callback=_make_progress_callback(pbar),
|
||||
)
|
||||
return loop.run_until_complete(_run_with_interrupt(_do()))
|
||||
finally:
|
||||
@@ -674,17 +537,14 @@ class NanoBanana(io.ComfyNode):
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_single)
|
||||
generated_images, first_token_ms, download_ms = future.result(timeout=_REQUEST_TIMEOUT)
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
generated_images, task_ms, parse_ms = future.result(timeout=_REQUEST_TIMEOUT)
|
||||
|
||||
output_tensor = _images_to_tensor_safe(generated_images, _NODE)
|
||||
elapsed = time.time() - start_time
|
||||
time_str = f"{elapsed:.3f}s" if elapsed < 1 else f"{elapsed:.2f}s"
|
||||
ft_str = f"{first_token_ms/1000:.2f}s"
|
||||
dl_str = f"{download_ms/1000:.2f}s"
|
||||
print(f"完成!总耗时 {time_str} | 首字 {ft_str} | 下载 {dl_str} | 成功 {len(generated_images)}张")
|
||||
task_str = f"{task_ms/1000:.2f}s"
|
||||
parse_str = f"{parse_ms/1000:.2f}s"
|
||||
print(f"完成!总耗时 {time_str} | 异步任务 {task_str} | 解析 {parse_str} | 成功 {len(generated_images)}张")
|
||||
|
||||
import gc; gc.collect()
|
||||
return io.NodeOutput(output_tensor)
|
||||
@@ -701,7 +561,7 @@ class NanoBanana(io.ComfyNode):
|
||||
except RuntimeError as e:
|
||||
raise RuntimeError(str(e)) from None
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
finally:
|
||||
if not was_interrupted:
|
||||
try:
|
||||
|
||||
+12
-17
@@ -31,7 +31,7 @@ from PIL import Image
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor, parse_batch_prompts
|
||||
from ..utils.file_utils import load_images_from_folder, pair_images_by_name, pair_images_cartesian
|
||||
from ..utils.config import get_api_key_or_raise, NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
from ..utils.http_error import async_request_with_retry, extract_structured_error_message, get_friendly_message
|
||||
from ..models_config import (
|
||||
get_enabled_async_models,
|
||||
get_model_provider,
|
||||
@@ -241,6 +241,10 @@ class NanoBananaV2:
|
||||
@staticmethod
|
||||
def _friendly_error(error_msg: str) -> str:
|
||||
"""将上游错误转化为用户友好的提示"""
|
||||
structured_message = extract_structured_error_message(error_msg)
|
||||
if structured_message and structured_message != error_msg:
|
||||
error_msg = structured_message
|
||||
|
||||
if "No available channel for model" in error_msg:
|
||||
return (
|
||||
"当前分组下模型不可用,请检查分组是否正确。"
|
||||
@@ -343,16 +347,16 @@ class NanoBananaV2:
|
||||
error_text = await response.text()
|
||||
if not error_text.strip():
|
||||
error_text = "(服务器未返回错误详情)"
|
||||
raise RuntimeError(f"查询任务失败 ({response.status}): {error_text}")
|
||||
raise RuntimeError(f"查询任务失败: {get_friendly_message(response.status, error_text)}")
|
||||
|
||||
result = await response.json()
|
||||
status = provider.extract_status(result)
|
||||
|
||||
# 提取进度并回调(封顶 1.0 防止异常值导致进度条溢出)
|
||||
if on_progress and status in ("SUBMITTED", "IN_PROGRESS"):
|
||||
# 提取进度并回调;运行中状态不显示 100%,只有 SUCCESS 才补满。
|
||||
if on_progress and status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
|
||||
p = provider.extract_progress(result)
|
||||
if p is not None:
|
||||
p = min(p, 1.0)
|
||||
p = min(p, 0.99)
|
||||
if p > last_progress:
|
||||
on_progress(p - last_progress)
|
||||
last_progress = p
|
||||
@@ -374,7 +378,7 @@ class NanoBananaV2:
|
||||
print(f"{self.NODE_LABEL}: [轮询] FAILURE 但无错误信息,原始响应: {json.dumps(result, ensure_ascii=False)[:500]}")
|
||||
friendly_msg = self._friendly_error(error_msg)
|
||||
raise RuntimeError(f"任务失败: {friendly_msg}")
|
||||
elif status in ("SUBMITTED", "IN_PROGRESS"):
|
||||
elif status in ("SUBMITTED", "QUEUED", "IN_PROGRESS"):
|
||||
# 分段 sleep,每 0.1 秒检查一次取消信号
|
||||
sleep_iterations = int(_POLL_INTERVAL / _INTERRUPT_CHECK_INTERVAL)
|
||||
for _ in range(sleep_iterations):
|
||||
@@ -406,10 +410,7 @@ class NanoBananaV2:
|
||||
"error": None,
|
||||
}
|
||||
|
||||
contributed = [0.0] # mutable container,追踪本任务已贡献的 pbar 进度
|
||||
|
||||
def _track_progress(delta):
|
||||
contributed[0] += delta
|
||||
if on_progress:
|
||||
on_progress(delta)
|
||||
|
||||
@@ -435,15 +436,9 @@ class NanoBananaV2:
|
||||
result["request_time"] = request_time
|
||||
result["download_time"] = download_time
|
||||
except InterruptProcessingException:
|
||||
# 用户取消:补齐进度后向上传播,不吞掉
|
||||
if contributed[0] < 1.0 and on_progress:
|
||||
on_progress(1.0 - contributed[0])
|
||||
raise
|
||||
except Exception as e:
|
||||
result["error"] = str(e) or f"{type(e).__name__}(无错误详情)"
|
||||
# 失败也补齐 1.0 进度,保证进度条总数正确
|
||||
if contributed[0] < 1.0 and on_progress:
|
||||
on_progress(1.0 - contributed[0])
|
||||
|
||||
return result
|
||||
|
||||
@@ -704,7 +699,7 @@ class NanoBananaV2:
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
finally:
|
||||
# 查询并打印余额
|
||||
@@ -1223,7 +1218,7 @@ class NanoBananaV2Batch(NanoBananaV2):
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
except Exception as e:
|
||||
raise type(e)(str(e)) from None
|
||||
raise RuntimeError(str(e)) from None
|
||||
|
||||
finally:
|
||||
try:
|
||||
|
||||
@@ -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
@@ -4,7 +4,9 @@ Seedance 视频生成节点
|
||||
- Seedance: 文生视频 / 图生视频 / 首尾帧生视频(根据图片输入自动切换模式)
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
@@ -13,7 +15,7 @@ import torch
|
||||
|
||||
from ..clients.seedance_client import SeedanceClient
|
||||
from ..clients.gemini_client import GeminiAPIClient
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64, pil_to_tensor
|
||||
from ..utils.image_utils import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.r2_uploader import upload_video, upload_audio
|
||||
from ..utils.config import NETWORK_ROUTE_OPTIONS, get_base_url_by_route
|
||||
|
||||
@@ -28,6 +30,9 @@ _MODELS = [
|
||||
|
||||
_RESOLUTIONS = ["720p", "1080p", "480p"]
|
||||
|
||||
_MAX_IMAGE_BYTES = 30 * 1024 * 1024
|
||||
_MAX_REQUEST_BODY_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
# ── 模型能力判断 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -38,13 +43,45 @@ def _supports_camera_fixed(model: str) -> bool:
|
||||
|
||||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _tensor_to_base64_url(tensor) -> str:
|
||||
def _format_mb(size_bytes: int) -> str:
|
||||
return f"{size_bytes / 1024 / 1024:.2f}MB"
|
||||
|
||||
|
||||
def _tensor_to_base64_url(tensor, label: str = "图片") -> str:
|
||||
"""ComfyUI IMAGE tensor → data:image/png;base64,xxx"""
|
||||
pil_images = tensor_to_pil(tensor)
|
||||
b64 = encode_image_to_base64(pil_images[0], format="PNG")
|
||||
image = pil_images[0]
|
||||
if image.mode == "RGBA":
|
||||
image = image.convert("RGB")
|
||||
|
||||
buffered = io.BytesIO()
|
||||
image.save(buffered, format="PNG")
|
||||
image_bytes = buffered.getvalue()
|
||||
image_size = len(image_bytes)
|
||||
|
||||
if image_size > _MAX_IMAGE_BYTES:
|
||||
raise ValueError(
|
||||
f"Seedance {label}大小 {_format_mb(image_size)} 超过单张图片 "
|
||||
f"{_format_mb(_MAX_IMAGE_BYTES)} 限制,请先压缩或缩小图片。"
|
||||
)
|
||||
|
||||
b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
return f"data:image/png;base64,{b64}"
|
||||
|
||||
|
||||
def _validate_request_body_size(body: dict, tag: str):
|
||||
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
|
||||
if body_size > _MAX_REQUEST_BODY_BYTES:
|
||||
raise ValueError(
|
||||
f"{tag} 请求体大小 {_format_mb(body_size)} 超过 "
|
||||
f"{_format_mb(_MAX_REQUEST_BODY_BYTES)} 限制,请减少参考图片数量或降低图片尺寸。"
|
||||
)
|
||||
print(
|
||||
f"[{tag}] 请求体大小: {_format_mb(body_size)} "
|
||||
f"(限制 {_format_mb(_MAX_REQUEST_BODY_BYTES)})"
|
||||
)
|
||||
|
||||
|
||||
|
||||
async def _url_to_tensor(url: str) -> torch.Tensor:
|
||||
"""从 URL 下载图片并转为 ComfyUI IMAGE tensor,失败时返回 None"""
|
||||
@@ -201,7 +238,7 @@ class Seedance:
|
||||
}
|
||||
|
||||
elif mode == "i2v":
|
||||
first_url = _tensor_to_base64_url(first_image)
|
||||
first_url = _tensor_to_base64_url(first_image, "首帧图片")
|
||||
metadata["content"] = [
|
||||
{
|
||||
"type": "image_url",
|
||||
@@ -218,8 +255,8 @@ class Seedance:
|
||||
}
|
||||
|
||||
else: # flipflop
|
||||
first_url = _tensor_to_base64_url(first_image)
|
||||
last_url = _tensor_to_base64_url(last_image)
|
||||
first_url = _tensor_to_base64_url(first_image, "首帧图片")
|
||||
last_url = _tensor_to_base64_url(last_image, "尾帧图片")
|
||||
metadata["content"] = [
|
||||
{
|
||||
"type": "image_url",
|
||||
@@ -240,6 +277,8 @@ class Seedance:
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
_validate_request_body_size(body, tag)
|
||||
|
||||
# 保存路径(临时文件,避免与下游保存节点重复落盘)
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix=f"{file_prefix}_")
|
||||
|
||||
@@ -315,6 +354,7 @@ class SeedanceMultiModal:
|
||||
web_search = _first(kwargs.get("联网搜索"), "关闭") == "打开"
|
||||
return_last = _first(kwargs.get("返回末帧图片"), "关闭") == "打开"
|
||||
seed = _first(kwargs.get("seed"), 0)
|
||||
network_route = _first(kwargs.get("网络线路"), "全球加速")
|
||||
|
||||
# 参考图片:INPUT_IS_LIST 时是 [tensor, tensor, ...] 列表,直接保留
|
||||
raw_images = kwargs.get("参考图片", None)
|
||||
@@ -344,11 +384,11 @@ class SeedanceMultiModal:
|
||||
imgs = ref_images[:9]
|
||||
if len(ref_images) > 9:
|
||||
print(f"[SeedanceMultiModal] 参考图片超过9张,仅取前9张(共{len(ref_images)}张)")
|
||||
for img_tensor in imgs:
|
||||
for idx, img_tensor in enumerate(imgs, start=1):
|
||||
# 每个 tensor 可能是 [1,H,W,C] 或 [H,W,C],统一确保有 batch 维
|
||||
if img_tensor.dim() == 3:
|
||||
img_tensor = img_tensor.unsqueeze(0)
|
||||
url = _tensor_to_base64_url(img_tensor)
|
||||
url = _tensor_to_base64_url(img_tensor, f"参考图片{idx}")
|
||||
content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": url},
|
||||
@@ -414,11 +454,13 @@ class SeedanceMultiModal:
|
||||
if first_image_url:
|
||||
body["image"] = first_image_url
|
||||
|
||||
_validate_request_body_size(body, "Seedance多模态")
|
||||
|
||||
# ── 保存路径(临时文件,避免与下游保存节点重复落盘)──────────────────
|
||||
_, save_path = tempfile.mkstemp(suffix=".mp4", prefix="seedance_mm_")
|
||||
|
||||
client = SeedanceClient()
|
||||
client.base_url = get_base_url_by_route(kwargs.get("网络线路", "全球加速"))
|
||||
client.base_url = get_base_url_by_route(network_route)
|
||||
pbar = _make_pbar()
|
||||
on_stage, on_prog = _make_callbacks("Seedance多模态", pbar)
|
||||
|
||||
|
||||
@@ -148,4 +148,6 @@ def get_async_api_base_url() -> str:
|
||||
|
||||
def get_base_url_by_route(route: str) -> str:
|
||||
"""根据网络线路选项返回对应域名,未匹配则走 config 垫底"""
|
||||
if isinstance(route, (list, tuple)):
|
||||
route = route[0] if route else None
|
||||
return NETWORK_ROUTES.get(route, get_api_base_url())
|
||||
|
||||
+54
-2
@@ -9,8 +9,9 @@
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from typing import Optional
|
||||
from typing import Any, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
@@ -45,15 +46,66 @@ DEFAULT_MAX_DELAY = 30.0 # 最大等待秒数
|
||||
DEFAULT_BACKOFF_FACTOR = 2.0 # 指数退避因子
|
||||
|
||||
|
||||
def _extract_message_from_payload(payload: Any) -> str:
|
||||
if isinstance(payload, str):
|
||||
text = payload.strip()
|
||||
if not text:
|
||||
return ""
|
||||
if text.startswith("{") or text.startswith("["):
|
||||
try:
|
||||
return _extract_message_from_payload(json.loads(text))
|
||||
except Exception:
|
||||
return text
|
||||
return text
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
return ""
|
||||
|
||||
error = payload.get("error")
|
||||
if isinstance(error, dict):
|
||||
for key in ("message", "msg", "detail", "reason"):
|
||||
value = error.get(key)
|
||||
if value:
|
||||
return _extract_message_from_payload(value)
|
||||
elif error:
|
||||
return _extract_message_from_payload(error)
|
||||
|
||||
for key in ("message", "msg", "detail", "reason", "error_message"):
|
||||
value = payload.get(key)
|
||||
if value:
|
||||
return _extract_message_from_payload(value)
|
||||
|
||||
for key in ("data", "result", "response", "output"):
|
||||
value = payload.get(key)
|
||||
nested = _extract_message_from_payload(value)
|
||||
if nested:
|
||||
return nested
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def extract_structured_error_message(raw_message: str) -> str:
|
||||
if not isinstance(raw_message, str):
|
||||
return ""
|
||||
text = raw_message.strip()
|
||||
if not (text.startswith("{") or text.startswith("[")):
|
||||
return ""
|
||||
return _extract_message_from_payload(text)
|
||||
|
||||
|
||||
def get_friendly_message(status_code: int, raw_message: str = "") -> str:
|
||||
"""根据状态码/错误内容返回友好文案,未匹配则返回原始信息"""
|
||||
if status_code == 524:
|
||||
return "Gateway timed out while waiting for upstream image generation. Please retry, lower resolution/count, or switch network route."
|
||||
if raw_message:
|
||||
raw_message_lower = raw_message.lower()
|
||||
structured_message = extract_structured_error_message(raw_message)
|
||||
message_for_matching = structured_message or raw_message
|
||||
raw_message_lower = message_for_matching.lower()
|
||||
for keyword, friendly_msg in ERROR_CONTENT_MESSAGES.items():
|
||||
if keyword.lower() in raw_message_lower:
|
||||
return friendly_msg
|
||||
if structured_message:
|
||||
return structured_message
|
||||
if status_code == 500:
|
||||
return "服务器返回 500:上游生成失败或服务端临时异常。请稍后重试;如果多次出现,请降低分辨率/数量,或调整提示词。"
|
||||
friendly = HTTP_ERROR_MESSAGES.get(status_code)
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
@@ -4,10 +4,16 @@ app.registerExtension({
|
||||
name: "o1key.hideSidebarItems",
|
||||
async setup() {
|
||||
const hide = () => {
|
||||
// 隐藏侧边栏的"说明"、"应用"、"模型"、"节点"按钮
|
||||
const hiddenLabels = ["说明", "帮助", "Help", "应用", "Apps", "模型", "Models", "节点", "Nodes"];
|
||||
// 隐藏侧边栏的"说明"、"应用"、"模型"、"节点"、"模板"按钮
|
||||
const hiddenLabels = ["说明", "帮助", "help", "应用", "apps", "模型", "models", "节点", "nodes", "模板", "templates", "template"];
|
||||
document.querySelectorAll(".side-bar-button, [class*='sidebar'] button, .p-togglebutton").forEach(btn => {
|
||||
const label = btn.getAttribute("aria-label") || btn.textContent || "";
|
||||
const label = [
|
||||
btn.getAttribute("aria-label"),
|
||||
btn.getAttribute("title"),
|
||||
btn.getAttribute("data-title"),
|
||||
btn.getAttribute("data-label"),
|
||||
btn.textContent,
|
||||
].filter(Boolean).join(" ").toLowerCase();
|
||||
if (hiddenLabels.some(k => label.includes(k))) {
|
||||
btn.style.display = "none";
|
||||
}
|
||||
|
||||
+77
-11
@@ -1,7 +1,9 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
import { api } from "../../../scripts/api.js";
|
||||
|
||||
const STORAGE_KEY = "o1key-notes";
|
||||
const SEEDED_KEY = "o1key-notes-seeded-v2";
|
||||
const NOTES_API = "/o1key/notes";
|
||||
const STYLE_ID = "o1key-notes-styles";
|
||||
|
||||
let notes = [];
|
||||
@@ -207,27 +209,91 @@ function injectStyles() {
|
||||
document.head.appendChild(el);
|
||||
}
|
||||
|
||||
function loadNotes() {
|
||||
function readCachedNotes() {
|
||||
try {
|
||||
const raw = localStorage.getItem(STORAGE_KEY);
|
||||
notes = raw ? JSON.parse(raw).map(makeNote) : [];
|
||||
if (raw === null) return { found: false, notes: [] };
|
||||
const parsed = JSON.parse(raw);
|
||||
return { found: true, notes: Array.isArray(parsed) ? parsed.map(makeNote) : [] };
|
||||
} catch {
|
||||
notes = [];
|
||||
return { found: false, notes: [] };
|
||||
}
|
||||
}
|
||||
|
||||
if (!localStorage.getItem(SEEDED_KEY)) {
|
||||
const existingTitles = new Set(notes.map(n => n.title));
|
||||
const samples = SAMPLE_NOTES.map(makeNote).filter(n => !existingTitles.has(n.title));
|
||||
notes = [...samples, ...notes];
|
||||
function writeCachedNotes() {
|
||||
try {
|
||||
localStorage.setItem(STORAGE_KEY, JSON.stringify(notes));
|
||||
} catch {}
|
||||
}
|
||||
|
||||
function hasSeededNotes() {
|
||||
try {
|
||||
return !!localStorage.getItem(SEEDED_KEY);
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
function markSeededNotes() {
|
||||
try {
|
||||
localStorage.setItem(SEEDED_KEY, "1");
|
||||
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() {
|
||||
writeCachedNotes();
|
||||
void persistNotesToFile();
|
||||
}
|
||||
|
||||
async function persistNotesToFile() {
|
||||
try {
|
||||
localStorage.setItem(STORAGE_KEY, JSON.stringify(notes));
|
||||
} catch {}
|
||||
const resp = await api.fetchApi(NOTES_API, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ notes }),
|
||||
});
|
||||
if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
|
||||
return true;
|
||||
} catch (e) {
|
||||
console.warn("[o1key notes] failed to save notes file", e);
|
||||
setPanelStatus("笔记文件保存失败,已保存在浏览器缓存");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
function allTags() {
|
||||
@@ -250,7 +316,7 @@ function filteredNotes() {
|
||||
app.registerExtension({
|
||||
name: "o1key.notePanel",
|
||||
async setup() {
|
||||
loadNotes();
|
||||
await loadNotes();
|
||||
app.extensionManager.registerSidebarTab({
|
||||
id: "o1key-notes",
|
||||
title: "笔记",
|
||||
|
||||
+60
-5
@@ -43,6 +43,7 @@ app.registerExtension({
|
||||
if (!confirm("确定要重启 ComfyUI 吗?")) return;
|
||||
btn.style.opacity = "0.5";
|
||||
btn.style.pointerEvents = "none";
|
||||
await disableExperimentalAssetApi();
|
||||
try { await fetch("/o1key/restart", { method: "POST" }); } catch {}
|
||||
pollUntilReady();
|
||||
});
|
||||
@@ -51,16 +52,70 @@ app.registerExtension({
|
||||
injected = true;
|
||||
}
|
||||
|
||||
async function disableExperimentalAssetApi() {
|
||||
if (!(await shouldDisableExperimentalAssetApi())) return;
|
||||
try {
|
||||
await fetch("/api/settings/Comfy.Assets.UseAssetAPI", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(false),
|
||||
signal: AbortSignal.timeout(2000),
|
||||
});
|
||||
} catch {}
|
||||
}
|
||||
|
||||
async function shouldDisableExperimentalAssetApi() {
|
||||
try {
|
||||
const r = await fetch("/api/settings/Comfy.Assets.UseAssetAPI", {
|
||||
cache: "no-store",
|
||||
signal: AbortSignal.timeout(2000),
|
||||
});
|
||||
if (!r.ok || !(await r.json())) return false;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
return !(await fetchOk("/api/assets/seed/status", 2000));
|
||||
}
|
||||
|
||||
async function fetchOk(url, timeout = 2500) {
|
||||
try {
|
||||
const r = await fetch(url, {
|
||||
cache: "no-store",
|
||||
signal: AbortSignal.timeout(timeout),
|
||||
});
|
||||
return r.ok;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
async function comfyReady() {
|
||||
const [statsOk, modelFoldersOk] = await Promise.all([
|
||||
fetchOk("/api/system_stats"),
|
||||
fetchOk("/api/experiment/models"),
|
||||
]);
|
||||
return statsOk && modelFoldersOk;
|
||||
}
|
||||
|
||||
function pollUntilReady() {
|
||||
let attempts = 0;
|
||||
const maxAttempts = 40;
|
||||
const maxAttempts = 80;
|
||||
const minRestartWaitMs = 5000;
|
||||
const startedAt = Date.now();
|
||||
let sawUnavailable = false;
|
||||
const interval = setInterval(async () => {
|
||||
attempts++;
|
||||
if (attempts > maxAttempts) { clearInterval(interval); forceReload(); return; }
|
||||
try {
|
||||
const r = await fetch("/api/system_stats", { signal: AbortSignal.timeout(2000) });
|
||||
if (r.ok) { clearInterval(interval); forceReload(); }
|
||||
} catch {}
|
||||
const ready = await comfyReady();
|
||||
if (!ready) {
|
||||
sawUnavailable = true;
|
||||
return;
|
||||
}
|
||||
if (!sawUnavailable && Date.now() - startedAt < minRestartWaitMs) return;
|
||||
|
||||
clearInterval(interval);
|
||||
await disableExperimentalAssetApi();
|
||||
setTimeout(forceReload, 800);
|
||||
}, 1500);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user