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,14 @@
|
||||
"""
|
||||
API 客户端模块
|
||||
包含与外部 API 通信的客户端实现
|
||||
"""
|
||||
|
||||
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 .openai_client import OpenAIAPIClient
|
||||
|
||||
__all__ = ['BaseAPIClient', 'GeminiAPIClient', 'GeminiFlashClient', 'SoraClient', 'KlingClient', 'VeoClient', 'OpenAIAPIClient']
|
||||
@@ -0,0 +1,545 @@
|
||||
"""
|
||||
API 客户端基类
|
||||
提供通用的 HTTP 请求、响应解析和错误处理功能
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
|
||||
class BaseAPIClient(ABC):
|
||||
"""
|
||||
API 客户端抽象基类
|
||||
|
||||
子类需要实现以下方法：
|
||||
- get_endpoint(): 获取 API 端点
|
||||
- build_request_body(): 构建请求体
|
||||
- parse_response(): 解析响应
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
max_request_size: int = 100 * 1024 * 1024
|
||||
):
|
||||
"""
|
||||
初始化客户端
|
||||
|
||||
Args:
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
max_request_size: 最大请求体大小（字节），默认 100MB
|
||||
"""
|
||||
self.base_url = base_url
|
||||
self.api_key = api_key
|
||||
self.max_request_size = max_request_size
|
||||
|
||||
@abstractmethod
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
"""
|
||||
获取 API 端点路径
|
||||
|
||||
Args:
|
||||
**kwargs: 额外参数（如模型名、分辨率等）
|
||||
|
||||
Returns:
|
||||
端点路径字符串
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
构建 API 请求体
|
||||
|
||||
Args:
|
||||
**kwargs: 请求参数
|
||||
|
||||
Returns:
|
||||
请求体字典
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
"""
|
||||
解析 API 响应
|
||||
|
||||
Args:
|
||||
response: API 响应字典
|
||||
|
||||
Returns:
|
||||
解析后的结果
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_headers(self, use_bearer_token: bool = False) -> Dict[str, str]:
|
||||
"""
|
||||
获取请求头
|
||||
|
||||
Args:
|
||||
use_bearer_token: 是否使用 Bearer Token 认证（默认为 False）
|
||||
|
||||
Returns:
|
||||
请求头字典
|
||||
"""
|
||||
if use_bearer_token:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"x-goog-api-key": self.api_key,
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def check_request_size(self, request_body: Dict[str, Any]) -> None:
|
||||
"""
|
||||
检查请求体大小是否超过限制
|
||||
|
||||
Args:
|
||||
request_body: 请求体字典
|
||||
|
||||
Raises:
|
||||
ValueError: 如果请求体超过限制
|
||||
"""
|
||||
request_json = json.dumps(request_body)
|
||||
request_size = len(request_json.encode('utf-8'))
|
||||
|
||||
if request_size > self.max_request_size:
|
||||
raise ValueError(
|
||||
"请求体积超过100MB限制，请调整分辨率或减少图片数量"
|
||||
)
|
||||
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""
|
||||
子类可重写：为指定 HTTP 状态码返回自定义错误文案。
|
||||
若返回 None，则使用基类默认拼接文案。
|
||||
|
||||
Args:
|
||||
status_code: HTTP 状态码（如 429、503）
|
||||
error_message: API 返回的原始错误信息
|
||||
|
||||
Returns:
|
||||
自定义完整错误文案，或 None 表示使用默认
|
||||
"""
|
||||
return None
|
||||
|
||||
async def request_async(
|
||||
self,
|
||||
endpoint: str,
|
||||
request_body: Dict[str, Any],
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
use_bearer_token: bool = False,
|
||||
timeout: Optional[int] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
发送异步 HTTP 请求（带详细计时）
|
||||
|
||||
Args:
|
||||
endpoint: API 端点
|
||||
request_body: 请求体
|
||||
session: aiohttp 会话（可选）
|
||||
use_bearer_token: 是否使用 Bearer Token 认证
|
||||
timeout: 超时时间（秒）- 已废弃，由服务器端控制
|
||||
|
||||
Returns:
|
||||
响应 JSON
|
||||
|
||||
Raises:
|
||||
RuntimeError: 请求失败时
|
||||
"""
|
||||
import time
|
||||
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token)
|
||||
|
||||
# 检查请求大小
|
||||
self.check_request_size(request_body)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
# 连接计时
|
||||
connect_start = time.time()
|
||||
|
||||
async with session.post(url, json=request_body, headers=headers) as response:
|
||||
connect_time = time.time() - connect_start
|
||||
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
|
||||
# 尝试解析 JSON 错误信息，提取关键内容
|
||||
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，使用原始文本
|
||||
pass
|
||||
|
||||
# 针对常见错误状态码提供友好提示
|
||||
if response.status == 400:
|
||||
raise RuntimeError(
|
||||
f"请求参数错误 (400 Bad Request)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 检查 API 密钥是否有效\n"
|
||||
f" - 确认请求参数格式正确"
|
||||
)
|
||||
elif response.status == 401:
|
||||
raise RuntimeError(
|
||||
f"认证失败 (401 Unauthorized)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 检查 API 密钥是否正确\n"
|
||||
f" - 确认 API 密钥是否过期"
|
||||
)
|
||||
elif response.status == 403:
|
||||
raise RuntimeError(
|
||||
f"权限不足 (403 Forbidden)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 检查 API 密钥权限\n"
|
||||
f" - 确认账户余额充足"
|
||||
)
|
||||
elif response.status == 404:
|
||||
raise RuntimeError(
|
||||
f"端点不存在 (404 Not Found)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 检查模型名称是否正确\n"
|
||||
f" - 使用其他可用模型"
|
||||
)
|
||||
elif response.status == 429:
|
||||
custom = self.get_http_error_message(429, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(
|
||||
f"请求频率超限 (429 Too Many Requests)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 等待一段时间后重试\n"
|
||||
f" - 检查 API 配额是否充足"
|
||||
)
|
||||
elif response.status == 503:
|
||||
custom = self.get_http_error_message(503, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(
|
||||
f"服务暂时不可用 (503 Service Unavailable)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 稍后重试\n"
|
||||
f" - 尝试使用其他模型"
|
||||
)
|
||||
elif response.status == 504:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：\n"
|
||||
f" - 尝试使用其他模型\n"
|
||||
f" - 稍后重试\n"
|
||||
f" - 降低分辨率或减少输入图像数量"
|
||||
)
|
||||
elif response.status == 502:
|
||||
raise RuntimeError(
|
||||
"糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status})\n"
|
||||
f"API 返回错误：{error_message}"
|
||||
)
|
||||
|
||||
# 接收响应体
|
||||
wait_start = time.time()
|
||||
response_data = await response.json()
|
||||
download_time = time.time() - wait_start
|
||||
|
||||
# 附加计时信息到响应数据（供上层使用）
|
||||
response_size = len(str(response_data))
|
||||
if not isinstance(response_data, dict):
|
||||
response_data = {"data": response_data}
|
||||
|
||||
# 将计时信息存储在响应的元数据中
|
||||
response_data["_timing"] = {
|
||||
"connect_time": connect_time,
|
||||
"download_time": download_time,
|
||||
"response_size": response_size
|
||||
}
|
||||
|
||||
return response_data
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def request_get_async(
|
||||
self,
|
||||
endpoint: str,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
use_bearer_token: bool = True,
|
||||
timeout: Optional[int] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
发送异步 HTTP GET 请求
|
||||
|
||||
Args:
|
||||
endpoint: API 端点
|
||||
session: aiohttp 会话（可选）
|
||||
use_bearer_token: 是否使用 Bearer Token 认证（默认为 True）
|
||||
timeout: 超时时间（秒）- 已废弃，由服务器端控制
|
||||
|
||||
Returns:
|
||||
响应 JSON
|
||||
|
||||
Raises:
|
||||
RuntimeError: 请求失败时
|
||||
"""
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
|
||||
# 尝试解析 JSON 错误信息，提取关键内容
|
||||
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，使用原始文本
|
||||
pass
|
||||
|
||||
# 针对常见错误状态码提供友好提示
|
||||
if response.status == 400:
|
||||
raise RuntimeError(
|
||||
f"请求参数错误 (400 Bad Request)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：检查请求参数"
|
||||
)
|
||||
elif response.status == 401:
|
||||
raise RuntimeError(
|
||||
f"认证失败 (401 Unauthorized)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：检查 API 密钥"
|
||||
)
|
||||
elif response.status == 429:
|
||||
custom = self.get_http_error_message(429, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(
|
||||
f"请求频率超限 (429 Too Many Requests)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：等待一段时间后重试"
|
||||
)
|
||||
elif response.status == 503:
|
||||
custom = self.get_http_error_message(503, error_message)
|
||||
if custom is not None:
|
||||
raise RuntimeError(custom)
|
||||
raise RuntimeError(
|
||||
f"服务暂时不可用 (503 Service Unavailable)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：稍后重试"
|
||||
)
|
||||
elif response.status == 504:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：稍后重试"
|
||||
)
|
||||
elif response.status == 502:
|
||||
raise RuntimeError(
|
||||
"糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status})\n"
|
||||
f"API 返回错误：{error_message}"
|
||||
)
|
||||
|
||||
return await response.json()
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def batch_request_async(
|
||||
self,
|
||||
requests: List[Dict[str, Any]],
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None
|
||||
) -> List[Any]:
|
||||
"""
|
||||
批量并发请求
|
||||
|
||||
Args:
|
||||
requests: 请求列表，每个元素包含 endpoint 和 request_body
|
||||
progress_callback: 进度回调函数 (current, total)
|
||||
|
||||
Returns:
|
||||
响应结果列表
|
||||
"""
|
||||
results = []
|
||||
completed = 0
|
||||
total = len(requests)
|
||||
|
||||
# 创建无限制的连接器
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
tasks = []
|
||||
|
||||
for req in requests:
|
||||
task = self.request_async(
|
||||
endpoint=req['endpoint'],
|
||||
request_body=req['request_body'],
|
||||
session=session
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# 并发执行
|
||||
responses = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for i, resp in enumerate(responses):
|
||||
if isinstance(resp, Exception):
|
||||
print(f"⚠️ 第 {i+1} 个请求失败: {str(resp)}")
|
||||
continue
|
||||
|
||||
try:
|
||||
parsed = self.parse_response(resp)
|
||||
results.append(parsed)
|
||||
completed += 1
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(completed, total)
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ 第 {i+1} 个响应解析失败: {str(e)}")
|
||||
|
||||
return results
|
||||
|
||||
def run_async_in_thread(self, coro) -> Any:
|
||||
"""
|
||||
在独立线程中运行异步代码（用于 ComfyUI 同步接口）
|
||||
|
||||
Args:
|
||||
coro: 协程对象
|
||||
|
||||
Returns:
|
||||
协程执行结果
|
||||
"""
|
||||
result_container = []
|
||||
error_container = []
|
||||
|
||||
def run_in_thread():
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
try:
|
||||
result = loop.run_until_complete(coro)
|
||||
result_container.append(result)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
except Exception as e:
|
||||
error_container.append(e)
|
||||
|
||||
thread = threading.Thread(target=run_in_thread)
|
||||
thread.start()
|
||||
thread.join()
|
||||
|
||||
if error_container:
|
||||
raise error_container[0]
|
||||
|
||||
if not result_container:
|
||||
raise RuntimeError("异步任务未返回结果")
|
||||
|
||||
return result_container[0]
|
||||
|
||||
async def query_balance_async(self) -> Dict[str, Any]:
|
||||
"""
|
||||
异步查询账户余额
|
||||
|
||||
Returns:
|
||||
余额信息字典，包含 name、total_available 等字段
|
||||
|
||||
Raises:
|
||||
RuntimeError: 查询失败时
|
||||
"""
|
||||
endpoint = "/api/usage/token"
|
||||
response = await self.request_get_async(endpoint, use_bearer_token=True)
|
||||
|
||||
if not response.get("code"):
|
||||
raise RuntimeError("余额查询响应格式错误")
|
||||
|
||||
data = response.get("data", {})
|
||||
return data
|
||||
|
||||
def query_balance_sync(self) -> Dict[str, Any]:
|
||||
"""
|
||||
同步查询账户余额（用于 ComfyUI 节点）
|
||||
|
||||
Returns:
|
||||
余额信息字典
|
||||
|
||||
Raises:
|
||||
RuntimeError: 查询失败时
|
||||
"""
|
||||
coro = self.query_balance_async()
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
def format_balance_info(self, balance_data: Dict[str, Any]) -> str:
|
||||
"""
|
||||
格式化余额信息为展示文本
|
||||
|
||||
Args:
|
||||
balance_data: 余额信息字典
|
||||
|
||||
Returns:
|
||||
格式化文本，如 "当前余额：100.00 | API：xxx"
|
||||
|
||||
Example:
|
||||
>>> data = {"name": "test-api", "total_available": 50000000}
|
||||
>>> client.format_balance_info(data)
|
||||
'当前余额：100.00 | API：test-api'
|
||||
"""
|
||||
api_name = balance_data.get("name", "未知")
|
||||
total_available = balance_data.get("total_available", 0)
|
||||
|
||||
# 实际显示余额 = total_available / 500000，单位：美元
|
||||
balance_in_dollars = total_available / 500000
|
||||
|
||||
return f"当前余额：{balance_in_dollars:.2f} | API：{api_name}"
|
||||
@@ -0,0 +1,198 @@
|
||||
"""
|
||||
Flux 图像编辑 API 客户端
|
||||
通过 vip.o1key.com 调用 Flux2 图像编辑 + SeedVR2 超分辨率服务
|
||||
|
||||
工作流程：
|
||||
1. submit_task → POST /v1/images/edits (multipart/form-data 提交主图+参考图+提示词)
|
||||
2. poll_result → GET /v1/images/edits/{task_id} (直连容器轮询)
|
||||
"""
|
||||
|
||||
import base64
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
|
||||
|
||||
# 显示名 → 实际请求值的映射
|
||||
SIZE_DISPLAY_MAP = {
|
||||
"2K": "2048",
|
||||
"4K": "4096",
|
||||
}
|
||||
|
||||
# 轮询直连容器地址，绕过代理层
|
||||
POLL_BASE_URL = "https://xrrh7tn08tfgwa8w-8188.container.x-gpu.com"
|
||||
|
||||
|
||||
class FluxEditClient:
|
||||
"""
|
||||
Flux 图像编辑客户端
|
||||
|
||||
对接 vip.o1key.com 上的 /v1/images/edits 接口，
|
||||
将图像编辑+超分辨率任务提交到远程服务器执行。
|
||||
"""
|
||||
|
||||
SUBMIT_ENDPOINT = "/v1/images/edits"
|
||||
STATUS_ENDPOINT = "/v1/images/edits/{task_id}"
|
||||
|
||||
DEFAULT_POLL_INTERVAL = 15 # 秒
|
||||
|
||||
def __init__(self):
|
||||
self.api_key = get_api_key_or_raise()
|
||||
self.base_url = get_api_base_url()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 同步方法（供 ComfyUI 节点调用）
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def submit_and_wait(
|
||||
self,
|
||||
image_bytes: bytes,
|
||||
mask_bytes: bytes,
|
||||
prompt: str,
|
||||
size: str = "4K",
|
||||
poll_interval: int = DEFAULT_POLL_INTERVAL,
|
||||
progress_callback=None,
|
||||
) -> bytes:
|
||||
"""
|
||||
提交任务并同步等待结果（阻塞直到完成）
|
||||
|
||||
Args:
|
||||
image_bytes: 主图二进制数据
|
||||
mask_bytes: 参考图二进制数据
|
||||
prompt: 编辑提示词
|
||||
size: 分辨率显示名 ("2K" 或 "4K")
|
||||
poll_interval: 轮询间隔（秒）
|
||||
progress_callback: 进度回调 fn(status_str)
|
||||
|
||||
Returns:
|
||||
结果图像的二进制数据
|
||||
|
||||
Raises:
|
||||
RuntimeError: 任务失败
|
||||
"""
|
||||
size_value = SIZE_DISPLAY_MAP.get(size, size)
|
||||
|
||||
# 1. 提交任务（走代理）
|
||||
task_id = self._submit_task_sync(image_bytes, mask_bytes, prompt, size_value)
|
||||
if progress_callback:
|
||||
progress_callback(f"任务已提交: {task_id[:8]}...")
|
||||
|
||||
# 2. 轮询等待（直连容器）
|
||||
return self._poll_result_sync(
|
||||
task_id, poll_interval, progress_callback
|
||||
)
|
||||
|
||||
def _submit_task_sync(
|
||||
self,
|
||||
image_bytes: bytes,
|
||||
mask_bytes: bytes,
|
||||
prompt: str,
|
||||
size: str,
|
||||
) -> str:
|
||||
"""同步提交任务，返回 task_id"""
|
||||
url = f"{self.base_url}{self.SUBMIT_ENDPOINT}"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
files = {
|
||||
"image": ("image.jpg", image_bytes, "image/jpeg"),
|
||||
"mask": ("mask.jpg", mask_bytes, "image/jpeg"),
|
||||
}
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
"model": "flux2-fp8-dualr",
|
||||
}
|
||||
|
||||
try:
|
||||
resp = requests.post(url, files=files, data=data, headers=headers, timeout=60)
|
||||
except requests.exceptions.Timeout:
|
||||
raise RuntimeError("提交任务超时，请检查网络连接")
|
||||
except requests.exceptions.ConnectionError:
|
||||
raise RuntimeError("无法连接到服务器，请检查网络或服务器地址")
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(
|
||||
f"提交任务失败 (HTTP {resp.status_code})\n"
|
||||
f"响应: {resp.text[:500]}"
|
||||
)
|
||||
|
||||
result = resp.json()
|
||||
task_id = result.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"服务器返回异常: 未获取到任务ID\n{result}")
|
||||
|
||||
return task_id
|
||||
|
||||
def _poll_result_sync(
|
||||
self,
|
||||
task_id: str,
|
||||
poll_interval: int,
|
||||
progress_callback=None,
|
||||
) -> bytes:
|
||||
"""同步轮询任务状态（直连容器），返回结果图像二进制"""
|
||||
url = f"{POLL_BASE_URL}{self.STATUS_ENDPOINT.format(task_id=task_id)}"
|
||||
|
||||
start_time = time.time()
|
||||
last_status = None
|
||||
|
||||
while True:
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
try:
|
||||
resp = requests.get(url, timeout=30)
|
||||
except requests.exceptions.ConnectionError:
|
||||
raise RuntimeError("轮询时无法连接到服务器，请检查网络")
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise RuntimeError(
|
||||
f"查询任务状态失败 (HTTP {resp.status_code})\n"
|
||||
f"响应: {resp.text[:500]}"
|
||||
)
|
||||
|
||||
result = resp.json()
|
||||
status = result.get("status", "unknown")
|
||||
|
||||
# 状态变化时打印日志
|
||||
if status != last_status:
|
||||
elapsed_str = f"{elapsed:.0f}s"
|
||||
print(f"Flux Edit: [{elapsed_str}] 任务 {task_id[:8]}... → {status}")
|
||||
last_status = status
|
||||
|
||||
if progress_callback:
|
||||
elapsed_str = f"{elapsed:.0f}s"
|
||||
status_desc = {
|
||||
"pending": "排队中",
|
||||
"processing": "处理中",
|
||||
"generating": "生图中，请耐心等待，预计耗时140s左右",
|
||||
}.get(status, status)
|
||||
progress_callback(f"{status_desc} (当前进度：{elapsed_str})")
|
||||
|
||||
if status == "completed":
|
||||
# 解码 base64 图像
|
||||
b64_data = result.get("result")
|
||||
if not b64_data:
|
||||
raise RuntimeError("任务完成但未返回图像数据")
|
||||
return base64.b64decode(b64_data)
|
||||
|
||||
elif status == "failed":
|
||||
error_msg = result.get("error", "未知错误")
|
||||
raise RuntimeError(
|
||||
f"图像编辑任务失败\n"
|
||||
f"错误: {error_msg}"
|
||||
)
|
||||
|
||||
elif status in ("not_found",):
|
||||
raise RuntimeError(
|
||||
f"任务未找到: {task_id}\n"
|
||||
f"可能已被清理或 ID 无效"
|
||||
)
|
||||
|
||||
# 继续等待
|
||||
time.sleep(poll_interval)
|
||||
|
||||
def query_balance_sync(self) -> dict:
|
||||
"""查询余额（兼容现有节点的 finally 块调用）"""
|
||||
return {"name": "flux-edit", "total_available": 0}
|
||||
@@ -0,0 +1,986 @@
|
||||
"""
|
||||
Gemini API 客户端
|
||||
处理与 api.o1key.com 的通信，用于图像生成
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
|
||||
class GeminiAPIClient(BaseAPIClient):
|
||||
"""
|
||||
Gemini API 客户端
|
||||
用于调用 Gemini 3 Pro 模型进行图像生成
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
"""
|
||||
初始化客户端
|
||||
|
||||
Args:
|
||||
api_key: API 密钥，如果为 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
|
||||
)
|
||||
|
||||
def get_endpoint(self, model: str = "", resolution: str = "2K", **kwargs) -> str:
|
||||
"""
|
||||
根据模型和分辨率获取 API 端点
|
||||
|
||||
Args:
|
||||
model: 模型名称
|
||||
resolution: 分辨率（1K, 2K, 4K）
|
||||
|
||||
Returns:
|
||||
API 端点路径
|
||||
"""
|
||||
from ..models_config import get_model_endpoint
|
||||
|
||||
# 特殊处理：动态端点模型（根据分辨率选择）
|
||||
if model == "nano-banana-pro-限时特价":
|
||||
if resolution == "1K":
|
||||
return "/v1beta/models/nano-banana-pro:generateContent"
|
||||
elif resolution == "2K":
|
||||
return "/v1beta/models/nano-banana-pro-2k:generateContent"
|
||||
elif resolution == "4K":
|
||||
return "/v1beta/models/nano-banana-pro-4k:generateContent"
|
||||
else:
|
||||
return "/v1beta/models/nano-banana-pro-2k:generateContent"
|
||||
|
||||
elif model == "nano-banana-2-官方计费":
|
||||
if resolution == "512":
|
||||
return "/v1beta/models/nano-banana-2-0.5k-official:generateContent"
|
||||
elif resolution == "1K":
|
||||
return "/v1beta/models/nano-banana-2-1k-official:generateContent"
|
||||
elif resolution == "2K":
|
||||
return "/v1beta/models/nano-banana-2-2k-official:generateContent"
|
||||
elif resolution == "4K":
|
||||
return "/v1beta/models/nano-banana-2-4k-official:generateContent"
|
||||
else:
|
||||
return "/v1beta/models/nano-banana-2-2k-official:generateContent"
|
||||
|
||||
elif model == "nano-banana-pro-官方计费":
|
||||
if resolution == "1K":
|
||||
return "/v1beta/models/nano-banana-pro-1k-official:generateContent"
|
||||
elif resolution == "2K":
|
||||
return "/v1beta/models/nano-banana-pro-2k-official:generateContent"
|
||||
elif resolution == "4K":
|
||||
return "/v1beta/models/nano-banana-pro-4k-official:generateContent"
|
||||
else:
|
||||
return "/v1beta/models/nano-banana-pro-2k-official:generateContent"
|
||||
|
||||
elif model == "gemini-3-pro-image-preview-url":
|
||||
if resolution == "1K":
|
||||
return "/v1beta/models/gemini-3-pro-image-preview-url:generateContent"
|
||||
elif resolution == "2K":
|
||||
return "/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent"
|
||||
elif resolution == "4K":
|
||||
return "/v1beta/models/gemini-3-pro-image-preview-4k-url:generateContent"
|
||||
else:
|
||||
return "/v1beta/models/gemini-3-pro-image-preview-2k-url:generateContent"
|
||||
|
||||
# 其他模型：从配置文件读取端点
|
||||
endpoint = get_model_endpoint(model)
|
||||
if endpoint:
|
||||
return endpoint
|
||||
|
||||
# 兜底：使用标准模式端点
|
||||
return "/v1beta/models/gemini-3-pro-image-preview:generateContent"
|
||||
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""Gemini 请求 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
|
||||
|
||||
def build_request_body(
|
||||
self,
|
||||
prompt: str = "",
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
aspect_ratio: str = "1:1",
|
||||
resolution: str = "2K",
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
candidate_count: int = 1,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
构建 API 请求体
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
images: 输入图像列表（可选）
|
||||
aspect_ratio: 宽高比
|
||||
resolution: 分辨率
|
||||
enable_grounding: 是否启用 Google Search Grounding
|
||||
enable_image_search: 是否同时启用 Google Image Search（仅 Gemini 3.1 Flash 支持）
|
||||
candidate_count: 单次请求返回的候选图数量，默认 1
|
||||
|
||||
Returns:
|
||||
请求体字典
|
||||
"""
|
||||
parts = []
|
||||
|
||||
# 添加文本部分
|
||||
parts.append({"text": prompt})
|
||||
|
||||
# 添加图像部分（如果有）
|
||||
if images:
|
||||
for img in images:
|
||||
img_base64 = encode_image_to_base64(img)
|
||||
parts.append({
|
||||
"inline_data": {
|
||||
"mime_type": "image/png",
|
||||
"data": img_base64
|
||||
}
|
||||
})
|
||||
|
||||
# 构建请求体
|
||||
request_body = {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": parts
|
||||
}
|
||||
],
|
||||
"generationConfig": {
|
||||
"candidateCount": candidate_count,
|
||||
"responseModalities": ["TEXT", "IMAGE"],
|
||||
"imageConfig": {
|
||||
"aspectRatio": aspect_ratio,
|
||||
"imageSize": resolution
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# 添加 Google Search Grounding 工具（如果启用）
|
||||
# 注意：enable_image_search=True 时会自动隐含 enable_grounding
|
||||
if enable_grounding or enable_image_search:
|
||||
if enable_image_search:
|
||||
# 同时启用网页搜索和图片搜索（仅 nano-banana-2 / gemini-3.1-flash-image-preview 支持）
|
||||
request_body["tools"] = [
|
||||
{
|
||||
"google_search": {
|
||||
"searchTypes": {
|
||||
"webSearch": {},
|
||||
"imageSearch": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
else:
|
||||
# 仅启用网页搜索（通用）
|
||||
request_body["tools"] = [{"google_search": {}}]
|
||||
|
||||
return request_body
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> List[Image.Image]:
|
||||
"""
|
||||
同步解析 API 响应（保留以满足抽象基类要求）
|
||||
|
||||
注意：此方法仅用于兼容基类接口，实际使用请调用 parse_response_async()
|
||||
|
||||
Args:
|
||||
response: API 响应字典
|
||||
|
||||
Returns:
|
||||
图像列表
|
||||
|
||||
Raises:
|
||||
RuntimeError: 此方法不应被直接调用
|
||||
"""
|
||||
raise RuntimeError(
|
||||
"parse_response() 不应被直接调用。"
|
||||
"请使用 generate_single_async() 或 generate_batch_async() 等高级方法。"
|
||||
)
|
||||
|
||||
async def parse_response_async(
|
||||
self,
|
||||
response: Dict[str, Any],
|
||||
session: Optional[aiohttp.ClientSession] = None
|
||||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||||
"""
|
||||
异步解析 API 响应，提取生成的图像
|
||||
|
||||
Args:
|
||||
response: API 响应字典
|
||||
session: aiohttp 会话（用于下载图片）
|
||||
|
||||
Returns:
|
||||
(图像列表, 格式信息字典)
|
||||
格式信息包含: type (base64/url), size, resolution, download_speed (仅URL)
|
||||
|
||||
Raises:
|
||||
RuntimeError: 解析失败或 API 拒绝时
|
||||
"""
|
||||
|
||||
# 初始化格式信息
|
||||
format_info = {
|
||||
"type": None, # "base64" or "url"
|
||||
"size": 0,
|
||||
"resolution": None,
|
||||
"download_speed": None
|
||||
}
|
||||
|
||||
candidates = response.get("candidates", [])
|
||||
|
||||
# ========== 错误检测（按优先级顺序）==========
|
||||
|
||||
# 1. 检查 candidatesTokenCount（最高优先级）
|
||||
usage_metadata = response.get("usageMetadata", {})
|
||||
candidates_token_count = usage_metadata.get("candidatesTokenCount", -1)
|
||||
|
||||
if candidates_token_count == 0:
|
||||
error_msg = (
|
||||
"Damn!你触发顶级风控啦!还没到生图阶段就被拒了。\n"
|
||||
"赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~"
|
||||
)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
# 2. 检查 finishReason（次优先级）
|
||||
candidates = response.get("candidates", [])
|
||||
if candidates:
|
||||
for candidate in candidates:
|
||||
finish_reason = candidate.get("finishReason", "")
|
||||
|
||||
if finish_reason and finish_reason != "STOP":
|
||||
error_msg = (
|
||||
"Ohh no! 生图过程触发风控,图片被拒绝生成!\n"
|
||||
"可能原因如下:\n"
|
||||
"1.违禁内容\n"
|
||||
"2.触发安全过滤器\n"
|
||||
"3.涉及版权问题\n"
|
||||
"4. Token超限\n"
|
||||
"赶紧调整一下图片或提示词吧!该情况不会返回图片且正常扣费!下次小心哦~"
|
||||
)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
# ========== 图像提取 ==========
|
||||
|
||||
images = []
|
||||
text_responses = [] # 收集文本响应
|
||||
|
||||
# 需要关闭 session 的标记
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
for candidate_idx, candidate in enumerate(candidates):
|
||||
content = candidate.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
|
||||
for part_idx, part in enumerate(parts):
|
||||
# 方式1: inline_data 或 inlineData (base64)
|
||||
# 兼容两种命名方式：蛇形（inline_data）和驼峰（inlineData）
|
||||
inline_data_key = None
|
||||
if "inline_data" in part:
|
||||
inline_data_key = "inline_data"
|
||||
elif "inlineData" in part:
|
||||
inline_data_key = "inlineData"
|
||||
|
||||
if inline_data_key:
|
||||
inline_data = part[inline_data_key]
|
||||
# 同样兼容 data/mimeType 的命名
|
||||
img_data = inline_data.get("data") or inline_data.get("data", "")
|
||||
|
||||
if img_data:
|
||||
img = decode_base64_to_pil(img_data)
|
||||
images.append(img)
|
||||
|
||||
# 记录格式信息
|
||||
if format_info["type"] is None:
|
||||
format_info["type"] = "base64"
|
||||
format_info["size"] = len(img_data) * 3 / 4 # Base64 解码后的字节数
|
||||
format_info["resolution"] = f"{img.size[0]}x{img.size[1]}"
|
||||
|
||||
# 方式2: text 中的 URL - 改为异步下载
|
||||
elif "text" in part:
|
||||
text = part["text"]
|
||||
|
||||
# 收集文本响应（用于后续错误检测）
|
||||
text_responses.append(text)
|
||||
|
||||
# 尝试 markdown 格式: 
|
||||
url_pattern_md = r'!\[.*?\]\((https?://[^\)]+)\)'
|
||||
urls = re.findall(url_pattern_md, text)
|
||||
|
||||
# 如果没找到，尝试纯 URL 格式
|
||||
if not urls:
|
||||
url_pattern_plain = r'https?://[^\s<>"{}|\\^`\[\]]+'
|
||||
urls = re.findall(url_pattern_plain, text)
|
||||
|
||||
if urls:
|
||||
for url_idx, url in enumerate(urls):
|
||||
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()
|
||||
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 # 静默失败
|
||||
|
||||
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:
|
||||
# API 返回了文本但没有图片，说明请求被拒绝
|
||||
combined_text = "\n".join(text_responses)
|
||||
error_msg = (
|
||||
f"API 拒绝响应\n\n"
|
||||
f"API 返回说明：\n{combined_text}\n\n"
|
||||
f"建议：\n"
|
||||
f" - 根据上述说明调整请求内容\n"
|
||||
f" - 确保提示词和参考图符合使用规范"
|
||||
)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
if not images:
|
||||
raise RuntimeError("API 响应中未找到生成的图像")
|
||||
|
||||
return images, format_info
|
||||
|
||||
async def generate_single_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
session=None,
|
||||
task_index: Optional[int] = None,
|
||||
total_tasks: Optional[int] = None,
|
||||
debug: bool = False,
|
||||
debug_request: bool = False,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
candidate_count: int = 1
|
||||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||||
"""
|
||||
单次异步生成请求（极简单行日志）
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
images: 输入图像列表
|
||||
session: aiohttp 会话
|
||||
task_index: 任务索引（用于批量任务）
|
||||
total_tasks: 总任务数（用于批量任务）
|
||||
debug: 是否打印完整 API 响应
|
||||
debug_request: 是否打印发送的请求体（base64 图片数据将被截断）
|
||||
enable_grounding: 是否启用 Google Search Grounding
|
||||
enable_image_search: 是否同时启用 Google Image Search
|
||||
|
||||
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()
|
||||
endpoint = self.get_endpoint(model=model, resolution=resolution)
|
||||
request_body = self.build_request_body(
|
||||
prompt=prompt,
|
||||
images=images,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
candidate_count=candidate_count
|
||||
)
|
||||
build_time = time.time() - build_start
|
||||
|
||||
# ========== 调试日志：打印请求体 ==========
|
||||
if debug_request:
|
||||
import json as _json
|
||||
|
||||
def _truncate_base64_req(obj, max_len=200):
|
||||
if isinstance(obj, dict):
|
||||
return {k: _truncate_base64_req(v, max_len) for k, v in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [_truncate_base64_req(item, max_len) for item in obj]
|
||||
elif isinstance(obj, str) and len(obj) > max_len:
|
||||
if all(c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' for c in obj[:50]):
|
||||
return f"<base64 data, {len(obj)} chars>"
|
||||
return obj
|
||||
return obj
|
||||
|
||||
safe_request = _truncate_base64_req(request_body)
|
||||
print(
|
||||
f"\n{'='*60}\n"
|
||||
f"[请求体日志] 任务 {task_prefix or '?'} 发送请求体：\n"
|
||||
f"端点: {endpoint}\n"
|
||||
f"{_json.dumps(safe_request, ensure_ascii=False, indent=2)}\n"
|
||||
f"{'='*60}\n"
|
||||
)
|
||||
|
||||
# 计算请求体大小
|
||||
request_size = len(json.dumps(request_body).encode('utf-8'))
|
||||
if request_size < 1024 * 1024:
|
||||
size_str = f"{request_size / 1024:.2f}KB"
|
||||
else:
|
||||
size_str = f"{request_size / (1024 * 1024):.2f}MB"
|
||||
|
||||
# ========== 2. 发送网络请求 ==========
|
||||
request_start = time.time()
|
||||
|
||||
try:
|
||||
response = await self.request_async(endpoint, request_body, session)
|
||||
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
|
||||
# 构建可安全序列化的响应副本（截断 base64 图片数据避免输出过长）
|
||||
def _truncate_base64(obj, max_len=200):
|
||||
if isinstance(obj, dict):
|
||||
return {k: _truncate_base64(v, max_len) for k, v in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [_truncate_base64(item, max_len) for item in obj]
|
||||
elif isinstance(obj, str) and len(obj) > max_len:
|
||||
# 判断是否为 base64 图片数据（不含空格/换行的长字符串）
|
||||
if all(c in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=' for c in obj[:50]):
|
||||
return f"<base64 data, {len(obj)} chars>"
|
||||
return obj
|
||||
return obj
|
||||
|
||||
safe_response = _truncate_base64(response)
|
||||
print(
|
||||
f"\n{'='*60}\n"
|
||||
f"[调试日志] 任务 {task_prefix or '?'} 完整 API 响应：\n"
|
||||
f"{_json.dumps(safe_response, ensure_ascii=False, indent=2)}\n"
|
||||
f"{'='*60}\n"
|
||||
)
|
||||
|
||||
# ========== 3. 解析响应 ==========
|
||||
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
|
||||
|
||||
# ========== 4. 格式化输出（单行） ==========
|
||||
# 格式化图像大小
|
||||
img_size = format_info.get("size", 0)
|
||||
if img_size < 1024 * 1024:
|
||||
img_size_str = f"{img_size / 1024:.2f}KB"
|
||||
else:
|
||||
img_size_str = 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)
|
||||
speed_str = f"{speed / (1024 * 1024):.1f}MB/s"
|
||||
download_info = f"URL {img_size_str} ({parse_time:.1f}s, {speed_str})"
|
||||
else:
|
||||
download_info = f"{img_size_str}"
|
||||
|
||||
# 单行输出
|
||||
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓")
|
||||
|
||||
# 返回结果和计时信息
|
||||
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
|
||||
|
||||
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,
|
||||
candidate_count: int = 1
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
批量全并发生成 - 改进版：支持分批处理和内存管理
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
batch_size: 批次大小
|
||||
images: 输入图像列表
|
||||
progress_callback: 进度回调，签名为 (completed, total, success, error_msg)
|
||||
debug: 是否打印完整 API 响应
|
||||
debug_request: 是否打印发送的请求体
|
||||
enable_grounding: 是否启用 Google Search Grounding
|
||||
enable_image_search: 是否同时启用 Google Image Search
|
||||
candidate_count: 单次请求返回的候选图数量
|
||||
|
||||
Returns:
|
||||
生成的图像列表
|
||||
"""
|
||||
import aiohttp
|
||||
import asyncio
|
||||
|
||||
all_images = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
first_error = None # 保存第一个错误
|
||||
|
||||
# 分批处理配置
|
||||
max_concurrent = 10 # 最大并发数
|
||||
save_batch_size = 10 # 分批保存大小
|
||||
|
||||
# 计算需要多少批次
|
||||
num_batches = (batch_size + max_concurrent - 1) // max_concurrent
|
||||
|
||||
print(f"GeminiClient: 批量生成 {batch_size} 张图片，并发数: {max_concurrent}，分 {num_batches} 批执行")
|
||||
|
||||
connector = aiohttp.TCPConnector(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_size_current = batch_end - batch_start
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"GeminiClient: 执行第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})...")
|
||||
|
||||
# 创建当前批次的任务
|
||||
tasks = []
|
||||
for i in range(batch_size_current):
|
||||
task_index = batch_start + i
|
||||
task = asyncio.create_task(
|
||||
self.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images,
|
||||
session=session,
|
||||
task_index=task_index + 1,
|
||||
total_tasks=batch_size,
|
||||
debug=debug,
|
||||
debug_request=debug_request,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
candidate_count=candidate_count
|
||||
),
|
||||
name=f"task_{task_index}"
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# 收集当前批次的结果
|
||||
batch_images = []
|
||||
batch_completed = 0
|
||||
|
||||
for coro in asyncio.as_completed(tasks):
|
||||
batch_completed += 1
|
||||
completed += 1
|
||||
|
||||
try:
|
||||
result_images, timing_info = await coro
|
||||
if result_images:
|
||||
# 立即处理生成的图片
|
||||
for img in result_images:
|
||||
batch_images.append(img)
|
||||
all_images.append(img)
|
||||
|
||||
success_count += 1
|
||||
|
||||
# 通知进度
|
||||
if progress_callback:
|
||||
progress_callback(completed, batch_size, True, None)
|
||||
|
||||
# 每成功生成一张图片就打印日志
|
||||
print(f"GeminiClient: 任务 {completed}/{batch_size} 成功生成图片 ✓")
|
||||
|
||||
except Exception as e:
|
||||
fail_count += 1
|
||||
# 保存第一个错误（用于后续抛出）
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
error_msg = str(e)
|
||||
|
||||
# 传递完整的错误信息（用于排查问题）
|
||||
if progress_callback:
|
||||
progress_callback(completed, batch_size, False, error_msg)
|
||||
|
||||
print(f"GeminiClient: 任务 {completed}/{batch_size} 失败 ✗")
|
||||
|
||||
# 当前批次完成后，立即清理内存
|
||||
if batch_images:
|
||||
print(f"GeminiClient: 第 {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"GeminiClient: 批量生成完成，成功 {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], None]] = None,
|
||||
debug: bool = False,
|
||||
debug_request: bool = False,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
candidate_count: int = 1
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
同步生成接口（用于 ComfyUI）
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
batch_size: 批次大小
|
||||
images: 输入图像列表
|
||||
progress_callback: 进度回调
|
||||
debug: 是否打印完整 API 响应
|
||||
debug_request: 是否打印发送的请求体
|
||||
enable_grounding: 是否启用 Google Search Grounding
|
||||
enable_image_search: 是否同时启用 Google Image Search
|
||||
candidate_count: 单次请求返回的候选图数量
|
||||
|
||||
Returns:
|
||||
生成的图像列表
|
||||
"""
|
||||
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,
|
||||
candidate_count=candidate_count
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
async def generate_multi_prompts_async(
|
||||
self,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: 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]:
|
||||
"""
|
||||
多提示词批量生成 - 改进版：支持分批处理和内存管理
|
||||
|
||||
为每个提示词生成指定数量的图像，分批并发执行。
|
||||
|
||||
Args:
|
||||
prompts: 提示词列表
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
images_per_prompt: 每个提示词生成的图像数量
|
||||
images: 输入图像列表（所有提示词共享）
|
||||
progress_callback: 进度回调，签名为 (completed, total, success, error_msg)
|
||||
debug: 是否打印完整 API 响应
|
||||
debug_request: 是否打印发送的请求体
|
||||
enable_grounding: 是否启用 Google Search Grounding
|
||||
enable_image_search: 是否同时启用 Google Image Search
|
||||
|
||||
Returns:
|
||||
生成的图像列表（长度 = len(prompts) * images_per_prompt）
|
||||
"""
|
||||
import aiohttp
|
||||
import asyncio
|
||||
|
||||
all_images = []
|
||||
completed = 0
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
first_error = None # 保存第一个错误
|
||||
total_tasks = len(prompts) * images_per_prompt
|
||||
|
||||
# 分批处理配置
|
||||
max_concurrent = 10 # 最大并发数
|
||||
|
||||
print(f"GeminiClient: 多提示词批量生成，共 {total_tasks} 个任务，{len(prompts)} 个提示词，每个 {images_per_prompt} 张")
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
# 创建所有任务
|
||||
tasks = []
|
||||
task_idx = 0
|
||||
for prompt in prompts:
|
||||
for _ in range(images_per_prompt):
|
||||
task = asyncio.create_task(
|
||||
self.generate_single_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images=images,
|
||||
session=session,
|
||||
task_index=task_idx + 1,
|
||||
total_tasks=total_tasks,
|
||||
debug=debug,
|
||||
debug_request=debug_request,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search
|
||||
),
|
||||
name=f"task_{task_idx}"
|
||||
)
|
||||
tasks.append(task)
|
||||
task_idx += 1
|
||||
|
||||
# 分批处理：每10个任务为一组
|
||||
batch_size = max_concurrent
|
||||
num_batches = (total_tasks + batch_size - 1) // batch_size
|
||||
|
||||
for batch_idx in range(num_batches):
|
||||
batch_start = batch_idx * batch_size
|
||||
batch_end = min(batch_start + batch_size, total_tasks)
|
||||
batch_tasks = tasks[batch_start:batch_end]
|
||||
|
||||
if num_batches > 1:
|
||||
print(f"GeminiClient: 执行第 {batch_idx + 1}/{num_batches} 批 ({batch_start + 1}-{batch_end})...")
|
||||
|
||||
# 收集当前批次的结果
|
||||
batch_images = []
|
||||
|
||||
for coro in asyncio.as_completed(batch_tasks):
|
||||
completed += 1
|
||||
|
||||
try:
|
||||
result_images, timing_info = await coro
|
||||
if result_images:
|
||||
# 立即处理生成的图片
|
||||
for img in result_images:
|
||||
batch_images.append(img)
|
||||
all_images.append(img)
|
||||
|
||||
success_count += 1
|
||||
|
||||
# 通知进度
|
||||
if progress_callback:
|
||||
progress_callback(completed, total_tasks, True, None)
|
||||
|
||||
# 每成功生成一张图片就打印日志
|
||||
print(f"GeminiClient: 任务 {completed}/{total_tasks} 成功生成图片 ✓")
|
||||
|
||||
except Exception as e:
|
||||
fail_count += 1
|
||||
# 保存第一个错误（用于后续抛出）
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
error_msg = str(e)
|
||||
|
||||
# 传递完整的错误信息（用于排查问题）
|
||||
if progress_callback:
|
||||
progress_callback(completed, total_tasks, False, error_msg)
|
||||
|
||||
print(f"GeminiClient: 任务 {completed}/{total_tasks} 失败 ✗")
|
||||
|
||||
# 当前批次完成后，立即清理内存
|
||||
if batch_images:
|
||||
print(f"GeminiClient: 第 {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"GeminiClient: 多提示词批量生成完成，成功 {success_count}/{total_tasks}，失败 {fail_count}")
|
||||
return all_images
|
||||
|
||||
def generate_multi_prompts_sync(
|
||||
self,
|
||||
prompts: List[str],
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images_per_prompt: int,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None,
|
||||
debug: bool = False,
|
||||
debug_request: bool = False,
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
多提示词批量生成（同步接口，用于 ComfyUI）
|
||||
|
||||
Args:
|
||||
prompts: 提示词列表
|
||||
model: 模型名称
|
||||
resolution: 分辨率
|
||||
aspect_ratio: 宽高比
|
||||
images_per_prompt: 每个提示词生成的图像数量
|
||||
images: 输入图像列表
|
||||
progress_callback: 进度回调
|
||||
debug: 是否打印完整 API 响应
|
||||
debug_request: 是否打印发送的请求体
|
||||
enable_grounding: 是否启用 Google Search Grounding
|
||||
enable_image_search: 是否同时启用 Google Image Search
|
||||
|
||||
Returns:
|
||||
生成的图像列表
|
||||
"""
|
||||
coro = self.generate_multi_prompts_async(
|
||||
prompts=prompts,
|
||||
model=model,
|
||||
resolution=resolution,
|
||||
aspect_ratio=aspect_ratio,
|
||||
images_per_prompt=images_per_prompt,
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,303 @@
|
||||
"""
|
||||
Gemini Flash API 客户端
|
||||
用于调用 Gemini 3 Flash 模型进行多模态文本生成
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..models_config import (
|
||||
get_flash_model_endpoint,
|
||||
get_enabled_flash_models,
|
||||
get_flash_model_thinking_level_value,
|
||||
)
|
||||
from .base_client import BaseAPIClient
|
||||
|
||||
|
||||
class GeminiFlashClient(BaseAPIClient):
|
||||
"""
|
||||
Gemini Flash API 客户端
|
||||
用于调用 Gemini 3 Flash 模型进行多模态文本生成
|
||||
|
||||
特点：
|
||||
- 支持图片和视频输入
|
||||
- 支持动态思考等级端点（不思考/低/中/高）
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
"""
|
||||
初始化客户端
|
||||
|
||||
Args:
|
||||
api_key: API 密钥，如果为 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 # 100MB
|
||||
)
|
||||
|
||||
def get_endpoint(
|
||||
self,
|
||||
model: str = "gemini-3-flash-preview",
|
||||
**kwargs
|
||||
) -> str:
|
||||
"""
|
||||
获取模型的 API 端点
|
||||
|
||||
Args:
|
||||
model: 模型名称
|
||||
|
||||
Returns:
|
||||
API 端点路径
|
||||
"""
|
||||
endpoint = get_flash_model_endpoint(model)
|
||||
|
||||
if endpoint is None:
|
||||
# 回退到第一个启用的模型端点
|
||||
default_models = get_enabled_flash_models()
|
||||
if default_models:
|
||||
endpoint = get_flash_model_endpoint(default_models[0])
|
||||
|
||||
if endpoint is None:
|
||||
raise ValueError(f"无法获取模型 '{model}' 的端点")
|
||||
|
||||
return endpoint
|
||||
|
||||
def get_http_error_message(self, status_code: int, error_message: str) -> Optional[str]:
|
||||
"""Gemini 请求 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
|
||||
|
||||
def build_request_body(
|
||||
self,
|
||||
prompt: str = "",
|
||||
model: str = "gemini-3-flash-preview",
|
||||
thinking_level: str = "不思考",
|
||||
image_data: Optional[List[Dict[str, str]]] = None,
|
||||
video_data: Optional[Dict[str, str]] = None,
|
||||
document_data: Optional[Dict[str, str]] = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
构建 API 请求体
|
||||
|
||||
Args:
|
||||
prompt: 用户提示词
|
||||
model: 模型名称
|
||||
thinking_level: 思考等级（不思考/低/中/高）- 通过动态端点控制，不需要在请求体中传递
|
||||
image_data: 图片数据列表，每个元素包含 mime_type 和 data
|
||||
video_data: 视频数据，包含 mime_type 和 data
|
||||
document_data: 文档数据，包含 mime_type 和 data
|
||||
|
||||
Returns:
|
||||
请求体字典
|
||||
"""
|
||||
parts = []
|
||||
|
||||
# 添加文本部分
|
||||
if prompt:
|
||||
parts.append({"text": prompt})
|
||||
|
||||
# 添加图片部分（如果有）
|
||||
if image_data:
|
||||
for img in image_data:
|
||||
parts.append({
|
||||
"inline_data": {
|
||||
"mime_type": img["mime_type"],
|
||||
"data": img["data"]
|
||||
}
|
||||
})
|
||||
|
||||
# 添加视频部分（如果有）
|
||||
if video_data:
|
||||
parts.append({
|
||||
"inline_data": {
|
||||
"mime_type": video_data["mime_type"],
|
||||
"data": video_data["data"]
|
||||
}
|
||||
})
|
||||
|
||||
# 添加文档部分（如果有）
|
||||
if document_data:
|
||||
parts.append({
|
||||
"inline_data": {
|
||||
"mime_type": document_data["mime_type"],
|
||||
"data": document_data["data"]
|
||||
}
|
||||
})
|
||||
|
||||
# 构建请求体
|
||||
request_body = {
|
||||
"contents": [
|
||||
{
|
||||
"parts": parts
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# 对于支持 thinkingConfig 的固定端点模型（如 gemini-3-pro-preview）
|
||||
# 通过请求体传递思考等级；动态端点模型（如 gemini-3-flash-preview）
|
||||
# 通过不同 URL 端点控制，无需此字段
|
||||
thinking_level_value = get_flash_model_thinking_level_value(model, thinking_level)
|
||||
if thinking_level_value is not None:
|
||||
request_body["generationConfig"] = {
|
||||
"thinkingConfig": {
|
||||
"thinkingLevel": thinking_level_value
|
||||
}
|
||||
}
|
||||
|
||||
return request_body
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> str:
|
||||
"""
|
||||
解析 API 响应，提取生成的文本
|
||||
|
||||
Args:
|
||||
response: API 响应字典
|
||||
|
||||
Returns:
|
||||
生成的文本内容
|
||||
|
||||
Raises:
|
||||
RuntimeError: 解析失败或 API 拒绝时
|
||||
"""
|
||||
# 检查 candidatesTokenCount
|
||||
usage_metadata = response.get("usageMetadata", {})
|
||||
candidates_token_count = usage_metadata.get("candidatesTokenCount", -1)
|
||||
|
||||
if candidates_token_count == 0:
|
||||
raise RuntimeError(
|
||||
"内容审核拒绝 - candidatesTokenCount = 0\n\n"
|
||||
"原因：提示词或输入内容包含不适当内容\n"
|
||||
"建议：检查并调整输入内容"
|
||||
)
|
||||
|
||||
# 检查 finishReason
|
||||
candidates = response.get("candidates", [])
|
||||
if candidates:
|
||||
for candidate in candidates:
|
||||
finish_reason = candidate.get("finishReason", "")
|
||||
|
||||
if finish_reason and finish_reason not in ["STOP", "MAX_TOKENS"]:
|
||||
reason_messages = {
|
||||
"PROHIBITED_CONTENT": "违禁内容拒绝",
|
||||
"SAFETY": "安全过滤器拒绝",
|
||||
"RECITATION": "版权问题"
|
||||
}
|
||||
error_title = reason_messages.get(finish_reason, f"生成异常 ({finish_reason})")
|
||||
raise RuntimeError(f"{error_title}\n建议：调整输入内容后重试")
|
||||
|
||||
# 提取文本内容
|
||||
text_parts = []
|
||||
|
||||
for candidate in candidates:
|
||||
content = candidate.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
|
||||
for part in parts:
|
||||
if "text" in part:
|
||||
text_parts.append(part["text"])
|
||||
|
||||
if not text_parts:
|
||||
raise RuntimeError("API 响应中未找到生成的文本")
|
||||
|
||||
# 合并所有文本部分
|
||||
return "\n".join(text_parts)
|
||||
|
||||
async def generate_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str = "gemini-3-flash-preview",
|
||||
thinking_level: str = "不思考",
|
||||
image_data: Optional[List[Dict[str, str]]] = None,
|
||||
video_data: Optional[Dict[str, str]] = None,
|
||||
document_data: Optional[Dict[str, str]] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None
|
||||
) -> str:
|
||||
"""
|
||||
异步生成文本
|
||||
|
||||
Args:
|
||||
prompt: 用户提示词
|
||||
model: 模型名称
|
||||
thinking_level: 思考等级（不思考/低/中/高）
|
||||
image_data: 图片数据列表
|
||||
video_data: 视频数据
|
||||
document_data: 文档数据
|
||||
session: aiohttp 会话
|
||||
|
||||
Returns:
|
||||
生成的文本内容
|
||||
"""
|
||||
endpoint = self.get_endpoint(model=model)
|
||||
request_body = self.build_request_body(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
thinking_level=thinking_level,
|
||||
image_data=image_data,
|
||||
video_data=video_data,
|
||||
document_data=document_data
|
||||
)
|
||||
|
||||
response = await self.request_async(
|
||||
endpoint,
|
||||
request_body,
|
||||
session
|
||||
)
|
||||
|
||||
return self.parse_response(response)
|
||||
|
||||
def generate_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str = "gemini-3-flash-preview",
|
||||
thinking_level: str = "不思考",
|
||||
image_data: Optional[List[Dict[str, str]]] = None,
|
||||
video_data: Optional[Dict[str, str]] = None,
|
||||
document_data: Optional[Dict[str, str]] = None
|
||||
) -> str:
|
||||
"""
|
||||
同步生成文本（用于 ComfyUI 节点）
|
||||
|
||||
Args:
|
||||
prompt: 用户提示词
|
||||
model: 模型名称
|
||||
thinking_level: 思考等级（不思考/低/中/高）
|
||||
image_data: 图片数据列表
|
||||
video_data: 视频数据
|
||||
document_data: 文档数据
|
||||
|
||||
Returns:
|
||||
生成的文本内容
|
||||
"""
|
||||
coro = self.generate_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
thinking_level=thinking_level,
|
||||
image_data=image_data,
|
||||
video_data=video_data,
|
||||
document_data=document_data
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
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
|
||||
|
||||
|
||||
class KlingClient:
|
||||
"""Kling 视频生成客户端"""
|
||||
|
||||
ENDPOINTS = {
|
||||
"image2video": "/kling/v1/videos/image2video",
|
||||
"text2video": "/kling/v1/videos/text2video",
|
||||
"motion_control": "/kling/v1/videos/motion-control",
|
||||
}
|
||||
|
||||
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]}"
|
||||
|
||||
async with session.post(url, json=body, headers=self._headers()) as resp:
|
||||
text = await resp.text()
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(f"提交失败 ({resp.status}): {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:
|
||||
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 = (
|
||||
data.get("status") or
|
||||
inner_data.get("task_status") or
|
||||
result.get("status") or
|
||||
""
|
||||
)
|
||||
status = status.lower() if status else ""
|
||||
|
||||
progress_str = data.get("progress", "0%")
|
||||
try:
|
||||
progress_pct = int(str(progress_str).replace("%", "").strip())
|
||||
except (ValueError, AttributeError):
|
||||
progress_pct = 0
|
||||
|
||||
print(f"[视频生成] 生成中 {progress_pct}%")
|
||||
|
||||
if on_progress:
|
||||
on_progress(progress_pct)
|
||||
|
||||
if status in ("success", "completed", "done", "finished", "succeed"):
|
||||
return result
|
||||
elif status in ("failed", "fail"):
|
||||
error_info = result.get("error", {})
|
||||
if isinstance(error_info, dict):
|
||||
error_msg = error_info.get("message", "未知错误")
|
||||
else:
|
||||
error_msg = str(error_info)
|
||||
raise RuntimeError(f"生成失败：{error_msg}")
|
||||
|
||||
await asyncio.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("[视频生成] 下载视频...")
|
||||
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):
|
||||
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(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
|
||||
@@ -0,0 +1,775 @@
|
||||
"""
|
||||
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": {
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"image_size": api_image_size
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 = aiohttp.ClientSession()
|
||||
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 图片链接：
|
||||
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
|
||||
|
||||
print(f"{task_prefix} 请求 {size_str} → API {request_time:.1f}s → {download_info} ✓")
|
||||
|
||||
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(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)
|
||||
@@ -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}"
|
||||
@@ -0,0 +1,510 @@
|
||||
"""
|
||||
Veo 视频生成 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
|
||||
|
||||
|
||||
class VeoClient(BaseAPIClient):
|
||||
"""
|
||||
Veo 视频生成客户端
|
||||
|
||||
工作流程：
|
||||
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 = 8,
|
||||
size: str = "720x1280",
|
||||
first_frame_bytes: Optional[bytes] = None,
|
||||
last_frame_bytes: Optional[bytes] = None,
|
||||
reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
提交视频生成任务
|
||||
|
||||
格式策略：
|
||||
- 无参考图片：application/json
|
||||
- 有参考图片：multipart/form-data，图片以 PNG 文件上传
|
||||
|
||||
Args:
|
||||
prompt: 提示词
|
||||
model: 模型名称
|
||||
seconds: 视频时长（秒）
|
||||
size: 分辨率
|
||||
first_frame_bytes: 首帧图片字节
|
||||
last_frame_bytes: 尾帧图片字节
|
||||
reference_bytes: 参考图片字节
|
||||
seed: 随机种子
|
||||
session: aiohttp 会话
|
||||
|
||||
Returns:
|
||||
API 响应 JSON，包含 video id 和初始状态
|
||||
"""
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
# 检查是否有图片
|
||||
has_images = any([first_frame_bytes, last_frame_bytes, reference_bytes])
|
||||
|
||||
if has_images:
|
||||
# 有图片：multipart/form-data + PNG 文件上传
|
||||
if first_frame_bytes and len(first_frame_bytes) > self.max_request_size:
|
||||
raise ValueError(f"首帧图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制")
|
||||
if last_frame_bytes and len(last_frame_bytes) > self.max_request_size:
|
||||
raise ValueError(f"尾帧图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制")
|
||||
if reference_bytes and len(reference_bytes) > self.max_request_size:
|
||||
raise ValueError(f"参考图片过大，超过 {self.max_request_size / 1024 / 1024:.0f}MB 限制")
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("model", model)
|
||||
form.add_field("seconds", str(seconds))
|
||||
form.add_field("size", size)
|
||||
# 注意：seed 不被上游 API 接受，仅在 ComfyUI 节点侧用于缓存刷新
|
||||
# if seed is not None:
|
||||
# form.add_field("seed", str(seed))
|
||||
|
||||
# 使用 input_reference 字段（OpenAI兼容格式）
|
||||
# 尝试支持多张图片：按顺序添加多个 input_reference 字段
|
||||
if first_frame_bytes:
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
first_frame_bytes,
|
||||
filename="first_frame.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
if last_frame_bytes:
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
last_frame_bytes,
|
||||
filename="last_frame.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
if reference_bytes:
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
reference_bytes,
|
||||
filename="reference.png",
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
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,
|
||||
}
|
||||
# 注意：seed 不被上游 API 接受，仅在 ComfyUI 节点侧用于缓存刷新
|
||||
# if seed is not None:
|
||||
# body["seed"] = str(seed)
|
||||
send_kwargs = {"json": body, "headers": headers}
|
||||
|
||||
# 打印请求调试信息
|
||||
import json
|
||||
if has_images:
|
||||
print(f"Veo: 使用 multipart/form-data 格式上传图片")
|
||||
else:
|
||||
print(f"Veo API 请求体: {json.dumps(body, ensure_ascii=False)}")
|
||||
|
||||
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 兼容大小写
|
||||
status = data.get("status", "").lower()
|
||||
|
||||
# progress 兼容整数和字符串
|
||||
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)
|
||||
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:
|
||||
"""
|
||||
下载生成的视频文件
|
||||
|
||||
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,
|
||||
first_frame_bytes: Optional[bytes] = None,
|
||||
last_frame_bytes: Optional[bytes] = None,
|
||||
reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
progress_callback: Optional[Callable[[int], None]] = None,
|
||||
on_stage: Optional[Callable[[str], None]] = None,
|
||||
) -> str:
|
||||
"""
|
||||
同步执行完整的视频生成流程（创建 → 轮询 → 下载）
|
||||
"""
|
||||
|
||||
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,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=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())
|
||||
|
||||
def generate_batch_videos_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
seconds: int,
|
||||
size: str,
|
||||
save_paths: List[str],
|
||||
first_frame_bytes: Optional[bytes] = None,
|
||||
last_frame_bytes: Optional[bytes] = None,
|
||||
reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
progress_callback: Optional[Callable[[int, int, bool, Optional[str]], None]] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
同步并发生成多个视频
|
||||
"""
|
||||
async def _run():
|
||||
batch_size = len(save_paths)
|
||||
connector = aiohttp.TCPConnector(limit=0)
|
||||
|
||||
async def generate_one(save_path: str):
|
||||
return await self._generate_one_video_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
save_path=save_path,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=reference_bytes,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
tasks = [generate_one(p) for p in save_paths]
|
||||
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"Veo: 第 {i + 1} 个视频生成失败")
|
||||
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
|
||||
|
||||
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,
|
||||
first_frame_bytes: Optional[bytes] = None,
|
||||
last_frame_bytes: Optional[bytes] = None,
|
||||
reference_bytes: Optional[bytes] = None,
|
||||
seed: Optional[int] = None,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
) -> str:
|
||||
"""异步生成单个视频"""
|
||||
result = await self.create_video_async(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
first_frame_bytes=first_frame_bytes,
|
||||
last_frame_bytes=last_frame_bytes,
|
||||
reference_bytes=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 _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