feat: GPT Image 节点支持点击取消立即中断请求

Co-Authored-By: Claude Opus 4.7 <[email protected]>

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,6 +328,7 @@ class GptImageClient:
         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:
                 t0 = time.time()
                 async with session.post(url, json=body, headers=self._json_headers()) as resp:
@@ -308,6 +356,8 @@ class GptImageClient:
                 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）──────────────────────────
 
     async def _edit_async(
@@ -384,6 +434,7 @@ class GptImageClient:
         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:
                 t0 = time.time()
                 async with session.post(
@@ -415,6 +466,8 @@ class GptImageClient:
                 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())
+
     # ── 同步统一入口（供节点调用）────────────────────────────────────────────
 
     def run_sync(

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}")