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,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:
+        async def _do_request():
-            t0 = time.time()
+            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
-            async with session.post(url, json=body, headers=self._json_headers()) as resp:
+                t0 = time.time()
-                elapsed = time.time() - t0
+                async with session.post(url, json=body, headers=self._json_headers()) as resp:
-                text = await resp.text()
+                    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)
+                        resp_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"响应 JSON 解析失败，原始内容：{text[:500]}")
-                    raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
 
-                try:
+                print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
-                    resp_json = json.loads(text)
+                return await self._parse_response(resp_json, session)
-                except Exception:
-                    raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
 
-            print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
+        return await self._run_with_interrupt(_do_request())
-            return await self._parse_response(resp_json, session)
 
     # ── 图像编辑（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:
+        async def _do_request():
-            t0 = time.time()
+            async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
-            async with session.post(
+                t0 = time.time()
-                url,
+                async with session.post(
-                data=form,
+                    url,
-                headers=self._auth_headers(),
+                    data=form,
-            ) as resp:
+                    headers=self._auth_headers(),
-                elapsed = time.time() - t0
+                ) as resp:
-                text = await resp.text()
+                    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)
+                        resp_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"响应 JSON 解析失败，原始内容：{text[:500]}")
-                    raise RuntimeError(f"请求失败 HTTP {resp.status}: {msg}")
 
-                try:
+                print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
-                    resp_json = json.loads(text)
+                return await self._parse_response(resp_json, session)
-                except Exception:
-                    raise RuntimeError(f"响应 JSON 解析失败，原始内容：{text[:500]}")
 
-            print(f"[o1key GPT Image] API 响应耗时 {elapsed:.1f}s")
+        return await self._run_with_interrupt(_do_request())
-            return await self._parse_response(resp_json, session)
 
     # ── 同步统一入口（供节点调用）────────────────────────────────────────────
 

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