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

Publish current ComfyUI O1Key code baseline

Browse Source 
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
2026-09-24 19:56:48 +08:00
parent 3e337722ab
commit ba920f2b66
183 changed files with 49496 additions and 9934 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clients/__init__.py
+31 -13
View File
@@ -1,16 +1,34 @@
"""
API 客户端模块
包含与外部 API 通信的客户端实现
"""API clients exposed through lazy imports.
Importing one client submodule no longer imports every provider client. This
keeps plugin startup lightweight and isolates optional provider dependencies.
"""
from .base_client import BaseAPIClient
from .gemini_client import GeminiAPIClient
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
from importlib import import_module
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'NewAPIVeoClient', 'GrokVideoClient', 'OpenAIAPIClient']
_EXPORTS = {
"BaseAPIClient": ("base_client", "BaseAPIClient"),
"GeminiAPIClient": ("gemini_client", "GeminiAPIClient"),
"GeminiFlashClient": ("gemini_flash_client", "GeminiFlashClient"),
"SoraClient": ("sora_client", "SoraClient"),
"VeoClient": ("veo_client", "VeoClient"),
"NewAPIVeoClient": ("newapi_veo_client", "NewAPIVeoClient"),
"MiniMaxH3Client": ("minimax_h3_client", "MiniMaxH3Client"),
"GrokVideoClient": ("grok_video_client", "GrokVideoClient"),
"OmniFlashClient": ("omni_flash_client", "OmniFlashClient"),
"SeedreamImageClient": ("seedream_image_client", "SeedreamImageClient"),
}
__all__ = list(_EXPORTS)
def __getattr__(name):
try:
module_name, attribute_name = _EXPORTS[name]
except KeyError as exc:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
value = getattr(import_module(f".{module_name}", __name__), attribute_name)
globals()[name] = value
return value
clients/base_async_provider.py
-183
View File
@@ -1,183 +0,0 @@
"""
异步生图 Provider 抽象基类
定义异步提交+轮询模式的统一接口，支持多种生图模型后端
每个 Provider 封装一种 API 后端的通信协议：
- 如何提交任务（端点、请求体格式）
- 如何轮询状态（端点、状态字段语义）
- 如何解析结果（响应格式、图片提取方式）
新增第三方生图模型时，只需实现此接口即可接入异步节点。
"""
from abc import ABC, abstractmethod
from typing import Dict, List, Optional
from PIL import Image
class BaseAsyncImageProvider(ABC):
"""异步生图 Provider 抽象基类"""
def __init__(self, api_key: str, proxy_url: Optional[str] = None):
self.api_key = api_key
self.proxy_url = proxy_url
# ========================================================================
# 必须实现的抽象方法
# ========================================================================
@property
@abstractmethod
def api_base_url(self) -> str:
"""异步 API 的基础 URL，如 https://cf-api.o1key.com"""
...
@abstractmethod
def get_submit_endpoint(self, model: str, resolution: str) -> str:
"""获取提交任务的 API 端点路径（不含 base_url）"""
...
@abstractmethod
def build_submit_body(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
**kwargs
) -> dict:
"""构建提交任务的请求体"""
...
@abstractmethod
def extract_task_id(self, response: dict) -> str:
"""从提交响应中提取 task_id"""
...
@abstractmethod
def extract_status(self, response: dict) -> str:
"""从轮询响应中提取任务状态（如 SUBMITTED / IN_PROGRESS / SUCCESS / FAILURE）"""
...
@abstractmethod
async def parse_result(
self,
result_data: dict,
session
) -> List[Image.Image]:
"""从任务完成后的 result data 中解析生成的图像列表"""
...
@abstractmethod
def get_models(self) -> List[str]:
"""获取此 Provider 支持的模型 ID 列表"""
...
@abstractmethod
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
"""获取指定模型支持的宽高比"""
...
@abstractmethod
def get_model_resolutions(self, model_id: str) -> List[str]:
"""获取指定模型支持的分辨率"""
...
# ========================================================================
# 可选的覆盖方法
# ========================================================================
def get_poll_endpoint(self, task_id: str) -> str:
"""获取轮询任务状态的 API 端点路径（默认实现适用于 o1key 异步 API）"""
return f"/async/v1/tasks/{task_id}"
def get_headers(self) -> dict:
"""获取 HTTP 请求头"""
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
def get_all_aspect_ratios(self) -> List[str]:
"""获取所有模型支持的宽高比（去重合并）"""
seen = set()
result = []
for model_id in self.get_models():
for ratio in self.get_model_aspect_ratios(model_id):
if ratio not in seen:
seen.add(ratio)
result.append(ratio)
return result
def get_all_resolutions(self) -> List[str]:
"""获取所有模型支持的分辨率（去重，按固定顺序排列）"""
_ORDER = ["512px", "1K", "2K", "4K"]
seen = set()
for model_id in self.get_models():
for res in self.get_model_resolutions(model_id):
seen.add(res)
return [r for r in _ORDER if r in seen]
def get_extra_inputs(self) -> dict:
"""
返回此 Provider 特有的额外 ComfyUI 输入参数。
子类重写以声明 Provider 专有的选项（如 Google Search Grounding）。
Returns:
dict，格式与 ComfyUI INPUT_TYPES 的 optional 字段一致
"""
return {}
def get_extra_kwargs(self, **kwargs) -> dict:
"""
从 ComfyUI kwargs 中提取此 Provider 特有的参数，
转换为 build_submit_body 可接收的 kwargs。
子类重写以处理 Provider 专有参数。
"""
return {}
def extract_progress(self, response: dict) -> Optional[float]:
"""
从轮询响应中提取生成进度。
Args:
response: 轮询接口返回的完整响应字典
Returns:
0.0-1.0 之间的进度值，或 None 表示该响应不含进度信息
"""
return None
def query_balance_sync(self) -> Optional[dict]:
"""
同步查询账户余额（可选）。
返回 None 表示不支持。
"""
return None
def format_balance_info(self, balance_data: dict) -> str:
"""格式化余额信息为展示文本"""
return ""
# ========================================================================
# 工具方法
# ========================================================================
@staticmethod
def build_proxy_url(port: str) -> Optional[str]:
"""
将端口号字符串转为 aiohttp 可用的 HTTP 代理 URL。
兼容 v2rayN (10808)、Clash Verge (7897) 等。
Args:
port: 用户填写的端口号，如 "7897"，空字符串返回 None
Returns:
代理 URL 或 None
"""
port = (port or "").strip()
if not port or not port.isdigit():
return None
return f"http://127.0.0.1:{port}"
clients/element_client.py
+74
View File
@@ -0,0 +1,74 @@
"""
可灵主体（Element）API 客户端
封装对 {base}/kling/v1/general/* 的调用，统一注入 Authorization。
两类调用方：
- 后端代理路由（__init__.py）：面板的上传/创建/刷新/列表/删除。
- 节点提交（K3_video.py）：生视频前查列表拿 名称→element_id 映射。
所有方法返回后端的响应信封 {"success", "message", "data"} 解出的 data，
失败抛 RuntimeError（带 message），由调用方决定如何呈现。
"""
import aiohttp
from ..utils.config import get_api_key_or_raise, get_base_url_by_route
_ELEMENT_PREFIX = "/kling/v1/general"
def _headers(api_key: str = None) -> dict:
key = api_key or get_api_key_or_raise()
return {"Authorization": f"Bearer {key}"}
def _resolve_base(route: str = None, base_url: str = None) -> str:
"""Prefer an explicit URL, otherwise use the global network route."""
if base_url:
return base_url.rstrip("/")
return get_base_url_by_route().rstrip("/")
def _unwrap(payload: dict):
"""从响应信封取 data；success=false 时抛 RuntimeError。"""
if not isinstance(payload, dict):
raise RuntimeError(f"主体接口返回异常：{payload!r}")
if not payload.get("success", False):
raise RuntimeError(payload.get("message") or "主体接口调用失败")
return payload.get("data")
async def list_elements(session: aiohttp.ClientSession, *, route=None, base_url=None,
include_all=False, api_key=None):
"""GET /advanced-custom-elements：返回主体列表（默认仅 succeed）。"""
base = _resolve_base(route, base_url)
url = f"{base}{_ELEMENT_PREFIX}/advanced-custom-elements"
params = {"pageNum": "1", "pageSize": "100"}
async with session.get(url, headers=_headers(api_key), params=params) as resp:
data = await resp.json()
result = _unwrap(data)
# 新API返回: {"code": 0, "data": [...], "total": N}
if isinstance(result, dict) and "data" in result:
elements = result.get("data", [])
else:
elements = result if isinstance(result, list) else []
# 过滤：默认只返回 succeed 状态
if not include_all:
elements = [e for e in elements if e.get("status") == "succeed"]
return elements
async def fetch_name_to_id_map(session: aiohttp.ClientSession, *, route=None,
base_url=None, api_key=None) -> dict:
"""生视频用：返回 {主体名称: element_id}，仅含已成功的主体。
注意：新API返回的 element_id 是 int64 数字类型，不是字符串。
"""
elements = await list_elements(session, route=route, base_url=base_url, api_key=api_key)
mapping = {}
for e in elements:
name = (e.get("name") or "").strip()
eid = e.get("element_id")
# element_id 可能是数字或字符串，统一保持原始类型（生视频时需要数字）
if name and eid is not None:
mapping[name] = eid
return mapping
clients/flux_edit_client.py
+2 -2
View File
@@ -1,6 +1,6 @@
"""
Flux 图像编辑 API 客户端
通过 vip.o1key.com 调用 Flux2 图像编辑 + SeedVR2 超分辨率服务
通过 api.o1key.cn 调用 Flux2 图像编辑 + SeedVR2 超分辨率服务
工作流程：
1. submit_task → POST /v1/images/edits (multipart/form-data 提交主图+参考图+提示词)
@@ -30,7 +30,7 @@ class FluxEditClient:
"""
Flux 图像编辑客户端
对接 vip.o1key.com 上的 /v1/images/edits 接口，
对接 api.o1key.cn 上的 /v1/images/edits 接口，
将图像编辑+超分辨率任务提交到远程服务器执行。
"""
clients/gemini_async_provider.py
-194
View File
@@ -1,194 +0,0 @@
"""
Gemini 异步生图 Provider
通过 cf-api.o1key.com 的异步提交+轮询接口调用 Gemini 图像生成模型
协议说明：
- 提交：POST {base}/async{gemini_endpoint}?image_format=url
- 轮询：GET {base}/async/v1/tasks/{task_id}
- 结果：可能直接返回 image_url，也可能返回 Gemini 标准 candidates 格式
"""
from io import BytesIO
from typing import Dict, List, Optional
from PIL import Image
from .base_async_provider import BaseAsyncImageProvider
from .gemini_client import GeminiAPIClient
from ..utils.config import get_async_api_base_url, get_api_key_or_raise
from ..models_config import (
get_enabled_models,
get_model_supported_aspect_ratios,
get_model_supported_resolutions,
)
class GeminiAsyncImageProvider(BaseAsyncImageProvider):
"""
Gemini 异步生图 Provider
委托 GeminiAPIClient 处理：
- 端点构造（get_endpoint）
- 请求体构建（build_request_body）
- 响应解析（parse_response_async）
"""
def __init__(self, api_key: str = None, proxy_url: str = None):
if api_key is None:
api_key = get_api_key_or_raise("O1KEY_API_KEY")
super().__init__(api_key=api_key, proxy_url=proxy_url)
self._client = GeminiAPIClient(api_key=api_key)
# ========================================================================
# 抽象方法实现
# ========================================================================
@property
def api_base_url(self) -> str:
return getattr(self, '_route_base_url', None) or get_async_api_base_url()
def get_submit_endpoint(self, model: str, resolution: str) -> str:
gemini_endpoint = self._client.get_endpoint(
model=model, resolution=resolution, image_format="url"
)
base = gemini_endpoint.split("?")[0]
async_endpoint = f"/async{base}"
if "?" in gemini_endpoint:
async_endpoint += "?" + gemini_endpoint.split("?", 1)[1]
return async_endpoint
def build_submit_body(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
**kwargs
) -> dict:
return self._client.build_request_body(
prompt=prompt,
images=images,
aspect_ratio=aspect_ratio,
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),
thinking_level=kwargs.get("thinking_level"),
request_log_enabled=False,
)
def extract_task_id(self, response: dict) -> str:
task_id = response.get("task_id")
if not task_id:
raise RuntimeError(f"提交响应中未找到 task_id: {response}")
return task_id
def extract_status(self, response: dict) -> str:
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:
async with session.get(image_url) as img_resp:
if img_resp.status == 200:
img_bytes = await img_resp.read()
return [Image.open(BytesIO(img_bytes))]
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
# 否则按 Gemini 标准格式解析
images_list, _ = await self._client.parse_response_async(result_data, session=session)
return images_list
def get_models(self) -> List[str]:
return get_enabled_models()
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
return get_model_supported_aspect_ratios(model_id)
def get_model_resolutions(self, model_id: str) -> List[str]:
return get_model_supported_resolutions(model_id)
# ========================================================================
# 可选方法覆盖
# ========================================================================
def get_extra_inputs(self) -> dict:
"""Gemini 专有：Google Search Grounding"""
return {
"联网功能": (["关闭", "打开"], {"default": "关闭"}),
}
def get_extra_kwargs(self, **kwargs) -> dict:
return {
"enable_grounding": kwargs.pop("联网功能", "关闭") == "打开",
}
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", "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", "percent"):
progress = _coerce(progress_info.get(field))
if progress is not None:
return progress
return None
def query_balance_sync(self) -> Optional[dict]:
try:
return self._client.query_balance_sync()
except Exception:
return None
def format_balance_info(self, balance_data: dict) -> str:
return self._client.format_balance_info(balance_data)
clients/gemini_client.py
+68 -50
View File
@@ -1,12 +1,13 @@
"""
Gemini API 客户端
处理与 api.o1key.com 的通信，用于图像生成
处理与 api.o1key.cn 的通信，用于图像生成
"""
import base64
import re
import time
from io import BytesIO
from typing import Any, Callable, Dict, List, Optional
from typing import Any, Awaitable, Callable, Dict, List, Optional
import aiohttp
from PIL import Image
@@ -80,7 +81,7 @@ class GeminiAPIClient(BaseAPIClient):
endpoint = "/v1beta/models/nano-banana-pro-2k:generateContent"
elif model == "nano-banana-2-次卡":
if resolution == "512px":
if resolution == "512":
endpoint = "/v1beta/models/nano-banana-2-0.5k:generateContent"
elif resolution == "1K":
endpoint = "/v1beta/models/nano-banana-2-1k:generateContent"
@@ -92,7 +93,7 @@ class GeminiAPIClient(BaseAPIClient):
endpoint = "/v1beta/models/nano-banana-2-2k:generateContent"
elif model == "nano-banana-2-官方计费":
if resolution == "512px":
if resolution == "512":
endpoint = "/v1beta/models/nano-banana-2-0.5k-official:generateContent"
elif resolution == "1K":
endpoint = "/v1beta/models/nano-banana-2-1k-official:generateContent"
@@ -214,7 +215,7 @@ class GeminiAPIClient(BaseAPIClient):
# 添加文本部分
parts.append({"text": prompt})
image_config = {"imageSize": resolution}
image_config = {"imageSize": {"512": "512px"}.get(resolution, resolution)}
if aspect_ratio and aspect_ratio != "智能":
image_config["aspectRatio"] = aspect_ratio
@@ -364,7 +365,8 @@ class GeminiAPIClient(BaseAPIClient):
async def parse_response_async(
self,
response: Dict[str, Any],
session: Optional[aiohttp.ClientSession] = None
session: Optional[aiohttp.ClientSession] = None,
image_downloader: Optional[Callable[[str], Awaitable[bytes]]] = None,
) -> tuple[List[Image.Image], Dict[str, Any]]:
"""
异步解析 API 响应，提取生成的图像
@@ -432,6 +434,24 @@ class GeminiAPIClient(BaseAPIClient):
if session is None:
session = self._make_session()
close_session = True
def _tag_image(img: Image.Image, raw_bytes: bytes) -> Image.Image:
fmt = (img.format or "").upper()
if fmt == "JPG":
fmt = "JPEG"
if fmt:
img.format = fmt
setattr(img, "_o1key_original_format", fmt)
setattr(img, "_o1key_original_bytes", raw_bytes)
return img
async def _download_url(url: str) -> bytes:
if image_downloader is not None:
return await image_downloader(url)
async with session.get(url) as img_response:
if img_response.status != 200:
raise RuntimeError(f"图片下载失败 ({img_response.status})")
return await img_response.read()
try:
for candidate_idx, candidate in enumerate(candidates):
@@ -446,7 +466,7 @@ class GeminiAPIClient(BaseAPIClient):
inline_data_key = "inline_data"
elif "inlineData" in part:
inline_data_key = "inlineData"
if inline_data_key:
# 跳过思考链草稿图（thought:true 标记的 inlineData 为模型自检用途，非最终输出）
if part.get("thought") is True:
@@ -458,6 +478,7 @@ class GeminiAPIClient(BaseAPIClient):
if img_data:
img = decode_base64_to_pil(img_data)
_tag_image(img, base64.b64decode(img_data))
images.append(img)
# 记录格式信息
@@ -475,21 +496,20 @@ class GeminiAPIClient(BaseAPIClient):
if url:
try:
download_start = time.time()
async with session.get(url) as img_response:
if img_response.status == 200:
img_bytes = await img_response.read()
download_time = time.time() - download_start
img_size = len(img_bytes)
speed = img_size / download_time if download_time > 0 else 0
img_bytes = await _download_url(url)
download_time = time.time() - download_start
img_size = len(img_bytes)
speed = img_size / download_time if download_time > 0 else 0
img = Image.open(BytesIO(img_bytes))
images.append(img)
img = Image.open(BytesIO(img_bytes))
_tag_image(img, img_bytes)
images.append(img)
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = img_size
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = img_size
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception:
pass # 静默失败
@@ -514,46 +534,44 @@ class GeminiAPIClient(BaseAPIClient):
try:
# 使用 aiohttp 异步下载
download_start = time.time()
async with session.get(url) as img_response:
if img_response.status == 200:
img_data = await img_response.read()
download_time = time.time() - download_start
img_size = len(img_data)
speed = img_size / download_time if download_time > 0 else 0
img = Image.open(BytesIO(img_data))
images.append(img)
# 记录格式信息（只记录第一张）
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = img_size
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception as e:
pass # 静默失败，继续尝试其他URL
# 方式3: 直接的 URL 字段 - 也改为异步
elif "imageUrl" in part or "url" in part:
url = part.get("imageUrl") or part.get("url")
try:
download_start = time.time()
async with session.get(url) as img_response:
if img_response.status == 200:
img_data = await img_response.read()
img_data = await _download_url(url)
download_time = time.time() - download_start
img_size = len(img_data)
speed = img_size / download_time if download_time > 0 else 0
img = Image.open(BytesIO(img_data))
_tag_image(img, img_data)
images.append(img)
# 记录格式信息
# 记录格式信息（只记录第一张）
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = img_size
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception as e:
pass # 静默失败，继续尝试其他URL
# 方式3: 直接的 URL 字段 - 也改为异步
elif "imageUrl" in part or "url" in part:
url = part.get("imageUrl") or part.get("url")
try:
download_start = time.time()
img_data = await _download_url(url)
download_time = time.time() - download_start
img_size = len(img_data)
speed = img_size / download_time if download_time > 0 else 0
img = Image.open(BytesIO(img_data))
_tag_image(img, img_data)
images.append(img)
# 记录格式信息
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = img_size
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception as e:
pass # 静默失败
clients/gpt_image_client.py
+892 -340
View File
File diff suppressed because it is too large Load Diff
clients/grok_image_client.py
+126 -23
View File
@@ -48,7 +48,7 @@ _RETRY_DELAY = 5
class GrokImageClient:
def __init__(self, route: str = "全球加速"):
def __init__(self, route: str = None):
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
self.base_url = get_base_url_by_route(route)
@@ -73,8 +73,13 @@ class GrokImageClient:
step = 0
while len(png_bytes) > max_bytes:
scale = 0.894
w = max(1, int(w * scale))
h = max(1, int(h * scale))
next_w = max(1, int(w * scale))
next_h = max(1, int(h * scale))
if (next_w, next_h) == (w, h):
raise RuntimeError(
f"Grok Image 无法将图像缩小到 {max_bytes} 字节以内"
)
w, h = next_w, next_h
img = img.resize((w, h), Image.LANCZOS)
buf = BytesIO()
img.save(buf, format="PNG")
@@ -99,6 +104,116 @@ class GrokImageClient:
tensors.append(torch.from_numpy(arr))
return torch.stack(tensors, dim=0)
@classmethod
def _build_edit_body(
cls,
prompt: str,
model: str,
aspect_ratio: str,
resolution: str,
image_list: List[torch.Tensor],
) -> dict:
"""构建最多三张参考图的编辑请求，并确保完整 JSON 不超过 20MB。"""
if len(image_list) > 3:
raise ValueError("Grok Image 最多支持 3 张参考图")
reference_images = []
for index, tensor in enumerate(image_list, start=1):
pil_images = tensor_to_pil(tensor)
if not pil_images:
raise ValueError(f"无法读取参考图{index}")
reference_images.append(pil_images[0])
if not reference_images:
raise ValueError("图像编辑至少需要 1 张参考图")
body: dict = {
"model": _MODEL_NAME_MAP.get(model, model),
"prompt": prompt,
"response_format": "b64_json",
}
if aspect_ratio and aspect_ratio != "auto":
body["aspect_ratio"] = aspect_ratio
if resolution:
body["resolution"] = resolution
# base64 大约是原始字节的 4/3。预留 1MB 给提示词和 JSON 字段，
# 同时保留旧逻辑的单图 PNG 最大 10MB 上限。
base_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
available = _MAX_BODY_BYTES - base_size - 1024 * 1024
if available <= 0:
raise ValueError("提示词和请求参数已超过 Grok Image 20MB 请求体限制")
per_image_limit = min(
_MAX_BODY_BYTES // 2,
max(1, int(available * 0.75) // len(reference_images)),
)
png_images = []
for index, image in enumerate(reference_images, start=1):
buffer = BytesIO()
image.save(buffer, format="PNG")
png_images.append(
cls._shrink_png_to_limit(
buffer.getvalue(),
per_image_limit,
label=f"参考图{index}",
)
)
def _set_images() -> None:
encoded = [base64.b64encode(data).decode("ascii") for data in png_images]
if len(encoded) == 1:
# 单图保持当前 o1key 兼容格式，不改变既有请求行为。
body.pop("images", None)
body["image"] = encoded[0]
else:
body.pop("image", None)
body["images"] = [
{
"type": "image_url",
"url": f"data:image/png;base64,{value}",
}
for value in encoded
]
_set_images()
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
for _ in range(8):
if body_size <= _MAX_BODY_BYTES:
return body
shrink_ratio = max(0.1, (_MAX_BODY_BYTES / body_size) * 0.95)
png_images = [
cls._shrink_png_to_limit(
data,
max(1, int(len(data) * shrink_ratio)),
label=f"参考图{index}",
)
for index, data in enumerate(png_images, start=1)
]
_set_images()
body_size = len(json.dumps(body, ensure_ascii=False).encode("utf-8"))
raise RuntimeError(
f"Grok Image 参考图缩放后请求体仍超过 20MB：{body_size / 1024 / 1024:.2f}MB"
)
@staticmethod
def _redact_edit_body(body: dict) -> dict:
result = dict(body)
image = result.get("image")
if isinstance(image, str) and len(image) > 50:
result["image"] = image[:50] + "..."
images = result.get("images")
if isinstance(images, list):
result["images"] = [
{
**item,
"url": item.get("url", "")[:50] + "...",
}
if isinstance(item, dict) else item
for item in images
]
return result
# ── 中断轮询 ──────────────────────────────────────────────────────────────
@staticmethod
@@ -202,28 +317,16 @@ class GrokImageClient:
n: int,
image_list: List[torch.Tensor],
) -> List[Image.Image]:
api_model = _MODEL_NAME_MAP.get(model, model)
body: dict = {
"model": api_model,
"prompt": prompt,
"response_format": "b64_json",
}
if aspect_ratio and aspect_ratio != "auto":
body["aspect_ratio"] = aspect_ratio
if resolution:
body["resolution"] = resolution
# 参考图转 base64 字符串
pil_images = tensor_to_pil(image_list[0])
img = pil_images[0]
buf = BytesIO()
img.save(buf, format="PNG")
png_bytes = buf.getvalue()
png_bytes = self._shrink_png_to_limit(png_bytes, _MAX_BODY_BYTES // 2)
body["image"] = base64.b64encode(png_bytes).decode("utf-8")
body = self._build_edit_body(
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
resolution=resolution,
image_list=image_list,
)
url = f"{self.base_url}{_ENDPOINT_EDITS}"
log_body = {k: (v[:50] + "..." if k == "image" and len(v) > 50 else v) for k, v in body.items()}
log_body = self._redact_edit_body(body)
print(f"[o1key Grok Image] 请求 URL: {url}")
print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")
clients/grok_video_client.py
+262 -384
View File
@@ -1,19 +1,12 @@
"""
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
"""
"""Client for the complete O1Key Grok Imagine Video API."""
import asyncio
import base64
import json
import os
import re
import time
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import quote
import aiohttp
@@ -21,44 +14,49 @@ 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 (
POLL_DEADLINE_SECONDS as VIDEO_POLL_DEADLINE_SECONDS,
check_interrupt,
download_video_to_file,
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"
"""Submit, poll, and download Grok video generation, edit, or extension tasks."""
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",
ENDPOINTS = {
"generate": "/grok/v1/videos/generations",
"edit": "/grok/v1/videos/edits",
"extend": "/grok/v1/videos/extensions",
}
STATUS_ENDPOINT = "/grok/v1/videos/{request_id}"
SUCCESS_STATUSES = {"complete", "completed", "succeed", "succeeded", "success", "done", "finished"}
FAILURE_STATUSES = {"fail", "failed", "failure", "error", "expired", "timeout", "cancelled", "canceled"}
BASE_MODEL = "grok-imagine-video"
LATEST_MODEL = "grok-imagine-video-1.5"
DEFAULT_MODEL = LATEST_MODEL
# Kept as a compatibility alias for callers that imported the old constant.
TEXT_TO_VIDEO_MODEL = BASE_MODEL
IMAGE_TO_VIDEO_MODELS = (LATEST_MODEL,)
MODEL_OPTIONS = (BASE_MODEL, LATEST_MODEL)
ASPECT_RATIO_OPTIONS = ("16:9", "9:16", "1:1", "4:3", "3:4", "3:2", "2:3")
RESOLUTION_OPTIONS = ("480p", "720p", "1080p")
SUCCESS_STATUSES = {"done"}
FAILURE_STATUSES = {"failed", "expired"}
POLL_DEADLINE_SECONDS = VIDEO_POLL_DEADLINE_SECONDS
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)
super().__init__(
base_url=(base_url or get_api_base_url()).rstrip("/"),
api_key=get_api_key_or_raise("O1KEY_API_KEY"),
)
def get_endpoint(self, **kwargs) -> str:
return self.CREATE_ENDPOINT
def get_endpoint(self, operation: str = "generate", **kwargs) -> str:
try:
return self.ENDPOINTS[operation]
except KeyError:
raise ValueError(f"不支持的 Grok 操作：{operation}。") from None
def build_request_body(self, **kwargs) -> Dict[str, Any]:
return self.build_video_body(**kwargs)
@@ -66,418 +64,298 @@ class GrokVideoClient(BaseAPIClient):
def parse_response(self, response: Dict[str, Any]) -> Any:
return response
@staticmethod
def _locator(
value: Optional[Dict[str, str]],
label: str,
allowed_keys: tuple[str, ...],
) -> Dict[str, str]:
if not isinstance(value, dict):
raise ValueError(f"{label}必须提供媒体定位对象。")
known_keys = ("url", "image_url", "file_id", "voice_id")
provided_keys = {
key
for key in known_keys
if value.get(key) is not None and str(value[key]).strip()
}
locator = {
key: str(value[key]).strip()
for key in allowed_keys
if value.get(key) is not None and str(value[key]).strip()
}
if len(locator) != 1 or provided_keys != set(locator):
supported = "、".join(allowed_keys)
raise ValueError(f"{label}必须且只能提供 {supported} 中的一项。")
return locator
@classmethod
def _image_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
return cls._locator(value, label, ("url", "image_url"))
@classmethod
def _audio_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
return cls._locator(value, label, ("url", "voice_id"))
@classmethod
def _video_locator(cls, value: Optional[Dict[str, str]], label: str) -> Dict[str, str]:
return cls._locator(value, label, ("url", "file_id"))
@classmethod
def _validate_common_generation(
cls, model: str, duration: int, aspect_ratio: str, resolution: str
) -> int:
if model not in cls.MODEL_OPTIONS:
raise ValueError(f"模型仅支持：{', '.join(cls.MODEL_OPTIONS)}。")
try:
duration = int(duration)
except (TypeError, ValueError):
raise ValueError("时长必须是整数。") from None
if not 1 <= duration <= 15:
raise ValueError("生成时长仅支持 1 到 15 秒。")
if aspect_ratio not in cls.ASPECT_RATIO_OPTIONS:
raise ValueError(f"宽高比仅支持：{', '.join(cls.ASPECT_RATIO_OPTIONS)}。")
if resolution not in cls.RESOLUTION_OPTIONS:
raise ValueError(f"分辨率仅支持：{', '.join(cls.RESOLUTION_OPTIONS)}。")
return duration
@classmethod
def build_video_body(
cls,
*,
operation: str,
prompt: str,
model: str,
aspect_ratio: str,
seconds: int,
quality: str = "720p",
images: Optional[List[str]] = None,
duration: Optional[int] = None,
aspect_ratio: str = "16:9",
resolution: str = "480p",
image: Optional[Dict[str, str]] = None,
reference_images: Optional[List[Dict[str, str]]] = None,
reference_audios: Optional[List[Dict[str, str]]] = None,
video: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
if operation not in cls.ENDPOINTS:
raise ValueError(f"不支持的 Grok 操作：{operation}。")
prompt = (prompt or "").strip()
if not prompt:
raise ValueError("提示词不能为空。")
if reference_images is not None and not isinstance(reference_images, (list, tuple)):
raise ValueError("reference_images 必须是数组。")
if reference_audios is not None and not isinstance(reference_audios, (list, tuple)):
raise ValueError("reference_audios 必须是数组。")
references = list(reference_images or [])
audios = list(reference_audios or [])
if operation == "generate":
duration = cls._validate_common_generation(model, duration, aspect_ratio, resolution)
normal_image = cls._image_locator(image, "图生视频参考图") if image else None
normal_references = [cls._image_locator(item, "参考图") for item in references]
normal_audios = [cls._audio_locator(item, "参考音频") for item in audios]
if normal_image and normal_references:
raise ValueError("image 和 reference_images 不能同时使用。")
if len(normal_references) > 7:
raise ValueError("参考生视频最多支持 7 张参考图。")
if len(normal_audios) > 3:
raise ValueError("参考生视频最多支持 3 个参考音频。")
has_reference_assets = bool(normal_references or normal_audios)
if has_reference_assets:
if not prompt:
raise ValueError("参考图/音频生视频必须填写提示词。")
if resolution == "1080p":
raise ValueError("参考图/音频生视频不支持 1080p。")
elif not normal_image and not prompt:
raise ValueError("文生视频必须填写提示词。")
if resolution == "1080p" and model != cls.LATEST_MODEL:
raise ValueError("1080p 仅支持 grok-imagine-video-1.5 的文生或图生视频。")
body: Dict[str, Any] = {
"model": model,
"duration": duration,
"aspect_ratio": aspect_ratio,
"resolution": resolution,
}
if prompt:
body["prompt"] = prompt
if normal_image:
body["image"] = normal_image
if normal_references:
body["reference_images"] = normal_references
if normal_audios:
body["reference_audios"] = normal_audios
return body
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)}。")
raise ValueError(f"模型仅支持：{', '.join(cls.MODEL_OPTIONS)}。")
if not prompt:
raise ValueError(f"{operation} 必须填写提示词。")
normal_video = cls._video_locator(video, "输入视频")
if operation == "edit":
return {"model": model, "prompt": prompt, "video": normal_video}
try:
seconds_value = int(seconds)
duration = int(duration)
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] = {
raise ValueError("续写时长必须是整数。") from None
if not 2 <= duration <= 10:
raise ValueError("视频续写时长仅支持 2 到 10 秒。")
return {
"model": model,
"prompt": prompt,
"aspect_ratio": aspect_ratio,
"seconds": str(seconds_value),
"quality": api_quality,
"video": normal_video,
"duration": duration,
}
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)
def _extract_request_id(payload: Dict[str, Any]) -> Optional[str]:
for source in (payload, payload.get("data")):
if isinstance(source, dict) and source.get("request_id"):
return str(source["request_id"])
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)
def _safe_filename(request_id: str) -> str:
return re.sub(r"[^A-Za-z0-9_.-]+", "_", request_id).strip("._") or "grok_video"
@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)}",
]
)
@staticmethod
def _safe_error_message(value: object) -> str:
message = str(value or "").strip()
message = re.sub(r"data:[^\s,;]+;base64,[A-Za-z0-9+/=_-]+", "<base64 omitted>", message)
message = re.sub(r"https?://[^\s\"'<>]+", "<temporary URL omitted>", message)
return message[:500]
async def _request_json_with_retry(
async def _request_json(
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,
request_id: Optional[str] = None,
) -> 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):
last_status, last_text = 0, ""
for attempt in range(4):
check_interrupt()
response = None
try:
response = await run_with_interrupt(
session.request(method, url, json=json_body, headers=headers, timeout=timeout)
session.request(
method, url, json=json_body,
headers=self.get_headers(use_bearer_token=True), timeout=timeout,
)
)
text = await run_with_interrupt(response.text())
last_status = response.status
last_text = text
last_status, last_text = response.status, 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
return json.loads(text) if text.strip() else {}
except json.JSONDecodeError:
raise RuntimeError("Grok Video 响应不是有效 JSON。") from None
if response.status not in RETRYABLE_STATUS_CODES or attempt == 3:
break
delay = min(2 ** attempt, 8)
print(f"Grok Video：HTTP {response.status}，{delay}s 后重试…")
await interruptible_sleep(delay)
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
if attempt == 3:
raise RuntimeError(
f"Grok Video 网络错误：{type(exc).__name__}"
) from None
delay = min(2 ** attempt, 8)
print(f"Grok Video：网络错误，{delay}s 后重试…")
await interruptible_sleep(delay)
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,
message = self._safe_error_message(
get_friendly_message(last_status, last_text) or "请求失败"
)
detail = f"Grok Video 请求失败：HTTP {last_status}，{message}"
if request_id:
detail += f"（request_id: {request_id}）"
raise RuntimeError(detail)
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,
async def _poll(
self, request_id: str, session: aiohttp.ClientSession, *, poll_interval: int,
timeout: int, progress_callback: Optional[Callable[[int, str, float], 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)
endpoint = self.STATUS_ENDPOINT.format(request_id=quote(request_id, safe=""))
started_at = time.monotonic()
while True:
data = await self._request_json_with_retry(
"GET",
endpoint,
session=session,
task_id=task_id,
timeout_seconds=60,
await interruptible_sleep(poll_interval)
response = await self._request_json(
"GET", endpoint, session, timeout_seconds=60, request_id=request_id
)
status = extract_status(data)
progress = extract_progress(data)
elapsed = time.time() - start
status = str(response.get("status", "")).strip().lower()
try:
progress = max(0, min(100, int(float(response.get("progress") or 0))))
except (TypeError, ValueError):
progress = 0
elapsed = time.monotonic() - started_at
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 status in self.SUCCESS_STATUSES:
return response
if status in self.FAILURE_STATUSES:
message = self._safe_error_message(
extract_error_message(response, default="未知错误")
)
raise RuntimeError(
f"Grok Video 任务{status}（request_id: {request_id}）："
f"{message}"
)
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"
f"Grok Video 轮询超时（request_id: {request_id}，状态：{status or 'unknown'}）。"
)
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,
def run_video_sync(
self, *, operation: str, prompt: str, model: str, duration: Optional[int] = None,
aspect_ratio: str = "16:9", resolution: str = "480p",
image: Optional[Dict[str, str]] = None,
reference_images: Optional[List[Dict[str, str]]] = None,
reference_audios: Optional[List[Dict[str, str]]] = None,
video: Optional[Dict[str, str]] = None, output_dir: Optional[str] = None,
poll_interval: int = 5, timeout: int = VIDEO_POLL_DEADLINE_SECONDS,
progress_callback: Optional[Callable[[int, str, float], None]] = None,
) -> Dict[str, Any]:
async def _run():
async def run_request() -> Dict[str, Any]:
async with self._make_session() as session:
endpoint = self.get_endpoint(operation)
body = self.build_video_body(
prompt=prompt,
model=model,
aspect_ratio=aspect_ratio,
seconds=seconds,
quality=quality,
images=images,
operation=operation, prompt=prompt, model=model, duration=duration,
aspect_ratio=aspect_ratio, resolution=resolution, image=image,
reference_images=reference_images, reference_audios=reference_audios,
video=video,
)
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,
print(f"Grok Video：正在提交{operation}任务…")
created = await self._request_json(
"POST", endpoint, session, json_body=body, timeout_seconds=180
)
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)
request_id = self._extract_request_id(created)
if not request_id:
raise RuntimeError("Grok Video 创建响应中没有 request_id。")
print(f"Grok Video：任务已提交，request_id：{request_id}")
completed = await self._poll(
request_id, session, poll_interval=max(1, int(poll_interval)),
timeout=timeout, progress_callback=progress_callback,
)
video_data = completed.get("video")
video_url = video_data.get("url") if isinstance(video_data, dict) else None
if not video_url:
raise RuntimeError(f"Grok Video 完成响应中没有 video.url（request_id: {request_id}）。")
directory = output_dir or os.getcwd()
os.makedirs(directory, exist_ok=True)
save_path = os.path.join(directory, f"{self._safe_filename(request_id)}.mp4")
print("Grok Video：视频生成完成，正在下载…")
video_path = await download_video_to_file(session, video_url, save_path, label="Grok Video")
return {
"task_id": task_id,
"status": extract_status(status_response),
"request_id": request_id,
"video_path": video_path,
"raw_json": {
"create": create_response,
"status": status_response,
},
"duration": video_data.get("duration"),
"raw_json": {"create": created, "status": completed},
}
return self.run_async_in_thread(_run())
return self.run_async_in_thread(run_request())
clients/kling_client.py
-285
View File
@@ -1,285 +0,0 @@
"""
Kling 视频生成 API 客户端
"""
import asyncio
import json
import os
from typing import Any, Callable, Dict, Optional
import aiohttp
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.http_error import async_request_with_retry
from ..utils.video_task import (
check_interrupt,
extract_error_message,
extract_progress,
extract_status,
interruptible_sleep,
is_failure_status,
is_success_status,
run_with_interrupt,
)
class KlingClient:
"""Kling 视频生成客户端"""
ENDPOINTS = {
"image2video": "/kling/v1/videos/image2video",
"text2video": "/kling/v1/videos/text2video",
"motion_control": "/kling/v1/videos/motion-control",
}
# new API 三段式端点（动作控制走这里）
NEW_API_CREATE = "/v1/videos"
NEW_API_STATUS = "/v1/videos/{video_id}"
NEW_API_CONTENT = "/v1/videos/{video_id}/content"
POLL_INITIAL_INTERVAL = 3
POLL_MAX_INTERVAL = 15
def __init__(self):
self.api_key = get_api_key_or_raise()
self.base_url = get_api_base_url()
def _headers(self) -> Dict[str, str]:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
# ── 提交任务 ──────────────────────────────────────────────────────
async def create_video_async(
self,
endpoint_type: str,
body: Dict[str, Any],
session: aiohttp.ClientSession,
) -> Dict[str, Any]:
url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}"
check_interrupt()
resp = await run_with_interrupt(async_request_with_retry(
session, "POST", url, json=body, headers=self._headers(), prefix="Kling 提交: "
))
check_interrupt()
text = await resp.text()
return json.loads(text)
# ── 轮询状态 ──────────────────────────────────────────────────────
async def poll_status_async(
self,
task_id: str,
endpoint_type: str,
session: aiohttp.ClientSession,
on_progress: Optional[Callable[[int], None]] = None,
) -> Dict[str, Any]:
url = f"{self.base_url}{self.ENDPOINTS[endpoint_type]}/{task_id}"
interval = self.POLL_INITIAL_INTERVAL
while True:
check_interrupt()
async with session.get(url, headers=self._headers()) as resp:
text = await resp.text()
if resp.status != 200:
raise RuntimeError(f"状态查询失败 ({resp.status}): {text}")
result = json.loads(text)
data = result.get("data", {})
inner_data = data.get("data", {}) if isinstance(data, dict) else {}
status = extract_status(result)
progress_pct = extract_progress(result)
print(f"[视频生成] 生成中 {progress_pct}%")
if on_progress:
on_progress(progress_pct)
if is_success_status(status):
return result
elif is_failure_status(status, result):
error_msg = extract_error_message(result)
raise RuntimeError(f"生成失败：{error_msg}")
await interruptible_sleep(interval)
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
# ── 下载视频 ──────────────────────────────────────────────────────
async def download_video_async(
self,
video_url: str,
save_path: str,
session: aiohttp.ClientSession,
) -> str:
print("[视频生成] 下载视频...")
check_interrupt()
async with session.get(video_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
check_interrupt()
f.write(chunk)
return save_path
# ── 异步入口（供节点调用）────────────────────────────────────────
async def generate_async(
self,
endpoint_type: str,
body: Dict[str, Any],
save_path: str,
on_stage: Optional[Callable[[str], None]] = None,
on_progress: Optional[Callable[[int], None]] = None,
) -> str:
"""提交 → 轮询 → 下载，返回本地文件路径"""
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
if on_stage:
on_stage("submitting")
result = await self.create_video_async(endpoint_type, body, session)
# 提交响应结构：result.data.task_id
task_id = result.get("task_id") or result.get("data", {}).get("task_id")
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{result}")
if on_stage:
on_stage(f"submitted:{task_id}")
if on_stage:
on_stage("polling")
final = await self.poll_status_async(
task_id, endpoint_type, session, on_progress=on_progress
)
# 兼容多种URL路径
# 响应结构：result.data.result_url 或 result.data.data.task_result.videos[0].url
data = final.get("data", {})
inner_data = data.get("data", {}) if isinstance(data, dict) else {}
video_url = (
data.get("result_url") or
final.get("url") or
final.get("video_url") or
(inner_data.get("task_result", {}).get("videos", [{}])[0].get("url")
if inner_data.get("task_result", {}).get("videos") else None)
)
if not video_url:
raise RuntimeError(f"API 未返回视频 URL，响应：{final}")
if on_stage:
on_stage("downloading")
path = await self.download_video_async(video_url, save_path, session)
if on_stage:
on_stage("done")
return path
# ── 动作控制：走 new API 三段式流程 ──────────────────────────────
async def motion_control_async(
self,
body: Dict[str, Any],
save_path: str,
on_stage: Optional[Callable[[str], None]] = None,
on_progress: Optional[Callable[[int], None]] = None,
) -> str:
"""
动作控制专用入口：
POST /v1/videos → GET /v1/videos/{id} → GET /v1/videos/{id}/content
body 字段与 Kling 官方动作控制接口一致（image_url/video_url/prompt/...）。
"""
headers = {"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"}
interval = self.POLL_INITIAL_INTERVAL
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
# 1. 提交
check_interrupt()
if on_stage:
on_stage("submitting")
create_url = f"{self.base_url}{self.NEW_API_CREATE}"
resp = await run_with_interrupt(async_request_with_retry(
session, "POST", create_url, json=body, headers=headers, prefix="Kling 动作控制提交: "
))
check_interrupt()
text = await resp.text()
create_resp = json.loads(text)
video_id = create_resp.get("id")
if not video_id:
raise RuntimeError(f"API 未返回视频 ID，响应：{create_resp}")
if on_stage:
on_stage(f"submitted:{video_id}")
# 2. 轮询
status_url = f"{self.base_url}{self.NEW_API_STATUS.format(video_id=video_id)}"
while True:
check_interrupt()
async with session.get(status_url, headers=headers) as resp:
text = await resp.text()
if resp.status != 200:
try:
err = json.loads(text)
msg = err.get("error", {}).get("message") or err.get("message") or text
except Exception:
msg = text
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
status_resp = json.loads(text)
status = extract_status(status_resp)
progress_pct = extract_progress(status_resp)
print(f"[动作控制] 生成中 {progress_pct}%")
if on_progress:
on_progress(progress_pct)
if is_success_status(status):
break
if is_failure_status(status, status_resp):
error_msg = extract_error_message(status_resp)
raise RuntimeError(f"动作控制生成失败：{error_msg}")
await interruptible_sleep(interval)
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
# 3. 下载
check_interrupt()
if on_stage:
on_stage("downloading")
content_url = f"{self.base_url}{self.NEW_API_CONTENT.format(video_id=video_id)}"
async with session.get(content_url, headers=headers,
allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
content_type = resp.headers.get("Content-Type", "")
if "application/json" in content_type:
data = await resp.json()
download_url = data.get("url") or data.get("download_url")
if not download_url:
raise RuntimeError("视频下载失败：响应中未找到下载链接")
async with session.get(download_url) as dl_resp:
if dl_resp.status != 200:
raise RuntimeError(f"从下载链接获取视频失败 ({dl_resp.status})")
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in dl_resp.content.iter_chunked(8192):
check_interrupt()
f.write(chunk)
else:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
check_interrupt()
f.write(chunk)
if on_stage:
on_stage("done")
return save_path
clients/minimax_h3_client.py
+242
View File
@@ -0,0 +1,242 @@
"""MiniMax H3 video client for the New API gateway."""
import json
from typing import Any, Callable, Dict, Optional
from urllib.parse import quote
import aiohttp
from ..utils.config import get_api_key_or_raise
from ..utils.http_error import async_request_with_retry
from ..utils.video_task import (
POLL_DEADLINE_SECONDS as VIDEO_POLL_DEADLINE_SECONDS,
PollDeadline,
check_interrupt,
download_video_to_file,
extract_error_message,
extract_progress,
interruptible_sleep,
run_with_interrupt,
)
PENDING_STATUSES = {
"NOT_START",
"SUBMITTED",
"QUEUED",
"IN_PROGRESS",
"RUNNING",
"UNKNOWN",
}
SUCCESS_STATUSES = {"SUCCESS", "COMPLETED", "SUCCEEDED"}
FAILURE_STATUSES = {"FAILURE", "FAILED", "CANCELLED", "CANCELED"}
def extract_public_task_id(payload: Dict[str, Any]) -> str:
"""Return the New API public task ID, preferring ``id`` as documented."""
task_id = payload.get("id") or payload.get("task_id")
if not task_id:
raise RuntimeError("MiniMax H3 创建成功但未返回任务 ID。")
return str(task_id)
def parse_task_snapshot(payload: Dict[str, Any]) -> Dict[str, Any]:
"""Normalize New API's wrapper and MiniMax's official V2 task shape."""
if not isinstance(payload, dict):
raise RuntimeError("MiniMax H3 查询响应不是 JSON 对象。")
raw_data = payload.get("data")
data = raw_data if isinstance(raw_data, dict) else payload
raw_task = payload.get("task")
task = raw_task if isinstance(raw_task, dict) else {}
task_content = task.get("content") if isinstance(task.get("content"), dict) else {}
task_error = task.get("error") if isinstance(task.get("error"), dict) else {}
data_error = data.get("error") if isinstance(data.get("error"), dict) else {}
root_error = payload.get("error") if isinstance(payload.get("error"), dict) else {}
metadata = data.get("metadata") if isinstance(data.get("metadata"), dict) else {}
status = str(
data.get("status") or task.get("status") or payload.get("status") or ""
).strip().upper()
error_message = str(
task_error.get("message")
or data_error.get("message")
or root_error.get("message")
or ""
).strip()
error_code = str(
task_error.get("code")
or data_error.get("code")
or root_error.get("code")
or ""
).strip()
if error_message and error_code:
error_message = f"{error_message}（错误码 {error_code}）"
elif error_code:
error_message = f"错误码 {error_code}"
return {
"status": status,
"progress": extract_progress(payload),
"result_url": str(
data.get("result_url")
or task_content.get("url")
or metadata.get("url")
or data.get("url")
or payload.get("result_url")
or payload.get("url")
or ""
).strip(),
"fail_reason": str(
data.get("fail_reason")
or error_message
or extract_error_message(payload, "视频生成失败")
).strip(),
}
class MiniMaxH3Client:
"""Create, poll, and immediately download a MiniMax-H3 video task."""
CREATE_ENDPOINT = "/v1/video/generations"
STATUS_ENDPOINT = "/v1/videos/{task_id}"
POLL_INTERVAL_SECONDS = 10.0
POLL_DEADLINE_SECONDS = VIDEO_POLL_DEADLINE_SECONDS
def __init__(self, base_url: str, api_key: Optional[str] = None):
self.base_url = (base_url or "").rstrip("/")
if not self.base_url:
raise ValueError("MiniMax H3 New API Base URL 不能为空。")
self.api_key = api_key or get_api_key_or_raise()
def _headers(self) -> Dict[str, str]:
# Model API requests use the application API token. New-Api-User is
# intentionally not sent because it belongs to management API auth.
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
@staticmethod
async def _read_json(response: aiohttp.ClientResponse, action: str) -> Dict[str, Any]:
raw = await response.text()
try:
payload = json.loads(raw)
except json.JSONDecodeError:
raise RuntimeError(f"MiniMax H3 {action}返回了无效 JSON。") from None
if not isinstance(payload, dict):
raise RuntimeError(f"MiniMax H3 {action}响应不是 JSON 对象。")
return payload
async def submit_async(
self,
body: Dict[str, Any],
session: aiohttp.ClientSession,
) -> str:
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
check_interrupt()
response = await run_with_interrupt(
async_request_with_retry(
session,
"POST",
url,
json=body,
headers=self._headers(),
prefix="MiniMax H3 创建任务：",
)
)
payload = await self._read_json(response, "创建任务")
return extract_public_task_id(payload)
async def poll_async(
self,
task_id: str,
session: aiohttp.ClientSession,
on_progress: Optional[Callable[[int], None]] = None,
) -> str:
encoded_task_id = quote(task_id, safe="")
url = f"{self.base_url}{self.STATUS_ENDPOINT.format(task_id=encoded_task_id)}"
deadline = PollDeadline(
seconds=self.POLL_DEADLINE_SECONDS,
label=f"MiniMax H3（任务 {task_id}）",
)
while True:
deadline.check()
check_interrupt()
response = await run_with_interrupt(
async_request_with_retry(
session,
"GET",
url,
headers=self._headers(),
prefix="MiniMax H3 查询任务：",
)
)
payload = await self._read_json(response, "查询任务")
snapshot = parse_task_snapshot(payload)
status = snapshot["status"]
progress = snapshot["progress"]
# A successful terminal state is authoritative even if an older
# gateway omits data.progress or returns a stale percentage.
if status in SUCCESS_STATUSES:
progress = 100
print(f"[MiniMax H3] 任务 {task_id}：{status or 'UNKNOWN'} {progress}%")
if on_progress:
on_progress(progress)
if status in SUCCESS_STATUSES:
result_url = snapshot["result_url"]
if not result_url:
raise RuntimeError(
f"MiniMax H3 任务 {task_id} 已成功，但响应缺少 data.result_url。"
)
return result_url
if status in FAILURE_STATUSES:
raise RuntimeError(
f"MiniMax H3 任务 {task_id} 生成失败：{snapshot['fail_reason']}"
)
if status not in PENDING_STATUSES:
raise RuntimeError(
f"MiniMax H3 任务 {task_id} 返回不支持的状态 {status or '<空>'}。"
)
await interruptible_sleep(self.POLL_INTERVAL_SECONDS)
async def generate_async(
self,
body: Dict[str, Any],
save_path: str,
on_stage: Optional[Callable[[str], None]] = None,
on_progress: Optional[Callable[[int], None]] = None,
) -> tuple[str, str]:
connector = aiohttp.TCPConnector(force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
if on_stage:
on_stage("submitting")
task_id = await self.submit_async(body, session)
print(f"[MiniMax H3] 已提交公开任务 ID：{task_id}")
if on_stage:
on_stage(f"submitted:{task_id}")
result_url = await self.poll_async(
task_id,
session,
on_progress=on_progress,
)
if on_stage:
on_stage("downloading")
await download_video_to_file(
session,
result_url,
save_path,
label=f"MiniMax H3 {task_id}",
)
if on_stage:
on_stage("done")
return save_path, task_id
clients/newapi_veo_client.py
+19 -44
View File
@@ -16,6 +16,10 @@ import aiohttp
from .base_client import BaseAPIClient
from ..utils.config import get_api_base_url, get_api_key_or_raise
from ..utils.video_task import (
POLL_DEADLINE_SECONDS as VIDEO_POLL_DEADLINE_SECONDS,
download_video_to_file,
)
class NewAPIVeoClient(BaseAPIClient):
@@ -26,6 +30,7 @@ class NewAPIVeoClient(BaseAPIClient):
RETRYABLE_STATUS_CODES = {408, 409, 425, 429, 500, 502, 503, 504}
COMPLETED_STATUSES = {"completed", "succeeded", "success", "done"}
FAILED_STATUSES = {"failed", "error", "cancelled", "canceled"}
POLL_DEADLINE_SECONDS = VIDEO_POLL_DEADLINE_SECONDS
def __init__(
self,
@@ -318,7 +323,7 @@ class NewAPIVeoClient(BaseAPIClient):
self,
task_id: str,
poll_interval: int = 5,
timeout: int = 900,
timeout: int = VIDEO_POLL_DEADLINE_SECONDS,
progress_callback: Optional[Callable[[int, str, float], None]] = None,
session: Optional[aiohttp.ClientSession] = None,
) -> Dict[str, Any]:
@@ -376,30 +381,8 @@ class NewAPIVeoClient(BaseAPIClient):
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)
)
# 抗超时 / 断点续传 / 无限重试 / 可取消
await download_video_to_file(session, url, save_path, label="VEO 视频")
async def download_video_async(
self,
@@ -410,7 +393,7 @@ class NewAPIVeoClient(BaseAPIClient):
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)
timeout = aiohttp.ClientTimeout(total=120, connect=30, sock_read=120)
close_session = False
if session is None:
@@ -420,6 +403,7 @@ class NewAPIVeoClient(BaseAPIClient):
try:
last_status = 0
last_error = ""
# 先探测 content 端点：JSON 则取真实下载链接，否则视为视频流交给健壮下载器。
for attempt in range(4):
async with session.get(url, headers=headers, timeout=timeout, allow_redirects=True) as response:
if response.status < 300:
@@ -440,30 +424,21 @@ class NewAPIVeoClient(BaseAPIClient):
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
return save_path
break # 非 JSON：content 端点即视频流（幂等 GET，可续传）
last_status = response.status
last_error = await response.text()
if response.status not in self.RETRYABLE_STATUS_CODES or attempt >= 3:
break
raise RuntimeError(
self._format_http_error(endpoint, last_status, last_error, task_id=task_id)
)
await asyncio.sleep(min(2 ** attempt, 8))
raise RuntimeError(
self._format_http_error(endpoint, last_status, last_error, task_id=task_id)
# 抗超时 / 断点续传 / 无限重试 / 可取消
return await download_video_to_file(
session, url, save_path, headers=headers, label="VEO 视频",
)
finally:
if close_session:
@@ -481,7 +456,7 @@ class NewAPIVeoClient(BaseAPIClient):
generate_audio: bool = True,
image_bytes: Optional[bytes] = None,
poll_interval: int = 5,
timeout: int = 900,
timeout: int = VIDEO_POLL_DEADLINE_SECONDS,
reuse_task_id: str = "",
progress_callback: Optional[Callable[[int, str, float], None]] = None,
) -> Dict[str, Any]:
clients/omni_flash_client.py
+358
View File
@@ -0,0 +1,358 @@
"""O1Key Omni Flash JSON video tasks."""
from __future__ import annotations
import asyncio
import json
import re
import time
from typing import Any, Callable
import aiohttp
from ..utils.config import get_api_key_or_raise, get_base_url_by_route
from ..utils.video_task import (
InterruptProcessingException, check_interrupt, download_video_to_file,
extract_error_message, extract_progress, extract_video_url,
interruptible_sleep, is_failure_status, is_success_status,
)
MODELS = {"omni_flash_8s", "omni_flash_10s", "omni_flash_abra_edit"}
RESOLUTIONS = {"720p", "1080p"}
RATIOS = {"16:9", "9:16"}
POLL_SECONDS = 7
POLL_DEADLINE_SECONDS = 2000
ERROR_HINTS = {
"invalid_request": "请求参数有误，请检查模型、分辨率、宽高比和素材",
"model_not_available": "当前模型不可用，请重新选择模型",
"image_url_required_for_i2v": "参考图地址缺失或无效，请连接图片后重试",
"invalid_api_key": "O1Key 令牌无效或已停用，请在令牌管理中更新",
"insufficient_balance": "O1Key 余额不足",
"task_not_found": "视频任务不存在或已失效",
"rate_limit_exceeded": "请求过于频繁，请稍后重试",
}
HTTP_HINTS = {
400: "请求参数错误", 401: "令牌验证失败", 402: "余额不足",
404: "任务不存在", 429: "请求过于频繁",
}
SENSITIVE_RESPONSE_KEYS = {
"authorization", "api_key", "apikey", "access_token", "refresh_token",
"token", "secret", "password", "b64_json", "base64", "image_base64",
"video_base64",
}
MAX_LOG_BODY = 16000
def build_video_body(
*, model: str, prompt: str, resolution: str, aspect_ratio: str,
mode: str, references: list[str] | None = None, source_video_url: str = "",
) -> dict[str, Any]:
"""Validate all scalar inputs before a paid request."""
if model not in MODELS:
raise ValueError("Omni Flash 模型无效")
prompt = str(prompt or "").strip()
if not prompt:
raise ValueError("提示词不能为空")
if len(prompt) > 20000:
raise ValueError("提示词过长")
if resolution not in RESOLUTIONS or aspect_ratio not in RATIOS:
raise ValueError("分辨率或宽高比无效")
references = list(references or [])
if any(not isinstance(value, str) or len(value) > 4096 or not value.startswith(("https://", "http://")) for value in references):
raise ValueError("参考图必须是 HTTP(S) 直链")
body: dict[str, Any] = {
"model": model, "prompt": prompt,
"resolution": resolution, "aspect_ratio": aspect_ratio,
}
if mode == "edit":
if model != "omni_flash_abra_edit" or len(source_video_url) > 4096 or not source_video_url.startswith(("https://", "http://")):
raise ValueError("视频编辑需要编辑模型和源视频直链")
if len(references) > 5:
raise ValueError("视频编辑最多支持 5 张参考图")
body["source_video_url"] = source_video_url
elif mode in {"text", "reference", "first_last_frame"}:
if model == "omni_flash_abra_edit" or source_video_url:
raise ValueError("生成模式不能使用编辑模型或源视频")
if mode == "text" and references:
raise ValueError("文生视频不能提供参考图")
if mode == "reference" and not references:
raise ValueError("参考图模式至少需要 1 张图片")
if mode == "first_last_frame":
if not 1 <= len(references) <= 2:
raise ValueError("首尾帧模式需要首帧图片，尾帧图片可选")
# The provider's frame-pair flag is for a transition between two
# frames. A lone first frame uses the documented single-image i2v
# request, avoiding a pair request with a missing end frame.
if len(references) == 2:
body["first_last_frame"] = True
else:
raise ValueError("Omni Flash 生成模式无效")
if references:
body["input_reference"] = references[0] if len(references) == 1 else references
return body
def _submission_payload(body: dict[str, Any]) -> dict[str, Any]:
"""Use a scalar JSON reference, or repeat the field in multipart for several."""
references = body.get("input_reference")
if not isinstance(references, list):
return {"json": body}
form = aiohttp.FormData()
for name, value in body.items():
values = value if name == "input_reference" else [value]
for item in values:
text = "true" if item is True else "false" if item is False else str(item)
form.add_field(name, text, content_type="text/plain")
return {"data": form}
def _redact_log_string(text: str) -> str:
text = re.sub(r"https?://[^\s\"'<>]+", "<URL 已隐藏>", text)
text = re.sub(r"(?i)bearer\s+[^\s\"']+", "Bearer <已隐藏>", text)
text = re.sub(r"(?i)(?:api[_-]?key|token|authorization)[\"']?\s*[:=]\s*[\"']?[^\s,;\"']+", "<凭据已隐藏>", text)
text = re.sub(r"(?i)data:[^,\s]+;base64,[A-Za-z0-9+/=]+", "<Base64 已隐藏>", text)
return re.sub(r"[A-Za-z0-9+/]{256,}={0,2}", "<长数据已隐藏>", text)
def _safe_error(value: Any) -> str:
text = _redact_log_string(str(value or "请求失败"))
return text[:400]
def _log_value(value: Any, key: str = "", depth: int = 0) -> Any:
if key.lower() in SENSITIVE_RESPONSE_KEYS:
return "<已隐藏>"
if depth >= 12:
return "<嵌套内容已省略>"
if isinstance(value, dict):
return {
_redact_log_string(str(name)[:200]): _log_value(item, str(name), depth + 1)
for name, item in value.items()
}
if isinstance(value, list):
return [_log_value(item, key, depth + 1) for item in value[:50]] + (
[f"<其余 {len(value) - 50} 项已省略>"] if len(value) > 50 else []
)
if isinstance(value, str):
if len(value) > 1200:
return f"<长文本 {len(value)} 字符已省略>"
return _redact_log_string(value)
return value
def _log_response_body(stage: str, status: int, raw_body: str) -> None:
try:
payload = json.loads(raw_body)
except (ValueError, TypeError):
safe_body = _safe_error(raw_body) if raw_body else "<空响应体>"
else:
safe_body = json.dumps(_log_value(payload), ensure_ascii=False, separators=(",", ":"))
if len(safe_body) > MAX_LOG_BODY:
safe_body = f"{safe_body[:MAX_LOG_BODY]}...<后续内容已省略>"
print(f"[Omni Flash] {stage} HTTP {status} 原始响应体（敏感值已隐藏）：{safe_body}")
async def _response_text(response: aiohttp.ClientResponse, stage: str) -> str:
raw_body = await response.text()
_log_response_body(stage, response.status, raw_body)
return raw_body
def _error_code(payload: Any) -> str:
if not isinstance(payload, dict):
return ""
data = payload.get("data")
inner = data.get("data") if isinstance(data, dict) else None
for source in (inner, data, payload):
if not isinstance(source, dict):
continue
error = source.get("error")
for value in (error.get("code") if isinstance(error, dict) else None, source.get("code")):
if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_-]{1,80}", value):
return value.lower()
return ""
def _response_error(payload: Any, status: int) -> str:
detail = extract_error_message(payload, default="") if isinstance(payload, dict) else payload
code = _error_code(payload)
hint = ERROR_HINTS.get(code) or HTTP_HINTS.get(status, "视频接口请求失败")
detail = _safe_error(detail) if detail else ""
if detail == code or detail == hint:
detail = ""
suffix = f"：{detail}" if detail else ""
code_note = f"，{code}" if code else ""
return f"{hint}（HTTP {status}{code_note}）{suffix}"
def _task_error(payload: dict[str, Any]) -> str:
code = _error_code(payload)
hint = ERROR_HINTS.get(code, "视频任务生成失败")
detail = extract_error_message(payload, default="")
detail = _safe_error(detail) if detail else ""
if detail == code or detail == hint:
detail = ""
code_note = f"（{code}）" if code else ""
return f"{hint}{code_note}{f'：{detail}' if detail else ''}"
def _task_id(payload: Any) -> str | None:
if not isinstance(payload, dict):
return None
data = payload.get("data")
inner = data.get("data") if isinstance(data, dict) else None
for source in (inner, data, payload):
if isinstance(source, dict):
for name in ("id", "task_id", "video_id"):
value = source.get(name)
if isinstance(value, str) and re.fullmatch(r"[A-Za-z0-9_-]{8,128}", value):
return value
return None
def _task_status(payload: dict[str, Any]) -> str:
data = payload.get("data")
inner = data.get("data") if isinstance(data, dict) else None
for source in (inner, data, payload):
if isinstance(source, dict):
for name in ("task_status", "task_state", "status", "state"):
value = source.get(name)
if value is not None and str(value).strip():
return str(value).strip().lower()
return ""
def _video_url(payload: Any) -> str | None:
value = extract_video_url(payload) if isinstance(payload, dict) else None
if not value and isinstance(payload, dict):
data = payload.get("data")
for source in (data, payload):
if isinstance(source, dict):
output = source.get("output")
if isinstance(output, dict):
value = output.get("video_url") or output.get("url")
if value:
break
return value if isinstance(value, str) and len(value) <= 8192 and value.startswith(("https://", "http://")) else None
class OmniFlashClient:
def __init__(self, *, base_url: str | None = None, api_key: str | None = None):
self.base_url = (base_url or get_base_url_by_route()).rstrip("/")
self.api_key = api_key or get_api_key_or_raise("O1KEY_API_KEY")
async def _download_completed(
self, session: aiohttp.ClientSession, task_id: str, save_path: str,
headers: dict[str, str], status_payload: dict[str, Any],
) -> None:
content_url = f"{self.base_url}/v1/videos/{task_id}/content"
result_url = _video_url(status_payload)
download_url, download_headers = content_url, headers
try:
async with session.get(content_url, headers=headers, allow_redirects=True) as response:
if response.status >= 300:
raw_body = await _response_text(response, "下载")
if not result_url:
try:
error = json.loads(raw_body)
except ValueError:
error = raw_body
raise RuntimeError(_response_error(error, response.status))
download_url, download_headers = result_url, {}
elif "json" in response.headers.get("Content-Type", "").lower():
raw_body = await _response_text(response, "下载")
try:
content_payload = json.loads(raw_body)
except ValueError:
raise RuntimeError("视频下载接口返回了无效 JSON") from None
if _error_code(content_payload) in ERROR_HINTS:
raise RuntimeError(_task_error(content_payload))
download_url = _video_url(content_payload) or result_url
if not download_url:
raise RuntimeError("任务已完成，但下载响应未提供视频地址")
download_headers = {}
else:
print(f"[Omni Flash] 下载 HTTP {response.status} 响应体：<视频二进制，未打印>")
except (aiohttp.ClientError, asyncio.TimeoutError):
# The streaming downloader handles transient connection failures and resumes.
if result_url:
download_url, download_headers = result_url, {}
await download_video_to_file(
session, download_url, save_path, headers=download_headers or None,
label="Omni Flash 视频",
)
async def generate(
self, body: dict[str, Any], save_path: str,
progress: Callable[[str, int, str], None] | None = None,
) -> str:
headers = {"Authorization": f"Bearer {self.api_key}"}
submit_headers = dict(headers)
if body.get("model") == "omni_flash_abra_edit":
submit_headers["X-No-Watermark"] = "video"
timeout = aiohttp.ClientTimeout(total=120, connect=30)
async with aiohttp.ClientSession(timeout=timeout) as session:
async with session.post(
f"{self.base_url}/v1/videos", headers=submit_headers,
**_submission_payload(body),
) as response:
raw_body = await _response_text(response, "提交")
try:
payload = json.loads(raw_body)
except ValueError:
if response.status >= 300:
raise RuntimeError(_response_error(raw_body, response.status)) from None
raise RuntimeError("提交接口未返回有效 JSON") from None
if response.status >= 300:
raise RuntimeError(_response_error(payload, response.status))
if _error_code(payload) in ERROR_HINTS:
raise RuntimeError(_task_error(payload))
task_id = _task_id(payload)
if not task_id:
raise RuntimeError("接口未返回有效任务 ID")
if progress:
progress("polling", 0, task_id)
deadline = time.monotonic() + POLL_DEADLINE_SECONDS
last_status = ""
while time.monotonic() < deadline:
await interruptible_sleep(POLL_SECONDS)
check_interrupt()
try:
async with session.get(f"{self.base_url}/v1/videos/{task_id}", headers=headers) as response:
raw_body = await _response_text(response, "查询")
if response.status in {408, 500, 502, 503, 504}:
continue
try:
status_payload = json.loads(raw_body)
except ValueError:
if response.status >= 300:
raise RuntimeError(_response_error(raw_body, response.status)) from None
raise
if response.status >= 300:
raise RuntimeError(_response_error(status_payload, response.status))
except (aiohttp.ClientError, asyncio.TimeoutError, json.JSONDecodeError, ValueError):
continue
if not isinstance(status_payload, dict):
continue
status = _task_status(status_payload)
last_status = status or last_status
code = _error_code(status_payload)
if is_failure_status(status) or code in ERROR_HINTS or (
not status and code and code not in {"ok", "success", "0"}
):
raise RuntimeError(_task_error(status_payload))
if is_success_status(status) or _video_url(status_payload):
if progress:
progress("downloading", 100, task_id)
try:
await self._download_completed(session, task_id, save_path, headers, status_payload)
except InterruptProcessingException:
raise
except Exception as exc:
raise RuntimeError(_safe_error(exc)) from None
return task_id
if progress:
progress("polling", extract_progress(status_payload), task_id)
status_note = f"，最后状态：{_safe_error(last_status)}" if last_status else ""
raise TimeoutError(f"Omni Flash 任务 {task_id} 等待超时{status_note}")
clients/openai_client.py
-784
View File
@@ -1,784 +0,0 @@
"""
OpenAI 兼容 API 客户端
端点固定为 /v1/chat/completions，模型名放入请求体 model 字段
"""
import re
import time
from io import BytesIO
from typing import Any, Callable, Dict, List, Optional
import aiohttp
from PIL import Image
from ..utils.image_utils import encode_image_to_base64, decode_base64_to_pil
from ..utils.config import get_api_key_or_raise, get_api_base_url
from .base_client import BaseAPIClient
# 固定端点
_ENDPOINT = "/v1/chat/completions"
class OpenAIAPIClient(BaseAPIClient):
"""
OpenAI 兼容格式的图像生成客户端
与 GeminiAPIClient 的主要区别：
- 端点固定为 /v1/chat/completions（不再动态拼模型名到 URL）
- 解析后的模型字符串放入请求体的 model 字段
- 请求体采用 messages 数组格式，图片以 data URI 内联
- 顶层追加 modalities 和 image_config 字段
- 响应解析对应 choices[0].message.content 结构
"""
def __init__(self, api_key: Optional[str] = None):
if api_key is None:
api_key = get_api_key_or_raise("O1KEY_API_KEY")
super().__init__(
base_url=get_api_base_url(),
api_key=api_key,
max_request_size=100 * 1024 * 1024
)
# ------------------------------------------------------------------ #
# 模型名解析 #
# 原 GeminiAPIClient.get_endpoint() 里动态拼 URL 的逻辑 #
# 现在改为：同样的输入 → 返回纯模型名字符串，放进请求体 #
# ------------------------------------------------------------------ #
def resolve_model_name(self, model: str, resolution: str) -> str:
"""
将「节点选中的模型 ID + 分辨率」解析为实际请求所用的模型名称。
对应关系与原 GeminiAPIClient.get_endpoint() 完全一致，
只是把拼在 URL 路径里的模型段提取出来单独返回。
Args:
model: 节点下拉框中的模型 ID，如 "nano-banana-pro-次卡"
resolution: 分辨率字符串，如 "1K" / "2K" / "4K" / "512"
Returns:
实际模型名，如 "nano-banana-pro-2k"
"""
# ── 动态端点模型 ──────────────────────────────────────────────────
if model == "nano-banana-pro-次卡":
if resolution == "1K":
return "nano-banana-pro"
elif resolution == "4K":
return "nano-banana-pro-4k"
else: # 2K（默认）
return "nano-banana-pro-2k"
elif model == "nano-banana-pro-官方计费":
if resolution == "1K":
return "nano-banana-pro-1k-official"
elif resolution == "4K":
return "nano-banana-pro-4k-official"
else: # 2K（默认）
return "nano-banana-pro-2k-official"
elif model == "nano-banana-2-官方计费":
if resolution == "512":
return "nano-banana-2-0.5k-official"
elif resolution == "1K":
return "nano-banana-2-1k-official"
elif resolution == "4K":
return "nano-banana-2-4k-official"
else: # 2K（默认）
return "nano-banana-2-2k-official"
elif model == "gemini-3-pro-image-preview-url":
if resolution == "1K":
return "gemini-3-pro-image-preview-url"
elif resolution == "4K":
return "gemini-3-pro-image-preview-4k-url"
else: # 2K（默认）
return "gemini-3-pro-image-preview-2k-url"
# ── 固定端点模型：从 models_config 里取端点，提取模型名段 ──────────
from ..models_config import get_model_endpoint
endpoint = get_model_endpoint(model)
if endpoint:
# 端点格式：/v1beta/models/<model-name>:generateContent
# 提取 <model-name> 部分
match = re.search(r"/models/([^:]+):", endpoint)
if match:
return match.group(1)
# ── 兜底：直接用 model ID ──────────────────────────────────────────
return model
# ------------------------------------------------------------------ #
# BaseAPIClient 抽象方法实现 #
# ------------------------------------------------------------------ #
def get_endpoint(self, **kwargs) -> str:
"""固定返回 /v1/chat/completions，模型信息已移入请求体。"""
return _ENDPOINT
def build_request_body(
self,
prompt: str = "",
images: Optional[List[Image.Image]] = None,
aspect_ratio: str = "1:1",
resolution: str = "2K",
model: str = "",
**kwargs
) -> Dict[str, Any]:
"""
构建 OpenAI /v1/chat/completions 格式请求体。
文生图示例输出：
{
"model": "nano-banana-pro-2k",
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "一个中国女子的OOTD"}
]
}
],
"modalities": ["image", "text"],
"stream": false,
"extra_body": {
"google": {
"image_config": {
"aspect_ratio": "16:9",
"image_size": "2K"
}
}
}
}
图生图时 content 数组追加若干 image_url 块：
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,<...>"}
}
Args:
prompt: 提示词
images: 参考图列表（可选，图生图时传入）
aspect_ratio: 宽高比，如 "16:9"
resolution: 分辨率，如 "2K"
model: 已解析好的模型名（由 resolve_model_name 返回）
"""
# ── 构建 content 数组 ─────────────────────────────────────────────
content: List[Dict[str, Any]] = []
# 1. 文本部分（始终在最前）
content.append({
"type": "text",
"text": prompt
})
# 2. 图片部分（图生图时追加，每张图一个 image_url block）
if images:
for img in images:
b64 = encode_image_to_base64(img)
content.append({
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{b64}"
}
})
# ── 分辨率映射（节点内部值 → API 所需值） ────────────────────────────
_resolution_map = {"512": "0.5K", "1K": "1K", "2K": "2K", "4K": "4K"}
api_image_size = _resolution_map.get(resolution, resolution)
# ── 组装完整请求体 ─────────────────────────────────────────────────
request_body: Dict[str, Any] = {
"model": model,
"messages": [
{
"role": "user",
"content": content
}
],
"modalities": ["image", "text"],
"stream": False,
"extra_body": {
"google": {
"image_config": {
"image_size": api_image_size
}
}
}
}
if aspect_ratio and aspect_ratio != "智能":
request_body["extra_body"]["google"]["image_config"]["aspect_ratio"] = aspect_ratio
return request_body
def parse_response(self, response: Dict[str, Any]) -> List[Image.Image]:
"""同步 parse_response，仅为满足抽象基类要求，实际不应被直接调用。"""
raise RuntimeError(
"parse_response() 不应被直接调用。"
"请使用 generate_single_async() 等高级方法。"
)
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
"""429 / 503 友好文案。"""
if status_code == 429:
return (
"莫慌!该模型暂时超出速率限制啦\n"
"解决方案如下(任意一种):\n"
"1.切换当前模型\n"
"2.前往后台,修改令牌分组"
)
if status_code == 503:
return (
"警报!服务器当前过载!\n"
"解决方案如下:\n"
"1.摸会儿鱼吧,稍后会恢复,嘿嘿~\n"
"2.切换其他模型\n"
"3.前往后台,修改令牌分组"
)
return None
# ------------------------------------------------------------------ #
# 响应解析 #
# ------------------------------------------------------------------ #
async def parse_response_async(
self,
response: Dict[str, Any],
session: Optional[aiohttp.ClientSession] = None
) -> tuple[List[Image.Image], Dict[str, Any]]:
"""
异步解析 /v1/chat/completions 格式响应，提取生成的图像。
响应结构（OpenAI 格式）：
{
"choices": [
{
"message": {
"role": "assistant",
"content": [
{"type": "text", "text": "..."},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,..."}}
// 或直接 inline_data / inlineData（兼容 Gemini 风格回包）
]
},
"finish_reason": "stop"
}
],
"usage": {...}
}
"""
format_info: Dict[str, Any] = {
"type": None, # "base64" | "url"
"size": 0,
"resolution": None,
"download_speed": None
}
# ── 错误前置检测 ───────────────────────────────────────────────────
# 1. usage.completion_tokens == 0 → 风控拦截（对齐 Gemini 的 candidatesTokenCount==0）
usage = response.get("usage", {})
completion_tokens = usage.get("completion_tokens", -1)
if completion_tokens == 0:
raise RuntimeError(
"Damn!你触发顶级风控啦!还没到生图阶段就被拒了。\n"
"赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~"
)
# 2. finish_reason 不是 "stop" → 安全过滤 / token 超限等
choices = response.get("choices", [])
if choices:
for choice in choices:
finish_reason = choice.get("finish_reason", "")
if finish_reason and finish_reason != "stop":
raise RuntimeError(
"Ohh no! 生图过程触发风控,图片被拒绝生成!\n"
"可能原因如下:\n"
"1.违禁内容\n"
"2.触发安全过滤器\n"
"3.涉及版权问题\n"
"4. Token超限\n"
"赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~"
)
# ── 图像提取 ───────────────────────────────────────────────────────
images: List[Image.Image] = []
text_responses: List[str] = []
close_session = False
if session is None:
session = self._make_session()
close_session = True
try:
for choice in choices:
message = choice.get("message", {})
# ── 优先从 message.images 提取（非标准扩展字段） ──────────────
# 部分服务端把图片放在独立的 images 字段，content 同时为 null
msg_images = message.get("images") or []
for img_part in msg_images:
part_type = img_part.get("type", "")
if part_type == "image_url":
url_obj = img_part.get("image_url", {})
url = url_obj.get("url", "")
if url.startswith("data:"):
try:
_, b64_data = url.split(",", 1)
img = decode_base64_to_pil(b64_data)
images.append(img)
if format_info["type"] is None:
format_info["type"] = "base64"
format_info["size"] = len(b64_data) * 3 / 4
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
except Exception:
pass
elif url.startswith("http"):
try:
dl_start = time.time()
async with session.get(url) as img_resp:
if img_resp.status == 200:
img_data = await img_resp.read()
dl_time = time.time() - dl_start
speed = len(img_data) / dl_time if dl_time > 0 else 0
img = Image.open(BytesIO(img_data))
images.append(img)
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = len(img_data)
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception:
pass
# ── 再从 message.content 提取（标准 OpenAI 格式） ─────────────
# content 为 null 时用空列表兜底，避免 for in None 崩溃
raw_content = message.get("content") or []
# content 可能是字符串（纯文本）或数组（多模态）
if isinstance(raw_content, str):
text_responses.append(raw_content)
continue
for part in raw_content:
part_type = part.get("type", "")
# ── 情况 A：OpenAI image_url 格式 ─────────────────────
if part_type == "image_url":
url_obj = part.get("image_url", {})
url = url_obj.get("url", "")
if url.startswith("data:"):
# data URI → 直接 base64 解码
# 格式：data:image/png;base64,<data>
try:
header, b64_data = url.split(",", 1)
img = decode_base64_to_pil(b64_data)
images.append(img)
if format_info["type"] is None:
format_info["type"] = "base64"
format_info["size"] = len(b64_data) * 3 / 4
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
except Exception:
pass
elif url.startswith("http"):
# 远程 URL → 异步下载
try:
dl_start = time.time()
async with session.get(url) as img_resp:
if img_resp.status == 200:
img_data = await img_resp.read()
dl_time = time.time() - dl_start
speed = len(img_data) / dl_time if dl_time > 0 else 0
img = Image.open(BytesIO(img_data))
images.append(img)
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = len(img_data)
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception:
pass
# ── 情况 B：Gemini 风格 inline_data / inlineData（兼容） ─
elif part_type in ("inline_data", "inlineData") or \
"inline_data" in part or "inlineData" in part:
inline_key = "inline_data" if "inline_data" in part else "inlineData"
inline = part.get(inline_key, {})
b64_data = inline.get("data", "")
if b64_data:
try:
img = decode_base64_to_pil(b64_data)
images.append(img)
if format_info["type"] is None:
format_info["type"] = "base64"
format_info["size"] = len(b64_data) * 3 / 4
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
except Exception:
pass
# ── 情况 C：text 中嵌套 URL（markdown 或纯链接） ─────────
elif part_type == "text":
text = part.get("text", "")
text_responses.append(text)
# markdown 图片链接：![alt](url)
urls = re.findall(r'!\[.*?\]\((https?://[^\)]+)\)', text)
if not urls:
urls = re.findall(r'https?://[^\s<>"{}|\\^`\[\]]+', text)
for url in urls:
try:
dl_start = time.time()
async with session.get(url) as img_resp:
if img_resp.status == 200:
img_data = await img_resp.read()
dl_time = time.time() - dl_start
speed = len(img_data) / dl_time if dl_time > 0 else 0
img = Image.open(BytesIO(img_data))
images.append(img)
if format_info["type"] is None:
format_info["type"] = "url"
format_info["size"] = len(img_data)
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
format_info["download_speed"] = speed
except Exception:
pass
except RuntimeError:
raise
except Exception as e:
raise RuntimeError(f"解析 API 响应失败: {str(e)}")
finally:
if close_session:
await session.close()
# ── 3. 无图像但有文本 → API 拒绝说明 ─────────────────────────────
if not images and text_responses:
combined = "\n".join(text_responses)
raise RuntimeError(
f"API 拒绝响应\n\n"
f"API 返回说明：\n{combined}\n\n"
f"建议：\n"
f" - 根据上述说明调整请求内容\n"
f" - 确保提示词和参考图符合使用规范"
)
if not images:
raise RuntimeError("API 响应中未找到生成的图像")
return images, format_info
# ------------------------------------------------------------------ #
# 核心生成方法（接口与 GeminiAPIClient 保持一致，节点可无缝切换） #
# ------------------------------------------------------------------ #
async def generate_single_async(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
images: Optional[List[Image.Image]] = None,
session: Optional[aiohttp.ClientSession] = None,
task_index: Optional[int] = None,
total_tasks: Optional[int] = None,
debug: bool = False,
debug_request: bool = False,
enable_grounding: bool = False, # 保留签名兼容，OpenAI 格式暂不使用
enable_image_search: bool = False # 保留签名兼容，OpenAI 格式暂不使用
) -> tuple[List[Image.Image], Dict[str, Any]]:
"""
单次异步生成请求（OpenAI /v1/chat/completions 格式）。
Args:
prompt: 提示词
model: 节点选中的模型 ID（将自动解析为实际模型名）
resolution: 分辨率
aspect_ratio: 宽高比
images: 参考图列表（图生图时传入）
session: 复用的 aiohttp 会话
task_index: 任务序号（批量时用于日志）
total_tasks: 总任务数（批量时用于日志）
debug: 打印完整 API 响应
debug_request: 打印请求体（base64 自动截断）
Returns:
(生成的图像列表, 计时信息字典)
"""
import json
total_start = time.time()
task_prefix = f"[{task_index}/{total_tasks}]" if task_index is not None and total_tasks else ""
# ── 1. 解析模型名 & 构建请求体 ────────────────────────────────────
build_start = time.time()
resolved_model = self.resolve_model_name(model, resolution)
endpoint = self.get_endpoint()
request_body = self.build_request_body(
prompt=prompt,
images=images,
aspect_ratio=aspect_ratio,
resolution=resolution,
model=resolved_model
)
build_time = time.time() - build_start
# ── 调试：打印请求体 ───────────────────────────────────────────────
if debug_request:
import json as _json
def _shorten_b64(obj):
if isinstance(obj, dict):
return {k: _shorten_b64(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_shorten_b64(i) for i in obj]
if isinstance(obj, str):
if obj.startswith("data:"):
header, _, data = obj.partition(",")
return f"{header},<base64 {len(data)} chars>"
if len(obj) > 200 and all(
c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
for c in obj[:64]
):
return f"<base64 {len(obj)} chars>"
return obj
print(
f"\n{'='*60}\n"
f"[请求体日志] 任务 {task_prefix or '?'}\n"
f"端点: {self.base_url}{endpoint}\n"
f"{_json.dumps(_shorten_b64(request_body), ensure_ascii=False, indent=2)}\n"
f"{'='*60}\n"
)
# ── 2. 计算请求体大小 ─────────────────────────────────────────────
request_size = len(json.dumps(request_body).encode("utf-8"))
size_str = (
f"{request_size / 1024:.2f}KB"
if request_size < 1024 * 1024
else f"{request_size / (1024 * 1024):.2f}MB"
)
# ── 3. 发送请求（Bearer Token 认证） ─────────────────────────────
request_start = time.time()
try:
response = await self.request_async(
endpoint,
request_body,
session,
use_bearer_token=True
)
except Exception as e:
request_time = time.time() - request_start
error_first_line = str(e).split("\n")[0]
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 失败: {error_first_line} ✗")
raise
request_time = time.time() - request_start
# ── 调试：打印完整响应 ─────────────────────────────────────────────
if debug:
import json as _json
def _shorten_b64(obj):
if isinstance(obj, dict):
return {k: _shorten_b64(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_shorten_b64(i) for i in obj]
if isinstance(obj, str):
if obj.startswith("data:"):
header, _, data = obj.partition(",")
return f"{header},<base64 {len(data)} chars>"
if len(obj) > 200 and all(
c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="
for c in obj[:64]
):
return f"<base64 {len(obj)} chars>"
return obj
print(
f"\n{'='*60}\n"
f"[调试日志] 任务 {task_prefix or '?'} 完整 API 响应：\n"
f"{_json.dumps(_shorten_b64(response), ensure_ascii=False, indent=2)}\n"
f"{'='*60}\n"
)
# ── 4. 解析响应 ───────────────────────────────────────────────────
parse_start = time.time()
try:
result_images, format_info = await self.parse_response_async(response, session)
except Exception as e:
parse_time = time.time() - parse_start
error_first_line = str(e).split("\n")[0]
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → 解析失败: {error_first_line} ✗")
raise
parse_time = time.time() - parse_start
# ── 5. 单行日志输出 ───────────────────────────────────────────────
img_size = format_info.get("size", 0)
img_size_str = (
f"{img_size / 1024:.2f}KB"
if img_size < 1024 * 1024
else f"{img_size / (1024 * 1024):.2f}MB"
)
if format_info.get("type") == "base64":
download_info = f"Base64 {img_size_str} ({parse_time:.1f}s)"
elif format_info.get("type") == "url":
speed = format_info.get("download_speed", 0)
download_info = f"URL {img_size_str} ({parse_time:.1f}s, {speed / (1024*1024):.1f}MB/s)"
else:
download_info = img_size_str
timing = response.get("_timing", {})
net_connect = timing.get("connect_time")
net_download = timing.get("download_time")
if net_connect is not None and net_download is not None:
net_str = f" | 连接 {net_connect:.2f}s | 下载 {net_download:.2f}s"
else:
net_str = ""
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓{net_str}")
total_time = time.time() - total_start
timing_info = {
"build_time": build_time,
"request_time": request_time,
"parse_time": parse_time,
"total_time": total_time,
"format_type": format_info.get("type", "unknown")
}
return result_images, timing_info
# ------------------------------------------------------------------ #
# 批量 & 同步接口（与 GeminiAPIClient 接口签名一致） #
# ------------------------------------------------------------------ #
async def generate_batch_async(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
batch_size: int,
images: Optional[List[Image.Image]] = None,
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
debug: bool = False,
debug_request: bool = False,
enable_grounding: bool = False,
enable_image_search: bool = False
) -> List[Image.Image]:
"""批量全并发生成（单提示词 × batch_size 张）。"""
import asyncio
all_images: List[Image.Image] = []
completed = 0
success_count = 0
fail_count = 0
first_error = None
max_concurrent = 10
num_batches = (batch_size + max_concurrent - 1) // max_concurrent
print(f"OpenAIClient: 批量生成 {batch_size} 张，并发数: {max_concurrent}，分 {num_batches} 批")
connector = aiohttp.TCPConnector(ssl=False, limit=0, limit_per_host=0)
async with aiohttp.ClientSession(connector=connector) as session:
for batch_idx in range(num_batches):
batch_start = batch_idx * max_concurrent
batch_end = min(batch_start + max_concurrent, batch_size)
batch_count = batch_end - batch_start
if num_batches > 1:
print(f"OpenAIClient: 第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})")
tasks = [
asyncio.create_task(
self.generate_single_async(
prompt=prompt,
model=model,
resolution=resolution,
aspect_ratio=aspect_ratio,
images=images,
session=session,
task_index=batch_start + i + 1,
total_tasks=batch_size,
debug=debug,
debug_request=debug_request
),
name=f"task_{batch_start + i}"
)
for i in range(batch_count)
]
batch_images: List[Image.Image] = []
for coro in asyncio.as_completed(tasks):
completed += 1
try:
result_imgs, _ = await coro
for img in result_imgs:
batch_images.append(img)
all_images.append(img)
success_count += 1
if progress_callback:
progress_callback(completed, batch_size, True, None)
print(f"OpenAIClient: 任务 {completed}/{batch_size} 成功 ✓")
except Exception as e:
fail_count += 1
if first_error is None:
first_error = e
if progress_callback:
progress_callback(completed, batch_size, False, str(e))
print(f"OpenAIClient: 任务 {completed}/{batch_size} 失败 ✗")
if batch_images:
print(f"OpenAIClient: 第 {batch_idx + 1} 批完成，生成 {len(batch_images)} 张")
import gc
gc.collect()
await asyncio.sleep(0.1)
batch_images = []
if not all_images:
if first_error:
raise first_error
raise RuntimeError(f"批量生成失败，{fail_count} 个请求全部失败")
print(f"OpenAIClient: 批量完成，成功 {success_count}/{batch_size}，失败 {fail_count}")
return all_images
def generate_sync(
self,
prompt: str,
model: str,
resolution: str,
aspect_ratio: str,
batch_size: int,
images: Optional[List[Image.Image]] = None,
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
debug: bool = False,
debug_request: bool = False,
enable_grounding: bool = False,
enable_image_search: bool = False
) -> List[Image.Image]:
"""同步生成接口（用于 ComfyUI 节点，接口与 GeminiAPIClient 完全一致）。"""
coro = self.generate_batch_async(
prompt=prompt,
model=model,
resolution=resolution,
aspect_ratio=aspect_ratio,
batch_size=batch_size,
images=images,
progress_callback=progress_callback,
debug=debug,
debug_request=debug_request,
enable_grounding=enable_grounding,
enable_image_search=enable_image_search
)
return self.run_async_in_thread(coro)
clients/seedance_client.py
+50 -29
View File
@@ -3,17 +3,17 @@ Seedance 视频生成客户端
使用 new-api 原生格式：POST /v1/video/generations → GET /v1/video/generations/{task_id}
"""
import asyncio
import json
import os
from typing import Any, Callable, Dict, Optional
import aiohttp
from ..utils.config import get_api_key_or_raise
from ..utils.config import get_api_key_or_raise, get_base_url_by_route
from ..utils.http_error import async_request_with_retry
from ..utils.video_task import (
PollDeadline,
check_interrupt,
download_video_to_file,
extract_error_message,
extract_progress,
extract_status,
@@ -25,11 +25,16 @@ from ..utils.video_task import (
class SeedanceClient:
"""Seedance 视频生成客户端（new-api 原生三段式）"""
"""Seedance 视频生成客户端（new-api 原生三段式）
# 提交任务
注意：新旧格式模型（seedance-2-0-260128-d 等）共用同一套端点，
区别仅在于请求体结构（顶层 content vs metadata.content），
由调用方（节点层）通过 use_new_format 控制请求体拼装方式。
"""
# 提交任务（新旧格式模型共用）
CREATE_ENDPOINT = "/v1/video/generations"
# 查询任务状态：{task_id} 占位
# 查询任务状态：{task_id} 占位（新旧格式模型共用）
STATUS_ENDPOINT = "/v1/video/generations/{task_id}"
POLL_INITIAL_INTERVAL = 4 # 首次轮询等待秒数
@@ -41,7 +46,7 @@ class SeedanceClient:
def __init__(self):
self.api_key = get_api_key_or_raise()
self.base_url = "https://api.o1key.com"
self.base_url = get_base_url_by_route()
def _headers(self) -> Dict[str, str]:
return {
@@ -55,9 +60,17 @@ class SeedanceClient:
self,
body: Dict[str, Any],
session: aiohttp.ClientSession,
use_new_format: bool = False,
) -> str:
"""提交视频生成任务，返回 task_id"""
"""提交视频生成任务，返回 task_id
use_new_format 仅用于调试日志标注请求体格式，不影响端点选择
（新旧格式模型统一走 CREATE_ENDPOINT）。
"""
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
print(f"[Seedance] 提交 → {url} (body格式: {'新' if use_new_format else '旧'})")
check_interrupt()
resp = await run_with_interrupt(async_request_with_retry(
session, "POST", url, json=body, headers=self._headers(), prefix="Seedance 提交: "
@@ -69,7 +82,7 @@ class SeedanceClient:
# new-api 返回字段：id / task_id
task_id = data.get("id") or data.get("task_id")
if not task_id:
raise RuntimeError(f"API 未返回任务 ID，响应：{data}")
raise RuntimeError("API 未返回任务 ID")
return task_id
# ── 2. 轮询状态 ────────────────────────────────────────────────────
@@ -79,12 +92,15 @@ class SeedanceClient:
task_id: str,
session: aiohttp.ClientSession,
on_progress: Optional[Callable[[int], None]] = None,
use_new_format: bool = False,
) -> str:
"""轮询任务状态，成功后返回视频 URL"""
"""轮询任务状态，成功后返回视频 URL（新旧格式模型统一走 STATUS_ENDPOINT）"""
url = f"{self.base_url}{self.STATUS_ENDPOINT.format(task_id=task_id)}"
interval = self.POLL_INITIAL_INTERVAL
deadline = PollDeadline(label="Seedance")
while True:
deadline.check()
check_interrupt()
async with session.get(url, headers=self._headers()) as resp:
text = await resp.text()
@@ -123,7 +139,7 @@ class SeedanceClient:
or inner.get("url")
)
if not video_url:
raise RuntimeError(f"任务成功但未找到视频 URL，响应：{result}")
raise RuntimeError("任务成功但未找到视频 URL")
# 末帧图片 URL 在 inner["data"]["content"]["last_frame_url"]
last_frame_url = (
content.get("last_frame_url")
@@ -149,16 +165,9 @@ class SeedanceClient:
) -> str:
"""下载视频到本地，返回本地路径"""
print(f"[Seedance] 下载视频...")
check_interrupt()
async with session.get(video_url, allow_redirects=True) as resp:
if resp.status != 200:
raise RuntimeError(f"视频下载失败 ({resp.status})")
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
async for chunk in resp.content.iter_chunked(8192):
check_interrupt()
f.write(chunk)
return save_path
return await download_video_to_file(
session, video_url, save_path, label="Seedance",
)
# ── 全流程入口（供节点调用）────────────────────────────────────────
@@ -168,6 +177,7 @@ class SeedanceClient:
save_path: str,
on_stage: Optional[Callable[[str], None]] = None,
on_progress: Optional[Callable[[int], None]] = None,
use_new_format: bool = False,
) -> tuple:
"""提交 → 轮询 → 下载，返回 (本地视频路径, 末帧图片URL或None)"""
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
@@ -177,19 +187,30 @@ class SeedanceClient:
check_interrupt()
if on_stage:
on_stage("submitting")
task_id = await self.submit_async(body, session)
task_id = await self.submit_async(body, session, use_new_format=use_new_format)
print(f"[Seedance] 任务已提交 → {task_id}")
if on_stage:
on_stage(f"submitted:{task_id}")
# 轮询
video_url, last_frame_url = await self.poll_async(task_id, session, on_progress=on_progress)
video_url, last_frame_url = await self.poll_async(task_id, session, on_progress=on_progress, use_new_format=use_new_format)
# 下载
# 下载（带"Video not ready"重试）
if on_stage:
on_stage("downloading")
path = await self.download_async(video_url, save_path, session)
if on_stage:
on_stage("done")
return path, last_frame_url
max_retries = 5
retry_delay = 3.0
for attempt in range(max_retries):
try:
path = await self.download_async(video_url, save_path, session)
if on_stage:
on_stage("done")
return path, last_frame_url
except Exception as e:
error_msg = str(e)
if "Video not ready" in error_msg and attempt < max_retries - 1:
print(f"[Seedance] 视频未就绪，{retry_delay}秒后重试 ({attempt + 1}/{max_retries})...")
await interruptible_sleep(retry_delay)
check_interrupt()
continue
raise
clients/seedance_element_client.py
+320
View File
@@ -0,0 +1,320 @@
"""
Seedance 2.0 真人素材（Element）API 客户端
封装标准素材接口与高并发素材接口
"""
import json
import aiohttp
from typing import Optional
from urllib.parse import quote
from ..utils.config import get_api_key_or_raise, get_base_url_by_route
from ..utils.video_task import PollDeadline, check_interrupt, interruptible_sleep
class SeedanceElementClient:
"""Seedance 2.0 真人素材客户端"""
_ASSET_REQUEST_TYPES = {"hc", "doubao"}
def __init__(self, base_url: str = None, api_key: str = None):
self.api_key = api_key or get_api_key_or_raise()
self.base_url = base_url or get_base_url_by_route()
def _headers(self) -> dict:
return {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
def _hc_headers(self) -> dict:
return {
**self._headers(),
"Accept": "application/json",
}
@classmethod
def _normalize_asset_request_type(cls, request_type: str) -> str:
normalized = str(request_type).strip().lower()
if normalized not in cls._ASSET_REQUEST_TYPES:
raise ValueError(f"不支持的素材请求类型：{request_type}")
return normalized
@staticmethod
def _hc_error_message(payload: dict, default: str) -> str:
error = payload.get("error")
if isinstance(error, dict):
error = error.get("message") or error.get("detail")
data = payload.get("data")
base_resp = data.get("base_resp", {}) if isinstance(data, dict) else {}
return str(
error
or payload.get("message")
or base_resp.get("status_msg")
or default
)
@staticmethod
async def _read_json_response(resp: aiohttp.ClientResponse) -> dict:
text = await resp.text()
try:
payload = json.loads(text)
except json.JSONDecodeError:
raise RuntimeError("素材接口返回了无效 JSON") from None
if not isinstance(payload, dict):
raise RuntimeError("素材接口返回格式错误")
return payload
async def create_element(
self,
name: str,
image_url: str,
description: Optional[str] = None,
channel_id: int = 0,
session: aiohttp.ClientSession = None,
) -> dict:
"""
创建素材
Args:
name: 素材名称
image_url: 图片 URL（必须是 http/https）
description: 素材描述（可选）
channel_id: 渠道ID，0表示自动选择
session: aiohttp会话，如果为None则创建临时会话
Returns:
{
"id": 123,
"name": "我的数字人",
"description": "真人形象描述",
"frontal_image": "https://xxx.jpg",
"element_id": "asset-abc123xyz", # 重要！上游Asset ID
"job_id": "group-xyz789",
"status": "succeed",
"created_at": 1719734400
}
"""
url = f"{self.base_url}/api/element/seedance"
body = {
"name": name,
"image_url": image_url,
"channel_id": channel_id,
}
if description:
body["description"] = description
should_close = session is None
if session is None:
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
session = aiohttp.ClientSession(connector=connector)
try:
async with session.post(url, json=body, headers=self._headers()) as resp:
result = await resp.json()
print(
"[Seedance素材][标准] 创建响应体:\n"
+ json.dumps(
{"http_status": resp.status, "success": bool(result.get("success"))},
ensure_ascii=False,
)
)
if resp.status != 200:
error_msg = result.get("message", result.get("error", str(result)))
raise RuntimeError(f"创建素材失败 ({resp.status}): {error_msg}")
if not result.get("success", False):
error_msg = result.get("message", "创建素材失败")
raise RuntimeError(error_msg)
data = result.get("data", result)
if isinstance(data, dict):
data = dict(data)
data["_create_response"] = result
return data
finally:
if should_close:
await session.close()
async def create_hc_asset(
self,
name: str,
asset_url: str,
asset_type: str,
session: aiohttp.ClientSession = None,
request_type: str = "hc",
) -> dict:
"""通过统一 Seedance 素材接口创建 HC 或 Doubao 素材。"""
normalized_asset_type = str(asset_type).strip().lower()
if normalized_asset_type not in {"image", "video", "audio"}:
raise ValueError(f"不支持的素材类型：{asset_type}")
normalized_request_type = self._normalize_asset_request_type(request_type)
request_label = "HC" if normalized_request_type == "hc" else "Doubao"
url = f"{self.base_url}/v1/seedance/assets"
body = {
"type": normalized_request_type,
"url": asset_url,
"asset_type": normalized_asset_type,
}
if name:
body["name"] = name
should_close = session is None
if session is None:
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
session = aiohttp.ClientSession(connector=connector)
try:
check_interrupt()
async with session.post(url, json=body, headers=self._hc_headers()) as resp:
result = await self._read_json_response(resp)
print(
f"[Seedance素材][{request_label}] 创建响应体:\n"
+ json.dumps(
{"http_status": resp.status, "success": bool(result.get("success"))},
ensure_ascii=False,
)
)
if resp.status < 200 or resp.status >= 300:
message = self._hc_error_message(result, "创建素材失败")
raise RuntimeError(f"创建 {request_label} 素材失败 ({resp.status}): {message}")
if not result.get("success", False):
raise RuntimeError(self._hc_error_message(result, f"创建 {request_label} 素材失败"))
data = result.get("data")
if not isinstance(data, dict) or not data.get("Id"):
raise RuntimeError(f"创建 {request_label} 素材成功但未返回 data.Id")
data = dict(data)
data["_create_response"] = result
return data
finally:
if should_close:
await session.close()
async def get_hc_asset(
self,
asset_id: str,
session: aiohttp.ClientSession = None,
request_type: str = "hc",
) -> dict:
"""通过统一 Seedance 素材接口查询 HC 或 Doubao 素材状态。"""
normalized_request_type = self._normalize_asset_request_type(request_type)
request_label = "HC" if normalized_request_type == "hc" else "Doubao"
encoded_id = quote(asset_id, safe="")
url = f"{self.base_url}/v1/seedance/assets/{encoded_id}"
should_close = session is None
if session is None:
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
session = aiohttp.ClientSession(connector=connector)
try:
check_interrupt()
async with session.get(
url,
params={"type": normalized_request_type},
headers=self._hc_headers(),
) as resp:
result = await self._read_json_response(resp)
if resp.status < 200 or resp.status >= 300:
message = self._hc_error_message(result, "查询素材状态失败")
raise RuntimeError(f"查询 {request_label} 素材失败 ({resp.status}): {message}")
if not result.get("success", False):
raise RuntimeError(self._hc_error_message(result, f"查询 {request_label} 素材失败"))
data = result.get("data")
if not isinstance(data, dict):
raise RuntimeError(f"查询 {request_label} 素材未返回 data")
return data
finally:
if should_close:
await session.close()
async def create_hc_asset_and_wait(
self,
name: str,
asset_url: str,
asset_type: str,
poll_interval: float = 3.0,
request_type: str = "hc",
) -> dict:
"""创建 HC 或 Doubao 素材并等待其进入 Active 状态。"""
normalized_request_type = self._normalize_asset_request_type(request_type)
request_label = "HC" if normalized_request_type == "hc" else "Doubao"
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
async with aiohttp.ClientSession(connector=connector) as session:
created = await self.create_hc_asset(
name=name,
asset_url=asset_url,
asset_type=asset_type,
session=session,
request_type=normalized_request_type,
)
asset_id = str(created["Id"])
deadline = PollDeadline(label=f"Seedance {request_label} 素材")
print(f"[Seedance素材][{request_label}] 已创建 {asset_id}，等待素材可用...")
while True:
deadline.check()
check_interrupt()
asset = await self.get_hc_asset(
asset_id,
session=session,
request_type=normalized_request_type,
)
status = str(asset.get("Status", "")).strip()
normalized_status = status.lower()
print(f"[Seedance素材][{request_label}] {asset_id} 状态: {status or '未知'}")
if normalized_status == "active":
asset = dict(asset)
asset["_create_response"] = created.get("_create_response", {})
return asset
if normalized_status == "failed":
message = self._hc_error_message({"data": asset}, "素材处理失败")
raise RuntimeError(f"{request_label} 素材处理失败：{message}")
if normalized_status != "processing":
raise RuntimeError(f"{request_label} 素材返回未知状态：{status or '空状态'}")
await interruptible_sleep(poll_interval)
async def delete_element(
self,
element_internal_id: int,
session: aiohttp.ClientSession = None,
) -> dict:
"""
删除素材记录（仅删除平台记录，不删除上游Asset）
Args:
element_internal_id: 平台内部记录ID（非element_id）
session: aiohttp会话
Returns:
{"message": "删除成功"}
"""
url = f"{self.base_url}/api/element/seedance/{element_internal_id}"
should_close = session is None
if session is None:
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
session = aiohttp.ClientSession(connector=connector)
try:
async with session.delete(url, headers=self._headers()) as resp:
result = await resp.json()
if resp.status != 200:
error_msg = result.get("message", result.get("error", str(result)))
raise RuntimeError(f"删除素材失败 ({resp.status}): {error_msg}")
if not result.get("success", False):
error_msg = result.get("message", "删除素材失败")
raise RuntimeError(error_msg)
return result.get("data", result)
finally:
if should_close:
await session.close()
clients/seedream_image_client.py
+409
View File
@@ -0,0 +1,409 @@
"""Seedream image client for O1Key's asynchronous image API."""
from __future__ import annotations
import os
import time
from typing import Any, Awaitable, Callable, Optional, Sequence
from PIL import Image
from ..utils.nano_banana_async import (
extract_async_image_result_urls,
image_to_upload_payload,
parse_completed_async_image_task,
poll_async_image_task,
submit_async_image_task,
upload_images_to_temp_urls,
)
from ..utils.o1key_image_catalog import (
MAX_UNIFIED_REFERENCE_IMAGES,
SEEDREAM_MODEL_OPTIONS,
SEEDREAM_LAYER_RESOLUTION_OPTIONS,
SEEDREAM_OUTPUT_FORMAT_OPTIONS,
SEEDREAM_SIZE_MATRIX,
UNIFIED_IMAGE_ROUTE_OPTIONS,
)
SEEDREAM_API_MODEL_ID = "dola-seedream-5-0-pro-260628-ep"
SEEDREAM_REFERENCE_MAX_BYTES = 30 * 1024 * 1024
SEEDREAM_REFERENCE_MAX_PIXELS = 6000 * 6000
SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE = 14
SEEDREAM_REFERENCE_MIN_ASPECT_RATIO = 1 / 16
SEEDREAM_REFERENCE_MAX_ASPECT_RATIO = 16
SEEDREAM_LAYER_REFERENCE_MIN_PIXELS = 512 * 512
def validate_seedream_reference_dimensions(
width: int,
height: int,
*,
label: str = "Seedream 参考图",
layer_decomposition: bool = False,
) -> None:
"""Validate the current Volcengine per-image reference-size contract."""
width = int(width)
height = int(height)
if width <= 0 or height <= 0:
raise ValueError(f"{label}尺寸无效：{width}×{height}")
pixels = width * height
if (
not layer_decomposition
and (
width <= SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE
or height <= SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE
)
):
raise ValueError(
f"{label}宽和高都必须大于 {SEEDREAM_REFERENCE_MIN_SIDE_EXCLUSIVE}px，"
f"当前为 {width}×{height}"
)
ratio = width / height
if (
ratio < SEEDREAM_REFERENCE_MIN_ASPECT_RATIO
or ratio > SEEDREAM_REFERENCE_MAX_ASPECT_RATIO
):
raise ValueError(
f"{label}宽高比必须在 1:16～16:1，当前为 {width}:{height}"
)
if layer_decomposition:
if not (
SEEDREAM_LAYER_REFERENCE_MIN_PIXELS
<= pixels
<= SEEDREAM_REFERENCE_MAX_PIXELS
):
raise ValueError(
f"{label}总像素必须在 512×512（262144）～6000×6000（36000000）之间，"
f"当前为 {width}×{height}（{pixels}）"
)
return
if pixels > SEEDREAM_REFERENCE_MAX_PIXELS:
raise ValueError(
f"{label}总像素不能超过 6000×6000（36000000），"
f"当前为 {width}×{height}（{pixels}）"
)
def validate_seedream_reference_image(
image: Image.Image,
*,
label: str = "Seedream 参考图",
layer_decomposition: bool = False,
) -> None:
"""Validate reference dimensions and the exact bytes sent to the uploader."""
validate_seedream_reference_dimensions(
image.width,
image.height,
label=label,
layer_decomposition=layer_decomposition,
)
payload, _extension, _content_type = image_to_upload_payload(image)
try:
payload_size = (
os.path.getsize(payload)
if isinstance(payload, (str, os.PathLike))
else len(payload)
)
except (OSError, TypeError) as exc:
raise ValueError(f"无法读取{label}文件大小") from exc
if payload_size > SEEDREAM_REFERENCE_MAX_BYTES:
raise ValueError(
f"{label}文件不能超过 30MB，当前为 {payload_size / 1024 / 1024:.2f}MB"
)
def validate_seedream_reference_images(
images: Sequence[Image.Image],
*,
layer_decomposition: bool = False,
) -> None:
for index, image in enumerate(images, start=1):
validate_seedream_reference_image(
image,
label=f"Seedream 参考图{index}",
layer_decomposition=layer_decomposition,
)
def resolve_seedream_model(model_name: str, route: str) -> str:
"""Map the stable workflow value to Seedream's API model identifier."""
if model_name == SEEDREAM_API_MODEL_ID:
return model_name
if model_name not in SEEDREAM_MODEL_OPTIONS:
raise ValueError(f"Seedream 模型无效：{model_name}")
if route not in UNIFIED_IMAGE_ROUTE_OPTIONS:
raise ValueError(f"Seedream 模型线路无效：{route}")
# O1Key currently exposes one Seedream endpoint for every displayed route.
return SEEDREAM_API_MODEL_ID
def build_seedream_submit_body(
*,
model: str,
prompt: str,
size: Optional[str],
output_format: str,
image_urls: Optional[Sequence[str]] = None,
layer_decomposition: bool = False,
) -> dict[str, Any]:
"""Build and validate the paid Seedream request without logging URLs."""
normalized_prompt = str(prompt or "").strip()
if not normalized_prompt and not layer_decomposition:
raise ValueError("请输入提示词")
if model != SEEDREAM_API_MODEL_ID:
raise ValueError(f"Seedream API 模型无效：{model}")
normalized_format = str(output_format or "").strip().lower()
if normalized_format not in SEEDREAM_OUTPUT_FORMAT_OPTIONS:
raise ValueError("Seedream 输出格式仅支持 png 或 jpeg")
if layer_decomposition and normalized_format != "png":
raise ValueError("Seedream 图层拆分仅支持 png 输出格式")
normalized_size = str(size or "").strip().lower().replace("*", "x").replace("×", "x")
if normalized_size:
if layer_decomposition:
normalized_size = "auto" if normalized_size == "auto" else normalized_size.upper()
if normalized_size not in SEEDREAM_LAYER_RESOLUTION_OPTIONS:
raise ValueError(f"Seedream 图层拆分分辨率无效：{size}")
elif normalized_size not in set(SEEDREAM_SIZE_MATRIX.values()):
raise ValueError(f"Seedream 图片尺寸无效：{size}")
urls = [str(url or "").strip() for url in (image_urls or ())]
if len(urls) > MAX_UNIFIED_REFERENCE_IMAGES:
raise ValueError(f"Seedream 参考图最多支持 {MAX_UNIFIED_REFERENCE_IMAGES} 张")
if any(not url.startswith("https://") for url in urls):
raise ValueError("Seedream 参考图必须使用临时素材 HTTPS URL")
if layer_decomposition and len(urls) != 1:
raise ValueError("Seedream 图层拆分必须且只能提供1张参考图")
body: dict[str, Any] = {
"model": model,
"n": 1,
"output_format": normalized_format,
"watermark": False,
}
if normalized_size:
body["size"] = normalized_size
if normalized_prompt:
body["prompt"] = normalized_prompt
if urls:
body["images"] = urls
if layer_decomposition:
body["layer_decomposition"] = True
return body
def _seedream_result_items(payload: Any) -> list[dict[str, Any]]:
"""Return the first documented image-item list without exposing its URLs."""
pending = [payload]
seen: set[int] = set()
while pending:
value = pending.pop(0)
if not isinstance(value, dict) or id(value) in seen:
continue
seen.add(id(value))
images = value.get("images")
if isinstance(images, list) and all(isinstance(item, dict) for item in images):
return images
for key in ("data", "result", "output"):
nested = value.get(key)
if isinstance(nested, dict):
pending.append(nested)
return []
def _bounded_int_list(value: Any, *, length: int) -> list[int] | None:
if not isinstance(value, (list, tuple)) or len(value) != length:
return None
try:
return [int(item) for item in value]
except (TypeError, ValueError):
return None
def extract_seedream_layer_metadata(payload: Any) -> list[dict[str, Any]]:
"""Sanitize layer metadata; result URLs are deliberately excluded."""
metadata: list[dict[str, Any]] = []
for offset, item in enumerate(_seedream_result_items(payload)):
try:
z_index = max(0, min(16, int(item.get("z_index", offset))))
except (TypeError, ValueError):
z_index = offset
safe: dict[str, Any] = {"z_index": z_index}
for key, limit in (("name", 200), ("description", 1000), ("size", 64), ("output_format", 16)):
value = item.get(key)
if isinstance(value, str) and value.strip():
safe[key] = value.strip()[:limit]
bounding_box = item.get("bounding_box")
if isinstance(bounding_box, dict):
absolute = _bounded_int_list(bounding_box.get("absolute"), length=4)
normalized = _bounded_int_list(bounding_box.get("normalized"), length=4)
safe_box = {}
if absolute is not None:
safe_box["absolute"] = absolute
if normalized is not None:
safe_box["normalized"] = normalized
if safe_box:
safe["bounding_box"] = safe_box
metadata.append(safe)
return metadata
class SeedreamImageClient:
"""Upload references, submit one Seedream task, poll it, and decode results."""
def __init__(self, *, base_url: str, api_key: str):
self.base_url = str(base_url).rstrip("/")
self.api_key = api_key
async def generate_async(
self,
*,
session: Any,
prompt: str,
model: str,
size: Optional[str],
output_format: str,
images: Optional[Sequence[Image.Image]] = None,
layer_decomposition: bool = False,
upload_cache: Optional[dict[int, Awaitable[str]]] = None,
check_interrupt: Optional[Callable[[], None]] = None,
progress_callback: Optional[Callable[[float], None]] = None,
result_url_callback: Optional[Callable[[str], None]] = None,
log_downloads: bool = True,
log_task_success: bool = True,
task_completed_callback: Optional[
Callable[[str, int, float, list[str]], None]
] = None,
) -> tuple[list[Image.Image], dict[str, Any]]:
if check_interrupt:
check_interrupt()
reference_images = list(images or ())
validate_seedream_reference_images(
reference_images,
layer_decomposition=layer_decomposition,
)
task_started = time.time()
image_urls = await upload_images_to_temp_urls(
session=session,
base_url=self.base_url,
api_key=self.api_key,
images=reference_images,
node_label="Seedream",
check_interrupt=check_interrupt,
upload_cache=upload_cache,
log_success=log_task_success,
)
body = build_seedream_submit_body(
model=model,
prompt=prompt,
size=size,
output_format=output_format,
image_urls=image_urls,
layer_decomposition=layer_decomposition,
)
task_id = await submit_async_image_task(
session,
self.base_url,
self.api_key,
body,
"Seedream",
log_body_enabled=False,
log_success=log_task_success,
)
task_payload = await poll_async_image_task(
session,
self.base_url,
self.api_key,
task_id,
"Seedream",
check_interrupt=check_interrupt,
log_body_enabled=False,
progress_callback=progress_callback,
log_success=log_task_success,
)
task_done = time.time()
parse_started = time.time()
task_payload, parsed = await parse_completed_async_image_task(
task_payload,
session,
self.base_url,
self.api_key,
task_id,
"Seedream",
check_interrupt=check_interrupt,
result_url_callback=None,
log_downloads=log_downloads,
)
if isinstance(parsed, tuple) and len(parsed) == 2:
result_images, metrics = parsed
else:
result_images = parsed
metrics = {
"download_bytes": 0,
"download_seconds": 0.0,
"download_wall_seconds": 0.0,
"inline_images": 0,
}
result_urls = extract_async_image_result_urls(task_payload)
result_metadata = extract_seedream_layer_metadata(task_payload)
if layer_decomposition:
paired = []
for index, image in enumerate(result_images):
metadata = (
result_metadata[index]
if index < len(result_metadata)
else {"z_index": index}
)
setattr(image, "_o1key_seedream_layer", metadata)
paired.append((metadata.get("z_index", index), index, image))
paired.sort(key=lambda item: (item[0], item[1]))
result_images = [item[2] for item in paired]
result_metadata = [
getattr(image, "_o1key_seedream_layer", {"z_index": index})
for index, image in enumerate(result_images)
]
if result_url_callback:
for url in result_urls:
result_url_callback(url)
if task_completed_callback:
task_completed_callback(
task_id,
len(result_images),
time.time() - task_started,
result_urls,
)
return result_images, {
"task_id": task_id,
"task_ids": [task_id],
"task_ms": (task_done - task_started) * 1000,
"parse_ms": (time.time() - parse_started) * 1000,
"download_ms": metrics["download_wall_seconds"] * 1000,
"download_total_ms": metrics["download_seconds"] * 1000,
"download_bytes": metrics["download_bytes"],
"inline_images": metrics["inline_images"],
"result_metadata": result_metadata,
}
__all__ = [
"SEEDREAM_API_MODEL_ID",
"SEEDREAM_LAYER_REFERENCE_MIN_PIXELS",
"SEEDREAM_REFERENCE_MAX_BYTES",
"SEEDREAM_REFERENCE_MAX_PIXELS",
"SeedreamImageClient",
"build_seedream_submit_body",
"extract_seedream_layer_metadata",
"resolve_seedream_model",
"validate_seedream_reference_dimensions",
"validate_seedream_reference_image",
"validate_seedream_reference_images",
]
clients/sora_client.py
+13 -14
View File
@@ -15,6 +15,7 @@ import aiohttp
from .base_client import BaseAPIClient
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import encode_image_to_base64
from ..utils.video_task import PollDeadline, download_video_to_file
def _translate_error_message(msg: str) -> str:
@@ -178,9 +179,11 @@ class SoraClient(BaseAPIClient):
close_session = True
interval = self.POLL_INITIAL_INTERVAL
deadline = PollDeadline(label="Sora 视频")
try:
while True:
deadline.check()
async with session.get(url, headers=headers) as response:
if response.status != 200:
error_text = await response.text()
@@ -253,18 +256,20 @@ class SoraClient(BaseAPIClient):
raise RuntimeError(f"视频下载失败: {error_message}")
content_type = response.headers.get("Content-Type", "")
download_url = None
if "application/json" in content_type:
data = await response.json()
download_url = data.get("url") or data.get("download_url")
if not download_url:
raise RuntimeError("视频下载失败: 响应中未找到下载链接")
await self._download_from_url(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(8192):
f.write(chunk)
if download_url:
await self._download_from_url(download_url, save_path, session)
else:
# content 端点直接返回视频流（幂等 GET，可安全重连续传）
await download_video_to_file(
session, url, save_path, headers=headers, label="Sora 视频",
)
return save_path
@@ -494,13 +499,7 @@ class SoraClient(BaseAPIClient):
session: aiohttp.ClientSession,
) -> None:
"""从给定 URL 下载文件到本地路径"""
os.makedirs(os.path.dirname(save_path), exist_ok=True)
async with session.get(url) as response:
if response.status != 200:
raise RuntimeError(f"从下载链接获取视频失败 (状态码: {response.status})")
with open(save_path, "wb") as f:
async for chunk in response.content.iter_chunked(8192):
f.write(chunk)
await download_video_to_file(session, url, save_path, label="Sora 视频")
@staticmethod
def _extract_error_message(error_text: str, status_code: int) -> str:
clients/veo_client.py
+13 -14
View File
@@ -15,6 +15,7 @@ import aiohttp
from .base_client import BaseAPIClient
from ..utils.config import get_api_key_or_raise, get_api_base_url
from ..utils.image_utils import encode_image_to_base64
from ..utils.video_task import PollDeadline, download_video_to_file
class VeoClient(BaseAPIClient):
@@ -206,9 +207,11 @@ class VeoClient(BaseAPIClient):
close_session = True
interval = self.POLL_INITIAL_INTERVAL
deadline = PollDeadline(label="Veo 视频")
try:
while True:
deadline.check()
async with session.get(url, headers=headers) as response:
if response.status != 200:
error_text = await response.text()
@@ -276,18 +279,20 @@ class VeoClient(BaseAPIClient):
raise RuntimeError(f"视频下载失败: {error_message}")
content_type = response.headers.get("Content-Type", "")
download_url = None
if "application/json" in content_type:
data = await response.json()
download_url = data.get("url") or data.get("download_url")
if not download_url:
raise RuntimeError("视频下载失败: 响应中未找到下载链接")
await self._download_from_url(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(8192):
f.write(chunk)
if download_url:
await self._download_from_url(download_url, save_path, session)
else:
# content 端点直接返回视频流（幂等 GET，可安全重连续传）
await download_video_to_file(
session, url, save_path, headers=headers, label="VEO 视频",
)
return save_path
@@ -474,13 +479,7 @@ class VeoClient(BaseAPIClient):
session: aiohttp.ClientSession,
) -> None:
"""从给定 URL 下载文件到本地路径"""
os.makedirs(os.path.dirname(save_path), exist_ok=True)
async with session.get(url) as response:
if response.status != 200:
raise RuntimeError(f"从下载链接获取视频失败 (状态码: {response.status})")
with open(save_path, "wb") as f:
async for chunk in response.content.iter_chunked(8192):
f.write(chunk)
await download_video_to_file(session, url, save_path, label="VEO 视频")
@staticmethod
def _extract_error_message(error_text: str, status_code: int) -> str: