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]>
2026-05-24 22:46:55 +08:00
parent c491731c99
commit 69279c654d
40 changed files with 3406 additions and 1599 deletions
@@ -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}"