feat: 新增 Grok 图像节点、前端 UI 增强、重构 nano-banana 系列
- 新增 Grok Image 节点及客户端 - 新增 save_image_format 节点 - 新增前端 JS 扩展：画笔工具、点阵网格、侧边栏隐藏、资源切换、重命名等 - 重构 nano-banana 节点，移除 pro 版本 - 移除 multi_res_preview 节点 - 新增 http_error 工具模块 - 各客户端和节点优化改进 Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Jony
+53
-40
@@ -14,6 +14,8 @@ import time
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ..utils.http_error import HTTP_ERROR_MESSAGES, RETRYABLE_STATUS_CODES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
|
||||
|
||||
|
||||
|
||||
class BaseAPIClient(ABC):
|
||||
@@ -206,7 +208,8 @@ class BaseAPIClient(ABC):
|
||||
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
raise RuntimeError(error_text)
|
||||
# 返回状态码和错误文本，由外层处理重试
|
||||
return {"_error": True, "_status": response.status, "_text": error_text}
|
||||
|
||||
wait_start = time.time()
|
||||
response_data = await response.json()
|
||||
@@ -231,31 +234,55 @@ class BaseAPIClient(ABC):
|
||||
return
|
||||
|
||||
try:
|
||||
if _interrupt_available:
|
||||
request_task = asyncio.ensure_future(_do_request())
|
||||
interrupt_task = asyncio.ensure_future(_poll_interrupt())
|
||||
last_error_status = None
|
||||
last_error_text = ""
|
||||
|
||||
done, pending = await asyncio.wait(
|
||||
[request_task, interrupt_task],
|
||||
return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
if _interrupt_available:
|
||||
request_task = asyncio.ensure_future(_do_request())
|
||||
interrupt_task = asyncio.ensure_future(_poll_interrupt())
|
||||
|
||||
# 取消未完成的任务
|
||||
for t in pending:
|
||||
t.cancel()
|
||||
try:
|
||||
await t
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
done, pending = await asyncio.wait(
|
||||
[request_task, interrupt_task],
|
||||
return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
|
||||
# 判断是哪个先完成
|
||||
if interrupt_task in done and request_task not in done:
|
||||
raise InterruptProcessingException()
|
||||
for t in pending:
|
||||
t.cancel()
|
||||
try:
|
||||
await t
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
# 请求完成，取出结果（可能含异常）
|
||||
return request_task.result()
|
||||
else:
|
||||
return await _do_request()
|
||||
if interrupt_task in done and request_task not in done:
|
||||
raise InterruptProcessingException()
|
||||
|
||||
result = request_task.result()
|
||||
else:
|
||||
result = await _do_request()
|
||||
|
||||
if isinstance(result, dict) and result.get("_error"):
|
||||
status = result["_status"]
|
||||
error_text = result["_text"]
|
||||
last_error_status = status
|
||||
last_error_text = error_text
|
||||
|
||||
if status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(status, f"请求失败 ({status})")
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
print(f"{friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
|
||||
if status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[status])
|
||||
raise RuntimeError(error_text)
|
||||
|
||||
return result
|
||||
|
||||
if last_error_status and last_error_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_error_status])
|
||||
raise RuntimeError(last_error_text)
|
||||
|
||||
except InterruptProcessingException:
|
||||
raise
|
||||
@@ -352,30 +379,16 @@ class BaseAPIClient(ABC):
|
||||
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"建议：等待一段时间后重试"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[429])
|
||||
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"建议：稍后重试"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[503])
|
||||
elif response.status == 504:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"API 返回错误：{error_message}\n"
|
||||
f"建议：稍后重试"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[504])
|
||||
elif response.status == 502:
|
||||
raise RuntimeError(
|
||||
"糟糕！请求到上游时遇到超时或过载！别担心，过会儿再次点击运行即可！"
|
||||
)
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[502])
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status})\n"
|
||||
|
||||
@@ -21,6 +21,7 @@ from PIL import Image
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
|
||||
|
||||
|
||||
# ── 固定端点 ──────────────────────────────────────────────────────────────────
|
||||
@@ -245,39 +246,55 @@ class DoubaoImageClient:
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
|
||||
# 3. 发送 POST 请求
|
||||
t0 = time.time()
|
||||
async with session.post(
|
||||
url,
|
||||
json=body,
|
||||
headers=self._headers(),
|
||||
) as resp:
|
||||
elapsed_req = time.time() - t0
|
||||
text = await resp.text()
|
||||
# 3. 发送 POST 请求（带退避重试）
|
||||
last_status = None
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
t0 = time.time()
|
||||
async with session.post(
|
||||
url,
|
||||
json=body,
|
||||
headers=self._headers(),
|
||||
) as resp:
|
||||
elapsed_req = time.time() - t0
|
||||
text = await resp.text()
|
||||
|
||||
if resp.status != 200:
|
||||
last_status = resp.status
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(resp.status)
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
print(f"[豆包生图] {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
if resp.status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
|
||||
try:
|
||||
err_json = json.loads(text)
|
||||
err_obj = err_json.get("error", {})
|
||||
if isinstance(err_obj, dict):
|
||||
msg = (
|
||||
err_obj.get("message")
|
||||
or err_obj.get("msg")
|
||||
or text
|
||||
)
|
||||
else:
|
||||
msg = str(err_obj) or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(
|
||||
f"请求失败 HTTP {resp.status}: {msg}"
|
||||
)
|
||||
|
||||
if resp.status != 200:
|
||||
# 尝试解析错误信息
|
||||
try:
|
||||
err_json = json.loads(text)
|
||||
err_obj = err_json.get("error", {})
|
||||
if isinstance(err_obj, dict):
|
||||
msg = (
|
||||
err_obj.get("message")
|
||||
or err_obj.get("msg")
|
||||
or text
|
||||
)
|
||||
else:
|
||||
msg = str(err_obj) or text
|
||||
resp_json = json.loads(text)
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(
|
||||
f"请求失败 HTTP {resp.status}: {msg}"
|
||||
)
|
||||
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
|
||||
|
||||
try:
|
||||
resp_json = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
|
||||
break
|
||||
else:
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
|
||||
raise RuntimeError(f"请求失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
|
||||
|
||||
print(f"[豆包生图] API 响应耗时 {elapsed_req:.1f}s，开始下载图像...")
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ from typing import Optional
|
||||
import requests
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.http_error import HTTP_ERROR_MESSAGES
|
||||
|
||||
|
||||
# 显示名 → 实际请求值的映射
|
||||
@@ -112,6 +113,8 @@ class FluxEditClient:
|
||||
raise RuntimeError("无法连接到服务器，请检查网络或服务器地址")
|
||||
|
||||
if resp.status_code != 200:
|
||||
if resp.status_code in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status_code])
|
||||
raise RuntimeError(
|
||||
f"提交任务失败 (HTTP {resp.status_code})\n"
|
||||
f"响应: {resp.text[:500]}"
|
||||
|
||||
@@ -44,7 +44,7 @@ class GeminiAsyncImageProvider(BaseAsyncImageProvider):
|
||||
|
||||
@property
|
||||
def api_base_url(self) -> str:
|
||||
return get_async_api_base_url()
|
||||
return getattr(self, '_route_base_url', None) or get_async_api_base_url()
|
||||
|
||||
def get_submit_endpoint(self, model: str, resolution: str) -> str:
|
||||
gemini_endpoint = self._client.get_endpoint(
|
||||
|
||||
@@ -187,6 +187,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_compression: str = None,
|
||||
thinking_level: str = None,
|
||||
**kwargs
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
@@ -227,15 +228,15 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
})
|
||||
|
||||
# 估算完整 body 大小（不含工具字段，工具字段很小可忽略）
|
||||
est_image_config = {"imageSize": resolution}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
est_image_config["aspectRatio"] = aspect_ratio
|
||||
estimated = self._estimate_body_size(
|
||||
parts + img_parts,
|
||||
{
|
||||
"generationConfig": {
|
||||
"responseModalities": ["IMAGE"],
|
||||
"imageConfig": {
|
||||
"aspectRatio": aspect_ratio,
|
||||
"imageSize": resolution
|
||||
}
|
||||
"imageConfig": est_image_config
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -267,6 +268,10 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
parts.extend(img_parts)
|
||||
|
||||
# 构建请求体
|
||||
image_config = {"imageSize": resolution}
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
image_config["aspectRatio"] = aspect_ratio
|
||||
|
||||
request_body = {
|
||||
"contents": [
|
||||
{
|
||||
@@ -276,13 +281,17 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
],
|
||||
"generationConfig": {
|
||||
"responseModalities": ["IMAGE"],
|
||||
"imageConfig": {
|
||||
"aspectRatio": aspect_ratio,
|
||||
"imageSize": resolution
|
||||
}
|
||||
"imageConfig": image_config
|
||||
}
|
||||
}
|
||||
|
||||
# 添加思考深度配置
|
||||
if thinking_level:
|
||||
request_body["generationConfig"]["thinkingConfig"] = {
|
||||
"thinkingLevel": thinking_level,
|
||||
"includeThoughts": True
|
||||
}
|
||||
|
||||
# 添加图片压缩参数
|
||||
if image_compression:
|
||||
request_body["image_compression"] = image_compression
|
||||
@@ -564,6 +573,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "base64",
|
||||
thinking_level: str = None,
|
||||
) -> tuple[List[Image.Image], Dict[str, Any]]:
|
||||
"""
|
||||
单次异步生成请求（极简单行日志）
|
||||
@@ -602,6 +612,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
resolution=resolution,
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
build_time = time.time() - build_start
|
||||
|
||||
@@ -737,6 +748,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "base64",
|
||||
thinking_level: str = None,
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
批量全并发生成 - 改进版：支持分批处理和内存管理
|
||||
@@ -801,6 +813,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
thinking_level=thinking_level,
|
||||
),
|
||||
name=f"task_{task_index}"
|
||||
)
|
||||
@@ -873,6 +886,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding: bool = False,
|
||||
enable_image_search: bool = False,
|
||||
image_format: str = "base64",
|
||||
thinking_level: str = None,
|
||||
) -> List[Image.Image]:
|
||||
"""
|
||||
同步生成接口（用于 ComfyUI）
|
||||
@@ -906,6 +920,7 @@ class GeminiAPIClient(BaseAPIClient):
|
||||
enable_grounding=enable_grounding,
|
||||
enable_image_search=enable_image_search,
|
||||
image_format=image_format,
|
||||
thinking_level=thinking_level,
|
||||
)
|
||||
|
||||
return self.run_async_in_thread(coro)
|
||||
|
||||
+38
-20
@@ -26,6 +26,7 @@ from PIL import Image
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.image_utils import tensor_to_pil, encode_image_to_base64
|
||||
from ..utils.http_error import RETRYABLE_STATUS_CODES, HTTP_ERROR_MESSAGES, _compute_delay, DEFAULT_MAX_RETRIES, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR
|
||||
|
||||
try:
|
||||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||
@@ -330,31 +331,46 @@ class GptImageClient:
|
||||
|
||||
async def _do_request():
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
t0 = time.time()
|
||||
async with session.post(url, json=body, headers=self._json_headers()) as resp:
|
||||
elapsed = time.time() - t0
|
||||
text = await resp.text()
|
||||
last_status = None
|
||||
for attempt in range(DEFAULT_MAX_RETRIES + 1):
|
||||
t0 = time.time()
|
||||
async with session.post(url, json=body, headers=self._json_headers()) as resp:
|
||||
elapsed = time.time() - t0
|
||||
text = await resp.text()
|
||||
|
||||
if resp.status != 200:
|
||||
last_status = resp.status
|
||||
if resp.status in RETRYABLE_STATUS_CODES and attempt < DEFAULT_MAX_RETRIES:
|
||||
friendly = HTTP_ERROR_MESSAGES.get(resp.status)
|
||||
delay = _compute_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY, DEFAULT_BACKOFF_FACTOR)
|
||||
print(f"[o1key GPT Image] {friendly} {delay:.1f}s 后重试 ({attempt+1}/{DEFAULT_MAX_RETRIES})...")
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
if resp.status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
|
||||
try:
|
||||
err_json = json.loads(text)
|
||||
err_obj = err_json.get("error", {})
|
||||
msg = (
|
||||
err_obj.get("message") or err_obj.get("msg") or text
|
||||
if isinstance(err_obj, dict)
|
||||
else str(err_obj) or text
|
||||
)
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
|
||||
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err_json = json.loads(text)
|
||||
err_obj = err_json.get("error", {})
|
||||
msg = (
|
||||
err_obj.get("message") or err_obj.get("msg") or text
|
||||
if isinstance(err_obj, dict)
|
||||
else str(err_obj) or text
|
||||
)
|
||||
resp_json = json.loads(text)
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
|
||||
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
|
||||
|
||||
try:
|
||||
resp_json = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
|
||||
print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
|
||||
return await self._parse_response(resp_json, session)
|
||||
|
||||
print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
|
||||
return await self._parse_response(resp_json, session)
|
||||
if last_status and last_status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[last_status])
|
||||
raise RuntimeError(f"请求失败: 重试 {DEFAULT_MAX_RETRIES} 次后仍然失败")
|
||||
|
||||
return await self._run_with_interrupt(_do_request())
|
||||
|
||||
@@ -446,6 +462,8 @@ class GptImageClient:
|
||||
text = await resp.text()
|
||||
|
||||
if resp.status != 200:
|
||||
if resp.status in HTTP_ERROR_MESSAGES:
|
||||
raise RuntimeError(HTTP_ERROR_MESSAGES[resp.status])
|
||||
try:
|
||||
err_json = json.loads(text)
|
||||
err_obj = err_json.get("error", {})
|
||||
|
||||
@@ -0,0 +1,355 @@
|
||||
"""
|
||||
Grok Image API 客户端
|
||||
支持两个接口：
|
||||
- POST /v1/images/generations 文生图
|
||||
- POST /v1/images/edits 图生图（带参考图）
|
||||
|
||||
上游 API 格式与 OpenAI Images API 兼容。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from io import BytesIO
|
||||
from typing import List, Optional
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_base_url_by_route
|
||||
from ..utils.image_utils import tensor_to_pil
|
||||
|
||||
try:
|
||||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||||
_INTERRUPT_AVAILABLE = True
|
||||
except ImportError:
|
||||
_INTERRUPT_AVAILABLE = False
|
||||
InterruptProcessingException = RuntimeError
|
||||
processing_interrupted = lambda: False
|
||||
|
||||
_ENDPOINT_GENERATIONS = "/v1/images/generations"
|
||||
_ENDPOINT_EDITS = "/v1/images/edits"
|
||||
|
||||
_MODEL_NAME_MAP = {
|
||||
"Grok Image": "grok-imagine-image",
|
||||
"Grok Image Pro": "grok-imagine-image-quality",
|
||||
}
|
||||
|
||||
_REQUEST_TIMEOUT = 900
|
||||
_MAX_BODY_BYTES = 20 * 1024 * 1024
|
||||
|
||||
_MAX_RETRIES = 3
|
||||
_RETRY_DELAY = 5
|
||||
|
||||
|
||||
class GrokImageClient:
|
||||
|
||||
def __init__(self, route: str = "全球加速"):
|
||||
self.api_key = get_api_key_or_raise("O1KEY_API_KEY")
|
||||
self.base_url = get_base_url_by_route(route)
|
||||
|
||||
def _json_headers(self) -> dict:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def _auth_headers(self) -> dict:
|
||||
return {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
# ── 图像工具 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
def _shrink_png_to_limit(png_bytes: bytes, max_bytes: int, label: str = "") -> bytes:
|
||||
if len(png_bytes) <= max_bytes:
|
||||
return png_bytes
|
||||
img = Image.open(BytesIO(png_bytes))
|
||||
w, h = img.size
|
||||
original_size = len(png_bytes)
|
||||
step = 0
|
||||
while len(png_bytes) > max_bytes:
|
||||
scale = 0.894
|
||||
w = max(1, int(w * scale))
|
||||
h = max(1, int(h * scale))
|
||||
img = img.resize((w, h), Image.LANCZOS)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
png_bytes = buf.getvalue()
|
||||
step += 1
|
||||
tag = f" ({label})" if label else ""
|
||||
print(
|
||||
f"[o1key Grok Image] 图像{tag}超出 {max_bytes // (1024*1024)}MB 限制，"
|
||||
f"已等比缩放 {step} 次：{original_size // 1024}KB → {len(png_bytes) // 1024}KB "
|
||||
f"（{w}×{h}）"
|
||||
)
|
||||
return png_bytes
|
||||
|
||||
@staticmethod
|
||||
def _pil_list_to_tensor(images: List[Image.Image]) -> torch.Tensor:
|
||||
if not images:
|
||||
placeholder = Image.new("RGB", (512, 512), (128, 128, 128))
|
||||
images = [placeholder]
|
||||
tensors = []
|
||||
for img in images:
|
||||
arr = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
||||
tensors.append(torch.from_numpy(arr))
|
||||
return torch.stack(tensors, dim=0)
|
||||
|
||||
# ── 中断轮询 ──────────────────────────────────────────────────────────────
|
||||
|
||||
@staticmethod
|
||||
async def _poll_interrupt():
|
||||
while True:
|
||||
await asyncio.sleep(0.5)
|
||||
if _INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
async def _run_with_interrupt(coro):
|
||||
if not _INTERRUPT_AVAILABLE:
|
||||
return await coro
|
||||
request_task = asyncio.ensure_future(coro)
|
||||
interrupt_task = asyncio.ensure_future(GrokImageClient._poll_interrupt())
|
||||
done, pending = await asyncio.wait(
|
||||
[request_task, interrupt_task],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for t in pending:
|
||||
t.cancel()
|
||||
try:
|
||||
await t
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
if interrupt_task in done and request_task not in done:
|
||||
raise InterruptProcessingException()
|
||||
return request_task.result()
|
||||
|
||||
# ── 响应解析 ──────────────────────────────────────────────────────────────
|
||||
|
||||
async def _parse_response(self, resp_json: dict, session: aiohttp.ClientSession) -> List[Image.Image]:
|
||||
if "error" in resp_json:
|
||||
err = resp_json["error"]
|
||||
msg = (
|
||||
err.get("message") or err.get("msg") or json.dumps(err, ensure_ascii=False)
|
||||
if isinstance(err, dict) else str(err)
|
||||
)
|
||||
raise RuntimeError(f"API 返回错误: {msg}")
|
||||
data_list = resp_json.get("data")
|
||||
if not data_list:
|
||||
raise RuntimeError(f"API 响应中未找到 data 字段")
|
||||
images: List[Image.Image] = []
|
||||
for idx, item in enumerate(data_list):
|
||||
b64 = item.get("b64_json", "")
|
||||
url = item.get("url", "")
|
||||
if b64:
|
||||
img_bytes = base64.b64decode(b64)
|
||||
img = Image.open(BytesIO(img_bytes))
|
||||
images.append(img)
|
||||
elif url and url.startswith("http"):
|
||||
async with session.get(url, allow_redirects=True) as r:
|
||||
if r.status != 200:
|
||||
raise RuntimeError(f"图像下载失败 HTTP {r.status}")
|
||||
img_bytes = await r.read()
|
||||
images.append(Image.open(BytesIO(img_bytes)))
|
||||
else:
|
||||
print(f"[o1key Grok Image] 警告：第 {idx + 1} 条数据无有效图像，已跳过")
|
||||
return images
|
||||
|
||||
# ── 文生图（generations 接口）─────────────────────────────────────────────
|
||||
|
||||
async def _generate_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
n: int,
|
||||
) -> List[Image.Image]:
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
body: dict = {
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio if aspect_ratio else "auto",
|
||||
"resolution": resolution if resolution else "1k",
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_GENERATIONS}"
|
||||
log_body = {k: v for k, v in body.items()}
|
||||
print(f"[o1key Grok Image] 请求 URL: {url}")
|
||||
print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")
|
||||
|
||||
results = []
|
||||
for i in range(n):
|
||||
images = await self._do_request_with_retry(url, body)
|
||||
results.extend(images)
|
||||
if n > 1:
|
||||
print(f"[o1key Grok Image] 第 {i+1}/{n} 张完成")
|
||||
return results
|
||||
|
||||
# ── 图生图（edits 接口）───────────────────────────────────────────────────
|
||||
|
||||
async def _edit_async(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
n: int,
|
||||
image_list: List[torch.Tensor],
|
||||
) -> List[Image.Image]:
|
||||
api_model = _MODEL_NAME_MAP.get(model, model)
|
||||
body: dict = {
|
||||
"model": api_model,
|
||||
"prompt": prompt,
|
||||
"response_format": "b64_json",
|
||||
}
|
||||
if aspect_ratio and aspect_ratio != "auto":
|
||||
body["aspect_ratio"] = aspect_ratio
|
||||
if resolution:
|
||||
body["resolution"] = resolution
|
||||
|
||||
# 参考图转 base64 字符串
|
||||
pil_images = tensor_to_pil(image_list[0])
|
||||
img = pil_images[0]
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
png_bytes = buf.getvalue()
|
||||
png_bytes = self._shrink_png_to_limit(png_bytes, _MAX_BODY_BYTES // 2)
|
||||
body["image"] = base64.b64encode(png_bytes).decode("utf-8")
|
||||
|
||||
url = f"{self.base_url}{_ENDPOINT_EDITS}"
|
||||
log_body = {k: (v[:50] + "..." if k == "image" and len(v) > 50 else v) for k, v in body.items()}
|
||||
print(f"[o1key Grok Image] 请求 URL: {url}")
|
||||
print(f"[o1key Grok Image] 请求体: {json.dumps(log_body, ensure_ascii=False)}")
|
||||
|
||||
results = []
|
||||
for i in range(n):
|
||||
images = await self._do_request_with_retry(url, body)
|
||||
results.extend(images)
|
||||
if n > 1:
|
||||
print(f"[o1key Grok Image] 第 {i+1}/{n} 张完成")
|
||||
return results
|
||||
|
||||
# ── 带重试的请求 ────────────────────────────────────────────────────────
|
||||
|
||||
async def _do_request_with_retry(self, url: str, body: dict) -> List[Image.Image]:
|
||||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||
timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT)
|
||||
|
||||
async def _do_request():
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
last_error = None
|
||||
for attempt in range(1, _MAX_RETRIES + 1):
|
||||
t0 = time.time()
|
||||
async with session.post(url, json=body, headers=self._json_headers()) as resp:
|
||||
elapsed = time.time() - t0
|
||||
text = await resp.text()
|
||||
|
||||
if resp.status == 429 or resp.status in (502, 503, 504):
|
||||
last_error = f"HTTP {resp.status}"
|
||||
print(f"[o1key Grok Image] 重试 {attempt}/{_MAX_RETRIES}（{last_error}）")
|
||||
await asyncio.sleep(_RETRY_DELAY * attempt)
|
||||
continue
|
||||
|
||||
if resp.status == 400 and "high load" in text.lower():
|
||||
last_error = "high load"
|
||||
print(f"[o1key Grok Image] 重试 {attempt}/{_MAX_RETRIES}（服务繁忙）")
|
||||
await asyncio.sleep(_RETRY_DELAY * attempt)
|
||||
continue
|
||||
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err_json = json.loads(text)
|
||||
err_obj = err_json.get("error", {})
|
||||
msg = (
|
||||
err_obj.get("message") or err_obj.get("msg") or text
|
||||
if isinstance(err_obj, dict) else str(err_obj) or text
|
||||
)
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
|
||||
|
||||
try:
|
||||
resp_json = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
|
||||
|
||||
print(f"[o1key Grok Image] API 响应耗时 {elapsed:.1f}s")
|
||||
return await self._parse_response(resp_json, session)
|
||||
|
||||
raise RuntimeError(f"重试 {_MAX_RETRIES} 次后仍失败: {last_error}")
|
||||
|
||||
return await self._run_with_interrupt(_do_request())
|
||||
|
||||
# ── 同步入口 ──────────────────────────────────────────────────────────────
|
||||
|
||||
def run_sync(
|
||||
self,
|
||||
prompt: str,
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
resolution: str,
|
||||
n: int,
|
||||
image_list: Optional[List[torch.Tensor]] = None,
|
||||
) -> List[Image.Image]:
|
||||
if image_list:
|
||||
coro = self._edit_async(
|
||||
prompt=prompt, model=model, aspect_ratio=aspect_ratio,
|
||||
resolution=resolution, n=n, image_list=image_list,
|
||||
)
|
||||
else:
|
||||
coro = self._generate_async(
|
||||
prompt=prompt, model=model, aspect_ratio=aspect_ratio,
|
||||
resolution=resolution, n=n,
|
||||
)
|
||||
|
||||
def _run():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(coro)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(_run)
|
||||
try:
|
||||
return future.result(timeout=_REQUEST_TIMEOUT + 30)
|
||||
except TimeoutError:
|
||||
raise RuntimeError("Grok Image 请求超时，请检查网络或稍后重试")
|
||||
|
||||
# ── 余额查询 ──────────────────────────────────────────────────────────────
|
||||
|
||||
async def _query_balance_async(self) -> dict:
|
||||
url = f"{self.base_url}/api/usage/token"
|
||||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||||
timeout = aiohttp.ClientTimeout(total=10)
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
async with session.get(url, headers=self._auth_headers()) as resp:
|
||||
if resp.status != 200:
|
||||
raise RuntimeError(f"余额查询失败 HTTP {resp.status}")
|
||||
return await resp.json()
|
||||
|
||||
def query_balance_sync(self) -> dict:
|
||||
def _run():
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
return loop.run_until_complete(self._query_balance_async())
|
||||
finally:
|
||||
loop.close()
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
return executor.submit(_run).result(timeout=15)
|
||||
|
||||
@staticmethod
|
||||
def format_balance_info(balance_data: dict) -> str:
|
||||
data = balance_data.get("data", {})
|
||||
api_name = data.get("name", "未知")
|
||||
total_available = data.get("total_available", 0)
|
||||
balance_in_dollars = total_available / 500000
|
||||
return f"当前余额：{balance_in_dollars:.2f} | API：{api_name}"
|
||||
+11
-16
@@ -10,6 +10,7 @@ from typing import Any, Callable, Dict, Optional
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise, get_api_base_url
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
|
||||
|
||||
class KlingClient:
|
||||
@@ -49,11 +50,11 @@ class KlingClient:
|
||||
) -> 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)
|
||||
resp = await async_request_with_retry(
|
||||
session, "POST", url, json=body, headers=self._headers(), prefix="Kling 提交: "
|
||||
)
|
||||
text = await resp.text()
|
||||
return json.loads(text)
|
||||
|
||||
# ── 轮询状态 ──────────────────────────────────────────────────────
|
||||
|
||||
@@ -203,17 +204,11 @@ class KlingClient:
|
||||
if on_stage:
|
||||
on_stage("submitting")
|
||||
create_url = f"{self.base_url}{self.NEW_API_CREATE}"
|
||||
async with session.post(create_url, json=body, headers=headers) as resp:
|
||||
text = await resp.text()
|
||||
if resp.status != 200:
|
||||
# 尝试提取友好错误信息
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = err.get("error", {}).get("message") or err.get("message") or text
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"动作控制提交失败 ({resp.status}): {msg}")
|
||||
create_resp = json.loads(text)
|
||||
resp = await async_request_with_retry(
|
||||
session, "POST", create_url, json=body, headers=headers, prefix="Kling 动作控制提交: "
|
||||
)
|
||||
text = await resp.text()
|
||||
create_resp = json.loads(text)
|
||||
|
||||
video_id = create_resp.get("id")
|
||||
if not video_id:
|
||||
|
||||
@@ -204,13 +204,15 @@ class OpenAIAPIClient(BaseAPIClient):
|
||||
"extra_body": {
|
||||
"google": {
|
||||
"image_config": {
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"image_size": api_image_size
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if aspect_ratio and aspect_ratio != "智能":
|
||||
request_body["extra_body"]["google"]["image_config"]["aspect_ratio"] = aspect_ratio
|
||||
|
||||
return request_body
|
||||
|
||||
def parse_response(self, response: Dict[str, Any]) -> List[Image.Image]:
|
||||
|
||||
@@ -11,6 +11,7 @@ from typing import Any, Callable, Dict, Optional
|
||||
import aiohttp
|
||||
|
||||
from ..utils.config import get_api_key_or_raise
|
||||
from ..utils.http_error import async_request_with_retry
|
||||
|
||||
|
||||
class SeedanceClient:
|
||||
@@ -47,18 +48,11 @@ class SeedanceClient:
|
||||
) -> str:
|
||||
"""提交视频生成任务，返回 task_id"""
|
||||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||||
async with session.post(url, json=body, headers=self._headers()) as resp:
|
||||
text = await resp.text()
|
||||
if resp.status != 200:
|
||||
try:
|
||||
err = json.loads(text)
|
||||
msg = (err.get("error", {}).get("message")
|
||||
or err.get("message")
|
||||
or text)
|
||||
except Exception:
|
||||
msg = text
|
||||
raise RuntimeError(f"提交失败 ({resp.status}): {msg}")
|
||||
data = json.loads(text)
|
||||
resp = await async_request_with_retry(
|
||||
session, "POST", url, json=body, headers=self._headers(), prefix="Seedance 提交: "
|
||||
)
|
||||
text = await resp.text()
|
||||
data = json.loads(text)
|
||||
|
||||
# new-api 返回字段：id / task_id
|
||||
task_id = data.get("id") or data.get("task_id")
|
||||
|
||||
Reference in New Issue
Block a user