publisher/comfyui_o1key
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Grok and VEO video workflow support

Browse Source
This commit is contained in:
o1key
2026-06-04 16:06:23 +08:00
parent 5d9aff9ca7
commit 30e9603f77
24 changed files with 3752 additions and 528 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clients/__init__.py
+3 -1
View File
@@ -9,6 +9,8 @@ from .gemini_flash_client import GeminiFlashClient
from .sora_client import SoraClient
from .kling_client import KlingClient
from .veo_client import VeoClient
from .newapi_veo_client import NewAPIVeoClient
from .grok_video_client import GrokVideoClient
from .openai_client import OpenAIAPIClient
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient']
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient']
clients/base_client.py
+2 -22
View File
@@ -40,7 +40,8 @@ class BaseAPIClient(ABC):
Args:
base_url: API 基础 URL
api_key: API 密钥
max_request_size: 最大请求体大小（字节），默认 100MB
max_request_size: 兼容参数；基类不再用它限制 JSON 请求体，
部分子类仍用它作为上传文件大小限制
"""
self.base_url = base_url
self.api_key = api_key
@@ -116,24 +117,6 @@ class BaseAPIClient(ABC):
"Content-Type": "application/json"
}
def check_request_size(self, request_body: Dict[str, Any]) -> None:
"""
检查请求体大小是否超过限制
Args:
request_body: 请求体字典
Raises:
ValueError: 如果请求体超过限制
"""
request_json = json.dumps(request_body)
request_size = len(request_json.encode('utf-8'))
if request_size > self.max_request_size:
raise ValueError(
"请求体积超过100MB限制，请调整分辨率或减少图片数量"
)
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
"""
子类可重写：为指定 HTTP 状态码返回自定义错误文案。
@@ -185,9 +168,6 @@ class BaseAPIClient(ABC):
url = f"{self.base_url}{endpoint}"
headers = self.get_headers(use_bearer_token)
# 检查请求大小
self.check_request_size(request_body)
close_session = False
if session is None:
session = self._make_session()
clients/gemini_async_provider.py
+53 -9
View File
@@ -72,7 +72,9 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
resolution=resolution,
enable_grounding=kwargs.get("enable_grounding", False),
enable_image_search=kwargs.get("enable_image_search", False),
image_compression=getattr(self, 'image_compression', None),
image_compression=getattr(self, "image_compression", None),
thinking_level=kwargs.get("thinking_level"),
request_log_enabled=False,
)
def extract_task_id(self, response: dict) -> str:
@@ -85,6 +87,23 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
return response.get("status", "UNKNOWN")
async def parse_result(self, result_data: dict, session) -> List[Image.Image]:
images = result_data.get("images") if isinstance(result_data, dict) else None
if isinstance(images, list) and images:
parsed = []
for item in images:
if not isinstance(item, dict):
continue
image_url = item.get("url") or item.get("image_url")
if image_url:
async with session.get(image_url) as img_resp:
if img_resp.status == 200:
img_bytes = await img_resp.read()
parsed.append(Image.open(BytesIO(img_bytes)).convert("RGB"))
else:
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
if parsed:
return parsed
# 异步接口可能直接返回 image_url
image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else ""
if image_url:
@@ -124,19 +143,44 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
def extract_progress(self, response: dict) -> Optional[float]:
"""从轮询响应中提取进度（0.0-1.0）"""
def _coerce(val) -> Optional[float]:
if val is None or isinstance(val, bool):
return None
if isinstance(val, (int, float)):
progress = float(val)
elif isinstance(val, str):
text = val.strip()
if not text:
return None
has_percent_suffix = text.endswith("%")
if has_percent_suffix:
text = text[:-1].strip()
try:
progress = float(text)
except ValueError:
return None
if has_percent_suffix:
progress /= 100.0
else:
return None
if progress > 1.0:
progress /= 100.0
return max(0.0, min(progress, 1.0))
# 直接字段：progress / percentage
for field in ("progress", "percentage"):
val = response.get(field)
if val is not None and isinstance(val, (int, float)):
return val / 100.0 if val > 1 else float(val)
for field in ("progress", "percentage", "percent"):
progress = _coerce(response.get(field))
if progress is not None:
return progress
# 嵌套字段：progressInfo / progress_info
progress_info = response.get("progressInfo") or response.get("progress_info")
if isinstance(progress_info, dict):
for field in ("progress", "percentage"):
val = progress_info.get(field)
if val is not None and isinstance(val, (int, float)):
return val / 100.0 if val > 1 else float(val)
for field in ("progress", "percentage", "percent"):
progress = _coerce(progress_info.get(field))
if progress is not None:
return progress
return None
clients/gemini_client.py
+116 -93
View File
@@ -34,8 +34,7 @@ class GeminiAPIClient(BaseAPIClient):
super().__init__(
base_url=get_api_base_url(),
api_key=api_key,
max_request_size=100 * 1024 * 1024
api_key=api_key
)
@staticmethod
@@ -188,6 +187,7 @@ class GeminiAPIClient(BaseAPIClient):
enable_image_search: bool = False,
image_compression: str = None,
thinking_level: str = None,
request_log_enabled: bool = True,
**kwargs
) -> Dict[str, Any]:
"""
@@ -204,21 +204,90 @@ class GeminiAPIClient(BaseAPIClient):
Returns:
请求体字典
"""
import json
_MAX_BODY_BYTES = 20 * 1024 * 1024 # 20 MB
_BODY_TARGET_BYTES = int(_MAX_BODY_BYTES * 0.9)
parts = []
# 添加文本部分
parts.append({"text": prompt})
image_config = {"imageSize": resolution}
if aspect_ratio and aspect_ratio != "智能":
image_config["aspectRatio"] = aspect_ratio
def _build_request_body(body_parts: List[dict]) -> Dict[str, Any]:
body = {
"contents": [
{
"role": "user",
"parts": body_parts
}
],
"generationConfig": {
"responseModalities": ["IMAGE"],
"imageConfig": image_config
}
}
# 添加思考深度配置
if thinking_level:
body["generationConfig"]["thinkingConfig"] = {
"thinkingLevel": thinking_level,
"includeThoughts": True
}
# 添加图片压缩参数
if image_compression:
body["image_compression"] = image_compression
# 添加 Google Search Grounding（如果启用）
# 新异步接口要求直接放在请求体顶层：{"google_search": true}
if enable_grounding or enable_image_search:
body["google_search"] = True
return body
def _request_size(body: Dict[str, Any]) -> int:
return len(json.dumps(body).encode("utf-8"))
def _format_size(size: int) -> str:
if size < 1024 * 1024:
return f"{size / 1024:.2f}KB"
return f"{size / 1024 / 1024:.2f}MB"
def _shorten_base64_for_log(obj, max_len: int = 200):
if isinstance(obj, dict):
result = {}
for key, value in obj.items():
if key == "data" and isinstance(value, str) and len(value) > max_len:
result[key] = f"<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:
working_images = list(images)
# 编码一次，估算大小，超限则迭代缩放
for _attempt in range(10):
def _build_image_parts(src_images: List[Image.Image]) -> List[dict]:
img_parts = []
for img in working_images:
for img in src_images:
img_base64 = encode_image_to_base64(img)
img_parts.append({
"inline_data": {
@@ -226,94 +295,49 @@ class GeminiAPIClient(BaseAPIClient):
"data": img_base64
}
})
return img_parts
# 估算完整 body 大小（不含工具字段，工具字段很小可忽略）
est_image_config = {"imageSize": resolution}
if aspect_ratio and aspect_ratio != "智能":
est_image_config["aspectRatio"] = aspect_ratio
estimated = self._estimate_body_size(
parts + img_parts,
{
"generationConfig": {
"responseModalities": ["IMAGE"],
"imageConfig": est_image_config
}
}
def _estimate_with_images(img_parts: List[dict]) -> int:
return _request_size(_build_request_body(parts + img_parts))
working_images = list(images)
img_parts = _build_image_parts(working_images)
original_request_body = _build_request_body(parts + img_parts)
if request_log_enabled:
_log_original_request_body(original_request_body)
original_request_logged = True
estimated = _estimate_with_images(img_parts)
if estimated > _MAX_BODY_BYTES:
ratio = _BODY_TARGET_BYTES / estimated
scale = ratio ** 0.5 # 面积比 → 线性比
working_images = self._scale_images_to_fit(working_images, scale)
img_parts = _build_image_parts(working_images)
estimated = _estimate_with_images(img_parts)
orig_sizes = ", ".join(f"{img.width}×{img.height}" for img in images)
new_sizes = ", ".join(f"{img.width}×{img.height}" for img in working_images)
size_mb = estimated / (1024 * 1024)
target_mb = _BODY_TARGET_BYTES / (1024 * 1024)
print(
f"Nano Banana Pro: 输入图片已按请求体目标大小自动缩放\n"
f" 原始尺寸: {orig_sizes}\n"
f" 缩放后: {new_sizes}\n"
f" 请求体积: {size_mb:.2f}MB（目标 {target_mb:.2f}MB，限制 20MB）"
)
if estimated <= _MAX_BODY_BYTES:
parts.extend(img_parts)
if _attempt > 0:
orig_sizes = ", ".join(
f"{img.width}×{img.height}" for img in images
)
new_sizes = ", ".join(
f"{img.width}×{img.height}" for img in working_images
)
size_mb = estimated / (1024 * 1024)
print(
f"Nano Banana Pro: 输入图片已自动缩放以控制请求体积\n"
f" 原始尺寸: {orig_sizes}\n"
f" 缩放后: {new_sizes}\n"
f" 请求体积: {size_mb:.2f}MB（限制 20MB）"
)
break
else:
# 按像素面积比推算需要的线性缩放系数，留 5% 余量
ratio = (_MAX_BODY_BYTES * 0.95) / estimated
scale = ratio ** 0.5 # 面积比 → 线性比
working_images = self._scale_images_to_fit(working_images, scale)
else:
# 10 轮后仍超限，使用最后一次结果（极端情况兜底）
parts.extend(img_parts)
parts.extend(img_parts)
# 构建请求体
image_config = {"imageSize": resolution}
if aspect_ratio and aspect_ratio != "智能":
image_config["aspectRatio"] = aspect_ratio
request_body = _build_request_body(parts)
if request_log_enabled and not original_request_logged:
_log_original_request_body(request_body)
request_body = {
"contents": [
{
"role": "user",
"parts": parts
}
],
"generationConfig": {
"responseModalities": ["IMAGE"],
"imageConfig": image_config
}
}
# 添加思考深度配置
if thinking_level:
request_body["generationConfig"]["thinkingConfig"] = {
"thinkingLevel": thinking_level,
"includeThoughts": True
}
# 添加图片压缩参数
if image_compression:
request_body["image_compression"] = image_compression
# 添加 Google Search Grounding 工具（如果启用）
# 注意：enable_image_search=True 时会自动隐含 enable_grounding
if enable_grounding or enable_image_search:
if enable_image_search:
# 同时启用网页搜索和图片搜索（仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持）
request_body["tools"] = [
{
"google_search": {
"searchTypes": {
"webSearch": {},
"imageSearch": {}
}
}
}
]
else:
# 仅启用网页搜索（通用）
request_body["tools"] = [{"google_search": {}}]
request_size = _request_size(request_body)
if request_size > _MAX_BODY_BYTES:
raise ValueError(
f"请求体超过 20MB 限制（当前 {request_size / 1024 / 1024:.2f}MB），"
"已停止提交；请减少参考图数量、降低图片复杂度或缩短提示词"
)
return request_body
@@ -924,7 +948,7 @@ class GeminiAPIClient(BaseAPIClient):
)
return self.run_async_in_thread(coro)
async def generate_multi_prompts_async(
self,
prompts: List[str],
@@ -1104,6 +1128,5 @@ class GeminiAPIClient(BaseAPIClient):
enable_grounding=enable_grounding,
enable_image_search=enable_image_search
)
return self.run_async_in_thread(coro)
clients/gpt_image_client.py
+636 -10
View File
@@ -1,13 +1,17 @@
"""
GPT Image API 客户端
支持两个接口：
- POST /v1/images/generations/ 文生图 / 图生图（gpt-image-1 / gpt-image-1.5）
- POST /v1/images/edits/ 图像编辑（带蒙版 inpainting）
新版节点请求走异步任务接口：
- POST /async/v1/generateImage
- GET /async/v1/tasks/{task_id}
旧同步接口保留兼容代码，但 GPT Image 节点不再使用：
- POST /v1/images/generations/
- POST /v1/images/edits
设计原则：
- 与 doubao_image_client.py 保持相同的异步 + 同步双入口模式
- generations / edits 接口均使用 multipart/form-data
- 响应兼容 SSE 流式、JSON、url 和 b64_json
- 对 ComfyUI 节点暴露同步入口，内部提交异步任务并轮询
- 图片和蒙版以 data:image/png;base64,... 放入 JSON 请求体
- 响应优先读取 data.images[].url
"""
import asyncio
@@ -16,7 +20,7 @@ import json
import time
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO
from typing import List, Optional
from typing import Any, Callable, Dict, List, Optional
import aiohttp
import numpy as np
@@ -38,6 +42,8 @@ except ImportError:
# ── 接口端点 ──────────────────────────────────────────────────────────────────
_ENDPOINT_GENERATIONS = "/v1/images/generations/"
_ENDPOINT_EDITS = "/v1/images/edits"
_ENDPOINT_ASYNC_GENERATE = "/async/v1/generateImage"
_ENDPOINT_ASYNC_TASK = "/async/v1/tasks/{task_id}"
# ── 模型名映射（UI 显示名 → API 实际参数名）─────────────────────────────────
_MODEL_NAME_MAP = {
@@ -47,6 +53,49 @@ _MODEL_NAME_MAP = {
# ── 超时 ──────────────────────────────────────────────────────────────────────
_REQUEST_TIMEOUT = 900 # 秒
_ASYNC_POLL_SCHEDULE = [5.0, 20.0]
_ASYNC_POLL_INTERVAL = 3.0
_ASYNC_MAX_WAIT = 600.0
_ASYNC_RETRY_DELAYS = [2.0, 5.0, 10.0]
_ASYNC_RETRYABLE_ERROR_CODES = {
"image_rate_limited",
"image_upstream_busy",
"image_timeout",
"image_storage_failed",
"image_empty_result",
"image_upstream_error",
"image_internal_error",
"image_unknown_error",
}
_ASYNC_RETRYABLE_ERROR_CATEGORIES = {
"rate_limit",
"upstream_busy",
"timeout",
"storage",
"upstream_error",
"internal_error",
"unknown",
}
_ASYNC_NON_RETRYABLE_ERROR_CODES = {
"image_invalid_size",
"image_payload_too_large",
"image_invalid_mask",
"image_invalid_parameter",
"image_safety_blocked",
"image_provider_quota_exceeded",
"image_provider_permission_required",
"image_model_unavailable",
"image_reference_download_failed",
}
REQUEST_LOG_ENABLED = False
POLL_LOG_ENABLED = False
class _AsyncImageTaskFailure(RuntimeError):
def __init__(self, message: str, error_detail: Optional[dict] = None):
super().__init__(message)
self.error_detail = error_detail or {}
class GptImageClient:
@@ -54,10 +103,10 @@ class GptImageClient:
GPT Image API 客户端
接口说明：
generations：multipart/form-data，支持 quality / size / n / model
edits：multipart/form-data，图片和 mask 使用 PNG 文件上传
async generateImage：JSON 提交，返回 task_id
tasks/{task_id}：轮询任务状态，成功后读取 data.images[].url
响应支持 JSON 和 SSE 流式格式。
旧 generations / edits 同步接口保留为兼容代码。
"""
def __init__(self):
@@ -132,6 +181,226 @@ class GptImageClient:
)
return png_bytes
@staticmethod
def _png_bytes_to_data_url(png_bytes: bytes) -> str:
b64 = base64.b64encode(png_bytes).decode("ascii")
return f"data:image/png;base64,{b64}"
@staticmethod
def _json_body_size(body: dict) -> int:
return len(json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
@staticmethod
def _format_body_size(size: int) -> str:
if size < 1024 * 1024:
return f"{size / 1024:.2f}KB"
return f"{size / 1024 / 1024:.2f}MB"
@staticmethod
def _shorten_data_urls_for_log(obj, max_len: int = 200):
if isinstance(obj, dict):
return {
key: GptImageClient._shorten_data_urls_for_log(value, max_len)
for key, value in obj.items()
}
if isinstance(obj, list):
return [GptImageClient._shorten_data_urls_for_log(item, max_len) for item in obj]
if isinstance(obj, str) and obj.startswith("data:image") and len(obj) > max_len:
header, _, data = obj.partition(",")
return f"{header},<base64 data, {len(data)} chars>"
return obj
def _log_original_request_body(self, label: str, body: dict) -> None:
if not REQUEST_LOG_ENABLED:
return
body_size = self._json_body_size(body)
print(
f"\n{'=' * 60}\n"
f"[o1key GPT Image] 原始请求体日志 | {label} | "
f"请求体积: {self._format_body_size(body_size)} "
f"(data URL 已折叠显示 base64 长度)\n"
f"{json.dumps(self._shorten_data_urls_for_log(body), ensure_ascii=False, indent=2)}\n"
f"{'=' * 60}\n"
)
def _log_original_response_body(self, label: str, text: str) -> None:
if not REQUEST_LOG_ENABLED:
return
size = len(text.encode("utf-8"))
print(
f"\n{'=' * 60}\n"
f"[o1key GPT Image] 原始返回响应体日志 | {label} | "
f"响应体积: {self._format_body_size(size)}\n"
f"{text}\n"
f"{'=' * 60}\n"
)
@staticmethod
def _extract_error_message(payload_or_text, status_code: int = 0) -> str:
payload = payload_or_text
if isinstance(payload_or_text, str):
try:
payload = json.loads(payload_or_text)
except Exception:
return get_friendly_message(status_code, payload_or_text)
if isinstance(payload, dict):
error = payload.get("error")
if isinstance(error, str) and error.strip():
return error.strip()
if isinstance(error, dict):
msg = error.get("message") or error.get("msg") or error.get("error")
if msg:
return str(msg)
return json.dumps(error, ensure_ascii=False)
msg = payload.get("message") or payload.get("msg")
if msg:
return str(msg)
return get_friendly_message(status_code, str(payload_or_text))
@staticmethod
def _extract_error_detail(payload: dict) -> dict:
if not isinstance(payload, dict):
return {}
detail = payload.get("error_detail")
return detail if isinstance(detail, dict) else {}
@staticmethod
def _should_retry_async_failure(error_detail: dict, retry_index: int) -> bool:
if not isinstance(error_detail, dict) or not error_detail:
return False
code = error_detail.get("code")
category = error_detail.get("category")
retryable = error_detail.get("retryable")
if code in _ASYNC_NON_RETRYABLE_ERROR_CODES:
return False
if code == "image_unknown_error":
return retry_index == 0
if retryable is True:
return True
if retryable is False:
return False
return (
code in _ASYNC_RETRYABLE_ERROR_CODES
or category in _ASYNC_RETRYABLE_ERROR_CATEGORIES
)
@staticmethod
def _coerce_progress_percent(value: Any) -> Optional[int]:
if value is None:
return None
if isinstance(value, str):
text = value.strip().rstrip("%")
if not text:
return None
try:
value = float(text)
except ValueError:
return None
elif isinstance(value, (int, float)):
value = float(value)
else:
return None
if 0 <= value <= 1:
value *= 100
return max(0, min(100, int(round(value))))
@staticmethod
def _emit_progress(progress_callback: Optional[Callable[[int], None]], pct: int) -> None:
if progress_callback is None:
return
try:
progress_callback(pct)
except Exception as error:
print(f"[o1key GPT Image] progress callback failed: {error}")
@staticmethod
def _resize_png_bytes(source_image: Image.Image, scale: float) -> bytes:
if scale < 0.999:
width, height = source_image.size
new_width = max(1, int(width * scale))
new_height = max(1, int(height * scale))
image = source_image.resize((new_width, new_height), Image.LANCZOS)
else:
image = source_image
buf = BytesIO()
image.save(buf, format="PNG")
return buf.getvalue()
@staticmethod
def _pil_to_png_bytes(image: Image.Image) -> bytes:
buf = BytesIO()
image.save(buf, format="PNG")
return buf.getvalue()
def _fit_png_assets_to_body_limit(self, assets: List[Dict[str, Any]], build_body) -> Dict[str, bytes]:
"""
根据完整 JSON 请求体大小压缩图片资产，保证最终 body 不超过 20MB。
使用同一个缩放比例二分搜索，让压缩结果尽量贴近上限而不是过度压缩。
"""
asset_bytes = {asset["key"]: asset["bytes"] for asset in assets}
initial_size = self._json_body_size(build_body(asset_bytes))
if initial_size <= self._MAX_BODY_BYTES:
return asset_bytes
if not assets:
raise RuntimeError(
f"请求体大小 {initial_size // 1024}KB 超过 20MB，且没有可压缩图片"
)
originals = []
for asset in assets:
image = Image.open(BytesIO(asset["bytes"]))
image.load()
originals.append((asset, image.copy()))
low = 0.001
high = 1.0
best_bytes = None
best_size = 0
best_scale = 0.0
for _ in range(16):
scale = (low + high) / 2
candidate = {}
for asset, image in originals:
candidate[asset["key"]] = self._resize_png_bytes(image, scale)
body_size = self._json_body_size(build_body(candidate))
if body_size <= self._MAX_BODY_BYTES:
best_bytes = candidate
best_size = body_size
best_scale = scale
low = scale
else:
high = scale
if best_bytes is None:
candidate = {}
for asset, image in originals:
candidate[asset["key"]] = self._resize_png_bytes(image, low)
body_size = self._json_body_size(build_body(candidate))
if body_size > self._MAX_BODY_BYTES:
raise RuntimeError(
f"图片已压缩到最小比例，但请求体仍超过 20MB：{body_size // 1024}KB"
)
best_bytes = candidate
best_size = body_size
best_scale = low
print(
f"[o1key GPT Image] 请求体超过 20MB，已等比压缩图片："
f"{initial_size // 1024}KB → {best_size // 1024}KB，scale={best_scale:.3f}"
)
return best_bytes
@staticmethod
def _tensor_to_png_bytes(tensor: torch.Tensor) -> bytes:
"""
@@ -500,6 +769,316 @@ class GptImageClient:
return request_task.result()
# ── 新版异步 GPT Image 接口 ──────────────────────────────────────────────
def _build_async_generate_body(
self,
prompt: str,
model: str,
quality: str,
size: str,
n: int,
image_list: Optional[List[torch.Tensor]] = None,
mask_tensor: Optional[torch.Tensor] = None,
output_format: str = "png",
) -> dict:
api_model = _MODEL_NAME_MAP.get(model, model)
assets: List[Dict[str, Any]] = []
image_keys: List[str] = []
mask_key = None
if image_list:
for idx_img, img_tensor in enumerate(image_list):
pil_images = tensor_to_pil(img_tensor)
if not pil_images:
continue
key = f"image_{idx_img}"
image_keys.append(key)
assets.append({
"key": key,
"label": f"参考图{idx_img + 1}",
"bytes": self._pil_to_png_bytes(pil_images[0]),
})
if mask_tensor is not None:
if not image_list:
raise ValueError("提供了蒙版但未提供图片，请同时提供图片和蒙版")
first_tensor = image_list[0]
if first_tensor.dim() == 3:
first_tensor = first_tensor.unsqueeze(0)
image_size = (first_tensor.shape[1], first_tensor.shape[2])
mask_key = "mask"
assets.append({
"key": mask_key,
"label": "蒙版",
"bytes": self._mask_tensor_to_rgba_png_bytes(mask_tensor, image_size),
})
def _make_body(asset_bytes: Dict[str, bytes]) -> dict:
body = {
"model": api_model,
"prompt": prompt,
"images": [
self._png_bytes_to_data_url(asset_bytes[key])
for key in image_keys
if key in asset_bytes
],
"size": size if size else "auto",
"quality": quality,
"n": int(n),
"output_format": output_format or "png",
}
if mask_key and mask_key in asset_bytes:
body["mask"] = {
"image_url": self._png_bytes_to_data_url(asset_bytes[mask_key])
}
return body
original_asset_bytes = {asset["key"]: asset["bytes"] for asset in assets}
original_body = _make_body(original_asset_bytes)
self._log_original_request_body("async generateImage", original_body)
asset_bytes = self._fit_png_assets_to_body_limit(assets, _make_body)
body = _make_body(asset_bytes)
body_size = self._json_body_size(body)
if body_size > self._MAX_BODY_BYTES:
raise RuntimeError(f"请求体超过 20MB：{body_size // 1024}KB")
return body
async def _submit_generate_image_task(
self,
session: aiohttp.ClientSession,
payload: dict,
) -> str:
url = f"{self.base_url}{_ENDPOINT_ASYNC_GENERATE}"
last_status = None
for attempt in range(DEFAULT_MAX_RETRIES + 1):
t0 = time.time()
async with session.post(
url,
json=payload,
headers=self._json_headers(),
) as resp:
elapsed = time.time() - t0
text = await resp.text()
self._log_original_response_body(
f"submit generateImage status={resp.status}",
text,
)
if resp.status not in (200, 201, 202):
last_status = resp.status
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
friendly = get_friendly_message(resp.status)
delay = _compute_delay(
attempt,
DEFAULT_BASE_DELAY,
DEFAULT_MAX_DELAY,
DEFAULT_BACKOFF_FACTOR,
)
print(f"[o1key GPT Image] {friendly} {delay:.1f}s 后重试提交 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
await asyncio.sleep(delay)
continue
raise RuntimeError(self._extract_error_message(text, resp.status))
try:
data = json.loads(text)
except Exception:
raise RuntimeError(f"提交响应 JSON 解析失败，原始内容：{text[:500]}") from None
task_id = data.get("task_id")
if not task_id:
raise RuntimeError(f"提交响应中未找到 task_id: {data}")
status = data.get("status", "")
print(f"[o1key GPT Image] 异步任务已提交 | task_id={task_id} | status={status} | 耗时 {elapsed:.1f}s")
return task_id
if last_status and last_status in HTTP_ERROR_MESSAGES:
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
raise RuntimeError(f"异步任务提交失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
async def _poll_generate_image_task(
self,
session: aiohttp.ClientSession,
task_id: str,
progress_callback: Optional[Callable[[int], None]] = None,
) -> dict:
url = f"{self.base_url}{_ENDPOINT_ASYNC_TASK.format(task_id=task_id)}"
start_time = time.time()
poll_count = 0
last_poll_at = start_time
while True:
if poll_count < len(_ASYNC_POLL_SCHEDULE):
next_poll_at = start_time + _ASYNC_POLL_SCHEDULE[poll_count]
else:
next_poll_at = last_poll_at + _ASYNC_POLL_INTERVAL
sleep_time = next_poll_at - time.time()
if sleep_time > 0:
await asyncio.sleep(sleep_time)
last_poll_at = time.time()
elapsed = last_poll_at - start_time
if elapsed > _ASYNC_MAX_WAIT:
raise RuntimeError(f"任务 {task_id} 超时（>{int(_ASYNC_MAX_WAIT)}秒），请稍后用 task_id 查询结果")
poll_count += 1
async with session.get(url, headers=self._auth_headers()) as resp:
text = await resp.text()
self._log_original_response_body(
f"poll task status={resp.status}",
text,
)
if resp.status != 200:
raise RuntimeError(self._extract_error_message(text, resp.status))
try:
task = json.loads(text)
except Exception:
raise RuntimeError(f"任务查询响应 JSON 解析失败，原始内容：{text[:500]}") from None
status = task.get("status", "UNKNOWN")
progress = task.get("progress")
progress_pct = self._coerce_progress_percent(progress)
if POLL_LOG_ENABLED:
progress_text = f" | progress={progress}" if progress is not None else ""
print(f"[o1key GPT Image] 查询任务 #{poll_count} | task_id={task_id} | status={status}{progress_text}")
if status == "SUCCESS":
self._emit_progress(progress_callback, 100)
return task
if progress_pct is not None and progress_pct < 100:
self._emit_progress(progress_callback, progress_pct)
if status == "FAILURE":
error_message = self._extract_error_message(task, 500) or "生成失败"
raise _AsyncImageTaskFailure(
error_message,
self._extract_error_detail(task),
)
if status not in ("SUBMITTED", "IN_PROGRESS"):
raise RuntimeError(f"未知任务状态 {status}: {task}")
async def _parse_async_task_images(
self,
task: dict,
session: aiohttp.ClientSession,
) -> List[Image.Image]:
data = task.get("data", {})
image_items = data.get("images") if isinstance(data, dict) else None
if not isinstance(image_items, list) or not image_items:
raise RuntimeError(f"任务结果中未找到 data.images: {task}")
images: List[Image.Image] = []
for idx, item in enumerate(image_items, 1):
if not isinstance(item, dict):
continue
url = item.get("url") or item.get("image_url")
b64 = item.get("b64_json", "")
if url and isinstance(url, str) and url.startswith("data:image"):
try:
_, b64_data = url.split(",", 1)
img = Image.open(BytesIO(base64.b64decode(b64_data)))
images.append(img)
print(f"[o1key GPT Image] 第 {idx} 张 data URL 解码完成 ({img.size[0]}×{img.size[1]})")
except Exception as e:
raise RuntimeError(f"第 {idx} 张 data URL 解码失败: {e}") from None
elif url and isinstance(url, str) and url.startswith("http"):
async with session.get(url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"图像下载失败 HTTP {resp.status}，URL: {url}")
img_bytes = await resp.read()
img = Image.open(BytesIO(img_bytes))
images.append(img)
print(f"[o1key GPT Image] 第 {idx} 张 URL 下载完成 ({img.size[0]}×{img.size[1]}) | {url}")
elif b64:
img = self._decode_b64_image(b64, f"第 {idx} 张")
images.append(img)
else:
print(f"[o1key GPT Image] 警告：第 {idx} 条结果既无 url 也无 b64_json，已跳过")
if not images:
raise RuntimeError("任务成功但没有可用图片结果")
return images
async def _generate_image_task_async(
self,
prompt: str,
model: str,
quality: str,
size: str,
n: int,
seed: int,
image_tensor: Optional[List[torch.Tensor]] = None,
mask_tensor: Optional[torch.Tensor] = None,
output_format: str = "png",
progress_callback: Optional[Callable[[int], None]] = None,
) -> List[Image.Image]:
body = self._build_async_generate_body(
prompt=prompt,
model=model,
quality=quality,
size=size,
n=n,
image_list=image_tensor,
mask_tensor=mask_tensor,
output_format=output_format,
)
mode = "图像编辑" if mask_tensor is not None else ("图生图" if image_tensor else "文生图")
body_size = self._json_body_size(body)
print(
f"[o1key GPT Image] {mode} | 新异步接口 | 模型={model} | "
f"quality={quality} | size={size} | n={n} | body={body_size // 1024}KB"
)
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
timeout = aiohttp.ClientTimeout(total=_ASYNC_MAX_WAIT + 120)
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
last_error = None
for retry_index in range(len(_ASYNC_RETRY_DELAYS) + 1):
try:
task_id = await self._submit_generate_image_task(session, body)
task = await self._poll_generate_image_task(
session,
task_id,
progress_callback=progress_callback,
)
return await self._parse_async_task_images(task, session)
except _AsyncImageTaskFailure as error:
last_error = error
detail = error.error_detail
if (
retry_index < len(_ASYNC_RETRY_DELAYS)
and self._should_retry_async_failure(detail, retry_index)
):
delay = _ASYNC_RETRY_DELAYS[retry_index]
code = detail.get("code", "unknown")
category = detail.get("category", "unknown")
failed_task_id = detail.get("task_id", "")
task_text = f" | failed_task_id={failed_task_id}" if failed_task_id else ""
print(
f"[o1key GPT Image] 任务失败但可重试 | code={code} | "
f"category={category}{task_text} | {delay:.0f}s 后重试 "
f"({retry_index + 1}/{len(_ASYNC_RETRY_DELAYS)})"
)
await asyncio.sleep(delay)
continue
raise RuntimeError(str(error)) from None
if last_error is not None:
raise RuntimeError(str(last_error)) from None
raise RuntimeError("生成失败")
# ── 文生图 / 图生图（generations 接口）───────────────────────────────────
async def _generate_async(
@@ -806,6 +1385,53 @@ class GptImageClient:
f"o1key GPT Image 请求超时（>{_REQUEST_TIMEOUT}s），请检查网络或稍后重试"
)
def generate_image_async_sync(
self,
prompt: str,
model: str,
quality: str,
size: str,
n: int,
seed: int,
image_tensor: Optional[List[torch.Tensor]] = None,
mask_tensor: Optional[torch.Tensor] = None,
output_format: str = "png",
progress_callback: Optional[Callable[[int], None]] = None,
) -> List[Image.Image]:
"""
新版异步任务入口，供节点调用。
旧 run_sync 保留兼容，但 GPT Image 节点不再使用旧同步接口。
"""
coro = self._generate_image_task_async(
prompt=prompt,
model=model,
quality=quality,
size=size,
n=n,
seed=seed,
image_tensor=image_tensor,
mask_tensor=mask_tensor,
output_format=output_format,
progress_callback=progress_callback,
)
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return loop.run_until_complete(self._run_with_interrupt(coro))
finally:
loop.close()
with ThreadPoolExecutor(max_workers=1) as executor:
future = executor.submit(_run)
try:
return future.result(timeout=_ASYNC_MAX_WAIT + 150)
except TimeoutError:
raise RuntimeError(
f"o1key GPT Image 异步任务超时（>{int(_ASYNC_MAX_WAIT)}秒），请稍后用 task_id 查询结果"
)
# ── 余额查询 ──────────────────────────────────────────────────────────────
async def _query_balance_async(self) -> dict:
clients/grok_video_client.py
+483
View File
@@ -0,0 +1,483 @@
"""
Grok Video API client.
Flow:
1. POST /v1/videos
2. GET /v1/videos/{task_id}
3. GET /v1/videos/{task_id}/content, or download a URL from the status body
"""
import asyncio
import base64
import json
import os
import re
import time
from typing import Any, Callable, Dict, List, Optional
import aiohttp
from .base_client import BaseAPIClient
from ..utils.config import get_api_base_url, get_api_key_or_raise
from ..utils.http_error import RETRYABLE_STATUS_CODES, get_friendly_message
from ..utils.video_task import (
check_interrupt,
extract_error_message,
extract_progress,
extract_status,
extract_video_url,
interruptible_sleep,
is_failure_status,
is_success_status,
run_with_interrupt,
)
class GrokVideoClient(BaseAPIClient):
CREATE_ENDPOINT = "/v1/videos"
STATUS_ENDPOINT = "/v1/videos/{task_id}"
CONTENT_ENDPOINT = "/v1/videos/{task_id}/content"
MODEL_OPTIONS = ["grok-imagine-video-1.5-preview", "grok-imagine-1.0-video"]
ASPECT_RATIO_OPTIONS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"]
QUALITY_OPTIONS = ["720p"]
MODEL_SECONDS_OPTIONS = {
"grok-imagine-1.0-video": [6, 10, 12, 16, 20],
}
QUALITY_API_MAP = {
"720p": "high",
"high": "high",
}
SUCCESS_STATUSES = {"complete", "completed", "succeed", "succeeded", "success", "done", "finished"}
FAILURE_STATUSES = {"fail", "failed", "failure", "error", "expired", "timeout", "cancelled", "canceled"}
def __init__(self, base_url: Optional[str] = None):
api_key = get_api_key_or_raise("O1KEY_API_KEY")
resolved_base_url = (base_url or "").strip() or get_api_base_url()
super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key)
def get_endpoint(self, **kwargs) -> str:
return self.CREATE_ENDPOINT
def build_request_body(self, **kwargs) -> Dict[str, Any]:
return self.build_video_body(**kwargs)
def parse_response(self, response: Dict[str, Any]) -> Any:
return response
@classmethod
def build_video_body(
cls,
prompt: str,
model: str,
aspect_ratio: str,
seconds: int,
quality: str = "720p",
images: Optional[List[str]] = None,
) -> Dict[str, Any]:
prompt = (prompt or "").strip()
if not prompt:
raise ValueError("提示词不能为空。")
if model not in cls.MODEL_OPTIONS:
raise ValueError(f"模型仅支持: {', '.join(cls.MODEL_OPTIONS)}。")
if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
raise ValueError(f"宽高比仅支持: {', '.join(cls.ASPECT_RATIO_OPTIONS)}。")
try:
seconds_value = int(seconds)
except (TypeError, ValueError):
raise ValueError("秒数必须是整数。") from None
allowed_seconds = cls.MODEL_SECONDS_OPTIONS.get(model)
if allowed_seconds is not None:
if seconds_value not in allowed_seconds:
raise ValueError(
f"模型 {model} 仅支持秒数: "
f"{', '.join(str(s) for s in allowed_seconds)}。"
"请修改为正确的秒数后再发起请求。"
)
elif seconds_value < 5 or seconds_value > 15:
raise ValueError("秒数仅支持 5 到 15。")
api_quality = cls.QUALITY_API_MAP.get(str(quality), str(quality))
if api_quality != "high":
raise ValueError("画质仅支持 720p。")
body: Dict[str, Any] = {
"model": model,
"prompt": prompt,
"aspect_ratio": aspect_ratio,
"seconds": str(seconds_value),
"quality": api_quality,
}
image_list = [img for img in (images or []) if img]
if image_list:
body["images"] = image_list[:3]
return body
@staticmethod
def _safe_task_filename(task_id: str) -> str:
safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._")
return safe or "grok_video"
@staticmethod
def _mask_body_for_log(body: Dict[str, Any]) -> Dict[str, Any]:
log_body = dict(body)
images = log_body.get("images")
if isinstance(images, list):
log_body["images"] = [f"<data-url chars={len(item)}>" for item in images]
return log_body
@staticmethod
def _extract_task_id(payload: Dict[str, Any]) -> Optional[str]:
sources = [payload]
data = payload.get("data")
if isinstance(data, dict):
sources.append(data)
for source in sources:
for key in ("id", "task_id", "video_id"):
value = source.get(key)
if value:
return str(value)
return None
@staticmethod
def _format_http_error(endpoint: str, status: int, error_text: str, task_id: Optional[str] = None) -> str:
message = get_friendly_message(status, error_text)
parts = [
"Grok Video 请求失败。",
f"endpoint: {endpoint}",
f"http_status: {status}",
]
if task_id:
parts.append(f"task_id: {task_id}")
if message:
parts.append(f"message: {message}")
return "\n".join(parts)
@classmethod
def _format_task_failure(cls, task_id: str, payload: Dict[str, Any]) -> str:
return "\n".join(
[
"Grok Video 任务失败。",
f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}",
f"task_id: {task_id}",
f"message: {extract_error_message(payload)}",
]
)
async def _request_json_with_retry(
self,
method: str,
endpoint: str,
session: aiohttp.ClientSession,
task_id: Optional[str] = None,
json_body: Optional[Dict[str, Any]] = None,
max_retries: int = 3,
timeout_seconds: int = 120,
) -> Dict[str, Any]:
url = f"{self.base_url}{endpoint}"
headers = self.get_headers(use_bearer_token=True)
timeout = aiohttp.ClientTimeout(total=timeout_seconds, connect=30, sock_read=timeout_seconds)
last_status = 0
last_text = ""
for attempt in range(max_retries + 1):
check_interrupt()
response = None
try:
response = await run_with_interrupt(
session.request(method, url, json=json_body, headers=headers, timeout=timeout)
)
text = await run_with_interrupt(response.text())
last_status = response.status
last_text = text
if 200 <= response.status < 300:
if not text.strip():
return {}
try:
return json.loads(text)
except Exception:
raise RuntimeError(f"Grok Video 响应 JSON 解析失败，原始内容：{text[:500]}") from None
if response.status in RETRYABLE_STATUS_CODES and attempt < max_retries:
delay = min(2 ** attempt, 8)
print(
f"Grok Video：{get_friendly_message(response.status)} "
f"{delay}s 后重试 ({attempt + 1}/{max_retries})..."
)
await interruptible_sleep(delay)
continue
break
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
if attempt < max_retries:
delay = min(2 ** attempt, 8)
print(f"Grok Video：网络错误，{delay}s 后重试 ({attempt + 1}/{max_retries})...")
await interruptible_sleep(delay)
continue
raise RuntimeError(f"Grok Video 网络错误: {e}") from None
finally:
if response is not None:
response.release()
raise RuntimeError(self._format_http_error(endpoint, last_status, last_text, task_id=task_id))
async def create_video_async(
self,
body: Dict[str, Any],
session: aiohttp.ClientSession,
) -> Dict[str, Any]:
print("Grok Video：正在提交任务...")
return await self._request_json_with_retry(
"POST",
self.CREATE_ENDPOINT,
session=session,
json_body=body,
timeout_seconds=180,
)
async def poll_video_status_async(
self,
task_id: str,
session: aiohttp.ClientSession,
poll_interval: int = 5,
timeout: int = 900,
progress_callback: Optional[Callable[[int, str, float], None]] = None,
) -> Dict[str, Any]:
endpoint = self.STATUS_ENDPOINT.format(task_id=task_id)
start = time.time()
interval = max(1, int(poll_interval))
await interruptible_sleep(interval)
while True:
data = await self._request_json_with_retry(
"GET",
endpoint,
session=session,
task_id=task_id,
timeout_seconds=60,
)
status = extract_status(data)
progress = extract_progress(data)
elapsed = time.time() - start
if progress_callback:
progress_callback(progress, status, elapsed)
if status in self.SUCCESS_STATUSES or is_success_status(status):
return data
if status in self.FAILURE_STATUSES or is_failure_status(status, data):
raise RuntimeError(self._format_task_failure(task_id, data))
if elapsed >= timeout:
raise TimeoutError(
"Grok Video 任务轮询超时；任务未被标记为失败，可用 task_id 继续查询。\n"
f"endpoint: {endpoint}\n"
f"task_id: {task_id}\n"
f"status: {status or 'unknown'}\n"
f"timeout: {timeout}s"
)
await interruptible_sleep(min(interval, max(0.0, timeout - elapsed)))
async def _download_url_to_file(
self,
url: str,
save_path: str,
session: aiohttp.ClientSession,
max_retries: int = 3,
) -> str:
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
last_status = 0
last_text = ""
headers = None
resolved_url = url
if url.startswith("data:"):
if "," not in url:
raise RuntimeError("Grok Video 下载失败：data URL 格式无效。")
_, b64_data = url.split(",", 1)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
f.write(base64.b64decode(b64_data))
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
raise RuntimeError("Grok Video 下载失败：保存后的文件为空。")
return save_path
if url.startswith("/"):
resolved_url = f"{self.base_url}{url}"
headers = self.get_headers(use_bearer_token=True)
for attempt in range(max_retries + 1):
check_interrupt()
async with session.get(
resolved_url,
headers=headers,
timeout=timeout,
allow_redirects=True,
) as response:
if 200 <= response.status < 300:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in response.content.iter_chunked(1024 * 1024):
check_interrupt()
if chunk:
f.write(chunk)
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
raise RuntimeError("Grok Video 下载失败：保存后的文件为空。")
return save_path
last_status = response.status
last_text = await response.text()
if response.status not in RETRYABLE_STATUS_CODES or attempt >= max_retries:
break
delay = min(2 ** attempt, 8)
print(f"Grok Video：下载重试 {attempt + 1}/{max_retries}，{delay}s 后继续...")
await interruptible_sleep(delay)
raise RuntimeError(self._format_http_error("download_url", last_status, last_text))
async def download_video_async(
self,
task_id: str,
save_path: str,
session: aiohttp.ClientSession,
) -> str:
endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id)
url = f"{self.base_url}{endpoint}"
headers = self.get_headers(use_bearer_token=True)
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
last_status = 0
last_text = ""
for attempt in range(4):
check_interrupt()
async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
if 200 <= response.status < 300:
content_type = response.headers.get("Content-Type", "").lower()
if "application/json" in content_type:
data = await response.json(content_type=None)
download_url = extract_video_url(data)
if not download_url:
raise RuntimeError(
"Grok Video 下载失败：content 响应为 JSON，但未包含视频 URL。\n"
f"endpoint: {endpoint}\n"
f"task_id: {task_id}"
)
return await self._download_url_to_file(download_url, save_path, session)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in response.content.iter_chunked(1024 * 1024):
check_interrupt()
if chunk:
f.write(chunk)
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
raise RuntimeError(
"Grok Video 下载失败：保存后的文件为空。\n"
f"endpoint: {endpoint}\n"
f"task_id: {task_id}"
)
return save_path
last_status = response.status
last_text = await response.text()
if response.status not in RETRYABLE_STATUS_CODES or attempt >= 3:
break
delay = min(2 ** attempt, 8)
print(f"Grok 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())
clients/newapi_veo_client.py
+542
View File
@@ -0,0 +1,542 @@
"""
new-api Veo 3.1 video client.
Implements the OpenAI-compatible /v1/videos task flow:
submit, poll, and stream-download video content.
"""
import asyncio
import json
import os
import re
import time
from typing import Any, Callable, Dict, Optional
import aiohttp
from .base_client import BaseAPIClient
from ..utils.config import get_api_base_url, get_api_key_or_raise
class NewAPIVeoClient(BaseAPIClient):
CREATE_ENDPOINT = "/v1/videos"
STATUS_ENDPOINT = "/v1/videos/{task_id}"
CONTENT_ENDPOINT = "/v1/videos/{task_id}/content"
RETRYABLE_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504}
COMPLETED_STATUSES = {"completed", "succeeded", "success", "done"}
FAILED_STATUSES = {"failed", "error", "cancelled", "canceled"}
def __init__(
self,
base_url: Optional[str] = None,
):
api_key = get_api_key_or_raise("O1KEY_API_KEY")
resolved_base_url = (base_url or "").strip() or get_api_base_url()
super().__init__(base_url=resolved_base_url.rstrip("/"), api_key=api_key)
def get_endpoint(self, **kwargs) -> str:
return self.CREATE_ENDPOINT
def build_request_body(self, **kwargs) -> Dict[str, Any]:
return self._build_video_body(**kwargs)
def parse_response(self, response: Dict[str, Any]) -> Any:
return response
@staticmethod
def _build_video_body(
prompt: str,
model: str,
duration: int,
aspect_ratio: str,
resolution: str,
negative_prompt: str = "",
generate_audio: bool = True,
) -> Dict[str, Any]:
metadata: Dict[str, Any] = {
"aspectRatio": aspect_ratio,
"resolution": resolution,
"generateAudio": bool(generate_audio),
}
negative_prompt = (negative_prompt or "").strip()
if negative_prompt:
metadata["negativePrompt"] = negative_prompt
body: Dict[str, Any] = {
"model": model,
"prompt": prompt,
"duration": int(duration),
"metadata": metadata,
}
return body
@staticmethod
def _print_request_body(body: Dict[str, Any], image_bytes: Optional[bytes] = None) -> None:
log_body = dict(body)
if image_bytes is not None:
log_body["input_reference"] = f"<PNG bytes: {len(image_bytes)}>"
print(
"NewAPI Veo request body:\n"
f"{json.dumps(log_body, ensure_ascii=False, indent=2)}"
)
@staticmethod
def _safe_task_filename(task_id: str) -> str:
safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", task_id).strip("._")
return safe or "newapi_veo"
@staticmethod
def _extract_task_id(data: Dict[str, Any]) -> Optional[str]:
for key in ("id", "task_id", "video_id"):
value = data.get(key)
if value:
return str(value)
nested = data.get("data")
if isinstance(nested, dict):
for key in ("id", "task_id", "video_id"):
value = nested.get(key)
if value:
return str(value)
return None
@staticmethod
def _extract_status(data: Dict[str, Any]) -> str:
for key in ("status", "state", "task_status"):
value = data.get(key)
if value:
return str(value).lower()
nested = data.get("data")
if isinstance(nested, dict):
for key in ("status", "state", "task_status"):
value = nested.get(key)
if value:
return str(value).lower()
return "unknown"
@staticmethod
def _extract_progress(data: Dict[str, Any]) -> int:
progress = data.get("progress")
if progress is None and isinstance(data.get("data"), dict):
progress = data["data"].get("progress")
if isinstance(progress, str):
progress = progress.rstrip("%").strip()
try:
return int(float(progress))
except ValueError:
return 0
if isinstance(progress, (int, float)):
return int(progress)
return 0
@classmethod
def _format_http_error(
cls,
endpoint: str,
status: int,
error_text: str,
task_id: Optional[str] = None,
) -> str:
code = ""
message = error_text
try:
payload = json.loads(error_text)
error = payload.get("error", payload)
if isinstance(error, dict):
code = str(error.get("code") or error.get("type") or "")
message = str(error.get("message") or payload.get("message") or error_text)
elif error is not None:
message = str(error)
except Exception:
pass
message = (message or "").strip()
if len(message) > 1200:
message = message[:1200] + "...(truncated)"
if status in (401, 403):
hint = "凭证或分组权限问题，请检查 new-api token、模型分组或渠道权限。"
elif status == 429:
hint = "频率或额度限制，请稍后重试或检查 new-api 额度。"
elif status in (502, 503, 504):
hint = "上游服务暂时不可用或超时，请稍后用 task_id 继续查询。"
elif status == 400:
hint = "请求参数错误，请检查 model、duration、metadata 和图片输入。"
else:
hint = "new-api 视频请求失败。"
parts = [
hint,
f"endpoint: {endpoint}",
f"http_status: {status}",
]
if task_id:
parts.append(f"task_id: {task_id}")
if code:
parts.append(f"error_code: {code}")
if message:
parts.append(f"message: {message}")
return "\n".join(parts)
@classmethod
def _format_task_failure(cls, task_id: str, data: Dict[str, Any]) -> str:
error = data.get("error")
if error is None and isinstance(data.get("data"), dict):
error = data["data"].get("error")
if isinstance(error, dict):
code = error.get("code") or error.get("type") or ""
message = error.get("message") or json.dumps(error, ensure_ascii=False)
else:
code = ""
message = str(error or "未知错误")
return "\n".join(
[
"Veo 视频任务失败。",
f"endpoint: {cls.STATUS_ENDPOINT.format(task_id=task_id)}",
f"task_id: {task_id}",
f"error_code: {code}",
f"message: {message}",
]
)
async def create_video_async(
self,
prompt: str,
model: str,
duration: int,
aspect_ratio: str,
resolution: str,
negative_prompt: str = "",
generate_audio: bool = True,
image_bytes: Optional[bytes] = None,
session: Optional[aiohttp.ClientSession] = None,
) -> Dict[str, Any]:
body = self._build_video_body(
prompt=prompt,
model=model,
duration=duration,
aspect_ratio=aspect_ratio,
resolution=resolution,
negative_prompt=negative_prompt,
generate_audio=generate_audio,
)
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
close_session = False
if session is None:
session = self._make_session()
close_session = True
try:
timeout = aiohttp.ClientTimeout(total=120, connect=30, sock_read=120)
headers = {"Authorization": f"Bearer {self.api_key}"}
if image_bytes is not None:
if len(image_bytes) > self.max_request_size:
raise ValueError(
f"输入图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制"
)
self._print_request_body(body, image_bytes=image_bytes)
form = aiohttp.FormData()
form.add_field("model", body["model"])
form.add_field("prompt", body["prompt"])
form.add_field("duration", str(body["duration"]))
form.add_field("metadata", json.dumps(body["metadata"], ensure_ascii=False))
form.add_field(
"input_reference",
image_bytes,
filename="input_reference.png",
content_type="image/png",
)
request_kwargs = {"data": form, "headers": headers}
print(
"NewAPI Veo: POST /v1/videos multipart "
f"| model={model} | duration={duration}s | {resolution} {aspect_ratio}"
)
else:
self._print_request_body(body)
headers["Content-Type"] = "application/json"
request_kwargs = {"json": body, "headers": headers}
print(
"NewAPI Veo: POST /v1/videos json "
f"| model={model} | duration={duration}s | {resolution} {aspect_ratio}"
)
async with session.post(url, timeout=timeout, **request_kwargs) as response:
if response.status >= 300:
error_text = await response.text()
raise RuntimeError(
self._format_http_error(
self.CREATE_ENDPOINT,
response.status,
error_text,
)
)
return await response.json()
finally:
if close_session:
await session.close()
async def _get_json_with_retry(
self,
endpoint: str,
session: aiohttp.ClientSession,
task_id: Optional[str] = None,
max_retries: int = 3,
) -> Dict[str, Any]:
url = f"{self.base_url}{endpoint}"
headers = self.get_headers(use_bearer_token=True)
timeout = aiohttp.ClientTimeout(total=60, connect=30, sock_read=60)
last_error = ""
last_status = 0
for attempt in range(max_retries + 1):
async with session.get(url, headers=headers, timeout=timeout) as response:
if response.status < 300:
return await response.json()
last_status = response.status
last_error = await response.text()
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= max_retries:
break
await asyncio.sleep(min(2 ** attempt, 8))
raise RuntimeError(
self._format_http_error(endpoint, last_status, last_error, task_id=task_id)
)
async def poll_video_status_async(
self,
task_id: str,
poll_interval: int = 5,
timeout: int = 900,
progress_callback: Optional[Callable[[int, str, float], None]] = None,
session: Optional[aiohttp.ClientSession] = None,
) -> Dict[str, Any]:
endpoint = self.STATUS_ENDPOINT.format(task_id=task_id)
close_session = False
if session is None:
session = self._make_session()
close_session = True
start = time.time()
try:
poll_interval = max(1, int(poll_interval))
await asyncio.sleep(poll_interval)
while True:
data = await self._get_json_with_retry(endpoint, session, task_id=task_id)
status = self._extract_status(data)
elapsed = time.time() - start
progress = self._extract_progress(data)
if status == "unknown":
print(
"NewAPI Veo status response did not include a recognized status field:\n"
f"{json.dumps(data, ensure_ascii=False, indent=2)[:1200]}"
)
if progress_callback:
progress_callback(progress, status, elapsed)
if status in self.COMPLETED_STATUSES:
return data
if status in self.FAILED_STATUSES:
raise RuntimeError(self._format_task_failure(task_id, data))
if elapsed >= timeout:
raise TimeoutError(
"Veo 视频任务轮询超时；任务未被标记为失败，可用 task_id 继续查询。\n"
f"endpoint: {endpoint}\n"
f"task_id: {task_id}\n"
f"status: {status}\n"
f"timeout: {timeout}s"
)
remaining = max(0.0, timeout - elapsed)
await asyncio.sleep(min(poll_interval, remaining))
finally:
if close_session:
await session.close()
async def _download_url_to_file(
self,
url: str,
save_path: str,
session: aiohttp.ClientSession,
max_retries: int = 3,
) -> None:
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
last_status = 0
last_error = ""
for attempt in range(max_retries + 1):
async with session.get(url, timeout=timeout, allow_redirects=True) as response:
if response.status < 300:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in response.content.iter_chunked(1024 * 1024):
if chunk:
f.write(chunk)
return
last_status = response.status
last_error = await response.text()
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= max_retries:
break
await asyncio.sleep(min(2 ** attempt, 8))
raise RuntimeError(
self._format_http_error("download_url", last_status, last_error)
)
async def download_video_async(
self,
task_id: str,
save_path: str,
session: Optional[aiohttp.ClientSession] = None,
) -> str:
endpoint = self.CONTENT_ENDPOINT.format(task_id=task_id)
url = f"{self.base_url}{endpoint}"
headers = self.get_headers(use_bearer_token=True)
timeout = aiohttp.ClientTimeout(total=900, connect=30, sock_read=900)
close_session = False
if session is None:
session = self._make_session()
close_session = True
try:
last_status = 0
last_error = ""
for attempt in range(4):
async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
if response.status < 300:
content_type = response.headers.get("Content-Type", "")
if "application/json" in content_type.lower():
data = await response.json()
nested = data.get("data") if isinstance(data.get("data"), dict) else {}
download_url = (
data.get("url")
or data.get("download_url")
or nested.get("url")
or nested.get("download_url")
)
if not download_url:
raise RuntimeError(
"视频下载失败: content 响应为 JSON，但未包含 url/download_url。\n"
f"endpoint: {endpoint}\n"
f"task_id: {task_id}"
)
await self._download_url_to_file(download_url, save_path, session)
else:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in response.content.iter_chunked(1024 * 1024):
if chunk:
f.write(chunk)
if not os.path.isfile(save_path) or os.path.getsize(save_path) <= 0:
raise RuntimeError(
"视频下载失败: 保存后的文件为空。\n"
f"endpoint: {endpoint}\n"
f"task_id: {task_id}"
)
return save_path
last_status = response.status
last_error = await response.text()
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= 3:
break
await asyncio.sleep(min(2 ** attempt, 8))
raise RuntimeError(
self._format_http_error(endpoint, last_status, last_error, task_id=task_id)
)
finally:
if close_session:
await session.close()
def generate_video_sync(
self,
prompt: str,
model: str,
duration: int,
aspect_ratio: str,
resolution: str,
output_dir: str,
negative_prompt: str = "",
generate_audio: bool = True,
image_bytes: Optional[bytes] = None,
poll_interval: int = 5,
timeout: int = 900,
reuse_task_id: str = "",
progress_callback: Optional[Callable[[int, str, float], None]] = None,
) -> Dict[str, Any]:
async def _run():
async with self._make_session() as session:
create_response: Dict[str, Any] = {}
task_id = (reuse_task_id or "").strip()
if task_id:
print(f"NewAPI Veo: reuse task_id={task_id}")
else:
create_response = await self.create_video_async(
prompt=prompt,
model=model,
duration=duration,
aspect_ratio=aspect_ratio,
resolution=resolution,
negative_prompt=negative_prompt,
generate_audio=generate_audio,
image_bytes=image_bytes,
session=session,
)
task_id = self._extract_task_id(create_response) or ""
if not task_id:
raise RuntimeError(
"new-api 未返回视频任务 ID。\n"
f"endpoint: {self.CREATE_ENDPOINT}\n"
f"response: {json.dumps(create_response, ensure_ascii=False)[:1200]}"
)
status_response = await self.poll_video_status_async(
task_id=task_id,
poll_interval=poll_interval,
timeout=timeout,
progress_callback=progress_callback,
session=session,
)
status = self._extract_status(status_response)
os.makedirs(output_dir, exist_ok=True)
filename = f"{self._safe_task_filename(task_id)}.mp4"
save_path = os.path.join(output_dir, filename)
video_path = await self.download_video_async(
task_id=task_id,
save_path=save_path,
session=session,
)
return {
"task_id": task_id,
"status": status,
"video_path": video_path,
"raw_json": {
"create": create_response,
"status": status_response,
},
}
return self.run_async_in_thread(_run())