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