fix: image gate reacts to ComfyUI Interrupt
GateBus.wait() only checked the gate's own Stop flag, so pressing ComfyUI's Interrupt left the image gate blocked. Add should_cancel to wait() (mirroring wait_payload) and pass mm.processing_interrupted from gate.py, matching the text gate. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+3
-2
@@ -47,13 +47,14 @@ class ImageGate:
|
||||
def run(self, image, routes, unique_id):
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
from . import gate_server
|
||||
import comfy.model_management as mm
|
||||
|
||||
gate_bus.GateBus.arm(unique_id)
|
||||
gate_server.send_preview(unique_id, image, routes)
|
||||
try:
|
||||
chosen_1 = gate_bus.GateBus.wait(unique_id)
|
||||
chosen_1 = gate_bus.GateBus.wait(
|
||||
unique_id, should_cancel=mm.processing_interrupted)
|
||||
except gate_bus.GateCancelled:
|
||||
import comfy.model_management as mm
|
||||
raise mm.InterruptProcessingException()
|
||||
|
||||
mask = mask_from_stash(gate_bus.GateBus.pop_mask(unique_id), image)
|
||||
|
||||
+2
-2
@@ -27,10 +27,10 @@ class GateBus:
|
||||
cls.messages[str(node_id)] = int(message)
|
||||
|
||||
@classmethod
|
||||
def wait(cls, node_id, period=0.1):
|
||||
def wait(cls, node_id, period=0.1, should_cancel=None):
|
||||
sid = str(node_id)
|
||||
while sid not in cls.messages:
|
||||
if cls.cancelled:
|
||||
if cls.cancelled or (should_cancel is not None and should_cancel()):
|
||||
cls.cancelled = False
|
||||
raise GateCancelled()
|
||||
time.sleep(period)
|
||||
|
||||
@@ -65,3 +65,10 @@ def test_wait_payload_should_cancel_raises():
|
||||
gb.GateBus.arm("p")
|
||||
with pytest.raises(gb.GateCancelled):
|
||||
gb.GateBus.wait_payload("p", should_cancel=lambda: True)
|
||||
|
||||
def test_wait_should_cancel_raises():
|
||||
# image gate: ComfyUI Interrupt (should_cancel) must abort the wait too
|
||||
gb.GateBus.arm("7")
|
||||
with pytest.raises(gb.GateCancelled):
|
||||
gb.GateBus.wait("7", should_cancel=lambda: True)
|
||||
assert gb.GateBus.cancelled is False
|
||||
|
||||
Reference in New Issue
Block a user