JonyandClaude Opus 4.6
69279c654d
- 新增 Grok Image 节点及客户端 - 新增 save_image_format 节点 - 新增前端 JS 扩展：画笔工具、点阵网格、侧边栏隐藏、资源切换、重命名等 - 重构 nano-banana 节点，移除 pro 版本 - 移除 multi_res_preview 节点 - 新增 http_error 工具模块 - 各客户端和节点优化改进 Co-Authored-By: Claude Opus 4.6 <[email protected]>
356 lines
14 KiB
Python
356 lines
14 KiB
Python
"""
|
||
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}"
|