feat: ImageGate node — pause, route via ExecutionBlocker, mask out

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-21 17:42:54 +02:00
parent ea3438567a
commit d8dbc4fb4b
2 changed files with 52 additions and 0 deletions
@@ -22,3 +22,45 @@ def mask_from_stash(data, image):
    m = Image.open(io.BytesIO(data)).convert("L")
    arr = np.array(m, dtype=np.float32) / 255.0
    return torch.from_numpy(arr).unsqueeze(0)
+
+
+class ImageGate:
+    CATEGORY = "Datasete Gates"
+    FUNCTION = "run"
+    RETURN_TYPES = ("MASK",) + ("IMAGE",) * MAX_ROUTES
+    RETURN_NAMES = ("mask",) + tuple(f"route_{i + 1}" for i in range(MAX_ROUTES))
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "routes": ("INT", {"default": 2, "min": 1, "max": MAX_ROUTES}),
+            },
+            "hidden": {"unique_id": "UNIQUE_ID"},
+        }
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("nan")            # always pause; never cached
+
+    def run(self, image, routes, unique_id):
+        from comfy_execution.graph_utils import ExecutionBlocker
+        from . import gate_server
+
+        gate_bus.GateBus.arm(unique_id)
+        gate_server.send_preview(unique_id, image, routes)
+        try:
+            chosen_1 = gate_bus.GateBus.wait(unique_id)
+        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)
+        chosen = max(0, min(chosen_1 - 1, routes - 1))
+        blocker = ExecutionBlocker(None)
+        return (mask,) + route_tuple(chosen, image, blocker, MAX_ROUTES)
+
+
+NODE_CLASS_MAPPINGS = {"ImageGate": ImageGate}
+NODE_DISPLAY_NAME_MAPPINGS = {"ImageGate": "Image Gate (Manual Router)"}