diff --git a/README.md b/README.md index 32e8b18..e62d29b 100644 --- a/README.md +++ b/README.md @@ -1,79 +1,20 @@ # ComfyUI Datasete Gates -Custom nodes for curating image datasets in ComfyUI. +A suite of custom nodes for **curating, loading, and gating image datasets** in +ComfyUI — built for human-in-the-loop inpaint/sort pipelines where you review +images, route them, and reuse them across workflows without rewiring. -## Image Pool (Grid) +All nodes appear under the **“Datasete Gates”** category. -A node that holds a curated **pool of images** — each with its own remembered -mask and editable label — shown as an in-node thumbnail grid. One image is -selectable as the node's output (image + mask + index + count + label), so you -can switch which image flows downstream **without rewiring**. +## Nodes at a glance -![category: Datasete Gates] - -### What it does - -- **In-node grid** of the pooled images. Ingest by **paste** (Ctrl+V), - **drag-and-drop**, or the **Upload** button. -- **Click a thumbnail** to make it the active output. No rewiring needed to - switch images. -- **Per-slot mask**: click the 🖌 button to paint a mask in ComfyUI's - MaskEditor. The mask is remembered per image and never redrawn when you switch - between images. -- **Per-slot label**: type a label under each thumbnail; it's saved with the - pool and exposed on the `label` output. -- **Drag to reorder** thumbnails; **✕** to delete. -- The pool is stored on disk, so it **survives a ComfyUI restart** and travels - with the workflow (via a per-node pool id). - -### Inputs - -| Input | Type | Description | -|-----------|--------|-------------| -| `index` | INT | `-1` (default) outputs the in-node **selected** image. `0+` forces that slot index (clamped to the pool size). | -| `pool_id` | STRING | Per-node pool identifier. Managed automatically by the UI (a UUID minted per node and hidden); you normally never touch it. | - -### Outputs - -| Output | Type | Description | -|---------|--------|-------------| -| `image` | IMAGE | The selected image, `[1, H, W, 3]` float 0..1. A 1×1 black image when the pool is empty. | -| `mask` | MASK | The selected image's mask, `[1, H, W]` float 0..1. **All zeros** when the slot has no mask. | -| `index` | INT | The resolved slot index that was output. | -| `count` | INT | Number of images in the pool. | -| `label` | STRING | The selected slot's label. | - -### Mask polarity - -A mask is a grayscale PNG where **white (1.0) = the painted region of interest** -(the area you painted in the MaskEditor — i.e. the area to inpaint). No mask file -means an all-zeros MASK output. The MaskEditor stores the painted region in the -image's alpha channel; the extension bakes that alpha into a grayscale mask on -save so that white = painted. - -### Managed pool folder - -Each pool lives under ComfyUI's input directory: - -``` -input/grid_pool// -├── manifest.json # {active, slots:[{image, mask, label, added}], next_seq} -├── img_0001.png # an image -├── img_0001.mask.png # its mask (sidecar; optional) -├── img_0002.png -└── ... -``` - -- Images are named monotonically (`img_0001.png`, `img_0002.png`, …). -- A mask is stored as a `*.mask.png` sidecar next to its image. -- `manifest.json` is written atomically. If it's missing or corrupt, it is - rebuilt from the files on disk. - -### Cloning nodes - -Copy/paste of a node shares the source node's `pool_id` (both show the same -pool). To give a clone its **own** independent pool, right-click it → -**“Detach pool (new id)”**. +| Node | Class | What it does | +|------|-------|--------------| +| **Image Pool (Grid)** | `GridImagePool` | Holds a curated pool of images (each with a remembered mask + label) as an in-node grid; outputs the selected one — switch images without rewiring. | +| **Pool Profile** | `PoolProfile` | Companion node: create/select/manage **named profiles** so a pool's images can be reused in any workflow and moved between machines. | +| **Folder Image Loader** | `FolderImageLoader` | Loads an image by index from a folder (fixed or auto-advancing), with its sidecar `.txt` caption and alpha mask. | +| **Image Gate (Manual Router)** | `ImageGate` | Pauses the run and lets you **click a button to route** the image down one of up to 10 outputs; optional gate-time mask; Stop cancels. | +| **Text Gate (Manual Pass)** | `TextGate` | Pauses the run, shows the incoming text in an **editable** box, and passes it on a click; any-type `signal` in/out for ordering. | ## Install @@ -85,15 +26,212 @@ git clone /path/to/ComfyUI/custom_nodes/ComfyUI-Datasete-Gates ln -sfn /media/p5/ComfyUI-Datasete-Gates /path/to/ComfyUI/custom_nodes/ComfyUI-Datasete-Gates ``` -Restart ComfyUI. The node appears under the **“Datasete Gates”** category as -**“Image Pool (Grid)”**. +Restart ComfyUI. Dependencies (torch, Pillow, numpy, aiohttp) are already +provided by ComfyUI. -Dependencies (torch, Pillow, numpy, aiohttp) are already provided by ComfyUI. +--- + +## Image Pool (Grid) + +Holds a curated **pool of images** — each with its own remembered mask and +editable label — shown as an in-node thumbnail grid. One image is selectable as +the node's output, so you can switch which image flows downstream **without +rewiring**. + +### What it does + +- **In-node grid** of pooled images. Ingest by **paste** (Ctrl+V), + **drag-and-drop**, or the **Upload** button. +- **Click a thumbnail** to make it the active output. No rewiring to switch. +- **Per-slot mask**: click 🖌 to paint a mask in ComfyUI's MaskEditor. The mask + is remembered per image and never redrawn when you switch between images. +- **Per-slot label**: type a label under each thumbnail; saved with the pool and + exposed on the `label` output. +- **Drag to reorder** thumbnails; **✕** to delete. +- Stored on disk, so the pool **survives a restart**. Pair with **Pool Profile** + to reuse it across workflows. + +### Inputs + +| Input | Type | Description | +|-----------|---------------|-------------| +| `index` | INT | `-1` (default) outputs the in-node **selected** image. `0+` forces that slot index (clamped to the pool size). | +| `pool_id` | STRING | Per-node pool identifier; managed by the UI (a hidden per-node UUID). You normally never touch it. | +| `profile` | POOL_PROFILE | *Optional.* When connected to a **Pool Profile** node, the pool uses that profile instead of its own `pool_id` (`profile or pool_id`). | + +### Outputs + +| Output | Type | Description | +|---------|--------|-------------| +| `image` | IMAGE | The selected image, `[1, H, W, 3]` float 0..1. A 1×1 black image when the pool is empty. | +| `mask` | MASK | The selected image's mask, `[1, H, W]` float 0..1. **All zeros** when the slot has no mask. | +| `index` | INT | The resolved slot index that was output. | +| `count` | INT | Number of images in the pool. | +| `label` | STRING | The selected slot's label. | + +### Cloning nodes + +Copy/paste shares the source node's `pool_id` (both show the same pool). To give +a clone its **own** independent pool, right-click → **“Detach pool (new id)”**. + +--- + +## Pool Profile + +Companion to the Image Pool. Turns a pool's fragile per-node UUID into a +**named, reusable profile**, so the same images (with masks/labels) can be loaded +in any workflow — and exported to another machine. + +Wire its `profile` output into a pool's `profile` input. Selecting a profile +**live-switches** the connected pool's grid to that profile's images, and any +adds/masks land in it. + +### Actions + +- **Create** a new named profile, **Select** an existing one (dropdown). +- **Rename**, **Delete** (removes its images), **Duplicate / Save-as** (snapshot + to a new profile). +- **Export** a profile to a `.zip`, **Import** one back — pools become portable. + +### Inputs / Outputs + +| Port | Type | Description | +|------|------|-------------| +| `profile` (widget) | STRING | The selected profile name (UI renders a dropdown). | +| `profile_id` (widget) | STRING | Hidden, UI-managed stable id. | +| `profile` (output) | POOL_PROFILE | The selected profile's id → connect to a pool's `profile` input. | + +### Registry + +`input/grid_pool/profiles.json` maps friendly **name → stable id**; each +profile's data lives in the existing `input/grid_pool//` layout. Existing +unnamed pools keep working unchanged — they're just unregistered ids. + +--- + +## Folder Image Loader + +Loads an image by index from a folder, plus its sidecar caption text and alpha +mask. Built for sequential, one-image-per-run dataset processing. + +### Inputs + +| Input | Type | Description | +|----------|------|-------------| +| `folder` | STRING | Absolute path to a folder of images. | +| `index` | INT | Which image (after natural sort). Has **`control_after_generate`** → set it to fixed / increment / decrement; auto-advances after each run. | +| `depth` | INT | `0` = top-level only; `N` = recurse up to N levels; `-1` = unlimited. | + +### Outputs + +| Output | Type | Description | +|------------|--------|-------------| +| `image` | IMAGE | The loaded image. | +| `text` | STRING | Contents of the sidecar `.txt`, or `""` if none. | +| `mask` | MASK | From the image's alpha channel (`1 - alpha`); zeros sized to the image if no alpha. | +| `filename` | STRING | The file **stem** (no extension). | +| `index` | INT | The resolved index actually loaded. | + +### Notes + +- Files are **natural-sorted** (`img2` before `img10`); extensions + `.png/.jpg/.jpeg/.webp/.bmp/.tif/.tiff`. +- Walking past the end (or below 0) **raises** — a clean end-of-batch stop signal + when running in increment mode. Empty folder / bad path raise too. + +--- + +## Image Gate (Manual Router) + +Pauses the running prompt and shows the image with a row of labeled **route +buttons**. Click a route to send the image down that output; every other route is +silently skipped. Built for manual dataset sorting. + +### Inputs + +| Input | Type | Description | +|----------|-------|-------------| +| `image` | IMAGE | The image (or batch, routed as one unit). | +| `routes` | INT | Number of visible route buttons/outputs (1–10). | + +### Outputs + +| Output | Type | Description | +|-------------------|------|-------------| +| `mask` | MASK | Painted at the gate (🖌), or zeros. Always emitted. | +| `route_1`…`route_10` | IMAGE | The chosen route carries the image; the rest return an `ExecutionBlocker` so only the chosen branch runs. The UI shows only `routes` of them, with editable labels. | + +### How it works + +- During execution the node **blocks** until you click. **Route K** → image to + output K (others blocked). **🖌 Edit mask** → opens the MaskEditor; the result + comes out on `mask`. **■ Stop** → cancels the whole run cleanly. +- Because any `ExecutionBlocker` input skips a node, a non-chosen route's + downstream never runs — as long as it also consumes the routed image (the + normal wiring). It re-pauses on every run (never cached). + +--- + +## Text Gate (Manual Pass) + +Pauses the run, shows the incoming text in an **editable** box, and emits it +(edited) when you click **Pass**. An optional any-type `signal` lets you force +execution order. + +### Inputs + +| Input | Type | Description | +|----------|------|-------------| +| `text` | STRING (`forceInput`) | The incoming text to review/edit. | +| `signal` | `*` (any) | *Optional.* Accepts any type; only used to sequence this node after its source. | + +### Outputs + +| Output | Type | Description | +|----------|------|-------------| +| `text` | STRING | The text you passed (possibly edited). | +| `signal` | `*` (any) | Passthrough of the input signal — chain gates in a fixed order. | + +Pauses every run; ComfyUI's global **Cancel** unblocks it cleanly (no deadlock). + +--- + +## Concepts + +### Human-in-the-loop gates + +Image/Text Gate **block the executor thread** during a run and wait for a click +(a small server-side waiter + a `/datasete_gate/*` route the UI posts to). Stop / +Cancel raise ComfyUI's `InterruptProcessingException`. These nodes always +re-execute (`IS_CHANGED = nan`) so they pause every time. + +### Mask polarity + +A mask is a grayscale PNG where **white (1.0) = the painted region of interest** +(the area to inpaint). No mask → an all-zeros MASK. The MaskEditor stores the +painted region in the image's alpha channel; the extension bakes that alpha into +a grayscale mask on save so that white = painted. + +### Storage layout + +``` +input/grid_pool/ +├── profiles.json # {profiles:[{id, name, created}]} (Pool Profile) +└── / + ├── manifest.json # {active, slots:[{image, mask, label, added}], next_seq} + ├── img_0001.png # an image (named monotonically) + ├── img_0001.mask.png # its mask (sidecar; optional) + └── ... +``` + +`manifest.json` is written atomically; if missing or corrupt it is rebuilt from +the files on disk. + +--- ## Development -The storage layer (`gates/pool.py`) is pure stdlib and fully unit-tested without -ComfyUI. Run the tests with: +The pure storage/scan layers are stdlib-only and unit-tested without ComfyUI: ```bash python -m pytest tests/ -v @@ -101,8 +239,12 @@ python -m pytest tests/ -v Layout: -- `gates/pool.py` — pure storage (manifest, add/remove/reorder/active/label/mask). Stdlib only. +- `gates/pool.py` — pure pool storage (manifest, add/remove/reorder/active/label/mask). +- `gates/profiles.py` — pure profile registry + dir ops + zip export/import. +- `gates/scan.py` — pure folder scan (natural sort, depth, sidecar, index). +- `gates/gate_bus.py` — pure blocking choice/text/mask waiter for the gates. - `gates/imaging.py` — torch/PIL tensor loaders. -- `gates/node.py` — the `GridImagePool` node. -- `gates/handlers.py` / `gates/routes.py` — pure handlers + aiohttp routes (`/grid_pool/*`). -- `web/grid_image_pool.js` — the in-node grid UI + MaskEditor integration. +- `gates/node.py` · `loader.py` · `gate.py` · `textgate.py` · `profile_node.py` — the nodes. +- `gates/handlers.py` · `routes.py` · `gate_server.py` · `profiles_routes.py` — aiohttp glue + (`/grid_pool/*`, `/datasete_gate/*`, `/grid_pool/profiles/*`). +- `web/*.js` — the in-node UIs (grid + MaskEditor, gate previews, profile dropdown). diff --git a/pyproject.toml b/pyproject.toml index adeba72..d5fe7b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-datasete-gates" version = "0.1.0" -description = "Dataset Gates — Image Pool (Grid) node for ComfyUI" +description = "Dataset Gates — image pool, folder loader, and manual routing/text gates for ComfyUI dataset curation" requires-python = ">=3.10" [tool.comfy]