From 95b3417ff614c2a71feefe1675f082384aa2187e Mon Sep 17 00:00:00 2001 From: Ethan Fel Date: Sun, 21 Jun 2026 20:45:15 +0200 Subject: [PATCH] Add Image Gate send/get bus design + implementation plan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Disk-backed image bus (input/gate_bus//): gates auto-publish image+mask to a named send_id on pass; when image input is empty they load from get_id (dropdown) — wireless, cycle-free "restart from the gate point" across runs. Making image optional implements ignore-on-normal-path. TDD plan with a pure stdlib imagebus + tensor savers; comfy imports stay lazy. Co-Authored-By: Claude Opus 4.8 --- .../plans/2026-06-21-image-gate-bus-design.md | 95 +++++ ...026-06-21-image-gate-bus-implementation.md | 376 ++++++++++++++++++ 2 files changed, 471 insertions(+) create mode 100644 docs/plans/2026-06-21-image-gate-bus-design.md create mode 100644 docs/plans/2026-06-21-image-gate-bus-implementation.md diff --git a/docs/plans/2026-06-21-image-gate-bus-design.md b/docs/plans/2026-06-21-image-gate-bus-design.md new file mode 100644 index 0000000..17e3c8e --- /dev/null +++ b/docs/plans/2026-06-21-image-gate-bus-design.md @@ -0,0 +1,95 @@ +# Image Gate — Send/Get Bus (teleport + checkpoint) — Design + +Date: 2026-06-21 +Status: Approved (brainstorming complete, ready for implementation plan) + +## 1. Purpose + +Let Image Gates pass images to each other by **name** through a disk-backed bus, so you can +**re-enter the pipeline at a gate** after manual editing/looking — without dragging wires and +without creating graph cycles. A gate **auto-publishes** its passed image (+ mask) to a named +slot; another gate (or a fresh workflow) **loads** that slot to resume from that point. + +This is an enhancement to the existing `Image Gate (Manual Router)` — no new node. + +## 2. Why no wire / no cycle + +ComfyUI graphs must be acyclic; a real wire from a downstream gate's output back into an +upstream gate is a cycle and is rejected at validation. The bus links sender↔receiver by a +**string id**, so there is no live wire and no cycle. "Ignore on the normal path" falls out +naturally from making `image` optional (see §4). + +## 3. Changes to the Image Gate + +New ports/widgets (all backward compatible): + +| Port | Type | Description | +|------|------|-------------| +| `image` | IMAGE | **now optional.** Wired → normal path. Empty → load from `get_id`. | +| `send_id` | STRING (widget) | If non-empty, on every **pass** the chosen image + mask are written to the bus slot `send_id` (latest-wins). Empty = don't publish. | +| `get_id` | STRING (widget, dropdown) | Used only when `image` is **not** connected: load the latest image + mask from this bus slot, then gate as usual. Dropdown lists existing bus ids. | + +Existing inputs (`routes`) and outputs (`mask`, `route_1..route_10`) are unchanged. + +## 4. Run logic + +``` +base = input/gate_bus +image, loaded_mask = resolve_source(base, image, get_id) + # image given -> (image, None) [normal path; get ignored] + # else get_id -> load (image, mask) from bus slot [re-entry] + # else -> nothing: block all routes silently, return zero mask +pause + wait (Stop -> InterruptProcessingException) [unchanged] +mask = painted-at-gate OR loaded_mask OR zeros [precedence] +if send_id: write image+mask to bus[send_id] [auto-publish on pass] +return (mask,) + route_tuple(chosen) [unchanged routing] +``` + +`IS_CHANGED` stays `nan` (always pauses). A gate with no image and no valid `get_id` is a +silent no-op (all routes `ExecutionBlocker`, zero mask) so it never breaks a graph. + +## 5. Bus storage + +``` +input/gate_bus// +├── image.png # latest passed image for this slot +└── mask.png # its mask (white = painted) +``` +Latest-wins (overwrite). `id` is the human-chosen name. Survives restart → cross-run resume. + +## 6. Frontend (`web/image_gate.js`) + +- Make the `image` input optional (litegraph) — the node works with it empty. +- `send_id`: a plain text widget. +- `get_id`: render as a **dropdown** populated from `GET /datasete_gate/bus/list` + (refresh when opened), plus free-text entry. +- Pause/preview UI unchanged — `send_preview` runs after the source is resolved, so + get-loaded images preview correctly. + +## 7. Code shape + +- `gates/imagebus.py` *(new, stdlib)* — slot paths, `has`, `ensure_dir`, `list_ids`, + `delete_id`. Unit-testable. +- `gates/imaging.py` *(additive)* — `save_image_tensor`, `save_mask_tensor` (mirror the + existing loaders). Unit-testable with torch. +- `gates/gate.py` *(additive)* — `bus_save`/`bus_load`, pure `resolve_source`, and the + `run()` wiring (optional image, publish on pass). comfy imports stay lazy. +- `gates/gates_compat.py` *(additive)* — `gate_bus_base()` → `input/gate_bus`. +- `gates/gate_server.py` *(additive)* — `GET /datasete_gate/bus/list`. + +## 8. Edge cases + +- `image` empty + `get_id` empty/missing → silent no-op (no pause, all blocked). +- Mask precedence: gate-painted > loaded-from-bus > zeros. +- Same `send_id` from multiple gates → latest pass wins (documented). +- `get_id` referencing a deleted slot → treated as missing (no-op). +- Cross-run: publish in run A, load in run B (even after restart) — that's the whole point. + +## 9. Testing + +- pytest: `imagebus` (paths/has/list/delete); `imaging` save→load round-trip (shapes, mask + polarity); `gate.resolve_source` (image wins / get loads / nothing → None); `bus_save`+ + `bus_load` round-trip. +- Manual (live): publish at gate A (`send_id=cp1`), then a gate with empty image + + `get_id=cp1` loads it (even in a new workflow), edit mask, route onward; dropdown lists ids; + normal wired path ignores the bus. diff --git a/docs/plans/2026-06-21-image-gate-bus-implementation.md b/docs/plans/2026-06-21-image-gate-bus-implementation.md new file mode 100644 index 0000000..5856ccd --- /dev/null +++ b/docs/plans/2026-06-21-image-gate-bus-implementation.md @@ -0,0 +1,376 @@ +# Image Gate Send/Get Bus Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Extend `Image Gate (Manual Router)` so it can auto-publish a passed image+mask to a named disk bus (`send_id`) and, when its `image` input is empty, load from a named slot (`get_id`) — enabling wireless, cycle-free "restart from the gate point" across runs. + +**Architecture:** A pure stdlib `gates/imagebus.py` manages slot dirs under `input/gate_bus//`. `gates/imaging.py` gains tensor PNG savers mirroring its loaders. `gates/gate.py` gains `bus_save`/`bus_load` + a pure `resolve_source`, and `run()` makes `image` optional, loads from `get_id` when absent, and publishes to `send_id` on pass. A `GET /datasete_gate/bus/list` route feeds the `get_id` dropdown. + +**Tech Stack:** Python 3.12, torch 2.8, Pillow, numpy, aiohttp; pytest 9; vanilla JS. + +--- + +## Conventions (read once) + +- **Test python:** `/media/p5/miniforge3/bin/python` (`PY=...`). +- **Run tests:** `cd /media/p5/ComfyUI-Datasete-Gates && $PY -m pytest tests/test_imagebus.py tests/test_gate.py tests/test_imaging.py -v` +- All edits to `gate.py`, `imaging.py`, `gates_compat.py`, `gate_server.py` are **additive** — + re-Read first, keep the existing Image Gate behavior, run full suite after. +- `gates/imagebus.py` stays stdlib-only. `gate.py` keeps comfy imports lazy (inside `run`). +- Bus base dir = `gates_compat.gate_bus_base()` = `input/gate_bus`. +- Commit style: Conventional Commits + repo Co-Authored-By trailer; stage only this feature's paths. + +--- + +### Task 1: `gates_compat.py` — `gate_bus_base()` + +**Files:** Modify `gates/gates_compat.py` + +**Step 1:** Re-Read the file, then append (mirrors `grid_pool_base`): + +```python +def gate_bus_base(): + import folder_paths + return os.path.join(folder_paths.get_input_directory(), "gate_bus") +``` + +**Step 2:** Verify import: `$PY -c "import gates.gates_compat as c; print(hasattr(c,'gate_bus_base'))"` → `True`. + +**Step 3: Commit** `feat: gate_bus_base() path helper` + +--- + +### Task 2: `imagebus.py` — slot paths + list/has/delete + +**Files:** Create `gates/imagebus.py`; Test `tests/test_imagebus.py` + +**Step 1: Failing test** + +```python +# tests/test_imagebus.py +from gates import imagebus as ib + +def test_paths(tmp_path): + base = str(tmp_path) + assert ib.image_path(base, "cp1").name == "image.png" + assert ib.mask_path(base, "cp1").name == "mask.png" + assert ib.bus_dir(base, "cp1").name == "cp1" + +def test_has_and_ensure(tmp_path): + base = str(tmp_path) + assert ib.has(base, "cp1") is False + ib.ensure_dir(base, "cp1") + ib.image_path(base, "cp1").write_bytes(b"x") + assert ib.has(base, "cp1") is True + +def test_list_ids_only_populated(tmp_path): + base = str(tmp_path) + ib.ensure_dir(base, "empty") # dir but no image.png + ib.ensure_dir(base, "cp1"); ib.image_path(base, "cp1").write_bytes(b"x") + ib.ensure_dir(base, "cp2"); ib.image_path(base, "cp2").write_bytes(b"y") + assert ib.list_ids(base) == ["cp1", "cp2"] + +def test_delete(tmp_path): + base = str(tmp_path) + ib.ensure_dir(base, "cp1"); ib.image_path(base, "cp1").write_bytes(b"x") + ib.delete_id(base, "cp1") + assert not ib.bus_dir(base, "cp1").exists() +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +# gates/imagebus.py +"""Disk-backed image bus for Image Gate send/get. Stdlib only.""" +import shutil +from pathlib import Path + + +def bus_dir(base, bus_id): + return Path(base) / bus_id + + +def image_path(base, bus_id): + return bus_dir(base, bus_id) / "image.png" + + +def mask_path(base, bus_id): + return bus_dir(base, bus_id) / "mask.png" + + +def has(base, bus_id): + return image_path(base, bus_id).exists() + + +def ensure_dir(base, bus_id): + d = bus_dir(base, bus_id) + d.mkdir(parents=True, exist_ok=True) + return d + + +def list_ids(base): + p = Path(base) + if not p.is_dir(): + return [] + return sorted(d.name for d in p.iterdir() if d.is_dir() and (d / "image.png").exists()) + + +def delete_id(base, bus_id): + d = bus_dir(base, bus_id) + if d.exists(): + shutil.rmtree(d) +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: imagebus slot store` + +--- + +### Task 3: `imaging.py` — tensor PNG savers + +**Files:** Modify `gates/imaging.py`; Test `tests/test_imaging.py` + +**Step 1: Failing test** (add) + +```python +import torch +from gates import imaging + +def test_save_load_image_roundtrip(tmp_path): + img = torch.zeros((1, 6, 4, 3), dtype=torch.float32) + img[0, 0, 0, 0] = 1.0 # red corner + p = str(tmp_path / "image.png") + imaging.save_image_tensor(p, img) + back = imaging.load_image_tensor(p) + assert back.shape == (1, 6, 4, 3) + assert float(back[0, 0, 0, 0]) > 0.99 + +def test_save_load_mask_roundtrip(tmp_path): + mask = torch.ones((1, 6, 4), dtype=torch.float32) + p = str(tmp_path / "mask.png") + imaging.save_mask_tensor(p, mask) + back = imaging.load_mask_tensor(p, 6, 4) + assert back.shape == (1, 6, 4) + assert float(back.min()) > 0.99 +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement (append to `imaging.py`)** + +```python +def save_image_tensor(path, image): + arr = (image[0].cpu().numpy() * 255.0).clip(0, 255).astype("uint8") + Image.fromarray(arr).save(path) + + +def save_mask_tensor(path, mask): + arr = (mask[0].cpu().numpy() * 255.0).clip(0, 255).astype("uint8") + Image.fromarray(arr, mode="L").save(path) +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: imaging tensor PNG savers` + +--- + +### Task 4: `gate.py` — `bus_save` / `bus_load` / `resolve_source` + +**Files:** Modify `gates/gate.py`, `tests/test_gate.py` + +**Step 1: Failing test** + +```python +import torch +from gates import gate + +def _img(r=1.0): + t = torch.zeros((1, 6, 4, 3), dtype=torch.float32) + t[0, 0, 0, 0] = r + return t + +def test_bus_save_load_roundtrip(tmp_path): + base = str(tmp_path) + gate.bus_save(base, "cp1", _img(1.0), torch.ones((1, 6, 4))) + img, mask = gate.bus_load(base, "cp1") + assert img.shape == (1, 6, 4, 3) and float(img[0, 0, 0, 0]) > 0.99 + assert mask.shape == (1, 6, 4) and float(mask.min()) > 0.99 + +def test_resolve_source_image_wins(tmp_path): + img = _img() + out_img, out_mask = gate.resolve_source(str(tmp_path), img, "cp1") + assert out_img is img and out_mask is None # given image ignores the bus + +def test_resolve_source_loads_from_get(tmp_path): + base = str(tmp_path) + gate.bus_save(base, "cp1", _img(1.0), torch.zeros((1, 6, 4))) + out_img, out_mask = gate.resolve_source(base, None, "cp1") + assert out_img.shape == (1, 6, 4, 3) and out_mask.shape == (1, 6, 4) + +def test_resolve_source_nothing(tmp_path): + assert gate.resolve_source(str(tmp_path), None, "") == (None, None) + assert gate.resolve_source(str(tmp_path), None, "missing") == (None, None) +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement (append to `gate.py`; add `from . import imagebus, imaging` at top)** + +```python +def bus_save(base, bus_id, image, mask): + imagebus.ensure_dir(base, bus_id) + imaging.save_image_tensor(str(imagebus.image_path(base, bus_id)), image) + imaging.save_mask_tensor(str(imagebus.mask_path(base, bus_id)), mask) + + +def bus_load(base, bus_id): + img = imaging.load_image_tensor(str(imagebus.image_path(base, bus_id))) + h, w = int(img.shape[1]), int(img.shape[2]) + mp = imagebus.mask_path(base, bus_id) + mask = imaging.load_mask_tensor(str(mp) if mp.exists() else None, h, w) + return img, mask + + +def resolve_source(base, image, get_id): + if image is not None: + return image, None + if get_id and imagebus.has(base, get_id): + return bus_load(base, get_id) + return None, None +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: gate bus_save/bus_load/resolve_source` + +--- + +### Task 5: `gate.py` — wire send/get into `ImageGate` (MERGE) + +**Files:** Modify `gates/gate.py`, `tests/test_gate.py` + +**Step 1: Failing test** (input shape) + +```python +def test_image_gate_optional_inputs(): + it = gate.ImageGate.INPUT_TYPES() + assert "image" in it["optional"] + assert "send_id" in it["optional"] and "get_id" in it["optional"] + assert "routes" in it["required"] +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** — re-Read `gate.py`, then: +- `INPUT_TYPES`: + ```python + return { + "required": {"routes": ("INT", {"default": 2, "min": 1, "max": MAX_ROUTES})}, + "optional": { + "image": ("IMAGE",), + "send_id": ("STRING", {"default": ""}), + "get_id": ("STRING", {"default": ""}), + }, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + ``` +- `run` signature + body: + ```python + def run(self, routes, unique_id, image=None, send_id="", get_id=""): + from comfy_execution.graph_utils import ExecutionBlocker + from . import gate_server + from .gates_compat import gate_bus_base + + base = gate_bus_base() + image, loaded_mask = resolve_source(base, image, get_id) + blocker = ExecutionBlocker(None) + if image is None: # nothing to gate -> silent no-op + return (torch.zeros((1, 1, 1), dtype=torch.float32),) + tuple( + blocker for _ in range(MAX_ROUTES)) + + 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() + + painted = gate_bus.GateBus.pop_mask(unique_id) + if painted: + mask = mask_from_stash(painted, image) + elif loaded_mask is not None: + mask = loaded_mask + else: + mask = mask_from_stash(None, image) + + if send_id: + bus_save(base, send_id, image, mask) + + chosen = max(0, min(chosen_1 - 1, routes - 1)) + return (mask,) + route_tuple(chosen, image, blocker, MAX_ROUTES) + ``` + +**Step 4: Run → PASS** (existing gate tests still pass). + +**Step 5: Commit** `feat: Image Gate send_id/get_id bus (optional image, publish on pass)` + +--- + +### Task 6: `gate_server.py` — bus list route + +**Files:** Modify `gates/gate_server.py` + +**Step 1:** Re-Read, then append (additive): + +```python +@routes.get("/datasete_gate/bus/list") +async def _bus_list(request): + from .gates_compat import gate_bus_base + from . import imagebus + return web.json_response({"ids": imagebus.list_ids(gate_bus_base())}) +``` + +**Step 2:** Full suite green: `$PY -m pytest tests/ -v`. + +**Step 3: Commit** `feat: gate bus/list route for get_id dropdown` + +--- + +### Task 7: `web/image_gate.js` — optional image + send/get widgets + +**Files:** Modify `web/image_gate.js` + +- Ensure the node tolerates an **empty `image` input** (it's optional now). +- `send_id`: leave as a plain text widget. +- `get_id`: turn into a **dropdown** populated from `GET /datasete_gate/bus/list` (fetch on + node create and when the widget is opened/clicked); allow free-text too. +- No change to the pause/preview flow — preview still arrives from the server after the + source is resolved (so get-loaded images preview fine). + +**Manual note:** verify the dropdown lists published ids and refreshes after a pass elsewhere. + +**Commit** `feat: image gate frontend — send_id widget + get_id dropdown` + +--- + +### Task 8: Live smoke test in ComfyUI + +Restart ComfyUI. Verify: +- [ ] Existing gate with a wired `image` works exactly as before (bus ignored). +- [ ] Set `send_id=cp1` on a gate, pass an image → `input/gate_bus/cp1/{image,mask}.png` appear. +- [ ] A second gate with **no image wired** and `get_id=cp1` → loads that image (+ mask), + pauses, and routes onward. +- [ ] Works in a **new workflow** / after a restart (cross-run resume). +- [ ] `get_id` dropdown lists existing bus ids. +- [ ] Gate with no image and no/invalid `get_id` → silent no-op (nothing downstream runs). +- [ ] Mask precedence: paint at the get-gate overrides the loaded mask. + +**Commit** (if fixes) `fix: image gate bus live-test adjustments` + +--- + +## Definition of done + +- `$PY -m pytest tests/test_imagebus.py tests/test_imaging.py tests/test_gate.py -v` green; + full `tests/` green (existing gate/pool/loader/text unaffected). +- Manual checklist passes: publish on pass, get-load (incl. cross-run), dropdown, optional + image, mask precedence, silent no-op.