From 7e8878badef143c9755a55344827e72430db4d16 Mon Sep 17 00:00:00 2001 From: Ethan Fel Date: Sun, 21 Jun 2026 17:09:26 +0200 Subject: [PATCH] Add Image Gate (Manual Router) design + implementation plan Interactive chooser/router: pauses execution, shows the image with up to 10 labeled route buttons + edit-mask + stop. Chosen route gets the image, others ExecutionBlocker-ed; gate-painted mask on a fixed output; stop raises InterruptProcessingException. TDD plan with a pure torch-free gate_bus; lazy comfy imports keep node logic unit-testable. Co-Authored-By: Claude Opus 4.8 --- docs/plans/2026-06-21-image-gate-design.md | 83 ++++ .../2026-06-21-image-gate-implementation.md | 431 ++++++++++++++++++ 2 files changed, 514 insertions(+) create mode 100644 docs/plans/2026-06-21-image-gate-design.md create mode 100644 docs/plans/2026-06-21-image-gate-implementation.md diff --git a/docs/plans/2026-06-21-image-gate-design.md b/docs/plans/2026-06-21-image-gate-design.md new file mode 100644 index 0000000..77a3e20 --- /dev/null +++ b/docs/plans/2026-06-21-image-gate-design.md @@ -0,0 +1,83 @@ +# Image Gate (Manual Router) — Design + +Date: 2026-06-21 +Status: Approved (brainstorming complete, ready for implementation plan) + +## 1. Purpose + +An interactive "image chooser on steroids": during a prompt run the node **pauses**, +shows the incoming image with a row of labeled **route buttons**, and waits for a human +click. Clicking **route K** sends the image down output K (all other route branches are +silently skipped). A **Stop** button cancels the whole run. Optionally, an **Edit mask** +button opens ComfyUI's MaskEditor on the image and the painted mask is emitted on a +single `mask` output. Built for manual dataset sorting/gating in the "Dataset Gates" suite. + +Third node in the `ComfyUI-Datasete-Gates` package. + +## 2. IO + +| dir | name | type | notes | +|---|---|---|---| +| in | `image` | IMAGE | the image (or batch, routed as one unit) | +| widget | `routes` | INT, default 2, 1..10 | number of visible route buttons/outputs | +| widget | per-route labels | (frontend) | editable, default `1..N`; rename the visible output slots | +| hidden | `unique_id` | UNIQUE_ID | node id, used to key the pause/choice | +| out | `mask` | MASK | **fixed slot 0**; painted at the gate, zeros (sized to image) if none | +| out | `route_1 … route_10` | IMAGE | dynamic; JS shows only `routes` of them, labeled | + +`RETURN_TYPES = ("MASK",) + ("IMAGE",)*10`. The node always returns all 11 outputs; the +chosen route carries the image, every other route returns `ExecutionBlocker(None)`. JS +hides the unused slots (>`routes`); their `ExecutionBlocker` returns are harmless. + +## 3. Behavior (the pause) + +On execute: +1. Push the image to the UI (`PromptServer.send_sync`, base64 or temp file) so the node + body shows the preview + the N labeled route buttons + **🖌 Edit mask** + **■ Stop**. +2. **Block** the executor thread on our own `GateBus.wait(unique_id)` (a `MessageHolder`- + style singleton in a `sleep(0.1)` loop; separate namespace from cg-image-picker). +3. Resolve: + - **route K** → image to output `K`, `ExecutionBlocker(None)` to the other routes; + `mask` = the painted mask (or zeros). + - **🖌 Edit mask** → opens MaskEditor (reuse the pool node's clipspace flow); the mask + is POSTed to `/datasete_gate/mask` keyed by `unique_id` and picked up on resume. + - **■ Stop** → cancel the prompt cleanly via + `comfy.model_management.InterruptProcessingException` (confirm exact symbol in plan). + +`IS_CHANGED` returns `nan` → the gate pauses on **every** run (never cached). + +## 4. Why the global mask is safe + +Verified in `execution.py:257-266` + `305-306`: if **any** input of a node is an +`ExecutionBlocker`, the node is skipped and the blocker propagates to all its outputs. +So a non-chosen route's downstream (which consumes the blocked routed image) never runs, +regardless of the live `mask` value. Caveat: a node wired to `mask` *only* (no routed +image) would run unconditionally — not the intended wiring. + +## 5. Code shape (same package) + +- `gates/gate.py` — `ImageGate` node: `INPUT_TYPES`, `IS_CHANGED=nan`, `run()` (push + preview → block → route via `ExecutionBlocker`). Pure helper `route_tuple(chosen, image, + blocker, max_routes)` for unit testing. +- `gates/gate_server.py` — `GateBus` (start/put/wait/cancel) + mask stash; aiohttp routes + `/datasete_gate/choice` and `/datasete_gate/mask`; `send_preview()` helper. +- `web/image_gate.js` — dynamic labeled outputs (show `routes` of 10), preview render, + route/stop/mask buttons, posts the choice; reuses the pool's MaskEditor helper. + +## 6. Edge cases + +- `routes` changed between runs → JS re-syncs visible slots; Python clamps `chosen` to + `routes`. +- Stop while no mask painted → clean interrupt, no output. +- Multiple gates in one graph → execute sequentially (single executor thread), so only one + blocks at a time; still keyed by `unique_id`. +- Batch input → previewed as the first image / small grid; routed as one unit. +- External queue-cancel → `GateBus` honors the cancel flag and raises. + +## 7. Testing + +- pytest: `route_tuple` (image at chosen, blocker elsewhere, correct length); `GateBus` + (pre-seeded message returns; cancel raises; `start` resets); mask zero-fallback sizing. +- Manual (live): pause appears, buttons labeled, click routes image to the right branch + and only that branch runs; Edit mask round-trips and feeds `mask`; Stop cancels cleanly; + changing `routes` adds/removes slots. diff --git a/docs/plans/2026-06-21-image-gate-implementation.md b/docs/plans/2026-06-21-image-gate-implementation.md new file mode 100644 index 0000000..19f4a1f --- /dev/null +++ b/docs/plans/2026-06-21-image-gate-implementation.md @@ -0,0 +1,431 @@ +# Image Gate (Manual Router) Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Build a ComfyUI custom node `Image Gate (Manual Router)` that pauses a running prompt, shows the image with up to 10 labeled route buttons + a mask-edit + a stop button, and routes the image down the clicked output (others `ExecutionBlocker`-ed), emitting any gate-painted mask on a fixed `mask` output. + +**Architecture:** A pure, torch-free `gates/gate_bus.py` (a `MessageHolder`-style blocking waiter + mask stash) is unit-testable without ComfyUI. `gates/gate.py` holds the node plus pure helpers (`route_tuple`, `mask_from_stash`); it imports `ExecutionBlocker`/`model_management` lazily so tests don't need comfy. `gates/gate_server.py` is the aiohttp glue (choice/mask routes + `send_preview`). `web/image_gate.js` renders preview + dynamic labeled outputs + buttons and posts the choice; it reuses the pool node's MaskEditor helper. + +**Tech Stack:** Python 3.12, torch 2.8, Pillow, numpy, aiohttp; pytest 9; vanilla JS frontend. + +--- + +## Conventions (read once) + +- **Test python:** `/media/p5/miniforge3/bin/python` (`PY=...`). +- **Run tests:** `cd /media/p5/ComfyUI-Datasete-Gates && $PY -m pytest tests/test_gate_bus.py tests/test_gate.py -v` +- **Concurrency:** other sessions may share this working tree. Stage only this node's paths + when committing; re-Read `__init__.py` before editing (Task 6) and *extend*, don't overwrite. +- `gates/gate_bus.py` MUST be import-safe without comfy/torch (stdlib only). +- `gates/gate.py` MUST import `ExecutionBlocker` and `comfy.model_management` **lazily inside + `run()`** (and `send_preview` lazily) so `import gates.gate` works under pytest. +- Mask convention: grayscale `L`, white = painted; zeros sized to the image if none. +- Commit style: Conventional Commits + repo Co-Authored-By trailer. +- `MAX_ROUTES = 10`. + +--- + +### Task 1: `gate_bus.py` — `GateBus` (arm/put/wait/cancel) + +**Files:** Create `gates/gate_bus.py`; Test `tests/test_gate_bus.py` + +**Step 1: Failing test** + +```python +# tests/test_gate_bus.py +import pytest +from gates import gate_bus as gb + +def test_put_and_wait_returns_choice(): + gb.GateBus.arm("7") + gb.GateBus.put("7", "3") + assert gb.GateBus.wait("7") == 3 + +def test_wait_consumes_message(): + gb.GateBus.arm("7") + gb.GateBus.put("7", "2") + gb.GateBus.wait("7") + assert "7" not in gb.GateBus.messages + +def test_cancel_raises_and_resets(): + gb.GateBus.arm("7") + gb.GateBus.put("7", "__cancel__") + with pytest.raises(gb.GateCancelled): + gb.GateBus.wait("7") + assert gb.GateBus.cancelled is False # reset after raising + +def test_arm_clears_stale_state(): + gb.GateBus.put("1", "5") + gb.GateBus.cancelled = True + gb.GateBus.arm("1") + assert "1" not in gb.GateBus.messages + assert gb.GateBus.cancelled is False +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +# gates/gate_bus.py +"""Blocking choice bus for the Image Gate node. Stdlib only — no comfy/torch.""" +import time + + +class GateCancelled(Exception): + pass + + +class GateBus: + messages = {} # node_id(str) -> chosen int (1-based) + masks = {} # node_id(str) -> PNG bytes + cancelled = False + + @classmethod + def arm(cls, node_id): + cls.messages.pop(str(node_id), None) + cls.masks.pop(str(node_id), None) + cls.cancelled = False + + @classmethod + def put(cls, node_id, message): + if message == "__cancel__": + cls.cancelled = True + else: + cls.messages[str(node_id)] = int(message) + + @classmethod + def wait(cls, node_id, period=0.1): + sid = str(node_id) + while sid not in cls.messages: + if cls.cancelled: + cls.cancelled = False + raise GateCancelled() + time.sleep(period) + return cls.messages.pop(sid) +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: gate_bus blocking choice waiter` + +--- + +### Task 2: `gate_bus.py` — mask stash + +**Files:** Modify `gates/gate_bus.py`, `tests/test_gate_bus.py` + +**Step 1: Failing test** + +```python +def test_mask_stash_roundtrip(): + gb.GateBus.put_mask("9", b"PNGDATA") + assert gb.GateBus.pop_mask("9") == b"PNGDATA" + assert gb.GateBus.pop_mask("9") is None # popped + +def test_arm_clears_mask(): + gb.GateBus.put_mask("9", b"x") + gb.GateBus.arm("9") + assert gb.GateBus.pop_mask("9") is None +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement (append to `GateBus`)** + +```python + @classmethod + def put_mask(cls, node_id, data): + cls.masks[str(node_id)] = data + + @classmethod + def pop_mask(cls, node_id): + return cls.masks.pop(str(node_id), None) +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: gate_bus mask stash` + +--- + +### Task 3: `gate.py` — `route_tuple` pure helper + +**Files:** Create `gates/gate.py`; Test `tests/test_gate.py` + +**Step 1: Failing test** + +```python +# tests/test_gate.py +from gates import gate + +def test_route_tuple_places_image_at_chosen(): + B = object() + t = gate.route_tuple(2, "IMG", B, max_routes=5) + assert t == (B, B, "IMG", B, B) + +def test_route_tuple_length_is_max(): + B = object() + assert len(gate.route_tuple(0, "IMG", B, max_routes=10)) == 10 +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +# gates/gate.py +import io +import math + +import numpy as np +import torch +from PIL import Image + +from . import gate_bus + +MAX_ROUTES = 10 + + +def route_tuple(chosen, image, blocker, max_routes=MAX_ROUTES): + return tuple(image if i == chosen else blocker for i in range(max_routes)) +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: gate route_tuple helper` + +--- + +### Task 4: `gate.py` — `mask_from_stash` + +**Files:** Modify `gates/gate.py`, `tests/test_gate.py` + +**Step 1: Failing test** + +```python +import io, torch +from PIL import Image + +def test_mask_from_stash_none_is_zeros(): + img = torch.zeros((1, 6, 4, 3)) + m = gate.mask_from_stash(None, img) + assert m.shape == (1, 6, 4) and float(m.max()) == 0.0 + +def test_mask_from_stash_decodes_png(): + buf = io.BytesIO(); Image.new("L", (4, 6), 255).save(buf, "PNG") + img = torch.zeros((1, 6, 4, 3)) + m = gate.mask_from_stash(buf.getvalue(), img) + assert m.shape == (1, 6, 4) and float(m.min()) > 0.99 +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement (append)** + +```python +def mask_from_stash(data, image): + b, h, w = image.shape[0], image.shape[1], image.shape[2] + if not data: + return torch.zeros((b, h, w), dtype=torch.float32) + m = Image.open(io.BytesIO(data)).convert("L") + arr = np.array(m, dtype=np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: gate mask_from_stash (paint or zeros)` + +--- + +### Task 5: `gate.py` — `ImageGate` node class + +**Files:** Modify `gates/gate.py`, `tests/test_gate.py` + +**Step 0: Verify the interrupt symbol** (so Stop cancels cleanly): +`grep -n "class InterruptProcessingException\|def interrupt_current_processing" /media/p5/Comfyui/comfy/model_management.py` +Use whatever exists (expected: `InterruptProcessingException`). + +**Step 1: Failing test** + +```python +import math + +def test_is_changed_always_nan(): + v = gate.ImageGate.IS_CHANGED(image=None, routes=2, unique_id="1") + assert math.isnan(v) + +def test_return_types_shape(): + assert gate.ImageGate.RETURN_TYPES[0] == "MASK" + assert len(gate.ImageGate.RETURN_TYPES) == gate.MAX_ROUTES + 1 + assert all(t == "IMAGE" for t in gate.ImageGate.RETURN_TYPES[1:]) +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement (append)** + +```python +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() # confirm symbol in Step 0 + + 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)"} +``` + +**Step 4: Run → PASS.** (`run()` itself is covered by the live smoke test, not unit tests.) + +**Step 5: Commit** `feat: ImageGate node — pause, route via ExecutionBlocker, mask out` + +--- + +### Task 6: `gate_server.py` — routes + preview, and register (MERGE) + +**Files:** Create `gates/gate_server.py`; Modify `__init__.py` + +**Step 1: Implement `gates/gate_server.py`** (aiohttp glue — verified live, not unit-tested) + +```python +# gates/gate_server.py +import base64 +import io + +import numpy as np +from aiohttp import web +from PIL import Image +from server import PromptServer + +from .gate_bus import GateBus + +routes = PromptServer.instance.routes + + +def send_preview(node_id, image, n_routes): + arr = (image[0].cpu().numpy() * 255.0).clip(0, 255).astype("uint8") + buf = io.BytesIO() + Image.fromarray(arr).save(buf, "PNG") + b64 = base64.b64encode(buf.getvalue()).decode() + PromptServer.instance.send_sync( + "datasete-gate-show", + {"id": str(node_id), "image": b64, "routes": int(n_routes)}, + ) + + +@routes.post("/datasete_gate/choice") +async def _choice(request): + post = await request.post() + GateBus.put(post.get("id"), post.get("message")) + return web.json_response({}) + + +@routes.post("/datasete_gate/mask") +async def _mask(request): + reader = await request.multipart() + node_id, data = None, None + async for part in reader: + if part.name == "id": + node_id = await part.text() + elif part.name == "mask": + data = await part.read(decode=False) + if node_id is not None: + GateBus.put_mask(node_id, data) + return web.json_response({}) +``` + +**Step 2: Re-Read `__init__.py`** and extend the `if __package__:` block to merge the gate +node and import its server (registers routes): + +```python + from .gates.gate import NODE_CLASS_MAPPINGS as _GATE_NODES, \ + NODE_DISPLAY_NAME_MAPPINGS as _GATE_NAMES + from .gates import gate_server # noqa: F401 (registers /datasete_gate/* routes) + NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **_GATE_NODES} + NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **_GATE_NAMES} +``` +(Adapt to the file's current merge structure; the only requirement is the gate node ends up +in the mappings and `gate_server` is imported.) + +**Step 3:** `$PY -c "import gates.gate; print(gates.gate.NODE_CLASS_MAPPINGS)"` → shows ImageGate. + +**Step 4:** Full suite green: `$PY -m pytest tests/ -v` + +**Step 5: Commit** `feat: gate server routes + preview + register ImageGate` + +--- + +### Task 7: `web/image_gate.js` — preview, dynamic outputs, buttons + +**Files:** Create `web/image_gate.js` + +Implement an `app.registerExtension` for `ImageGate`: +- **Dynamic outputs:** on `nodeCreated` and when the `routes` widget changes, show only the + first `routes` of the 10 `route_*` outputs (hide/remove the rest); give each visible output + an editable label (default `1..N`) persisted in `widgets_values`; keep the `mask` output + (slot 0) always visible. (Reuse your existing dynamic-slot pattern.) +- **Preview + buttons:** listen for the `datasete-gate-show` socket event + (`api.addEventListener`); when it fires for this node's id, render the image in a DOM widget + with: one button per visible route (labeled), an **🖌 Edit mask** button, and a **■ Stop** + button. +- **Choice:** route button → POST `/datasete_gate/choice` `{id, message: <1-based index>}`. + Stop → POST `{id, message: "__cancel__"}`. +- **Mask:** 🖌 → open MaskEditor on the previewed image (reuse the pool node's clipspace + helper); on save, export the grayscale mask PNG and POST it to `/datasete_gate/mask` + (multipart `id`, `mask`) **before** clicking a route. + +**Manual verification (live, Task 8 covers the run):** node shows N labeled outputs that +track the `routes` widget; labels persist across reload. + +**Commit** `feat: image gate frontend — preview, dynamic outputs, route/stop/mask` + +--- + +### Task 8: Live smoke test in ComfyUI + +Restart ComfyUI (repo already symlinked into `custom_nodes`). Build: `Folder Image Loader → +Image Gate`, wire `route_1`/`route_2` to two `PreviewImage`/`SaveImage` nodes, `mask` to a +`MaskPreview`. Verify: +- [ ] "Image Gate (Manual Router)" appears under "Datasete Gates". +- [ ] Queue → execution **pauses**, image preview + labeled buttons + 🖌 + ■ appear. +- [ ] Click route 1 → only route-1's downstream runs; route-2's does not. +- [ ] Click route 2 → only route-2's downstream runs. +- [ ] 🖌 Edit mask → MaskEditor opens; paint, save; then click a route → `mask` output carries the painted mask; no mask painted → zeros. +- [ ] ■ Stop → the run cancels cleanly (no scary traceback; queue stops). +- [ ] Change `routes` from 2→4 → two more labeled outputs appear; reload keeps labels. +- [ ] Run twice in a row → it pauses **both** times (not cached). + +**Commit** (if fixes) `fix: image gate live-test adjustments` + +--- + +## Definition of done + +- `$PY -m pytest tests/test_gate_bus.py tests/test_gate.py -v` green; full `tests/` green. +- Manual checklist passes: pause, route isolation (ExecutionBlocker), mask round-trip, clean Stop, dynamic labeled outputs.