From 19e670ced664d402cb410a5417fd6b723c929bf7 Mon Sep 17 00:00:00 2001 From: Ethan Fel Date: Sun, 21 Jun 2026 12:46:27 +0200 Subject: [PATCH] Add Image Pool (Grid) implementation plan TDD, bite-sized tasks across 3 phases: pure pool storage layer + atomic manifest, tensor/imaging helpers, GridImagePool node with IS_CHANGED, aiohttp routes, in-node grid UI, and MaskEditor clipspace round-trip. Co-Authored-By: Claude Opus 4.8 --- ...26-06-21-image-pool-grid-implementation.md | 989 ++++++++++++++++++ 1 file changed, 989 insertions(+) create mode 100644 docs/plans/2026-06-21-image-pool-grid-implementation.md diff --git a/docs/plans/2026-06-21-image-pool-grid-implementation.md b/docs/plans/2026-06-21-image-pool-grid-implementation.md new file mode 100644 index 0000000..58a3b9b --- /dev/null +++ b/docs/plans/2026-06-21-image-pool-grid-implementation.md @@ -0,0 +1,989 @@ +# Image Pool (Grid) 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 Pool (Grid)` that holds a curated pool of images — each with its own remembered mask and editable label — displayed as an in-node grid, with one selectable as the output (image + mask + index + count + label). + +**Architecture:** Pure Python storage layer (`gates/pool.py`, stdlib only — fully unit-testable) manages a managed pool folder under `input/grid_pool//` with an atomic `manifest.json`. A thin tensor/imaging layer (`gates/imaging.py`, torch+PIL) loads slots into ComfyUI tensors. The node class (`gates/node.py`) wires them together. aiohttp routes (`gates/routes.py`) let the frontend mutate the pool. A JS extension (`web/grid_image_pool.js`) renders the grid, ingests images (paste/drop/upload), and reuses ComfyUI's MaskEditor via the clipspace mechanism. + +**Tech Stack:** Python 3.12, torch 2.8, Pillow, numpy, aiohttp (ComfyUI's `PromptServer`), pytest 9; vanilla JS frontend extension (LiteGraph DOM widget + ComfyUI `app`/`api`). + +--- + +## Conventions (read once) + +- **Test python:** `/media/p5/miniforge3/bin/python` (call as `PY=/media/p5/miniforge3/bin/python`). +- **Run tests:** `cd /media/p5/ComfyUI-Datasete-Gates && $PY -m pytest tests/ -v` +- **Repo root:** `/media/p5/ComfyUI-Datasete-Gates` (already a git repo with the design doc committed). +- **Install for live testing:** ComfyUI loads custom nodes from `/media/p5/Comfyui/custom_nodes/`. After phase 1 we symlink the repo in: + `ln -s /media/p5/ComfyUI-Datasete-Gates /media/p5/Comfyui/custom_nodes/ComfyUI-Datasete-Gates` +- **Mask convention:** a mask PNG is grayscale `L`; white (1.0) = the region of interest (area to inpaint). MASK output is `[1,H,W]` float 0..1. No mask file → all-zeros. +- **Image tensor:** `[1,H,W,3]` float 0..1 (ComfyUI IMAGE). +- **Commit style:** Conventional Commits, end body with the Co-Authored-By trailer used in the design-doc commit. +- `gates/pool.py` MUST stay stdlib-only (no torch / no `folder_paths`) so it tests without ComfyUI. +- `gates/node.py` MUST resolve the pool base dir through `_grid_pool_base()` so tests can monkeypatch it (never import `folder_paths` at module top level). + +--- + +# PHASE 1 — Pool storage, node output (no masking), grid UI + +### Task 1: Scaffold the package so ComfyUI can load it + +**Files:** +- Create: `gates/__init__.py` (empty) +- Create: `pyproject.toml` +- Create: `__init__.py` (repo root — mappings + WEB_DIRECTORY) +- Create: `web/grid_image_pool.js` (placeholder) +- Create: `tests/__init__.py` (empty) +- Create: `requirements.txt` (empty — deps already in comfy env) + +**Step 1: Write `pyproject.toml`** + +```toml +[project] +name = "comfyui-datasete-gates" +version = "0.1.0" +description = "Dataset Gates — Image Pool (Grid) node for ComfyUI" +requires-python = ">=3.10" + +[tool.comfy] +PublisherId = "ethanfel" +DisplayName = "ComfyUI Datasete Gates" +``` + +**Step 2: Write repo-root `__init__.py`** + +```python +"""ComfyUI-Datasete-Gates — custom nodes.""" +from .gates.node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .gates import routes # noqa: F401 (registers aiohttp routes on import) + +WEB_DIRECTORY = "./web" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] +``` + +> Note: `gates.node` and `gates.routes` are created in later tasks. Until Task 8/9 exist, temporarily comment those imports OR do Task 1 last. Recommended: create the files as empty stubs now (`NODE_CLASS_MAPPINGS = {}` etc.) and fill them in their tasks. + +**Step 3: Create stub `web/grid_image_pool.js`** + +```javascript +import { app } from "../../scripts/app.js"; +app.registerExtension({ name: "datasete.gates.imagepool" }); +``` + +**Step 4: Verify package imports (pure)** + +Run: `cd /media/p5/ComfyUI-Datasete-Gates && /media/p5/miniforge3/bin/python -c "import gates"` +Expected: no error. + +**Step 5: Commit** + +```bash +git add pyproject.toml __init__.py gates/ web/ tests/ requirements.txt +git commit -m "chore: scaffold ComfyUI-Datasete-Gates package" +``` + +--- + +### Task 2: `pool.py` — empty manifest + atomic read/write + +**Files:** +- Create: `gates/pool.py` +- Test: `tests/test_pool.py` + +**Step 1: Write the failing test** + +```python +# tests/test_pool.py +import json +from pathlib import Path +from gates import pool + +def test_empty_manifest_shape(): + m = pool.empty_manifest() + assert m == {"active": 0, "slots": [], "next_seq": 1} + +def test_read_missing_creates_empty(tmp_path): + m = pool.read_manifest(str(tmp_path), "p1") + assert m == pool.empty_manifest() + +def test_write_then_read_roundtrip(tmp_path): + m = pool.empty_manifest() + m["active"] = 2 + pool.write_manifest(str(tmp_path), "p1", m) + # file lives at /p1/manifest.json + assert (tmp_path / "p1" / "manifest.json").exists() + assert pool.read_manifest(str(tmp_path), "p1") == m + +def test_write_is_atomic_no_partial_temp_left(tmp_path): + pool.write_manifest(str(tmp_path), "p1", pool.empty_manifest()) + leftovers = list((tmp_path / "p1").glob("*.tmp")) + assert leftovers == [] +``` + +**Step 2: Run to verify fail** + +Run: `$PY -m pytest tests/test_pool.py -v` +Expected: FAIL (module/functions missing). + +**Step 3: Implement** + +```python +# gates/pool.py +"""Pure storage layer for the Image Pool node. Stdlib only — no torch, no comfy.""" +import json +import os +from pathlib import Path + + +def empty_manifest(): + return {"active": 0, "slots": [], "next_seq": 1} + + +def pool_dir(base_dir, pool_id): + return Path(base_dir) / pool_id + + +def manifest_path(base_dir, pool_id): + return pool_dir(base_dir, pool_id) / "manifest.json" + + +def read_manifest(base_dir, pool_id): + p = manifest_path(base_dir, pool_id) + if not p.exists(): + return empty_manifest() + try: + with open(p, "r", encoding="utf-8") as f: + m = json.load(f) + # minimal shape guard + if not isinstance(m, dict) or "slots" not in m: + raise ValueError("bad manifest") + m.setdefault("active", 0) + m.setdefault("next_seq", len(m.get("slots", [])) + 1) + return m + except (ValueError, json.JSONDecodeError): + return rebuild_manifest(base_dir, pool_id) + + +def write_manifest(base_dir, pool_id, manifest): + d = pool_dir(base_dir, pool_id) + d.mkdir(parents=True, exist_ok=True) + final = d / "manifest.json" + tmp = d / "manifest.json.tmp" + with open(tmp, "w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2) + os.replace(tmp, final) # atomic on same filesystem + return manifest +``` + +> `rebuild_manifest` is referenced here but implemented in Task 7. Add a temporary stub returning `empty_manifest()` now; replace it in Task 7. (Its test arrives in Task 7.) + +**Step 4: Run to verify pass** + +Run: `$PY -m pytest tests/test_pool.py -v` +Expected: PASS (4 tests). + +**Step 5: Commit** + +```bash +git add gates/pool.py tests/test_pool.py +git commit -m "feat: pool manifest read/write with atomic save" +``` + +--- + +### Task 3: `pool.py` — `next_image_name` + `add_image` + +**Files:** Modify `gates/pool.py`; Modify `tests/test_pool.py` + +**Step 1: Failing test** + +```python +def test_next_image_name_uses_next_seq(): + m = pool.empty_manifest() + assert pool.next_image_name(m) == "img_0001.png" + m["next_seq"] = 42 + assert pool.next_image_name(m) == "img_0042.png" + +def test_add_image_writes_file_and_appends_slot(tmp_path): + data = b"\x89PNG\r\n\x1a\n" + b"fake" # bytes are written verbatim + m = pool.add_image(str(tmp_path), "p1", data, ts=123) + assert len(m["slots"]) == 1 + slot = m["slots"][0] + assert slot == {"image": "img_0001.png", "mask": None, "label": "", "added": 123} + assert m["next_seq"] == 2 + assert (tmp_path / "p1" / "img_0001.png").read_bytes() == data + +def test_add_image_monotonic_after_growth(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + m = pool.add_image(str(tmp_path), "p1", b"b", ts=2) + assert [s["image"] for s in m["slots"]] == ["img_0001.png", "img_0002.png"] +``` + +**Step 2: Run → FAIL.** `$PY -m pytest tests/test_pool.py -v` + +**Step 3: Implement (append to `pool.py`)** + +```python +def next_image_name(manifest): + return f"img_{manifest.get('next_seq', 1):04d}.png" + + +def add_image(base_dir, pool_id, data, ts=0): + m = read_manifest(base_dir, pool_id) + name = next_image_name(m) + d = pool_dir(base_dir, pool_id) + d.mkdir(parents=True, exist_ok=True) + with open(d / name, "wb") as f: + f.write(data) + m["slots"].append({"image": name, "mask": None, "label": "", "added": ts}) + m["next_seq"] = m.get("next_seq", 1) + 1 + write_manifest(base_dir, pool_id, m) + return m +``` + +**Step 4: Run → PASS.** + +**Step 5: Commit** + +```bash +git add gates/pool.py tests/test_pool.py +git commit -m "feat: pool add_image + monotonic naming" +``` + +--- + +### Task 4: `pool.py` — `remove_slot` (deletes files, fixes active) + +**Files:** Modify `gates/pool.py`; Modify `tests/test_pool.py` + +**Step 1: Failing test** + +```python +def test_remove_slot_deletes_files_and_reindexes(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + pool.add_image(str(tmp_path), "p1", b"b", ts=2) + pool.add_image(str(tmp_path), "p1", b"c", ts=3) + m = pool.set_active(str(tmp_path), "p1", 2) # active=2 + m = pool.remove_slot(str(tmp_path), "p1", 0) # drop first + assert [s["image"] for s in m["slots"]] == ["img_0002.png", "img_0003.png"] + assert not (tmp_path / "p1" / "img_0001.png").exists() + assert m["active"] == 1 # shifted down + +def test_remove_active_clamps(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + pool.add_image(str(tmp_path), "p1", b"b", ts=2) + pool.set_active(str(tmp_path), "p1", 1) + m = pool.remove_slot(str(tmp_path), "p1", 1) # removed the active last one + assert m["active"] == 0 + assert len(m["slots"]) == 1 +``` + +> `set_active` is needed here; implement it in Task 5 first OR move Task 5 before Task 4. The plan keeps numbering but executor may reorder 4/5 freely. Simplest: do Task 5 then Task 4. (Marked.) + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +def remove_slot(base_dir, pool_id, index): + m = read_manifest(base_dir, pool_id) + if index < 0 or index >= len(m["slots"]): + return m + slot = m["slots"].pop(index) + d = pool_dir(base_dir, pool_id) + for key in ("image", "mask"): + name = slot.get(key) + if name: + f = d / name + if f.exists(): + f.unlink() + if index < m["active"]: + m["active"] -= 1 + m["active"] = _clamp_active(m) + write_manifest(base_dir, pool_id, m) + return m + + +def _clamp_active(m): + n = len(m["slots"]) + if n == 0: + return 0 + return max(0, min(m.get("active", 0), n - 1)) +``` + +**Step 4: Run → PASS.** + +**Step 5: Commit** `feat: pool remove_slot with file cleanup` + +--- + +### Task 5: `pool.py` — `set_active` + `resolve_slot` (the -1 / clamp rule) + +**Files:** Modify `gates/pool.py`; Modify `tests/test_pool.py` + +**Step 1: Failing test** + +```python +def test_set_active_clamps(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + pool.add_image(str(tmp_path), "p1", b"b", ts=2) + assert pool.set_active(str(tmp_path), "p1", 1)["active"] == 1 + assert pool.set_active(str(tmp_path), "p1", 9)["active"] == 1 # clamp high + assert pool.set_active(str(tmp_path), "p1", -5)["active"] == 0 # clamp low + +def test_resolve_slot_rules(): + m = {"active": 1, "slots": [0, 1, 2], "next_seq": 4} # 3 slots + assert pool.resolve_slot(m, -1) == 1 # manual -> active + assert pool.resolve_slot(m, 0) == 0 # forced + assert pool.resolve_slot(m, 9) == 2 # clamp high + assert pool.resolve_slot({"active": 0, "slots": [], "next_seq": 1}, -1) == -1 # empty +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +def set_active(base_dir, pool_id, index): + m = read_manifest(base_dir, pool_id) + m["active"] = index + m["active"] = _clamp_active(m) + write_manifest(base_dir, pool_id, m) + return m + + +def resolve_slot(manifest, index_widget): + n = len(manifest["slots"]) + if n == 0: + return -1 + idx = manifest.get("active", 0) if index_widget == -1 else index_widget + return max(0, min(idx, n - 1)) +``` + +**Step 4: Run → PASS.** + +**Step 5: Commit** `feat: pool set_active + resolve_slot selection rule` + +--- + +### Task 6: `pool.py` — `set_label` + +**Files:** Modify `gates/pool.py`; Modify `tests/test_pool.py` + +**Step 1: Failing test** + +```python +def test_set_label(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + m = pool.set_label(str(tmp_path), "p1", 0, "front view") + assert m["slots"][0]["label"] == "front view" + +def test_set_label_out_of_range_noop(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + m = pool.set_label(str(tmp_path), "p1", 5, "x") + assert m["slots"][0]["label"] == "" +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +def set_label(base_dir, pool_id, index, label): + m = read_manifest(base_dir, pool_id) + if 0 <= index < len(m["slots"]): + m["slots"][index]["label"] = str(label) + write_manifest(base_dir, pool_id, m) + return m +``` + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: pool set_label` + +--- + +### Task 7: `pool.py` — `rebuild_manifest` (corrupt/missing recovery) + +**Files:** Modify `gates/pool.py`; Modify `tests/test_pool.py` + +**Step 1: Failing test** + +```python +def test_rebuild_from_files(tmp_path): + d = tmp_path / "p1" + d.mkdir() + (d / "img_0001.png").write_bytes(b"a") + (d / "img_0001.mask.png").write_bytes(b"m") + (d / "img_0003.png").write_bytes(b"c") # gap on purpose + m = pool.rebuild_manifest(str(tmp_path), "p1") + assert [s["image"] for s in m["slots"]] == ["img_0001.png", "img_0003.png"] + assert m["slots"][0]["mask"] == "img_0001.mask.png" + assert m["slots"][1]["mask"] is None + assert m["next_seq"] == 4 # max seq 3 + 1 + assert m["active"] == 0 + +def test_read_corrupt_manifest_triggers_rebuild(tmp_path): + d = tmp_path / "p1"; d.mkdir() + (d / "img_0001.png").write_bytes(b"a") + (d / "manifest.json").write_text("{ not json") + m = pool.read_manifest(str(tmp_path), "p1") + assert [s["image"] for s in m["slots"]] == ["img_0001.png"] +``` + +**Step 2: Run → FAIL** (rebuild stub returns empty). + +**Step 3: Replace the stub** + +```python +import re + +def rebuild_manifest(base_dir, pool_id): + d = pool_dir(base_dir, pool_id) + m = empty_manifest() + if not d.exists(): + return m + imgs = sorted(p.name for p in d.glob("img_*.png") if not p.name.endswith(".mask.png")) + max_seq = 0 + for name in imgs: + match = re.match(r"img_(\d+)\.png$", name) + seq = int(match.group(1)) if match else 0 + max_seq = max(max_seq, seq) + mask_name = name.replace(".png", ".mask.png") + mask = mask_name if (d / mask_name).exists() else None + m["slots"].append({"image": name, "mask": mask, "label": "", "added": 0}) + m["next_seq"] = max_seq + 1 + return m +``` + +**Step 4: Run → PASS** (run the whole file: `$PY -m pytest tests/test_pool.py -v`). + +**Step 5: Commit** `feat: pool rebuild_manifest recovery` + +--- + +### Task 8: `imaging.py` — tensor loaders + change hash (torch) + +**Files:** +- Create: `gates/imaging.py` +- Test: `tests/test_imaging.py` + +**Step 1: Failing test** + +```python +# tests/test_imaging.py +import numpy as np, torch +from PIL import Image +from gates import imaging + +def _png(tmp_path, name, color, size=(4, 6)): # size = (w, h) + p = tmp_path / name + Image.new("RGB", size, color).save(p) + return str(p) + +def test_load_image_tensor_shape_and_range(tmp_path): + t = imaging.load_image_tensor(_png(tmp_path, "a.png", (255, 0, 0))) + assert t.shape == (1, 6, 4, 3) # [B,H,W,C] + assert t.dtype == torch.float32 + assert 0.0 <= float(t.min()) and float(t.max()) <= 1.0 + assert float(t[0, 0, 0, 0]) > 0.99 # red channel + +def test_load_mask_none_is_zeros(): + m = imaging.load_mask_tensor(None, h=6, w=4) + assert m.shape == (1, 6, 4) + assert float(m.max()) == 0.0 + +def test_load_mask_from_file(tmp_path): + p = tmp_path / "m.png" + Image.new("L", (4, 6), 255).save(p) + m = imaging.load_mask_tensor(str(p), h=6, w=4) + assert m.shape == (1, 6, 4) + assert float(m.min()) > 0.99 + +def test_empty_image_is_1x1_black(): + img, mask = imaging.empty_outputs() + assert img.shape == (1, 1, 1, 3) and float(img.max()) == 0.0 + assert mask.shape == (1, 1, 1) + +def test_change_hash_changes_with_mtime(): + h1 = imaging.change_hash("p", 0, [1000.0]) + h2 = imaging.change_hash("p", 0, [1001.0]) + assert h1 != h2 +``` + +**Step 2: Run → FAIL.** `$PY -m pytest tests/test_imaging.py -v` + +**Step 3: Implement** + +```python +# gates/imaging.py +"""Tensor/imaging helpers (torch + PIL). No comfy imports.""" +import hashlib +import numpy as np +import torch +from PIL import Image, ImageOps + + +def load_image_tensor(path): + img = Image.open(path) + img = ImageOps.exif_transpose(img).convert("RGB") + arr = np.array(img, dtype=np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) # [1,H,W,3] + + +def load_mask_tensor(path, h, w): + if not path: + return torch.zeros((1, h, w), dtype=torch.float32) + m = Image.open(path).convert("L") + arr = np.array(m, dtype=np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) # [1,H,W] + + +def empty_outputs(): + return (torch.zeros((1, 1, 1, 3), dtype=torch.float32), + torch.zeros((1, 1, 1), dtype=torch.float32)) + + +def change_hash(pool_id, index, mtimes): + key = f"{pool_id}|{index}|" + "|".join(f"{t:.3f}" for t in mtimes) + return hashlib.sha256(key.encode()).hexdigest() +``` + +**Step 4: Run → PASS.** + +**Step 5: Commit** `feat: imaging tensor loaders + change hash` + +--- + +### Task 9: `node.py` — the `GridImagePool` node + +**Files:** +- Create: `gates/node.py` +- Test: `tests/test_node.py` + +**Step 1: Failing test** + +```python +# tests/test_node.py +import numpy as np, torch +from PIL import Image +from gates import node, pool + +def _seed_pool(tmp_path, monkeypatch): + base = str(tmp_path / "grid_pool") + monkeypatch.setattr(node, "_grid_pool_base", lambda: base) + return base + +def _add_png(base, pid, name_bytes_color, ts): + # write a real PNG via pool.add_image + import io + buf = io.BytesIO(); Image.new("RGB", (4, 6), name_bytes_color).save(buf, "PNG") + return pool.add_image(base, pid, buf.getvalue(), ts=ts) + +def test_execute_empty_pool_returns_blank(tmp_path, monkeypatch): + _seed_pool(tmp_path, monkeypatch) + n = node.GridImagePool() + img, mask, idx, count, label = n.run(index=-1, pool_id="p1") + assert img.shape == (1, 1, 1, 3) + assert count == 0 and idx == 0 and label == "" + +def test_execute_selects_active(tmp_path, monkeypatch): + base = _seed_pool(tmp_path, monkeypatch) + _add_png(base, "p1", (255, 0, 0), 1) + _add_png(base, "p1", (0, 255, 0), 2) + pool.set_active(base, "p1", 1) + pool.set_label(base, "p1", 1, "green") + n = node.GridImagePool() + img, mask, idx, count, label = n.run(index=-1, pool_id="p1") + assert img.shape == (1, 6, 4, 3) + assert idx == 1 and count == 2 and label == "green" + assert float(img[0, 0, 0, 1]) > 0.99 # green channel + assert float(mask.max()) == 0.0 # no mask yet + +def test_execute_forced_index_clamps(tmp_path, monkeypatch): + base = _seed_pool(tmp_path, monkeypatch) + _add_png(base, "p1", (255, 0, 0), 1) + n = node.GridImagePool() + _, _, idx, count, _ = n.run(index=9, pool_id="p1") + assert idx == 0 and count == 1 + +def test_is_changed_differs_after_active_change(tmp_path, monkeypatch): + base = _seed_pool(tmp_path, monkeypatch) + _add_png(base, "p1", (255, 0, 0), 1) + _add_png(base, "p1", (0, 255, 0), 2) + h1 = node.GridImagePool.IS_CHANGED(index=-1, pool_id="p1") + pool.set_active(base, "p1", 1) + h2 = node.GridImagePool.IS_CHANGED(index=-1, pool_id="p1") + assert h1 != h2 +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement** + +```python +# gates/node.py +import os +from .gates_compat import grid_pool_base as _grid_pool_base # see note below +from . import pool, imaging + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + + +class GridImagePool: + CATEGORY = "Datasete Gates" + FUNCTION = "run" + RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING") + RETURN_NAMES = ("image", "mask", "index", "count", "label") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "index": ("INT", {"default": -1, "min": -1, "max": 9999}), + }, + "hidden": {"pool_id": "POOL_ID"}, + } + + @staticmethod + def _resolve(index, pool_id): + base = _grid_pool_base() + m = pool.read_manifest(base, pool_id) + idx = pool.resolve_slot(m, index) + return base, m, idx + + def run(self, index, pool_id="default"): + base, m, idx = self._resolve(index, pool_id) + if idx < 0: + img, mask = imaging.empty_outputs() + return (img, mask, 0, 0, "") + slot = m["slots"][idx] + d = pool.pool_dir(base, pool_id) + img = imaging.load_image_tensor(str(d / slot["image"])) + h, w = int(img.shape[1]), int(img.shape[2]) + mask_name = slot.get("mask") + mask = imaging.load_mask_tensor(str(d / mask_name) if mask_name else None, h, w) + return (img, mask, idx, len(m["slots"]), slot.get("label", "")) + + @classmethod + def IS_CHANGED(cls, index, pool_id="default", **kwargs): + base, m, idx = cls._resolve(index, pool_id) + if idx < 0: + return imaging.change_hash(pool_id, -1, []) + slot = m["slots"][idx] + d = pool.pool_dir(base, pool_id) + mtimes = [] + for key in ("image", "mask"): + name = slot.get(key) + p = d / name if name else None + mtimes.append(os.path.getmtime(p) if p and p.exists() else 0.0) + # include active so manual selection changes invalidate cache + return imaging.change_hash(pool_id, f"{idx}:{m.get('active')}", mtimes) + + +NODE_CLASS_MAPPINGS = {"GridImagePool": GridImagePool} +NODE_DISPLAY_NAME_MAPPINGS = {"GridImagePool": "Image Pool (Grid)"} +``` + +**Step 3b: Create `gates/gates_compat.py`** (isolates the comfy dependency for testability) + +```python +# gates/gates_compat.py +import os + +def grid_pool_base(): + import folder_paths # imported lazily; only available inside ComfyUI + return os.path.join(folder_paths.get_input_directory(), "grid_pool") +``` + +> The test monkeypatches `node._grid_pool_base`. Because `node.py` does +> `from .gates_compat import grid_pool_base as _grid_pool_base`, the name +> `node._grid_pool_base` exists at module scope and is patchable. ✅ + +**Step 4: Run → PASS.** `$PY -m pytest tests/test_node.py -v` + +**Step 5: Commit** `feat: GridImagePool node (image/mask/index/count/label + IS_CHANGED)` + +--- + +### Task 10: `routes.py` — aiohttp routes wired to `pool.py` + +**Files:** +- Create: `gates/routes.py` +- Test: `tests/test_routes_logic.py` (test the handler *logic* without a live server) + +> Routes are thin: parse request → call a pure `_handlers.py` function → JSON. We TDD the pure handler funcs; the aiohttp wrapper is verified live in Task 12. + +**Step 1: Failing test** + +```python +# tests/test_routes_logic.py +import io +from PIL import Image +from gates import handlers + +def _png_bytes(color=(1, 2, 3)): + b = io.BytesIO(); Image.new("RGB", (4, 4), color).save(b, "PNG"); return b.getvalue() + +def test_handle_add_then_list(tmp_path): + base = str(tmp_path) + m = handlers.handle_add(base, "p1", _png_bytes(), "png", ts=5) + assert len(m["slots"]) == 1 + assert handlers.handle_list(base, "p1")["slots"][0]["image"] == "img_0001.png" + +def test_handle_active_label_remove(tmp_path): + base = str(tmp_path) + handlers.handle_add(base, "p1", _png_bytes(), "png", ts=1) + handlers.handle_add(base, "p1", _png_bytes(), "png", ts=2) + assert handlers.handle_active(base, "p1", 1)["active"] == 1 + assert handlers.handle_label(base, "p1", 0, "hi")["slots"][0]["label"] == "hi" + assert len(handlers.handle_remove(base, "p1", 0)["slots"]) == 1 +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement `gates/handlers.py`** + +```python +# gates/handlers.py +"""Pure request handlers — no aiohttp. Each returns the updated manifest dict.""" +from . import pool + + +def handle_add(base, pool_id, data, ext, ts=0): + return pool.add_image(base, pool_id, data, ts=ts) + +def handle_remove(base, pool_id, index): + return pool.remove_slot(base, pool_id, index) + +def handle_active(base, pool_id, index): + return pool.set_active(base, pool_id, index) + +def handle_label(base, pool_id, index, label): + return pool.set_label(base, pool_id, index, label) + +def handle_list(base, pool_id): + return pool.read_manifest(base, pool_id) + +def handle_set_mask(base, pool_id, index, mask_png_bytes): + return pool.set_mask(base, pool_id, index, mask_png_bytes) # Task 11 +``` + +**Step 3b: Implement `gates/routes.py`** (aiohttp glue — not unit-tested, verified live) + +```python +# gates/routes.py +import json +from aiohttp import web +from server import PromptServer +from . import handlers +from .gates_compat import grid_pool_base + +routes = PromptServer.instance.routes + + +def _base(): + return grid_pool_base() + + +@routes.post("/grid_pool/add") +async def _add(request): + reader = await request.multipart() + pool_id, ts, data = "default", 0, None + async for part in reader: + if part.name == "pool_id": + pool_id = (await part.text()) + elif part.name == "ts": + ts = int(await part.text()) + elif part.name == "image": + data = await part.read(decode=False) + m = handlers.handle_add(_base(), pool_id, data, "png", ts=ts) + return web.json_response(m) + + +@routes.post("/grid_pool/remove") +async def _remove(request): + body = await request.json() + return web.json_response(handlers.handle_remove(_base(), body["pool_id"], int(body["index"]))) + + +@routes.post("/grid_pool/active") +async def _active(request): + body = await request.json() + return web.json_response(handlers.handle_active(_base(), body["pool_id"], int(body["index"]))) + + +@routes.post("/grid_pool/label") +async def _label(request): + body = await request.json() + return web.json_response(handlers.handle_label(_base(), body["pool_id"], int(body["index"]), body["label"])) + + +@routes.get("/grid_pool/list") +async def _list(request): + pool_id = request.query.get("pool_id", "default") + return web.json_response(handlers.handle_list(_base(), pool_id)) +``` + +> Update repo-root `__init__.py` to `from .gates import routes` (already there from Task 1). +> `handle_set_mask` route is added in Phase 2 (Task 11/12). + +**Step 4: Run → PASS** (`tests/test_routes_logic.py`). The aiohttp module import requires comfy, so do NOT import `gates.routes` in tests — only `gates.handlers`. + +**Step 5: Commit** `feat: pool handlers + aiohttp routes` + +--- + +### Task 11: Live smoke test — node loads + grid renders + ingest/select/delete/label + +**Files:** Modify `web/grid_image_pool.js` (full Phase-1 UI) + +**Step 1: Symlink into ComfyUI** + +```bash +ln -sfn /media/p5/ComfyUI-Datasete-Gates /media/p5/Comfyui/custom_nodes/ComfyUI-Datasete-Gates +``` + +**Step 2: Implement the grid widget JS** + +Write `web/grid_image_pool.js` with this structure (complete code): +- `app.registerExtension({ name, beforeRegisterNodeDef })` — for `GridImagePool` only. +- In `nodeCreated`: ensure a `pool_id` exists; if the hidden widget is empty, generate `crypto.randomUUID()` and store it on a hidden widget so it serializes into the workflow. +- `addDOMWidget("grid", "div", el, {})` — a scrollable flex-wrap container. +- `refresh()` → `api.fetchApi('/grid_pool/list?pool_id=' + id)` → render thumbnails. Each thumb ``, active border, a `` label, a ✕ delete button, a 🖌 mask button (wired in Phase 2). +- **Paste:** on `paste` event when node selected → read clipboard image → `FormData` → POST `/grid_pool/add` → refresh. +- **Drop:** `el.ondrop` → for each image file → POST `/grid_pool/add` → refresh. +- **Upload:** a button → hidden `` → POST each → refresh. +- **Select:** click thumb → POST `/grid_pool/active` `{pool_id,index}` → refresh; also set the node dirty so IS_CHANGED re-fires. +- **Label edit:** `change` on label input → POST `/grid_pool/label`. +- **Delete:** ✕ → POST `/grid_pool/remove` → refresh. + +(Provide the full JS in implementation; keep DOM minimal and dependency-free.) + +**Step 3: Restart ComfyUI, manual verification checklist** + +- [ ] Node "Image Pool (Grid)" appears under "Datasete Gates". +- [ ] Paste an image (Ctrl+V) → thumbnail appears. +- [ ] Drag 2 files onto node → both appear. +- [ ] Click a thumb → active border moves. +- [ ] Edit a label → reload workflow → label persists. +- [ ] Delete a thumb → it disappears + file removed from `input/grid_pool//`. +- [ ] Connect `IMAGE`/`MASK` to a PreviewImage / preview → run → selected image shows, mask all black. +- [ ] Set `index` widget to `0` → forces first regardless of active. +- [ ] Restart ComfyUI, reload workflow → pool still there. + +**Step 4: Fix any issues found, re-verify.** + +**Step 5: Commit** `feat: in-node grid UI — ingest/select/delete/label + Phase 1 complete` + +--- + +# PHASE 2 — MaskEditor integration + per-slot mask persistence + +### Task 12: `pool.py` — `set_mask` + handler + route + +**Files:** Modify `gates/pool.py`, `gates/handlers.py`, `gates/routes.py`; Modify `tests/test_pool.py`, `tests/test_routes_logic.py` + +**Step 1: Failing test** + +```python +def test_set_mask_writes_sidecar(tmp_path): + pool.add_image(str(tmp_path), "p1", b"a", ts=1) + m = pool.set_mask(str(tmp_path), "p1", 0, b"MASKBYTES") + assert m["slots"][0]["mask"] == "img_0001.mask.png" + assert (tmp_path / "p1" / "img_0001.mask.png").read_bytes() == b"MASKBYTES" + +def test_set_mask_out_of_range_noop(tmp_path): + m = pool.set_mask(str(tmp_path), "p1", 0, b"x") + assert m["slots"] == [] +``` + +**Step 2: Run → FAIL.** + +**Step 3: Implement (append to `pool.py`)** + +```python +def set_mask(base_dir, pool_id, index, mask_bytes): + m = read_manifest(base_dir, pool_id) + if not (0 <= index < len(m["slots"])): + return m + img_name = m["slots"][index]["image"] + mask_name = img_name.replace(".png", ".mask.png") + with open(pool_dir(base_dir, pool_id) / mask_name, "wb") as f: + f.write(mask_bytes) + m["slots"][index]["mask"] = mask_name + write_manifest(base_dir, pool_id, m) + return m +``` + +Add `/grid_pool/set_mask` route (multipart: `pool_id`, `index`, `mask` file) calling `handlers.handle_set_mask`. + +**Step 4: Run → PASS.** **Step 5: Commit** `feat: pool set_mask + route` + +--- + +### Task 13: MaskEditor round-trip in JS (clipspace integration) + +**Files:** Modify `web/grid_image_pool.js` + +**Background (verified against installed frontend):** the editor is opened via the clipspace mechanism. The flow per the legacy `maskeditor.js`: +1. Build a clipspace payload: `ComfyApp.clipspace = { imgs:[Image], images:[{filename,subfolder,type}], selectedIndex:0, ... }`. +2. Set `ComfyApp.clipspace_return_node = node`. +3. Call the editor open API. In the installed frontend this is exposed (`openMaskEditor`). Confirm the exact accessor at implement time: + `grep -rho "openMaskEditor[^,;]*" ` and check `app.extensionManager`/`ComfyApp` for the callable. +4. The editor saves the painted mask to `input/clipspace/...` via `/upload/mask` and calls `node.pasteFromClipspace(clipspace)` on close. + +**Implementation:** +- Add `🖌` button per thumbnail → `openMaskEditorForSlot(node, index)`: + - fetch the slot image as an `Image` from `/view?...`, + - set up `ComfyApp.clipspace` + `clipspace_return_node = node`, + - open the editor. +- Implement `node.pasteFromClipspace = async (clipspace) => {...}`: + - read the saved masked image (`clipspace.imgs[selectedIndex].src`) with `channel=a` to get the alpha, + - draw alpha to a canvas, export grayscale PNG blob (white = masked), + - POST to `/grid_pool/set_mask` (multipart) with `pool_id`, `index` (the slot being edited), `mask`, + - `refresh()` and mark node dirty. +- Track "which slot is being edited" on the node (`node._editingSlot = index`) so `pasteFromClipspace` knows the target. + +**Manual verification checklist:** +- [ ] Click 🖌 on a slot → MaskEditor opens with that image. +- [ ] Paint, save → returns to graph; thumbnail shows a "has-mask" dot. +- [ ] `input/grid_pool//img_XXXX.mask.png` exists. +- [ ] Run graph → MASK output matches the painted region (white = painted). +- [ ] Switch active to another image and back → mask still there (no redraw). +- [ ] Edit the mask again → MASK output updates on next run (IS_CHANGED via mtime). +- [ ] Verify mask orientation/scale matches the image (no flip / off-by-resize). + +**Commit** `feat: MaskEditor round-trip — per-slot mask persistence (Phase 2 complete)` + +> **Mask polarity check:** confirm whether the alpha from the editor needs inverting so that "painted area" == 1.0 in the MASK output. Adjust the canvas export accordingly and note the decision in code comments + README. + +--- + +# PHASE 3 — Polish (optional, do as needed) + +### Task 14: Right-click "Detach pool (new id)" +Add a node context-menu entry that assigns a fresh `crypto.randomUUID()` to `pool_id`, clears the displayed grid, and refreshes — so a cloned node can get its own pool. Manual verify: copy node → both share pool → detach one → independent. +**Commit** `feat: detach-pool context menu` + +### Task 15: Drag-to-reorder thumbnails +HTML5 drag-and-drop within the grid → POST a new order to a `/grid_pool/reorder` route that reorders `manifest.slots` (and fixes `active`). Add `pool.reorder(base, pool_id, order)` with a unit test first. +**Commit** `feat: drag-reorder slots` + +### Task 16: Badges + empty-state polish +Slot index badge, has-mask dot styling, count display, friendly empty-pool message. Manual verify. +**Commit** `feat: grid badges + empty state` + +### Task 17: README +Write `README.md`: what the node does, install (symlink/clone into `custom_nodes`), the IO table, mask polarity note, and the managed-pool-folder layout. +**Commit** `docs: README for Image Pool (Grid)` + +--- + +## Definition of done (Phase 1+2) + +- `$PY -m pytest tests/ -v` → all green. +- Manual checklists in Tasks 11 and 13 pass. +- Pool + masks + labels survive a ComfyUI restart. +- No rewiring needed to switch images; masks are never redrawn when switching.