feat: 新增AI生图（批量版）节点，模型改名，超时优化与友好报错
- 新增 BatchAsyncImageGenerator 节点（全并发+即时落盘，不怕中途失败丢图） - 原版 AsyncImageGenerator 移除批量提示词功能，单节点只处理单提示词 - 模型改名：限时特价→次卡，gemini→nano-banana-官方 - 异步节点过滤 官方计费 渠道，仅保留次卡和官方模型 - 单任务超时提升至900s，批量超时改为动态计算（批次数×900s） - No available channel for model 错误转化为中文友好提示 - 新增 base_async_provider / gemini_async_provider 异步客户端基类 Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
o1key
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
异步生图 Provider 抽象基类
|
||||
定义异步提交+轮询模式的统一接口，支持多种生图模型后端
|
||||
|
||||
每个 Provider 封装一种 API 后端的通信协议：
|
||||
- 如何提交任务（端点、请求体格式）
|
||||
- 如何轮询状态（端点、状态字段语义）
|
||||
- 如何解析结果（响应格式、图片提取方式）
|
||||
|
||||
新增第三方生图模型时，只需实现此接口即可接入异步节点。
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Optional
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class BaseAsyncImageProvider(ABC):
|
||||
"""异步生图 Provider 抽象基类"""
|
||||
|
||||
def __init__(self, api_key: str, proxy_url: Optional[str] = None):
|
||||
self.api_key = api_key
|
||||
self.proxy_url = proxy_url
|
||||
|
||||
# ========================================================================
|
||||
# 必须实现的抽象方法
|
||||
# ========================================================================
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def api_base_url(self) -> str:
|
||||
"""异步 API 的基础 URL，如 https://cf-api.o1key.com"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_submit_endpoint(self, model: str, resolution: str) -> str:
|
||||
"""获取提交任务的 API 端点路径（不含 base_url）"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def build_submit_body(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
**kwargs
|
||||
) -> dict:
|
||||
"""构建提交任务的请求体"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def extract_task_id(self, response: dict) -> str:
|
||||
"""从提交响应中提取 task_id"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def extract_status(self, response: dict) -> str:
|
||||
"""从轮询响应中提取任务状态（如 SUBMITTED / IN_PROGRESS / SUCCESS / FAILURE）"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def parse_result(
|
||||
self,
|
||||
result_data: dict,
|
||||
session
|
||||
) -> List[Image.Image]:
|
||||
"""从任务完成后的 result data 中解析生成的图像列表"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_models(self) -> List[str]:
|
||||
"""获取此 Provider 支持的模型 ID 列表"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
|
||||
"""获取指定模型支持的宽高比"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_model_resolutions(self, model_id: str) -> List[str]:
|
||||
"""获取指定模型支持的分辨率"""
|
||||
...
|
||||
|
||||
# ========================================================================
|
||||
# 可选的覆盖方法
|
||||
# ========================================================================
|
||||
|
||||
def get_poll_endpoint(self, task_id: str) -> str:
|
||||
"""获取轮询任务状态的 API 端点路径（默认实现适用于 o1key 异步 API）"""
|
||||
return f"/async/v1/tasks/{task_id}"
|
||||
|
||||
def get_headers(self) -> dict:
|
||||
"""获取 HTTP 请求头"""
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def get_all_aspect_ratios(self) -> List[str]:
|
||||
"""获取所有模型支持的宽高比（去重合并）"""
|
||||
seen = set()
|
||||
result = []
|
||||
for model_id in self.get_models():
|
||||
for ratio in self.get_model_aspect_ratios(model_id):
|
||||
if ratio not in seen:
|
||||
seen.add(ratio)
|
||||
result.append(ratio)
|
||||
return result
|
||||
|
||||
def get_all_resolutions(self) -> List[str]:
|
||||
"""获取所有模型支持的分辨率（去重，按固定顺序排列）"""
|
||||
_ORDER = ["512px", "1K", "2K", "4K"]
|
||||
seen = set()
|
||||
for model_id in self.get_models():
|
||||
for res in self.get_model_resolutions(model_id):
|
||||
seen.add(res)
|
||||
return [r for r in _ORDER if r in seen]
|
||||
|
||||
def get_extra_inputs(self) -> dict:
|
||||
"""
|
||||
返回此 Provider 特有的额外 ComfyUI 输入参数。
|
||||
子类重写以声明 Provider 专有的选项（如 Google Search Grounding）。
|
||||
|
||||
Returns:
|
||||
dict，格式与 ComfyUI INPUT_TYPES 的 optional 字段一致
|
||||
"""
|
||||
return {}
|
||||
|
||||
def get_extra_kwargs(self, **kwargs) -> dict:
|
||||
"""
|
||||
从 ComfyUI kwargs 中提取此 Provider 特有的参数，
|
||||
转换为 build_submit_body 可接收的 kwargs。
|
||||
|
||||
子类重写以处理 Provider 专有参数。
|
||||
"""
|
||||
return {}
|
||||
|
||||
def extract_progress(self, response: dict) -> Optional[float]:
|
||||
"""
|
||||
从轮询响应中提取生成进度。
|
||||
|
||||
Args:
|
||||
response: 轮询接口返回的完整响应字典
|
||||
|
||||
Returns:
|
||||
0.0-1.0 之间的进度值，或 None 表示该响应不含进度信息
|
||||
"""
|
||||
return None
|
||||
|
||||
def query_balance_sync(self) -> Optional[dict]:
|
||||
"""
|
||||
同步查询账户余额（可选）。
|
||||
返回 None 表示不支持。
|
||||
"""
|
||||
return None
|
||||
|
||||
def format_balance_info(self, balance_data: dict) -> str:
|
||||
"""格式化余额信息为展示文本"""
|
||||
return ""
|
||||
|
||||
# ========================================================================
|
||||
# 工具方法
|
||||
# ========================================================================
|
||||
|
||||
@staticmethod
|
||||
def build_proxy_url(port: str) -> Optional[str]:
|
||||
"""
|
||||
将端口号字符串转为 aiohttp 可用的 HTTP 代理 URL。
|
||||
兼容 v2rayN (10808)、Clash Verge (7897) 等。
|
||||
|
||||
Args:
|
||||
port: 用户填写的端口号，如 "7897"，空字符串返回 None
|
||||
|
||||
Returns:
|
||||
代理 URL 或 None
|
||||
"""
|
||||
port = (port or "").strip()
|
||||
if not port or not port.isdigit():
|
||||
return None
|
||||
return f"http://127.0.0.1:{port}"
|
||||
@@ -0,0 +1,149 @@
|
||||
"""
|
||||
Gemini 异步生图 Provider
|
||||
通过 cf-api.o1key.com 的异步提交+轮询接口调用 Gemini 图像生成模型
|
||||
|
||||
协议说明：
|
||||
- 提交：POST {base}/async{gemini_endpoint}?image_format=url
|
||||
- 轮询：GET {base}/async/v1/tasks/{task_id}
|
||||
- 结果：可能直接返回 image_url，也可能返回 Gemini 标准 candidates 格式
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
from typing import Dict, List, Optional
|
||||
from PIL import Image
|
||||
|
||||
from .base_async_provider import BaseAsyncImageProvider
|
||||
from .gemini_client import GeminiAPIClient
|
||||
from ..utils.config import get_async_api_base_url, get_api_key_or_raise
|
||||
from ..models_config import (
|
||||
get_enabled_models,
|
||||
get_model_supported_aspect_ratios,
|
||||
get_model_supported_resolutions,
|
||||
)
|
||||
|
||||
|
||||
class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
"""
|
||||
Gemini 异步生图 Provider
|
||||
|
||||
委托 GeminiAPIClient 处理：
|
||||
- 端点构造（get_endpoint）
|
||||
- 请求体构建（build_request_body）
|
||||
- 响应解析（parse_response_async）
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str = None, proxy_url: str = None):
|
||||
if api_key is None:
|
||||
api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
super().__init__(api_key=api_key, proxy_url=proxy_url)
|
||||
self._client = GeminiAPIClient(api_key=api_key)
|
||||
|
||||
# ========================================================================
|
||||
# 抽象方法实现
|
||||
# ========================================================================
|
||||
|
||||
@property
|
||||
def api_base_url(self) -> str:
|
||||
return get_async_api_base_url()
|
||||
|
||||
def get_submit_endpoint(self, model: str, resolution: str) -> str:
|
||||
gemini_endpoint = self._client.get_endpoint(
|
||||
model=model, resolution=resolution, image_format="url"
|
||||
)
|
||||
base = gemini_endpoint.split("?")[0]
|
||||
async_endpoint = f"/async{base}"
|
||||
if "?" in gemini_endpoint:
|
||||
async_endpoint += "?" + gemini_endpoint.split("?", 1)[1]
|
||||
return async_endpoint
|
||||
|
||||
def build_submit_body(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
resolution: str,
|
||||
aspect_ratio: str,
|
||||
images: Optional[List[Image.Image]] = None,
|
||||
**kwargs
|
||||
) -> dict:
|
||||
return self._client.build_request_body(
|
||||
prompt=prompt,
|
||||
images=images,
|
||||
aspect_ratio=aspect_ratio,
|
||||
resolution=resolution,
|
||||
enable_grounding=kwargs.get("enable_grounding", False),
|
||||
enable_image_search=kwargs.get("enable_image_search", False),
|
||||
)
|
||||
|
||||
def extract_task_id(self, response: dict) -> str:
|
||||
task_id = response.get("task_id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"提交响应中未找到 task_id: {response}")
|
||||
return task_id
|
||||
|
||||
def extract_status(self, response: dict) -> str:
|
||||
return response.get("status", "UNKNOWN")
|
||||
|
||||
async def parse_result(self, result_data: dict, session) -> List[Image.Image]:
|
||||
# 异步接口可能直接返回 image_url
|
||||
image_url = result_data.get("image_url", "") if isinstance(result_data, dict) else ""
|
||||
if image_url:
|
||||
async with session.get(image_url) as img_resp:
|
||||
if img_resp.status == 200:
|
||||
img_bytes = await img_resp.read()
|
||||
return [Image.open(BytesIO(img_bytes))]
|
||||
raise RuntimeError(f"下载图片失败 ({img_resp.status}): {image_url}")
|
||||
|
||||
# 否则按 Gemini 标准格式解析
|
||||
images_list, _ = await self._client.parse_response_async(result_data, session=session)
|
||||
return images_list
|
||||
|
||||
def get_models(self) -> List[str]:
|
||||
return get_enabled_models()
|
||||
|
||||
def get_model_aspect_ratios(self, model_id: str) -> List[str]:
|
||||
return get_model_supported_aspect_ratios(model_id)
|
||||
|
||||
def get_model_resolutions(self, model_id: str) -> List[str]:
|
||||
return get_model_supported_resolutions(model_id)
|
||||
|
||||
# ========================================================================
|
||||
# 可选方法覆盖
|
||||
# ========================================================================
|
||||
|
||||
def get_extra_inputs(self) -> dict:
|
||||
"""Gemini 专有：Google Search Grounding"""
|
||||
return {
|
||||
"联网功能": (["关闭", "打开"], {"default": "关闭"}),
|
||||
}
|
||||
|
||||
def get_extra_kwargs(self, **kwargs) -> dict:
|
||||
return {
|
||||
"enable_grounding": kwargs.pop("联网功能", "关闭") == "打开",
|
||||
}
|
||||
|
||||
def extract_progress(self, response: dict) -> Optional[float]:
|
||||
"""从轮询响应中提取进度（0.0-1.0）"""
|
||||
# 直接字段：progress / percentage
|
||||
for field in ("progress", "percentage"):
|
||||
val = response.get(field)
|
||||
if val is not None and isinstance(val, (int, float)):
|
||||
return val / 100.0 if val > 1 else float(val)
|
||||
|
||||
# 嵌套字段：progressInfo / progress_info
|
||||
progress_info = response.get("progressInfo") or response.get("progress_info")
|
||||
if isinstance(progress_info, dict):
|
||||
for field in ("progress", "percentage"):
|
||||
val = progress_info.get(field)
|
||||
if val is not None and isinstance(val, (int, float)):
|
||||
return val / 100.0 if val > 1 else float(val)
|
||||
|
||||
return None
|
||||
|
||||
def query_balance_sync(self) -> Optional[dict]:
|
||||
try:
|
||||
return self._client.query_balance_sync()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def format_balance_info(self, balance_data: dict) -> str:
|
||||
return self._client.format_balance_info(balance_data)
|
||||
@@ -70,7 +70,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
from ..models_config import get_model_endpoint
|
||||
|
||||
# 特殊处理：动态端点模型（根据分辨率选择）
|
||||
if model == "nano-banana-pro-限时特价":
|
||||
if model == "nano-banana-pro-次卡":
|
||||
if resolution == "1K":
|
||||
endpoint = "/v1beta/models/nano-banana-pro:generateContent"
|
||||
elif resolution == "2K":
|
||||
@@ -80,7 +80,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
else:
|
||||
endpoint = "/v1beta/models/nano-banana-pro-2k:generateContent"
|
||||
|
||||
elif model == "nano-banana-2-限时特价":
|
||||
elif model == "nano-banana-2-次卡":
|
||||
if resolution == "512px":
|
||||
endpoint = "/v1beta/models/nano-banana-2-0.5k:generateContent"
|
||||
elif resolution == "1K":
|
||||
|
||||
+11
-23
@@ -33,8 +33,7 @@ _ENDPOINT_EDITS = "/v1/images/edits/"
|
||||
|
||||
# ── 模型名映射（UI 显示名 → API 实际参数名）─────────────────────────────────
|
||||
_MODEL_NAME_MAP = {
|
||||
"gpt-image-1.5-特价": "gpt-image-1.5-special",
|
||||
"gpt-image-2-特价": "gpt-image-2-special",
|
||||
"gpt-image-2-次卡": "gpt-image-2-special",
|
||||
}
|
||||
|
||||
# ── 超时 ──────────────────────────────────────────────────────────────────────
|
||||
@@ -46,7 +45,7 @@ class GptImageClient:
|
||||
GPT Image API 客户端
|
||||
|
||||
接口说明：
|
||||
generations：JSON body，支持 background / quality / size / n / model
|
||||
generations：JSON body，支持 quality / size / n / model
|
||||
edits：multipart/form-data，必须包含 image（PNG），可选 mask（PNG）
|
||||
|
||||
两个接口的响应格式相同：
|
||||
@@ -227,7 +226,6 @@ class GptImageClient:
|
||||
prompt: str,
|
||||
model: str,
|
||||
quality: str,
|
||||
background: str,
|
||||
size: str,
|
||||
n: int,
|
||||
seed: int,
|
||||
@@ -244,7 +242,6 @@ class GptImageClient:
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"quality": quality,
|
||||
"background": background,
|
||||
"n": n,
|
||||
"moderation": "low",
|
||||
}
|
||||
@@ -278,7 +275,7 @@ class GptImageClient:
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_GENERATIONS}"
|
||||
print(f"[o1key GPT Image] {mode} | 模型={model} | quality={quality} | "
|
||||
f"background={background} | size={size} | n={n}")
|
||||
f"size={size} | n={n}")
|
||||
|
||||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
@@ -311,15 +308,12 @@ class GptImageClient:
|
||||
return await self._parse_response(resp_json, session)
|
||||
|
||||
# ── 图像编辑（edits 接口，multipart/form-data）──────────────────────────
|
||||
# 注意：o1key 中转服务的 edits 接口暂不支持 quality / background / moderation 参数，
|
||||
# 这些字段暂时不传递，待服务方更新后可恢复。
|
||||
|
||||
async def _edit_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
quality: str,
|
||||
background: str,
|
||||
size: str,
|
||||
n: int,
|
||||
seed: int,
|
||||
@@ -328,8 +322,6 @@ class GptImageClient:
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
调用 /v1/images/edits/ 接口（multipart/form-data）。
|
||||
当前仅传递 model / prompt / n / size / image / mask，
|
||||
quality / background / moderation 暂不支持（o1key 服务端限制）。
|
||||
"""
|
||||
# 模型名映射：UI 显示名 → API 参数名
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
@@ -342,13 +334,11 @@ class GptImageClient:
|
||||
normalized_tensors.append(t)
|
||||
num_images = len(normalized_tensors)
|
||||
|
||||
# o1key 中转服务的 edits 接口暂不支持 background / moderation / seed，
|
||||
# 待服务方更新后可重新加入。
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", api_model)
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("n", str(n))
|
||||
form.add_field("quality", quality)
|
||||
form.add_field("model", api_model)
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("n", str(n))
|
||||
form.add_field("quality", quality)
|
||||
|
||||
form.add_field("size", size if size else "auto")
|
||||
|
||||
@@ -387,7 +377,8 @@ class GptImageClient:
|
||||
mode = "图像编辑（无蒙版）"
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_EDITS}"
|
||||
print(f"[o1key GPT Image] {mode} | 模型={model} | 参考图={num_images}张 | size={size} | n={n}")
|
||||
print(f"[o1key GPT Image] {mode} | 模型={model} | 参考图={num_images}张 | "
|
||||
f"quality={quality} | size={size} | n={n}")
|
||||
|
||||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
@@ -430,7 +421,6 @@ class GptImageClient:
|
||||
prompt: str,
|
||||
model: str,
|
||||
quality: str,
|
||||
background: str,
|
||||
size: str,
|
||||
n: int,
|
||||
seed: int,
|
||||
@@ -443,21 +433,19 @@ class GptImageClient:
|
||||
路由逻辑：
|
||||
- 无 image_tensor → generations 接口（文生图，JSON body）
|
||||
- 有 image_tensor → edits 接口（图生图/编辑，multipart/form-data）
|
||||
所有模型统一走 multipart，quality/background 通过表单字段传递，
|
||||
new-api 开启"透传请求体"后原样转发给上游。
|
||||
"""
|
||||
use_edits = (image_tensor is not None)
|
||||
|
||||
if use_edits:
|
||||
coro = self._edit_async(
|
||||
prompt=prompt, model=model, quality=quality,
|
||||
background=background, size=size, n=n, seed=seed,
|
||||
size=size, n=n, seed=seed,
|
||||
image_list=image_tensor, mask_tensor=mask_tensor,
|
||||
)
|
||||
else:
|
||||
coro = self._generate_async(
|
||||
prompt=prompt, model=model, quality=quality,
|
||||
background=background, size=size, n=n, seed=seed,
|
||||
size=size, n=n, seed=seed,
|
||||
image_list=image_tensor,
|
||||
)
|
||||
|
||||
|
||||
@@ -56,14 +56,14 @@ class OpenAIAPIClient(BaseAPIClient):
|
||||
只是把拼在 URL 路径里的模型段提取出来单独返回。
|
||||
|
||||
Args:
|
||||
model: 节点下拉框中的模型 ID，如 "nano-banana-pro-限时特价"
|
||||
model: 节点下拉框中的模型 ID，如 "nano-banana-pro-次卡"
|
||||
resolution: 分辨率字符串，如 "1K" / "2K" / "4K" / "512"
|
||||
|
||||
Returns:
|
||||
实际模型名，如 "nano-banana-pro-2k"
|
||||
"""
|
||||
# ── 动态端点模型 ──────────────────────────────────────────────────
|
||||
if model == "nano-banana-pro-限时特价":
|
||||
if model == "nano-banana-pro-次卡":
|
||||
if resolution == "1K":
|
||||
return "nano-banana-pro"
|
||||
elif resolution == "4K":
|
||||
|
||||
Reference in New Issue
Block a user