Compare commits
2 Commits
d9134b4e9b
...
78b1b85a11
| Author | SHA1 | Date | |
|---|---|---|---|
| 78b1b85a11 | |||
| b50718f7fb |
+3
-2
@@ -47,13 +47,14 @@ class ImageGate:
|
|||||||
def run(self, image, routes, unique_id):
|
def run(self, image, routes, unique_id):
|
||||||
from comfy_execution.graph_utils import ExecutionBlocker
|
from comfy_execution.graph_utils import ExecutionBlocker
|
||||||
from . import gate_server
|
from . import gate_server
|
||||||
|
import comfy.model_management as mm
|
||||||
|
|
||||||
gate_bus.GateBus.arm(unique_id)
|
gate_bus.GateBus.arm(unique_id)
|
||||||
gate_server.send_preview(unique_id, image, routes)
|
gate_server.send_preview(unique_id, image, routes)
|
||||||
try:
|
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:
|
except gate_bus.GateCancelled:
|
||||||
import comfy.model_management as mm
|
|
||||||
raise mm.InterruptProcessingException()
|
raise mm.InterruptProcessingException()
|
||||||
|
|
||||||
mask = mask_from_stash(gate_bus.GateBus.pop_mask(unique_id), image)
|
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)
|
cls.messages[str(node_id)] = int(message)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def wait(cls, node_id, period=0.1):
|
def wait(cls, node_id, period=0.1, should_cancel=None):
|
||||||
sid = str(node_id)
|
sid = str(node_id)
|
||||||
while sid not in cls.messages:
|
while sid not in cls.messages:
|
||||||
if cls.cancelled:
|
if cls.cancelled or (should_cancel is not None and should_cancel()):
|
||||||
cls.cancelled = False
|
cls.cancelled = False
|
||||||
raise GateCancelled()
|
raise GateCancelled()
|
||||||
time.sleep(period)
|
time.sleep(period)
|
||||||
|
|||||||
@@ -65,3 +65,10 @@ def test_wait_payload_should_cancel_raises():
|
|||||||
gb.GateBus.arm("p")
|
gb.GateBus.arm("p")
|
||||||
with pytest.raises(gb.GateCancelled):
|
with pytest.raises(gb.GateCancelled):
|
||||||
gb.GateBus.wait_payload("p", should_cancel=lambda: True)
|
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