From aae98c0f891e6dcd15c7af2cb6a31c6d9f4a5bc3 Mon Sep 17 00:00:00 2001 From: o1key <951565127@qq.com> Date: Tue, 12 May 2026 14:55:18 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20GPT=20Image=20=E8=8A=82=E7=82=B9?= =?UTF-8?q?=E6=94=AF=E6=8C=81=E7=82=B9=E5=87=BB=E5=8F=96=E6=B6=88=E7=AB=8B?= =?UTF-8?q?=E5=8D=B3=E4=B8=AD=E6=96=AD=E8=AF=B7=E6=B1=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.7 --- clients/gpt_image_client.py | 145 ++++++++++++++++++++++++------------ nodes/gpt_image.py | 15 ++++ 2 files changed, 114 insertions(+), 46 deletions(-) diff --git a/clients/gpt_image_client.py b/clients/gpt_image_client.py index 6a69005..c18d351 100644 --- a/clients/gpt_image_client.py +++ b/clients/gpt_image_client.py @@ -27,6 +27,14 @@ 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 +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/" @@ -220,6 +228,45 @@ class GptImageClient: return images + # ── 中断轮询 ────────────────────────────────────────────────────────────── + + @staticmethod + async def _poll_interrupt(): + """每 0.5s 轮询一次 ComfyUI 中断标志""" + while True: + await asyncio.sleep(0.5) + if _INTERRUPT_AVAILABLE and processing_interrupted(): + return + + @staticmethod + async def _run_with_interrupt(coro): + """ + 将异步任务与中断轮询并发执行。 + 如果用户点击取消，cancel 掉 coro 并抛出 InterruptProcessingException。 + """ + if not _INTERRUPT_AVAILABLE: + return await coro + + request_task = asyncio.ensure_future(coro) + interrupt_task = asyncio.ensure_future(GptImageClient._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() + # ── 文生图 / 图生图（generations 接口）─────────────────────────────────── async def _generate_async( @@ -281,32 +328,35 @@ class GptImageClient: connector = aiohttp.TCPConnector(ssl=False, force_close=True) timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT) - 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() + 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() + + 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}") - 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) + return await self._run_with_interrupt(_do_request()) # ── 图像编辑（edits 接口，multipart/form-data）────────────────────────── @@ -384,36 +434,39 @@ class GptImageClient: connector = aiohttp.TCPConnector(ssl=False, force_close=True) timeout = aiohttp.ClientTimeout(total=_REQUEST_TIMEOUT) - async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: - t0 = time.time() - async with session.post( - url, - data=form, - headers=self._auth_headers(), - ) as resp: - elapsed = time.time() - t0 - text = await resp.text() + async def _do_request(): + async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: + t0 = time.time() + async with session.post( + url, + data=form, + headers=self._auth_headers(), + ) as resp: + elapsed = time.time() - t0 + text = await resp.text() + + 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}") - 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) + return await self._run_with_interrupt(_do_request()) # ── 同步统一入口（供节点调用）──────────────────────────────────────────── diff --git a/nodes/gpt_image.py b/nodes/gpt_image.py index 4435b4a..aaae936 100644 --- a/nodes/gpt_image.py +++ b/nodes/gpt_image.py @@ -7,6 +7,14 @@ import time from ..clients.gpt_image_client import GptImageClient from ..utils.image_utils import parse_batch_prompts +try: + from comfy.model_management import processing_interrupted, InterruptProcessingException + _INTERRUPT_AVAILABLE = True +except ImportError: + _INTERRUPT_AVAILABLE = False + processing_interrupted = lambda: False + InterruptProcessingException = RuntimeError + class O1keyGPTImage: """ @@ -163,6 +171,9 @@ class O1keyGPTImage: total = len(batch_prompts) print(f"[o1key GPT Image] 批量模式 | {total} 条提示词 | 每条生成 {生图数量} 张") for idx, p in enumerate(batch_prompts, 1): + if _INTERRUPT_AVAILABLE and processing_interrupted(): + print("[o1key GPT Image] 用户取消，已中断批量生成") + raise InterruptProcessingException() try: pil_images = client.run_sync( prompt=p, @@ -177,6 +188,8 @@ class O1keyGPTImage: all_pil_images.extend(pil_images) snippet = p[:30] + ("..." if len(p) >= 30 else "") print(f"[o1key GPT Image] [{idx}/{total}] ✓ {snippet}") + except InterruptProcessingException: + raise except Exception as e: error_msg = str(e).split('\n')[0] snippet = p[:30] + ("..." if len(p) >= 30 else "") @@ -197,6 +210,8 @@ class O1keyGPTImage: mask_tensor=遮罩, ) all_pil_images.extend(pil_images) + except InterruptProcessingException: + raise except Exception as e: error_msg = str(e).split('\n')[0] print(f"[o1key GPT Image] ❌ {error_msg}")