feat: sync latest local version as authoritative codebase
Complete rewrite/sync of comfyui_o1key custom nodes. Treat this commit as the current canonical version. Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
This commit is contained in:
o1key
@@ -0,0 +1,530 @@
|
||||
"""
|
||||
Sora 视频生成 API 客户端
|
||||
提供视频创建、状态轮询、视频下载功能
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _translate_error_message(msg: str) -> str:
|
||||
"""将 API 返回的已知英文错误信息翻译为中文友好提示"""
|
||||
if "people-in-user-uploads" in msg or (
|
||||
"moderation" in msg and "inputs" in msg
|
||||
):
|
||||
return "上传的参考图片中包含了真实人物【官方风控】，请尝试使用其他办法绕开。"
|
||||
return msg
|
||||
|
||||
|
||||
class SoraClient(BaseAPIClient):
|
||||
"""
|
||||
Sora 视频生成客户端
|
||||
|
||||
工作流程：
|
||||
1. create_video → POST /v1/videos (提交生成任务)
|
||||
2. poll_status → GET /v1/videos/{id} (轮询直到完成/失败)
|
||||
3. download_video→ GET /v1/videos/{id}/content (下载视频文件)
|
||||
"""
|
||||
|
||||
CREATE_ENDPOINT = "/v1/videos"
|
||||
STATUS_ENDPOINT = "/v1/videos/{video_id}"
|
||||
CONTENT_ENDPOINT = "/v1/videos/{video_id}/content"
|
||||
|
||||
POLL_INITIAL_INTERVAL = 3
|
||||
POLL_MAX_INTERVAL = 15
|
||||
|
||||
def __init__(self):
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_api_base_url()
|
||||
super().__init__(base_url=base_url, api_key=api_key)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# BaseAPIClient 抽象方法实现（本客户端主要使用自定义方法）
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
return self.CREATE_ENDPOINT
|
||||
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
return response
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 核心异步方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def create_video_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int = 4,
|
||||
size: str = "720x1280",
|
||||
input_reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
提交视频生成任务
|
||||
|
||||
格式策略（根据抓包确认）：
|
||||
- 无参考图片：application/json
|
||||
- 有参考图片：multipart/form-data，input_reference 以 PNG 文件上传
|
||||
|
||||
注意：seed 不被上游 API 接受，仅在 ComfyUI 节点侧用于缓存刷新
|
||||
|
||||
Returns:
|
||||
API 响应 JSON，包含 video id 和初始状态
|
||||
"""
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
# ============================================================
|
||||
# ⚠️ 已验证可用的标准请求方案，请勿随意修改！（2026-02-28）
|
||||
# ============================================================
|
||||
# 经多轮调试确认：
|
||||
# - 有图片：必须使用 multipart/form-data，input_reference 以 PNG 文件上传
|
||||
# · filename="reference.png", content_type="image/png"（与抓包一致）
|
||||
# · 不可改为 application/json + base64 → 400 "expected a file, got a string"
|
||||
# · 不可改为 application/json + data URI → 500 upstream error
|
||||
# · 不可改为 multipart + image/jpeg → 400 "Inpaint image must match..."（尺寸校验失败）
|
||||
# - 无图片：使用 application/json，已验证成功
|
||||
# ============================================================
|
||||
if input_reference_bytes:
|
||||
if len(input_reference_bytes) > self.max_request_size:
|
||||
raise ValueError(
|
||||
f"参考图片约 {len(input_reference_bytes) / 1024 / 1024:.1f}MB，"
|
||||
f"超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制，请使用较小的图片"
|
||||
)
|
||||
# ⚠️ 有图片：multipart/form-data + PNG 文件上传（唯一验证成功的方案）
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("model", model)
|
||||
form.add_field("seconds", str(seconds))
|
||||
form.add_field("size", size)
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
input_reference_bytes,
|
||||
filename="reference.png", # ⚠️ 不可改文件名/扩展名
|
||||
content_type="image/png", # ⚠️ 不可改为 image/jpeg
|
||||
)
|
||||
send_kwargs: Dict[str, Any] = {"data": form, "headers": headers}
|
||||
else:
|
||||
# ⚠️ 无图片：application/json（已验证成功）
|
||||
body: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"seconds": str(seconds),
|
||||
"size": size,
|
||||
}
|
||||
send_kwargs = {"json": body, "headers": headers}
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
async with session.post(url, **send_kwargs) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
error_message = self._extract_error_message(error_text, response.status)
|
||||
raise RuntimeError(error_message)
|
||||
|
||||
resp_json = await response.json()
|
||||
return resp_json
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def poll_video_status_async(
|
||||
self,
|
||||
video_id: str,
|
||||
progress_callback: Optional[Callable[[int, float], None]] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
轮询视频生成状态，直到完成或失败
|
||||
|
||||
Args:
|
||||
video_id: 视频任务 ID
|
||||
progress_callback: 进度回调 (progress_percent, elapsed_seconds)
|
||||
session: aiohttp 会话
|
||||
|
||||
Returns:
|
||||
最终状态的 API 响应
|
||||
|
||||
Raises:
|
||||
RuntimeError: 生成失败
|
||||
"""
|
||||
url = f"{self.base_url}{self.STATUS_ENDPOINT.format(video_id=video_id)}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
interval = self.POLL_INITIAL_INTERVAL
|
||||
|
||||
try:
|
||||
while True:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
error_message = self._extract_error_message(error_text, response.status)
|
||||
raise RuntimeError(error_message)
|
||||
|
||||
data = await response.json()
|
||||
|
||||
# status 兼容大小写：queued / in_progress / IN_PROGRESS / completed / COMPLETED
|
||||
status = data.get("status", "").lower()
|
||||
|
||||
# progress 兼容整数 (30) 和字符串 ("30%") 两种格式
|
||||
progress_raw = data.get("progress", 0)
|
||||
if isinstance(progress_raw, str):
|
||||
try:
|
||||
progress = int(progress_raw.rstrip("%").strip())
|
||||
except ValueError:
|
||||
progress = 0
|
||||
else:
|
||||
progress = int(progress_raw) if progress_raw else 0
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(progress)
|
||||
|
||||
if status == "completed":
|
||||
return data
|
||||
|
||||
if status == "failed":
|
||||
error_info = data.get("error", {})
|
||||
error_msg = error_info.get("message", "未知错误") if isinstance(error_info, dict) else str(error_info)
|
||||
error_msg = _translate_error_message(error_msg)
|
||||
raise RuntimeError(f"视频生成失败: {error_msg}")
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def download_video_async(
|
||||
self,
|
||||
video_id: str,
|
||||
save_path: str,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> str:
|
||||
"""
|
||||
下载生成的视频文件
|
||||
|
||||
处理两种情况：
|
||||
1. 响应为重定向或 JSON 含下载 URL → 跟随下载
|
||||
2. 响应为二进制视频流 → 直接保存
|
||||
|
||||
Returns:
|
||||
保存的文件路径
|
||||
"""
|
||||
url = f"{self.base_url}{self.CONTENT_ENDPOINT.format(video_id=video_id)}"
|
||||
headers = self.get_headers(use_bearer_token=True)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
async with session.get(url, headers=headers, allow_redirects=True) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
error_message = self._extract_error_message(error_text, response.status)
|
||||
raise RuntimeError(f"视频下载失败: {error_message}")
|
||||
|
||||
content_type = response.headers.get("Content-Type", "")
|
||||
|
||||
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)
|
||||
|
||||
return save_path
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 同步包装
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def generate_video_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int,
|
||||
size: str,
|
||||
save_path: str,
|
||||
input_reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
progress_callback: Optional[Callable[[int, float], None]] = None,
|
||||
on_stage: Optional[Callable[[str], None]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
同步执行完整的视频生成流程（创建 → 轮询 → 下载）
|
||||
|
||||
Args:
|
||||
on_stage: 阶段回调，用于打印状态切换信息
|
||||
|
||||
Returns:
|
||||
保存的视频文件路径
|
||||
"""
|
||||
|
||||
async def _run():
|
||||
connector = aiohttp.TCPConnector(limit=0)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
# 1. 提交任务
|
||||
if on_stage:
|
||||
on_stage("submitting")
|
||||
result = await self.create_video_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
input_reference_bytes=input_reference_bytes,
|
||||
seed=seed,
|
||||
session=session,
|
||||
)
|
||||
video_id = result.get("id")
|
||||
if not video_id:
|
||||
raise RuntimeError("API 未返回视频任务 ID")
|
||||
|
||||
if on_stage:
|
||||
on_stage(f"submitted:{video_id}")
|
||||
|
||||
# 2. 轮询状态
|
||||
if on_stage:
|
||||
on_stage("polling")
|
||||
await self.poll_video_status_async(
|
||||
video_id=video_id,
|
||||
progress_callback=progress_callback,
|
||||
session=session,
|
||||
)
|
||||
|
||||
# 3. 下载视频
|
||||
if on_stage:
|
||||
on_stage("downloading")
|
||||
path = await self.download_video_async(
|
||||
video_id=video_id,
|
||||
save_path=save_path,
|
||||
session=session,
|
||||
)
|
||||
|
||||
if on_stage:
|
||||
on_stage("done")
|
||||
return path
|
||||
|
||||
return self.run_async_in_thread(_run())
|
||||
|
||||
async def _generate_one_video_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int,
|
||||
size: str,
|
||||
save_path: str,
|
||||
input_reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> str:
|
||||
"""
|
||||
异步生成单个视频（创建 → 轮询 → 下载）
|
||||
|
||||
Returns:
|
||||
保存的视频文件路径
|
||||
"""
|
||||
result = await self.create_video_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
input_reference_bytes=input_reference_bytes,
|
||||
seed=seed,
|
||||
session=session,
|
||||
)
|
||||
video_id = result.get("id")
|
||||
if not video_id:
|
||||
raise RuntimeError("API 未返回视频任务 ID")
|
||||
|
||||
await self.poll_video_status_async(video_id=video_id, session=session)
|
||||
path = await self.download_video_async(
|
||||
video_id=video_id, save_path=save_path, session=session
|
||||
)
|
||||
return path
|
||||
|
||||
async def generate_batch_videos_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int,
|
||||
size: str,
|
||||
save_paths: List[str],
|
||||
input_reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
并发生成多个视频
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
seconds: 视频时长（秒）
|
||||
size: 分辨率
|
||||
save_paths: 各视频的保存路径列表，长度决定并发数量
|
||||
input_reference_bytes: 参考图片字节（可选）
|
||||
seed: 随机种子（仅节点侧使用）
|
||||
progress_callback: 进度回调 (current, total, success, error_msg)
|
||||
|
||||
Returns:
|
||||
成功生成的视频路径列表
|
||||
"""
|
||||
batch_size = len(save_paths)
|
||||
connector = aiohttp.TCPConnector(limit=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
tasks = [
|
||||
self._generate_one_video_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
save_path=save_paths[i],
|
||||
input_reference_bytes=input_reference_bytes,
|
||||
seed=seed,
|
||||
session=session,
|
||||
)
|
||||
for i in range(batch_size)
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
completed = 0
|
||||
paths: List[str] = []
|
||||
first_error = None
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
error_msg = str(result)
|
||||
print(f"Sora: 第 {i + 1} 个视频生成失败")
|
||||
print(f"原始错误详情:\n{error_msg}")
|
||||
if first_error is None:
|
||||
first_error = result
|
||||
if progress_callback:
|
||||
progress_callback(i + 1, batch_size, False, error_msg)
|
||||
else:
|
||||
completed += 1
|
||||
paths.append(result)
|
||||
if progress_callback:
|
||||
progress_callback(completed, batch_size, True, None)
|
||||
|
||||
if not paths:
|
||||
if first_error:
|
||||
raise first_error
|
||||
raise RuntimeError(f"批量视频生成失败，{batch_size} 个任务全部失败")
|
||||
|
||||
return paths
|
||||
|
||||
def generate_batch_videos_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int,
|
||||
size: str,
|
||||
save_paths: List[str],
|
||||
input_reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
同步并发生成多个视频（用于 ComfyUI 节点）
|
||||
|
||||
Args:
|
||||
save_paths: 各视频的保存路径列表，长度决定并发数量
|
||||
|
||||
Returns:
|
||||
成功生成的视频路径列表
|
||||
"""
|
||||
coro = self.generate_batch_videos_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
save_paths=save_paths,
|
||||
input_reference_bytes=input_reference_bytes,
|
||||
seed=seed,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 内部辅助方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _download_from_url(
|
||||
self,
|
||||
url: str,
|
||||
save_path: str,
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def _extract_error_message(error_text: str, status_code: int) -> str:
|
||||
"""从错误响应中提取可读的错误信息"""
|
||||
error_message = error_text
|
||||
try:
|
||||
error_json = json.loads(error_text)
|
||||
if "error" in error_json:
|
||||
if isinstance(error_json["error"], dict):
|
||||
error_message = error_json["error"].get("message", error_text)
|
||||
else:
|
||||
error_message = str(error_json["error"])
|
||||
elif "message" in error_json:
|
||||
error_message = error_json["message"]
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
pass
|
||||
|
||||
status_hints = {
|
||||
400: "请求参数错误 (400)",
|
||||
401: "认证失败 (401)，请检查 API 密钥",
|
||||
403: "权限不足 (403)，请检查账户权限或余额",
|
||||
429: "请求频率超限 (429)，请稍后重试",
|
||||
503: "服务暂时不可用 (503)，请稍后重试",
|
||||
504: "请求超时 (504)，请稍后重试",
|
||||
}
|
||||
hint = status_hints.get(status_code, f"API 请求失败 (状态码: {status_code})")
|
||||
return f"{hint}\nAPI 返回: {error_message}"
|
||||
Reference in New Issue
Block a user