Compare commits
38 Commits
e304b39009
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| 690278b592 | |||
| 3ee14819b7 | |||
| d6d2c98a58 | |||
| 36dd5c91ee | |||
| 954b9ec2e6 | |||
| 1881aa727f | |||
| 78b1b85a11 | |||
| b50718f7fb | |||
| d9134b4e9b | |||
| 3fb63e44a3 | |||
| 00c8c6a790 | |||
| 726cd7bf17 | |||
| 5b92e9b338 | |||
| b2f5850b46 | |||
| 31a7112052 | |||
| 66e664247c | |||
| 5419366bde | |||
| d0dafa1d39 | |||
| b4639a73d3 | |||
| 84fc4f1cf1 | |||
| 58a48d67e5 | |||
| c869ecee2a | |||
| 8d785f5ca2 | |||
| 0ace20a1bc | |||
| b90d1befe6 | |||
| fe95a9af3a | |||
| 99a5ccac82 | |||
| f2ac5e37f3 | |||
| ce371ffe13 | |||
| aa909448d7 | |||
| 037cbf27db | |||
| 969463a4e9 | |||
| 7f90b6878f | |||
| 0413e25571 | |||
| cdd742c950 | |||
| 628a945514 | |||
| 95b3417ff6 | |||
| 0ba3d81fbf |
@@ -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/<pool_id>/
|
||||
├── 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 <repo-url> /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/<id>/` 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 `<stem>.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)
|
||||
└── <pool_id or profile_id>/
|
||||
├── 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).
|
||||
|
||||
+8
-2
@@ -18,14 +18,20 @@ if __package__:
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES
|
||||
from .gates.profile_node import NODE_CLASS_MAPPINGS as _PROF_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _PROF_NAMES
|
||||
from .gates.bucket_node import NODE_CLASS_MAPPINGS as _BUCKET_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _BUCKET_NAMES
|
||||
from .gates.sidecar_node import NODE_CLASS_MAPPINGS as _SC_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _SC_NAMES
|
||||
from .gates import routes # noqa: F401 (registers aiohttp routes on import)
|
||||
from .gates import gate_server # noqa: F401 (registers /datasete_gate/* + text routes)
|
||||
from .gates import profiles_routes # noqa: F401 (registers /grid_pool/profiles/*)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES,
|
||||
**_TEXT_NODES, **_PROF_NODES}
|
||||
**_TEXT_NODES, **_PROF_NODES, **_BUCKET_NODES,
|
||||
**_SC_NODES}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES,
|
||||
**_TEXT_NAMES, **_PROF_NAMES}
|
||||
**_TEXT_NAMES, **_PROF_NAMES, **_BUCKET_NAMES,
|
||||
**_SC_NAMES}
|
||||
else: # pragma: no cover - exercised only under pytest collection
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# Bucket Resize (Klein 9B) — Design
|
||||
|
||||
Date: 2026-06-21
|
||||
Status: Approved (brainstorming complete, ready for implementation plan)
|
||||
Spec: `/media/unraid/davinci/comics-lora/dataset/KLEIN_BUCKET_SIZES.md`
|
||||
|
||||
## 1. Purpose
|
||||
|
||||
Automatically resize any image so it lands **exactly on its training bucket** — W×H
|
||||
multiples of 64 within a ~1.64 MP area budget (FLUX.2 [klein] 9B, ai-toolkit
|
||||
`resolution: [1280]`). Resize-to-cover + center-crop, with slight Lanczos upscale only when
|
||||
needed. Outputs the bucketed image (+ identically transformed mask) and the chosen size.
|
||||
|
||||
Sixth node in the `ComfyUI-Datasete-Gates` suite. **No custom frontend** — standard widgets.
|
||||
|
||||
## 2. Bucket selection (generated grid)
|
||||
|
||||
Budget = `resolution²` (default 1280 → 1,638,400 px). For an image of aspect `a = iw/ih`:
|
||||
|
||||
- Enumerate widths `w` in multiples of `divisible` (default 64). For each, take the **largest**
|
||||
on-grid height within budget: `h = floor(budget / w / divisible) * divisible` (skip if
|
||||
`h < divisible`). This is the max-area frontier per width.
|
||||
- Pick the candidate minimizing **log-aspect distance** `|ln(w/h) − ln(a)|`; tie-break by
|
||||
larger area. This reproduces the doc's 13 rows for normal aspects (square→1280×1280,
|
||||
0.5→896×1792, 2.0→1792×896, …) and extends to extreme aspects (≈0.09–2.67).
|
||||
|
||||
## 3. Fit: cover + center-crop
|
||||
|
||||
For chosen bucket `(W, H)` and image `(iw, ih)`:
|
||||
- `scale = max(W/iw, H/ih)` (cover). `new = (round(iw*scale), round(ih*scale))`.
|
||||
- Resize with **Lanczos** (good for up- and down-scale), then **center-crop** to exactly
|
||||
`W×H`: `left=(new_w−W)//2`, `top=(new_h−H)//2`.
|
||||
- If `scale > max_upscale` (default 1.5), still fit but **log a warning** (the doc warns big
|
||||
upscales soften texture).
|
||||
|
||||
The optional **mask** gets the identical scale+crop (so it stays aligned); absent → zeros
|
||||
sized to the bucket.
|
||||
|
||||
## 4. IO
|
||||
|
||||
| dir | name | type | notes |
|
||||
|-----|------|------|-------|
|
||||
| in | `image` | IMAGE | required |
|
||||
| in (opt) | `mask` | MASK | transformed identically; zeros if absent |
|
||||
| widget | `resolution` | INT (default 1280, min 64) | area budget = `resolution²` |
|
||||
| widget | `divisible` | INT (default 64, min 8) | grid step |
|
||||
| widget | `max_upscale` | FLOAT (default 1.5, min 1.0) | warn above this cover-scale |
|
||||
| out | `image` | IMAGE | exactly bucket `W×H`, `[1,H,W,3]` |
|
||||
| out | `mask` | MASK | `[1,H,W]` |
|
||||
| out | `width` | INT | chosen bucket width |
|
||||
| out | `height` | INT | chosen bucket height |
|
||||
| out | `label` | STRING | `"WxH"` (e.g. `1280x1280`) |
|
||||
|
||||
## 5. Code shape
|
||||
|
||||
- `gates/buckets.py` *(new, pure stdlib + math)* — `pick_bucket(iw, ih, resolution, divisible)`
|
||||
→ `(W, H)`; `cover_crop_params(iw, ih, W, H)` → `(new_w, new_h, left, top, scale)`.
|
||||
Fully unit-testable; **tested against the doc's table**.
|
||||
- `gates/bucket_node.py` *(new, torch/PIL)* — tensor↔PIL resize/crop using `buckets`, the
|
||||
`BucketResize` node. `run()` is pure compute (no comfy, no blocking) → fully unit-testable.
|
||||
- root `__init__.py` — additive merge of the node mapping.
|
||||
|
||||
## 6. Edge cases
|
||||
|
||||
- Batch `B>1`: bucket is chosen from the **first** image's aspect and applied to all (keeps a
|
||||
uniform output tensor); documented. (Dataset flow is typically one image per run.)
|
||||
- Image already exactly on a bucket → `scale≈1`, no crop.
|
||||
- Tiny/extreme aspect → handled by the generated grid (nearest of the frontier).
|
||||
- `max_upscale` only warns; it never refuses (the node always returns an on-grid image).
|
||||
- Mask resized with the same geometry (Lanczos), then clamped to [0,1].
|
||||
|
||||
## 7. Testing
|
||||
|
||||
- pytest `tests/test_buckets.py`: `pick_bucket` reproduces the doc rows for a set of aspects
|
||||
(1.0→1280×1280, 0.5→896×1792, 0.58→960×1664, 2.0→1792×896, …); all outputs are ÷divisible
|
||||
and ≤ budget; `cover_crop_params` math (cover scale, centered crop, exact target).
|
||||
- pytest `tests/test_bucket_node.py`: feed known tensor sizes → output is exactly the bucket
|
||||
shape; mask aligned; `label`/`width`/`height` correct; no-mask → zeros.
|
||||
- Manual (live): drop node after a loader, confirm odd-sized inputs come out on-grid and the
|
||||
label matches the table.
|
||||
@@ -0,0 +1,300 @@
|
||||
# Bucket Resize (Klein 9B) Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** A `BucketResize` node that snaps any image onto its ai-toolkit training bucket (W×H ÷64, ≤ ~1.64 MP) via cover-scale + center-crop (Lanczos), transforms an optional mask identically, and outputs the bucketed image + chosen `width`/`height`/`label`.
|
||||
|
||||
**Architecture:** A pure stdlib+math `gates/buckets.py` selects the bucket and computes the cover-crop geometry — fully unit-testable against the spec's table. `gates/bucket_node.py` (torch/PIL) does the actual tensor resize/crop; its `run()` is pure compute (no comfy, no blocking) so it unit-tests end-to-end. No custom frontend.
|
||||
|
||||
**Spec:** `/media/unraid/davinci/comics-lora/dataset/KLEIN_BUCKET_SIZES.md`
|
||||
|
||||
**Tech Stack:** Python 3.12, torch 2.8, Pillow, numpy; pytest 9.
|
||||
|
||||
---
|
||||
|
||||
## Conventions (read once)
|
||||
|
||||
- **Test python:** `/media/p5/miniforge3/bin/python` (`PY=...`).
|
||||
- **Run tests:** `cd /media/p5/ComfyUI-Datasete-Gates && $PY -m pytest tests/test_buckets.py tests/test_bucket_node.py -v`
|
||||
- `gates/buckets.py` is pure (stdlib + `math`); no torch/comfy.
|
||||
- IMAGE tensors are `[B,H,W,3]` float 0..1; MASK is `[B,H,W]`.
|
||||
- `__init__.py` edit is **additive** — re-Read first, extend the mappings.
|
||||
- Commit style: Conventional Commits + repo Co-Authored-By; stage only this node's paths.
|
||||
|
||||
---
|
||||
|
||||
### Task 1: `buckets.py` — `pick_bucket` (reproduce the spec table)
|
||||
|
||||
**Files:** Create `gates/buckets.py`; Test `tests/test_buckets.py`
|
||||
|
||||
**Step 1: Failing test**
|
||||
|
||||
```python
|
||||
# tests/test_buckets.py
|
||||
from gates import buckets
|
||||
|
||||
# (iw, ih) -> expected (W, H) from KLEIN_BUCKET_SIZES.md, budget 1280, ÷64
|
||||
CASES = [
|
||||
(1000, 1000, 1280, 1280), # square
|
||||
(1000, 2000, 896, 1792), # a=0.50 portrait
|
||||
(1000, 1730, 960, 1664), # a≈0.58
|
||||
(1000, 1100, 1216, 1344), # a≈0.90 -> portrait-leaning
|
||||
(2000, 1000, 1792, 896), # a=2.00 landscape
|
||||
(1500, 1000, 1536, 1024), # a=1.50
|
||||
]
|
||||
|
||||
def test_pick_bucket_matches_table():
|
||||
for iw, ih, W, H in CASES:
|
||||
assert buckets.pick_bucket(iw, ih, 1280, 64) == (W, H)
|
||||
|
||||
def test_buckets_are_on_grid_and_within_budget():
|
||||
for iw, ih, *_ in CASES:
|
||||
W, H = buckets.pick_bucket(iw, ih, 1280, 64)
|
||||
assert W % 64 == 0 and H % 64 == 0
|
||||
assert W * H <= 1280 * 1280
|
||||
|
||||
def test_square_is_exactly_1280():
|
||||
assert buckets.pick_bucket(512, 512, 1280, 64) == (1280, 1280)
|
||||
```
|
||||
|
||||
**Step 2: Run → FAIL.**
|
||||
|
||||
**Step 3: Implement**
|
||||
|
||||
```python
|
||||
# gates/buckets.py
|
||||
"""Pure bucket math for KLEIN_BUCKET_SIZES.md. Stdlib only."""
|
||||
import math
|
||||
|
||||
|
||||
def pick_bucket(iw, ih, resolution=1280, divisible=64):
|
||||
"""Choose the on-grid bucket (W,H), area <= resolution^2, nearest to the
|
||||
image aspect (log distance; tie-break larger area)."""
|
||||
budget = resolution * resolution
|
||||
target = iw / ih
|
||||
best = None
|
||||
w = divisible
|
||||
w_max = budget // divisible
|
||||
while w <= w_max:
|
||||
h = (budget // w // divisible) * divisible # largest on-grid h within budget
|
||||
if h >= divisible:
|
||||
err = abs(math.log(w / h) - math.log(target))
|
||||
cand = (err, -(w * h), w, h) # min err, then max area
|
||||
if best is None or cand < best:
|
||||
best = cand
|
||||
w += divisible
|
||||
return best[2], best[3]
|
||||
|
||||
|
||||
def cover_crop_params(iw, ih, W, H):
|
||||
"""Cover-scale + centered crop to land (iw,ih) exactly on (W,H)."""
|
||||
scale = max(W / iw, H / ih)
|
||||
new_w = max(W, round(iw * scale))
|
||||
new_h = max(H, round(ih * scale))
|
||||
left = (new_w - W) // 2
|
||||
top = (new_h - H) // 2
|
||||
return new_w, new_h, left, top, scale
|
||||
```
|
||||
|
||||
**Step 4: Run → PASS.** **Step 5: Commit** `feat: bucket selection matching Klein 9B table`
|
||||
|
||||
---
|
||||
|
||||
### Task 2: `buckets.py` — `cover_crop_params`
|
||||
|
||||
**Files:** Modify `tests/test_buckets.py`
|
||||
|
||||
**Step 1: Failing test**
|
||||
|
||||
```python
|
||||
def test_cover_crop_exact_aspect_no_crop():
|
||||
# a=2.0 image onto 1792x896 bucket -> scale 0.896, no crop
|
||||
new_w, new_h, left, top, scale = buckets.cover_crop_params(2000, 1000, 1792, 896)
|
||||
assert (new_w, new_h) == (1792, 896)
|
||||
assert (left, top) == (0, 0)
|
||||
assert round(scale, 3) == 0.896
|
||||
|
||||
def test_cover_crop_square_into_landscape_crops_height():
|
||||
new_w, new_h, left, top, scale = buckets.cover_crop_params(1000, 1000, 1792, 896)
|
||||
assert new_w == 1792 and new_h >= 896
|
||||
assert left == 0 and top == (new_h - 896) // 2 # centered vertical crop
|
||||
assert scale > 1.0 # upscaled to cover width
|
||||
|
||||
def test_cover_crop_upscale_square():
|
||||
*_, scale = buckets.cover_crop_params(1000, 1000, 1280, 1280)
|
||||
assert round(scale, 2) == 1.28
|
||||
```
|
||||
|
||||
**Step 2: Run → PASS** (implemented in Task 1). If it fails, fix `cover_crop_params`.
|
||||
|
||||
**Step 3:** (no new code — locks the geometry with tests.)
|
||||
|
||||
**Step 4: Commit** `test: bucket cover_crop_params geometry`
|
||||
|
||||
---
|
||||
|
||||
### Task 3: `bucket_node.py` — fit helpers + `BucketResize` node
|
||||
|
||||
**Files:** Create `gates/bucket_node.py`; Test `tests/test_bucket_node.py`
|
||||
|
||||
**Step 1: Failing test**
|
||||
|
||||
```python
|
||||
# tests/test_bucket_node.py
|
||||
import torch
|
||||
from gates import bucket_node as bn
|
||||
|
||||
def test_square_to_1280():
|
||||
out, m, w, h, label = bn.BucketResize().run(image=torch.rand((1, 1000, 1000, 3)))
|
||||
assert (w, h) == (1280, 1280)
|
||||
assert out.shape == (1, 1280, 1280, 3)
|
||||
assert m.shape == (1, 1280, 1280) and float(m.max()) == 0.0 # no mask -> zeros
|
||||
assert label == "1280x1280"
|
||||
|
||||
def test_landscape_bucket_shapes():
|
||||
# tensor [B,H,W,3] with H=1000,W=2000 -> aspect 2.0 -> 1792x896
|
||||
out, m, w, h, label = bn.BucketResize().run(image=torch.rand((1, 1000, 2000, 3)))
|
||||
assert (w, h) == (1792, 896)
|
||||
assert out.shape == (1, 896, 1792, 3)
|
||||
assert label == "1792x896"
|
||||
|
||||
def test_mask_resized_and_aligned():
|
||||
out, m, w, h, _ = bn.BucketResize().run(
|
||||
image=torch.rand((1, 1000, 1000, 3)), mask=torch.ones((1, 1000, 1000)))
|
||||
assert m.shape == (1, 1280, 1280) and float(m.min()) > 0.9
|
||||
|
||||
def test_outputs_are_on_grid():
|
||||
out, m, w, h, _ = bn.BucketResize().run(
|
||||
image=torch.rand((1, 777, 1333, 3)), resolution=1280, divisible=64)
|
||||
assert w % 64 == 0 and h % 64 == 0
|
||||
assert out.shape[1] == h and out.shape[2] == w
|
||||
```
|
||||
|
||||
**Step 2: Run → FAIL.**
|
||||
|
||||
**Step 3: Implement**
|
||||
|
||||
```python
|
||||
# gates/bucket_node.py
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from . import buckets
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
def _resize_crop_pil(pil, new_w, new_h, left, top, W, H):
|
||||
pil = pil.resize((new_w, new_h), Image.LANCZOS)
|
||||
return pil.crop((left, top, left + W, top + H))
|
||||
|
||||
|
||||
def fit_image(image, W, H):
|
||||
"""image [B,H,W,3] -> [B,H,W,3] at (W,H) using the first image's geometry."""
|
||||
b, ih, iw = image.shape[0], image.shape[1], image.shape[2]
|
||||
new_w, new_h, left, top, scale = buckets.cover_crop_params(iw, ih, W, H)
|
||||
out = []
|
||||
for i in range(b):
|
||||
arr = (image[i].cpu().numpy() * 255.0).clip(0, 255).astype("uint8")
|
||||
pil = _resize_crop_pil(Image.fromarray(arr), new_w, new_h, left, top, W, H)
|
||||
out.append(torch.from_numpy(np.array(pil, dtype=np.float32) / 255.0))
|
||||
return torch.stack(out, 0), scale
|
||||
|
||||
|
||||
def fit_mask(mask, W, H):
|
||||
b, ih, iw = mask.shape[0], mask.shape[1], mask.shape[2]
|
||||
new_w, new_h, left, top, _ = buckets.cover_crop_params(iw, ih, W, H)
|
||||
out = []
|
||||
for i in range(b):
|
||||
arr = (mask[i].cpu().numpy() * 255.0).clip(0, 255).astype("uint8")
|
||||
pil = _resize_crop_pil(Image.fromarray(arr, mode="L"), new_w, new_h, left, top, W, H)
|
||||
out.append(torch.from_numpy(np.array(pil, dtype=np.float32) / 255.0))
|
||||
return torch.stack(out, 0)
|
||||
|
||||
|
||||
class BucketResize:
|
||||
CATEGORY = "Datasete Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("image", "mask", "width", "height", "label")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"resolution": ("INT", {"default": 1280, "min": 64, "max": 8192}),
|
||||
"divisible": ("INT", {"default": 64, "min": 8, "max": 256}),
|
||||
"max_upscale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 8.0, "step": 0.1}),
|
||||
},
|
||||
"optional": {"mask": ("MASK",)},
|
||||
}
|
||||
|
||||
def run(self, image, resolution=1280, divisible=64, max_upscale=1.5, mask=None):
|
||||
ih, iw = int(image.shape[1]), int(image.shape[2])
|
||||
W, H = buckets.pick_bucket(iw, ih, resolution, divisible)
|
||||
out_img, scale = fit_image(image, W, H)
|
||||
if scale > max_upscale:
|
||||
print(f"[BucketResize] cover scale {scale:.2f}x exceeds max_upscale "
|
||||
f"{max_upscale} for {iw}x{ih} -> {W}x{H}")
|
||||
out_mask = fit_mask(mask, W, H) if mask is not None \
|
||||
else torch.zeros((out_img.shape[0], H, W), dtype=torch.float32)
|
||||
return (out_img, out_mask, W, H, f"{W}x{H}")
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"BucketResize": BucketResize}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"BucketResize": "Bucket Resize (Klein 9B)"}
|
||||
```
|
||||
|
||||
**Step 4: Run → PASS.** **Step 5: Commit** `feat: BucketResize node (cover-crop onto Klein buckets)`
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Register in `__init__.py` (MERGE)
|
||||
|
||||
**Files:** Modify `__init__.py`
|
||||
|
||||
**Step 1:** Re-Read `__init__.py`, then add inside the `if __package__:` block:
|
||||
|
||||
```python
|
||||
from .gates.bucket_node import NODE_CLASS_MAPPINGS as _BUCKET_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _BUCKET_NAMES
|
||||
```
|
||||
and merge:
|
||||
```python
|
||||
NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **_BUCKET_NODES}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **_BUCKET_NAMES}
|
||||
```
|
||||
(No routes/web — standard widgets only.)
|
||||
|
||||
**Step 2:** `$PY -c "import gates.bucket_node; print(gates.bucket_node.NODE_CLASS_MAPPINGS)"`.
|
||||
|
||||
**Step 3:** Full suite green: `$PY -m pytest tests/ -v`.
|
||||
|
||||
**Step 4: Commit** `feat: register BucketResize`
|
||||
|
||||
---
|
||||
|
||||
### Task 5: Live smoke test in ComfyUI
|
||||
|
||||
Restart ComfyUI. Build: `Folder Image Loader → Bucket Resize → PreviewImage` (+ a SaveImage
|
||||
using `label` for the filename). Verify:
|
||||
- [ ] "Bucket Resize (Klein 9B)" appears under "Datasete Gates".
|
||||
- [ ] A square-ish image → `1280x1280`; a 2:1 image → `1792x896`; a tall image → a portrait
|
||||
bucket — all ÷64, output exactly bucket-sized.
|
||||
- [ ] An odd size (e.g. 1333×777) lands on-grid with a clean center-crop.
|
||||
- [ ] Feeding a mask (e.g. from the loader's alpha) → mask comes out aligned at bucket size.
|
||||
- [ ] `width`/`height`/`label` outputs match the preview.
|
||||
- [ ] A small input triggers the console `max_upscale` warning but still outputs on-grid.
|
||||
|
||||
**Commit** (if fixes) `fix: bucket resize live-test adjustments`
|
||||
|
||||
---
|
||||
|
||||
## Definition of done
|
||||
|
||||
- `$PY -m pytest tests/test_buckets.py tests/test_bucket_node.py -v` green; full `tests/` green.
|
||||
- `pick_bucket` reproduces the spec table; outputs are always ÷divisible and ≤ budget.
|
||||
- Manual checklist passes: on-grid output, aligned mask, correct label, upscale warning.
|
||||
@@ -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/<id>/
|
||||
├── 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.
|
||||
@@ -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/<id>/`. `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.
|
||||
@@ -0,0 +1,88 @@
|
||||
# Multi-Reroute (Rail) — Design
|
||||
|
||||
Date: 2026-06-21
|
||||
Status: Approved (brainstorming complete, ready for implementation plan)
|
||||
|
||||
## 1. Purpose
|
||||
|
||||
A single node holding **N parallel pass-through lanes** (a "rail"), so you can run tidy
|
||||
bundles of wires across the graph instead of dropping many separate Reroute nodes. Each lane
|
||||
forwards any type; you grow/shrink the rail with +/− at either end.
|
||||
|
||||
Seventh node in the `ComfyUI-Datasete-Gates` suite.
|
||||
|
||||
## 2. Approach
|
||||
|
||||
A **real pass-through node** with **any-type lanes** (`AnyType("*")`). Lane `i`'s input is
|
||||
forwarded to lane `i`'s output. An unconnected lane outputs an `ExecutionBlocker` so nothing
|
||||
downstream of an unused lane runs. (Not the frontend-only virtual-reroute trick — simpler and
|
||||
robust across all types; the trade-off is slots read as `*` instead of adapting to the wired
|
||||
type.)
|
||||
|
||||
## 3. IO
|
||||
|
||||
- Up to `MAX_LANES` (32) lanes, each: optional input `in_<i>` (`*`) → output `out_<i>` (`*`).
|
||||
- The node always returns a length-`MAX_LANES` tuple; the frontend shows only the active
|
||||
lanes (default 4). Wired output indices are stable, so unshown trailing outputs are simply
|
||||
unconnected.
|
||||
|
||||
```
|
||||
RETURN_TYPES = (ANY,) * MAX_LANES RETURN_NAMES = ("out_1", …, "out_32")
|
||||
INPUT_TYPES = {"optional": {"in_1": (ANY,), …}}
|
||||
```
|
||||
|
||||
No `IS_CHANGED` override — a reroute should be transparent/cacheable (re-runs only when an
|
||||
input value actually changes).
|
||||
|
||||
## 4. Run logic
|
||||
|
||||
```python
|
||||
def run(self, **kwargs):
|
||||
blocker = ExecutionBlocker(None)
|
||||
return tuple(
|
||||
kwargs.get(f"in_{i+1}") if kwargs.get(f"in_{i+1}") is not None else blocker
|
||||
for i in range(MAX_LANES)
|
||||
)
|
||||
```
|
||||
|
||||
Lane-count-agnostic: connected lanes forward their value; empty lanes block. The visible lane
|
||||
count is purely a frontend concern.
|
||||
|
||||
## 5. Frontend (`web/multi_reroute.js`)
|
||||
|
||||
- Render `lanes` lane rows (input + output pair), default 4; persist the count in a hidden
|
||||
widget so reload restores the rail (the "use raw widgets_values to add slots before link
|
||||
rewiring" pattern already used in this repo).
|
||||
- **+/− buttons**:
|
||||
- **Bottom add/remove** (Phase 1): reveal/hide the next/last lane pair — trivial and
|
||||
wiring-safe (only the end moves).
|
||||
- **Top add/remove** (Phase 2): insert/remove a lane at the top while **preserving the
|
||||
other lanes' wiring** — requires capturing links and re-mapping slot indices
|
||||
(rgthree-style). Kept separate so a bug here can't scramble existing rails.
|
||||
- Lanes use the shared `AnyType` so any wire connects.
|
||||
- (Phase 3 polish) compact reroute-pill look / optional per-lane labels.
|
||||
|
||||
## 6. Edge cases
|
||||
|
||||
- Empty lane → `ExecutionBlocker` (downstream skipped). A legitimate `None` value is treated
|
||||
as empty (reroute values are objects/tensors, effectively never `None`).
|
||||
- Removing a lane is from the **end** in Phase 1 (indices stay stable → links intact).
|
||||
Mid/top removal is Phase 2 with remap.
|
||||
- More than `MAX_LANES` requested → capped (logged in UI).
|
||||
- Mixed types across lanes is fine — each lane is independent `*`.
|
||||
|
||||
## 7. Code shape
|
||||
|
||||
- `gates/anytype.py` *(new)* — shared `AnyType("*")` + `ANY` (textgate can dedupe onto this
|
||||
later; not touched now).
|
||||
- `gates/reroute_node.py` *(new)* — pure `build_outputs(values, max_lanes, blocker)` +
|
||||
`MultiReroute` node (lazy `ExecutionBlocker` import for testability).
|
||||
- `web/multi_reroute.js` *(new)* — dynamic lane slots + +/− buttons + persistence.
|
||||
- root `__init__.py` — additive merge of the node mapping.
|
||||
|
||||
## 8. Testing
|
||||
|
||||
- pytest: `anytype` equals-everything; `build_outputs` forwards connected lanes and blocks
|
||||
empty ones (length == MAX_LANES); node `RETURN_TYPES` length + all-`*`.
|
||||
- Manual (live): add/remove lanes (bottom, then top), wire mixed types through, confirm values
|
||||
pass and reload restores the rail; empty lanes don't trigger downstream.
|
||||
@@ -0,0 +1,237 @@
|
||||
# Multi-Reroute (Rail) Implementation Plan
|
||||
|
||||
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
||||
|
||||
**Goal:** A `MultiReroute` node — N parallel any-type pass-through lanes ("rail") with +/− to add/remove lanes — so wire bundles stay tidy without many separate Reroute nodes.
|
||||
|
||||
**Architecture:** A real pass-through node with `AnyType("*")` lanes: `in_i → out_i`, empty lane → `ExecutionBlocker`. Pure `build_outputs` is unit-tested; `ExecutionBlocker` is imported lazily. The frontend manages dynamic lane slots and persists the lane count. Bottom add/remove is Phase 1 (wiring-safe); top add/remove (link-preserving) is Phase 2.
|
||||
|
||||
**Tech Stack:** Python 3.12 (stdlib), 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_anytype.py tests/test_reroute.py -v`
|
||||
- `gates/anytype.py` and `gates/reroute_node.py` import-safe without comfy (lazy
|
||||
`ExecutionBlocker` inside `run`).
|
||||
- `__init__.py` edit is **additive** — re-Read first, extend the mappings.
|
||||
- `MAX_LANES = 32`, default visible lanes = 4.
|
||||
- Commit style: Conventional Commits + repo Co-Authored-By; stage only this node's paths.
|
||||
|
||||
---
|
||||
|
||||
### Task 1: `anytype.py` — shared wildcard
|
||||
|
||||
**Files:** Create `gates/anytype.py`; Test `tests/test_anytype.py`
|
||||
|
||||
**Step 1: Failing test**
|
||||
|
||||
```python
|
||||
# tests/test_anytype.py
|
||||
from gates import anytype
|
||||
|
||||
def test_any_equals_everything():
|
||||
assert (anytype.ANY != "IMAGE") is False
|
||||
assert (anytype.ANY != "LATENT") is False
|
||||
assert isinstance(anytype.ANY, str)
|
||||
```
|
||||
|
||||
**Step 2: Run → FAIL.**
|
||||
|
||||
**Step 3: Implement**
|
||||
|
||||
```python
|
||||
# gates/anytype.py
|
||||
"""Shared ComfyUI wildcard type."""
|
||||
|
||||
|
||||
class AnyType(str):
|
||||
def __ne__(self, other):
|
||||
return False
|
||||
|
||||
|
||||
ANY = AnyType("*")
|
||||
```
|
||||
|
||||
**Step 4: Run → PASS.** **Step 5: Commit** `feat: shared AnyType wildcard`
|
||||
|
||||
---
|
||||
|
||||
### Task 2: `reroute_node.py` — `build_outputs`
|
||||
|
||||
**Files:** Create `gates/reroute_node.py`; Test `tests/test_reroute.py`
|
||||
|
||||
**Step 1: Failing test**
|
||||
|
||||
```python
|
||||
# tests/test_reroute.py
|
||||
from gates import reroute_node as rr
|
||||
|
||||
def test_build_outputs_forwards_and_blocks():
|
||||
B = object() # blocker sentinel
|
||||
vals = {"in_1": "A", "in_3": "C"}
|
||||
out = rr.build_outputs(vals, max_lanes=4, blocker=B)
|
||||
assert out == ("A", B, "C", B)
|
||||
|
||||
def test_build_outputs_length():
|
||||
B = object()
|
||||
assert len(rr.build_outputs({}, max_lanes=rr.MAX_LANES, blocker=B)) == rr.MAX_LANES
|
||||
```
|
||||
|
||||
**Step 2: Run → FAIL.**
|
||||
|
||||
**Step 3: Implement**
|
||||
|
||||
```python
|
||||
# gates/reroute_node.py
|
||||
from .anytype import ANY
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
MAX_LANES = 32
|
||||
|
||||
|
||||
def build_outputs(values, max_lanes, blocker):
|
||||
out = []
|
||||
for i in range(max_lanes):
|
||||
v = values.get(f"in_{i + 1}")
|
||||
out.append(v if v is not None else blocker)
|
||||
return tuple(out)
|
||||
```
|
||||
|
||||
**Step 4: Run → PASS.** **Step 5: Commit** `feat: multi-reroute build_outputs`
|
||||
|
||||
---
|
||||
|
||||
### Task 3: `reroute_node.py` — `MultiReroute` node
|
||||
|
||||
**Files:** Modify `gates/reroute_node.py`, `tests/test_reroute.py`
|
||||
|
||||
**Step 1: Failing test**
|
||||
|
||||
```python
|
||||
def test_io_shape():
|
||||
assert len(rr.MultiReroute.RETURN_TYPES) == rr.MAX_LANES
|
||||
assert all(t == "*" for t in rr.MultiReroute.RETURN_TYPES)
|
||||
assert rr.MultiReroute.RETURN_NAMES[0] == "out_1"
|
||||
it = rr.MultiReroute.INPUT_TYPES()
|
||||
assert "in_1" in it["optional"] and f"in_{rr.MAX_LANES}" in it["optional"]
|
||||
```
|
||||
|
||||
**Step 2: Run → FAIL.**
|
||||
|
||||
**Step 3: Implement (append)**
|
||||
|
||||
```python
|
||||
class MultiReroute:
|
||||
CATEGORY = "Datasete Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = (ANY,) * MAX_LANES
|
||||
RETURN_NAMES = tuple(f"out_{i + 1}" for i in range(MAX_LANES))
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"optional": {f"in_{i + 1}": (ANY,) for i in range(MAX_LANES)}}
|
||||
|
||||
def run(self, **kwargs):
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
return build_outputs(kwargs, MAX_LANES, ExecutionBlocker(None))
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"MultiReroute": MultiReroute}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"MultiReroute": "Multi Reroute (Rail)"}
|
||||
```
|
||||
|
||||
> `RETURN_TYPES` entries are the `AnyType` instance (`== "*"`), so the test's `t == "*"`
|
||||
> holds. No `IS_CHANGED` (transparent/cacheable passthrough).
|
||||
|
||||
**Step 4: Run → PASS.** **Step 5: Commit** `feat: MultiReroute node (any-type pass-through lanes)`
|
||||
|
||||
---
|
||||
|
||||
### Task 4: Register in `__init__.py` (MERGE)
|
||||
|
||||
**Files:** Modify `__init__.py`
|
||||
|
||||
**Step 1:** Re-Read `__init__.py`; add inside `if __package__:`:
|
||||
|
||||
```python
|
||||
from .gates.reroute_node import NODE_CLASS_MAPPINGS as _RR_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _RR_NAMES
|
||||
```
|
||||
and merge into the final dicts:
|
||||
```python
|
||||
NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **_RR_NODES}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **_RR_NAMES}
|
||||
```
|
||||
|
||||
**Step 2:** `$PY -c "import gates.reroute_node; print(gates.reroute_node.NODE_CLASS_MAPPINGS)"`.
|
||||
|
||||
**Step 3:** Full suite green: `$PY -m pytest tests/ -v`.
|
||||
|
||||
**Step 4: Commit** `feat: register MultiReroute`
|
||||
|
||||
---
|
||||
|
||||
### Task 5: `web/multi_reroute.js` — dynamic lanes + bottom +/− (Phase 1)
|
||||
|
||||
**Files:** Create `web/multi_reroute.js`
|
||||
|
||||
`app.registerExtension` for `MultiReroute`:
|
||||
- On `nodeCreated`: ensure the node shows `lanes` lane pairs (default 4). Keep a hidden
|
||||
`lanes` widget (count) so the rail count **persists** across save/reload — restore by
|
||||
re-adding that many slot pairs (mirror the repo's existing dynamic-slot restore that reads
|
||||
raw `widgets_values` before link rewiring).
|
||||
- Maintain exactly `lanes` visible input slots `in_1..in_lanes` and output slots
|
||||
`out_1..out_lanes` (add/remove the trailing pair as the count changes).
|
||||
- **Bottom +**: `lanes++` → add `in_{n}`/`out_{n}`. **Bottom −**: `lanes--` → remove the last
|
||||
pair (only the end moves, so existing wiring is untouched). Clamp `1..MAX_LANES`.
|
||||
- Buttons via on-node widgets (or the node context menu) labeled `+ lane` / `− lane`.
|
||||
- Slots are any-type (`*`) so any wire connects.
|
||||
|
||||
**Manual verify:** add/remove lanes from the bottom; wire IMAGE through lane 1 and LATENT
|
||||
through lane 2; values pass; save+reload restores the lane count and wiring.
|
||||
|
||||
**Commit** `feat: multi-reroute frontend — dynamic lanes + bottom add/remove`
|
||||
|
||||
---
|
||||
|
||||
### Task 6: `web/multi_reroute.js` — top +/− with wiring preservation (Phase 2)
|
||||
|
||||
**Files:** Modify `web/multi_reroute.js`
|
||||
|
||||
- **Top +**: insert a new lane at the top. Because slots are positional and links bind to
|
||||
slot index, capture all current input/output links, add a pair, and **re-map links** so the
|
||||
existing lanes keep their connections and the new empty lane is visually first.
|
||||
- **Top −**: remove the top lane, remap the rest up.
|
||||
- Implement against a small lane-model (list of logical lanes ↔ slot indices); rebuild slots
|
||||
and reconnect from the model so a failure can't silently drop wires.
|
||||
|
||||
**Manual verify:** with lanes 1–3 wired, Top + adds an empty lane at top and 1–3 stay wired;
|
||||
Top − removes it cleanly. Save/reload still consistent.
|
||||
|
||||
**Commit** `feat: multi-reroute top add/remove (wiring-preserving)`
|
||||
|
||||
---
|
||||
|
||||
### Task 7: Live smoke test in ComfyUI
|
||||
|
||||
Restart ComfyUI. Verify:
|
||||
- [ ] "Multi Reroute (Rail)" appears under "Datasete Gates" with 4 lanes.
|
||||
- [ ] Bottom +/− add/remove lanes; Top +/− (Phase 2) keep existing wiring.
|
||||
- [ ] Route mixed types (IMAGE, MASK, LATENT, STRING) through separate lanes → all pass intact.
|
||||
- [ ] An unconnected lane doesn't trigger its downstream (ExecutionBlocker).
|
||||
- [ ] Save + reload restores lane count and all connections.
|
||||
|
||||
**Commit** (if fixes) `fix: multi-reroute live-test adjustments`
|
||||
|
||||
---
|
||||
|
||||
## Definition of done
|
||||
|
||||
- `$PY -m pytest tests/test_anytype.py tests/test_reroute.py -v` green; full `tests/` green.
|
||||
- Manual checklist passes: lanes add/remove (bottom; top in P2), mixed-type pass-through,
|
||||
empty-lane blocking, persistence.
|
||||
@@ -0,0 +1,64 @@
|
||||
# Text Gate — "Run from here" + sticky edit (design)
|
||||
|
||||
**Goal:** Bring the Text Gate to parity with the Image Gate's "Run from here"
|
||||
affordance, plus a text-specific touch: keep the user's edited text across
|
||||
re-runs ("start from there").
|
||||
|
||||
**Scope:** Frontend only — `web/text_gate.js`. No changes to `gates/textgate.py`,
|
||||
`gates/gate_bus.py`, or `gates/gate_server.py`. The gate already re-arms and
|
||||
re-pauses on every run (`GateBus.arm` → `wait_payload`) and `IS_CHANGED` returns
|
||||
`NaN`, so re-queuing the prompt is enough to "resume": cached upstream means the
|
||||
gate re-pauses near-instantly.
|
||||
|
||||
## State machine
|
||||
|
||||
The node currently has no explicit state. Add three:
|
||||
|
||||
- **idle** — before the first run. Pass shown, Run-from-here hidden.
|
||||
- **paused** — socket `datasete-textgate-show` arrived. Textarea editable &
|
||||
populated, **▶ Pass** shown, **Run from here** hidden, status `edit, then Pass`.
|
||||
- **passed** — after Pass click. Textarea keeps the edited text, **Pass** hidden,
|
||||
**▶ Run from here** shown, status `passed — Run from here to re-run`.
|
||||
|
||||
**Run from here** click → executes the `Comfy.QueuePrompt` command via
|
||||
`app.extensionManager.command.execute(...)` — the same path the Run button and
|
||||
Ctrl+Enter use, so the prompt actually starts. A bare `app.queuePrompt(0, 1)`
|
||||
enqueues but skips the command's run setup, so the 1.47 frontend doesn't kick off
|
||||
execution (you'd have to press Run yourself). `app.queuePrompt` remains a fallback
|
||||
for older frontends without the command registry.
|
||||
|
||||
## Sticky edited text (by intent, not text comparison)
|
||||
|
||||
The Image Gate keeps its mask sticky; the Text Gate keeps its text. Stickiness is
|
||||
keyed off **which action triggered the run**, not a text comparison — because the
|
||||
upstream feeding `text` is often non-deterministic (random/seeded prompts), so a
|
||||
text comparison would wrongly clobber the edit on every Run-from-here.
|
||||
|
||||
- The "Run from here" button sets `node._tgKeepEdit = true` before re-queuing.
|
||||
- On the next re-pause (`datasete-textgate-show`):
|
||||
- if `node._tgKeepEdit` → **keep** the current textarea value and clear the
|
||||
flag, so the gate re-emits *your* edited text downstream.
|
||||
- else (a normal toolbar Queue) → overwrite the textarea with the incoming
|
||||
upstream text.
|
||||
|
||||
Net: Run-from-here always preserves your edit; a deliberate full Queue shows the
|
||||
fresh upstream text. `_tgKeepEdit` is per-session (not serialized).
|
||||
|
||||
**Out of scope:** re-queuing still recomputes non-cacheable upstream nodes — that
|
||||
is inherent to ComfyUI and identical for the Image Gate. With intent-based
|
||||
stickiness the regenerated text is simply ignored, so it can't change the result;
|
||||
to skip the compute, Bypass (Ctrl+B) the upstream node manually.
|
||||
|
||||
## Verification
|
||||
|
||||
- `node --check web/text_gate.js` (no JS test harness in the repo — consistent
|
||||
with the other `web/*.js`).
|
||||
- Manual: pause → edit → Pass → button appears → Run-from-here re-pauses showing
|
||||
your edited text → downstream re-runs; change something upstream → new input
|
||||
shows.
|
||||
|
||||
## Dropped (YAGNI)
|
||||
|
||||
- A separate "↺ reset to input" button — the upstream-change detection covers the
|
||||
stale-edit footgun.
|
||||
- Any backend auto-pass / bypass mode — not requested.
|
||||
@@ -0,0 +1,60 @@
|
||||
# Save Image + chainable sidecars (design)
|
||||
|
||||
**Goal:** A save-image node (like KJ's `SaveImageKJ`) that, instead of a single
|
||||
caption, writes any number of **sidecar** text/JSON files alongside the image,
|
||||
each sharing the image's base name so associations never break.
|
||||
|
||||
Decisions (from brainstorming):
|
||||
- **One unified `Sidecar` node** (content + name + extension), not per-type nodes.
|
||||
- **JSON is just a string** — content is a STRING written verbatim; the extension
|
||||
decides `.txt` vs `.json`.
|
||||
|
||||
## Nodes (backend-only — standard widgets/slots, no web JS)
|
||||
|
||||
### `Sidecar` — one link in the chain
|
||||
- Inputs: `content` (STRING, forceInput) · `name` (STRING, default `""`) ·
|
||||
`extension` (STRING, default `.txt`) · `sidecar` (optional `SIDECAR` chain-in).
|
||||
- Output: `sidecar` (`SIDECAR`) — a list; appends `{content, name, ext}` to the
|
||||
incoming chain and passes it on. Pure, no comfy imports.
|
||||
- `SIDECAR` is a custom type so only sidecar/save ports interconnect.
|
||||
|
||||
### `Save Image (Sidecars)` — end of the chain
|
||||
- Inputs: `images` (IMAGE) · `filename_prefix` (default `ComfyUI`) ·
|
||||
`output_folder` (default `output`; absolute or under the ComfyUI output dir) ·
|
||||
`sidecar` (optional `SIDECAR`). `OUTPUT_NODE`, returns the image preview.
|
||||
- `folder_paths.get_save_image_path()` → `base = f"{filename}_{counter:05}_"`
|
||||
(mirrors `SaveImageKJ`). Saves `base.png`, then each sidecar as `base + name + ext`.
|
||||
|
||||
## Filename rule
|
||||
|
||||
`base` already ends in `_`, so it is the separator:
|
||||
|
||||
| name | ext | file |
|
||||
|---|---|---|
|
||||
| `""` | `.txt` | `ComfyUI_00001_.txt` (caption, shares the image base) |
|
||||
| `""` | `.json` | `ComfyUI_00001_.json` |
|
||||
| `variant_a` | `.txt` | `ComfyUI_00001_variant_a.txt` |
|
||||
|
||||
Image: `ComfyUI_00001_.png`. Batch > 1 writes each sidecar per image.
|
||||
|
||||
## Validation (all before any file is written → no partial output)
|
||||
|
||||
- **Duplicate → error:** two sidecars resolving to the same `name+ext`
|
||||
(two empty-name `.txt`, or two `variant_a.json`) raise a clear `ValueError`.
|
||||
- **Extension allowlist:** `.txt .caption .json .yaml .yml .md .csv .tsv .xml .log
|
||||
.ini .toml`. `name`/`extension` sanitized to a basename; per-file `commonpath`
|
||||
path-traversal guard. (All copied from `SaveImageKJ`.)
|
||||
|
||||
## Code layout / testing
|
||||
|
||||
- `gates/sidecar.py` — pure logic: `ALLOWED_EXTENSIONS`, `normalize_ext`,
|
||||
`sanitize_name`, `append_spec`, `build_plan`. Unit-tested (chain build, filename
|
||||
resolution, duplicate raises, bad-ext raises) — no torch/comfy.
|
||||
- `gates/sidecar_node.py` — the two node classes; torch/PIL/`folder_paths`
|
||||
imported lazily inside `save()`. `build_plan` runs before any I/O.
|
||||
- `__init__.py` — additive registration.
|
||||
|
||||
## Rejected / deferred
|
||||
|
||||
- Per-type text/json nodes (unified node covers both).
|
||||
- Structured JSON/dict input (string-JSON is enough).
|
||||
@@ -0,0 +1,60 @@
|
||||
# Text Gate — "Protected" mode (standalone text node)
|
||||
|
||||
**Goal:** A `protected` switch that turns the Text Gate into a standalone text
|
||||
node: no pause, it outputs the text you typed every run, ignoring upstream. Toggle
|
||||
off → back to the normal pause/edit/Pass gate using the upstream text.
|
||||
|
||||
Decisions (from brainstorming):
|
||||
- Protected = **plain-text-node behavior** (no pause), not a "still-pause-but-lock".
|
||||
- The upstream wire is **kept but its value ignored** while protected (toggle off
|
||||
resumes upstream seamlessly — no reconnecting).
|
||||
|
||||
## Backend (`gates/textgate.py`)
|
||||
|
||||
The authored text and the flag must reach `run()`, so:
|
||||
|
||||
- `text` input: `required` → **`optional`** (`forceInput` kept), so the node runs
|
||||
standalone. Existing connections still work.
|
||||
- New serializing widgets:
|
||||
- `protected` (BOOLEAN, default `False`) — the switch.
|
||||
- `stored_text` (STRING) — the authored text, hidden in the UI behind the DOM
|
||||
editor; the textarea syncs into it.
|
||||
- `run(self, unique_id=None, text=None, signal=None, protected=False, stored_text="")`:
|
||||
- `protected` → `return (stored_text, signal)` immediately — no `GateBus`, no
|
||||
pause, upstream ignored. (Returns early *before* importing comfy, so it stays
|
||||
import-safe/unit-testable.)
|
||||
- else → current pause flow, guarding an unconnected input with `text or ""`.
|
||||
- `IS_CHANGED`: `protected` → return `stored_text` (cache-friendly like a real text
|
||||
node; downstream only re-runs when the text changes). Else → `float("nan")` (so
|
||||
the existing NaN test still passes).
|
||||
|
||||
## Frontend (`web/text_gate.js`)
|
||||
|
||||
- Hide the auto-created `stored_text` widget (`computeSize → [0,-4]`, the pool
|
||||
node's trick); the DOM textarea stays the single editor and writes its value into
|
||||
`stored_text` on every edit (persists + reaches the backend).
|
||||
- Read the `protected` boolean toggle (label "🔒 Protected (text node)"). On **ON**:
|
||||
snapshot the current textarea into `stored_text`, hide Pass / Run-from-here, show
|
||||
status "🔒 protected — outputs this text, upstream ignored", keep the textarea
|
||||
editable. On **OFF**: revert to the normal pause UI.
|
||||
- Ignore the `datasete-textgate-show` socket while protected. On load, populate the
|
||||
textarea from `stored_text`.
|
||||
|
||||
## Persistence & compat
|
||||
|
||||
`protected` + `stored_text` are real widgets → save/reload restores mode + text.
|
||||
Old saved TextGates get `protected=false`, `stored_text=""` defaults (the DOM editor
|
||||
is `serialize:false`, so old nodes carry no conflicting widgets_values).
|
||||
|
||||
## Testing
|
||||
|
||||
- Unit: `run(protected=True, stored_text="hi")` → `("hi", signal)` without touching
|
||||
`GateBus`; `IS_CHANGED(protected=True, stored_text="hi")` → `"hi"`;
|
||||
`IS_CHANGED(protected=False)` → `NaN`; `text` is in `INPUT_TYPES()["optional"]`.
|
||||
- Frontend: `node --check`; manual — toggle protect, edit freely, Run doesn't
|
||||
overwrite, save/reload keeps text, toggle off resumes upstream.
|
||||
|
||||
## Rejected
|
||||
|
||||
A frontend-only "lock" that still pauses — doesn't give true text-node behavior
|
||||
(you'd still click Pass each run), which is the point of the switch.
|
||||
@@ -0,0 +1,72 @@
|
||||
"""BucketResize node: cover-crop an image (and optional mask) onto a Klein
|
||||
training bucket. Pure compute (torch + PIL); no comfy imports in run()."""
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from . import buckets
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
def _resize_crop_pil(pil, new_w, new_h, left, top, W, H):
|
||||
pil = pil.resize((new_w, new_h), Image.LANCZOS)
|
||||
return pil.crop((left, top, left + W, top + H))
|
||||
|
||||
|
||||
def fit_image(image, W, H):
|
||||
"""image [B,H,W,3] -> [B,H,W,3] at (W,H) using the first image's geometry."""
|
||||
b, ih, iw = image.shape[0], image.shape[1], image.shape[2]
|
||||
new_w, new_h, left, top, scale = buckets.cover_crop_params(iw, ih, W, H)
|
||||
out = []
|
||||
for i in range(b):
|
||||
arr = (image[i].cpu().numpy() * 255.0).clip(0, 255).astype("uint8")
|
||||
pil = _resize_crop_pil(Image.fromarray(arr), new_w, new_h, left, top, W, H)
|
||||
out.append(torch.from_numpy(np.array(pil, dtype=np.float32) / 255.0))
|
||||
return torch.stack(out, 0), scale
|
||||
|
||||
|
||||
def fit_mask(mask, W, H):
|
||||
b, ih, iw = mask.shape[0], mask.shape[1], mask.shape[2]
|
||||
new_w, new_h, left, top, _ = buckets.cover_crop_params(iw, ih, W, H)
|
||||
out = []
|
||||
for i in range(b):
|
||||
arr = (mask[i].cpu().numpy() * 255.0).clip(0, 255).astype("uint8")
|
||||
pil = _resize_crop_pil(Image.fromarray(arr), new_w, new_h, left, top, W, H)
|
||||
out.append(torch.from_numpy(np.array(pil, dtype=np.float32) / 255.0))
|
||||
return torch.stack(out, 0)
|
||||
|
||||
|
||||
class BucketResize:
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("image", "mask", "width", "height", "label")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"resolution": ("INT", {"default": 1280, "min": 64, "max": 8192}),
|
||||
"divisible": ("INT", {"default": 64, "min": 8, "max": 256}),
|
||||
"max_upscale": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 8.0, "step": 0.1}),
|
||||
},
|
||||
"optional": {"mask": ("MASK",)},
|
||||
}
|
||||
|
||||
def run(self, image, resolution=1280, divisible=64, max_upscale=1.5, mask=None):
|
||||
ih, iw = int(image.shape[1]), int(image.shape[2])
|
||||
W, H = buckets.pick_bucket(iw, ih, resolution, divisible)
|
||||
out_img, scale = fit_image(image, W, H)
|
||||
if scale > max_upscale:
|
||||
print(f"[BucketResize] cover scale {scale:.2f}x exceeds max_upscale "
|
||||
f"{max_upscale} for {iw}x{ih} -> {W}x{H}")
|
||||
out_mask = fit_mask(mask, W, H) if mask is not None \
|
||||
else torch.zeros((out_img.shape[0], H, W), dtype=torch.float32)
|
||||
return (out_img, out_mask, W, H, f"{W}x{H}")
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"BucketResize": BucketResize}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"BucketResize": "Bucket Resize (Klein 9B)"}
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Pure bucket math for KLEIN_BUCKET_SIZES.md. Stdlib only."""
|
||||
import math
|
||||
|
||||
|
||||
def pick_bucket(iw, ih, resolution=1280, divisible=64):
|
||||
"""Choose the on-grid bucket (W,H), area <= resolution^2, nearest to the
|
||||
image aspect (log distance; tie-break larger area)."""
|
||||
budget = resolution * resolution
|
||||
target = iw / ih
|
||||
best = None
|
||||
w = divisible
|
||||
w_max = budget // divisible
|
||||
while w <= w_max:
|
||||
h = (budget // w // divisible) * divisible # largest on-grid h within budget
|
||||
if h >= divisible:
|
||||
err = abs(math.log(w / h) - math.log(target))
|
||||
cand = (err, -(w * h), w, h) # min err, then max area
|
||||
if best is None or cand < best:
|
||||
best = cand
|
||||
w += divisible
|
||||
return best[2], best[3]
|
||||
|
||||
|
||||
def cover_crop_params(iw, ih, W, H):
|
||||
"""Cover-scale + centered crop to land (iw,ih) exactly on (W,H)."""
|
||||
scale = max(W / iw, H / ih)
|
||||
new_w = max(W, round(iw * scale))
|
||||
new_h = max(H, round(ih * scale))
|
||||
left = (new_w - W) // 2
|
||||
top = (new_h - H) // 2
|
||||
return new_w, new_h, left, top, scale
|
||||
+4
-3
@@ -25,7 +25,7 @@ def mask_from_stash(data, image):
|
||||
|
||||
|
||||
class ImageGate:
|
||||
CATEGORY = "Datasete Gates"
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("MASK",) + ("IMAGE",) * MAX_ROUTES
|
||||
RETURN_NAMES = ("mask",) + tuple(f"route_{i + 1}" for i in range(MAX_ROUTES))
|
||||
@@ -47,13 +47,14 @@ class ImageGate:
|
||||
def run(self, image, routes, unique_id):
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
from . import gate_server
|
||||
import comfy.model_management as mm
|
||||
|
||||
gate_bus.GateBus.arm(unique_id)
|
||||
gate_server.send_preview(unique_id, image, routes)
|
||||
try:
|
||||
chosen_1 = gate_bus.GateBus.wait(unique_id)
|
||||
chosen_1 = gate_bus.GateBus.wait(
|
||||
unique_id, should_cancel=mm.processing_interrupted)
|
||||
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)
|
||||
|
||||
+2
-2
@@ -27,10 +27,10 @@ class GateBus:
|
||||
cls.messages[str(node_id)] = int(message)
|
||||
|
||||
@classmethod
|
||||
def wait(cls, node_id, period=0.1):
|
||||
def wait(cls, node_id, period=0.1, should_cancel=None):
|
||||
sid = str(node_id)
|
||||
while sid not in cls.messages:
|
||||
if cls.cancelled:
|
||||
if cls.cancelled or (should_cancel is not None and should_cancel()):
|
||||
cls.cancelled = False
|
||||
raise GateCancelled()
|
||||
time.sleep(period)
|
||||
|
||||
+1
-1
@@ -27,7 +27,7 @@ def load_image_and_mask(path):
|
||||
|
||||
|
||||
class FolderImageLoader:
|
||||
CATEGORY = "Datasete Gates"
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "STRING", "MASK", "STRING", "INT")
|
||||
RETURN_NAMES = ("image", "text", "mask", "filename", "index")
|
||||
|
||||
+1
-1
@@ -8,7 +8,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
class GridImagePool:
|
||||
CATEGORY = "Datasete Gates"
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("image", "mask", "index", "count", "label")
|
||||
|
||||
@@ -4,7 +4,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
class PoolProfile:
|
||||
CATEGORY = "Datasete Gates"
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("POOL_PROFILE",)
|
||||
RETURN_NAMES = ("profile",)
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Pure sidecar-file planning: filename resolution, extension allowlist, and
|
||||
duplicate detection. No comfy/torch imports so it stays unit-testable."""
|
||||
import os
|
||||
|
||||
# Plain-text / data formats only — never image or executable extensions.
|
||||
ALLOWED_EXTENSIONS = {
|
||||
".txt", ".caption", ".json", ".yaml", ".yml", ".md",
|
||||
".csv", ".tsv", ".xml", ".log", ".ini", ".toml",
|
||||
}
|
||||
|
||||
|
||||
def normalize_ext(ext):
|
||||
"""Sanitize an extension to a leading-dot basename and allowlist it."""
|
||||
ext = os.path.basename((ext or "").strip())
|
||||
if ext and not ext.startswith("."):
|
||||
ext = "." + ext
|
||||
if ext.lower() not in ALLOWED_EXTENSIONS:
|
||||
raise ValueError(
|
||||
f"Disallowed sidecar extension {ext!r}. "
|
||||
f"Allowed: {', '.join(sorted(ALLOWED_EXTENSIONS))}")
|
||||
return ext
|
||||
|
||||
|
||||
def sanitize_name(name):
|
||||
"""Reduce a name field to a bare filename token (no path traversal)."""
|
||||
return os.path.basename((name or "").strip())
|
||||
|
||||
|
||||
def append_spec(chain, content, name, ext):
|
||||
"""Return a new chain list with this sidecar spec appended (no mutation)."""
|
||||
out = list(chain) if chain else []
|
||||
out.append({"content": content, "name": name, "ext": ext})
|
||||
return out
|
||||
|
||||
|
||||
def build_plan(specs):
|
||||
"""Resolve specs to a list of (suffix, content), where the file written is
|
||||
`<image_base> + suffix` and suffix is `name + ext`. Validates extensions and
|
||||
rejects duplicate filenames, raising ValueError *before* any I/O so a bad
|
||||
chain writes nothing."""
|
||||
seen = set()
|
||||
plan = []
|
||||
for s in specs or []:
|
||||
ext = normalize_ext(s.get("ext"))
|
||||
name = sanitize_name(s.get("name"))
|
||||
suffix = f"{name}{ext}"
|
||||
if suffix in seen:
|
||||
raise ValueError(
|
||||
f"Duplicate sidecar file '<base>{suffix}': two sidecars resolve "
|
||||
f"to the same name (name={name!r}, ext={ext}). "
|
||||
f"Give one a distinct name.")
|
||||
seen.add(suffix)
|
||||
plan.append((suffix, s.get("content", "")))
|
||||
return plan
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Save Image + chainable sidecar text/JSON files.
|
||||
|
||||
`Sidecar` nodes chain a list of {content, name, ext} specs (SIDECAR type);
|
||||
`SaveImageSidecars` saves the image and writes each sidecar next to it sharing
|
||||
the image's base name. Heavy deps (torch/PIL/folder_paths) are imported lazily
|
||||
inside save() so this module imports without comfy for unit tests."""
|
||||
from . import sidecar as sc
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
class Sidecar:
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("SIDECAR",)
|
||||
RETURN_NAMES = ("sidecar",)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"content": ("STRING", {"forceInput": True}),
|
||||
"name": ("STRING", {"default": ""}),
|
||||
"extension": ("STRING", {"default": ".txt"}),
|
||||
},
|
||||
"optional": {
|
||||
"sidecar": ("SIDECAR",), # chain-in from a previous Sidecar
|
||||
},
|
||||
}
|
||||
|
||||
def run(self, content, name, extension, sidecar=None):
|
||||
return (sc.append_spec(sidecar, content, name, extension),)
|
||||
|
||||
|
||||
class SaveImageSidecars:
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "save"
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"filename_prefix": ("STRING", {"default": "ComfyUI"}),
|
||||
"output_folder": ("STRING", {"default": "output"}),
|
||||
},
|
||||
"optional": {
|
||||
"sidecar": ("SIDECAR",),
|
||||
},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _safe_path(folder, filename):
|
||||
import os
|
||||
path = os.path.join(folder, filename)
|
||||
root = os.path.abspath(folder)
|
||||
if os.path.commonpath((root, os.path.abspath(path))) != root:
|
||||
raise ValueError(f"Refusing to write outside the target folder: {path}")
|
||||
return path
|
||||
|
||||
def save(self, images, filename_prefix, output_folder, sidecar=None,
|
||||
prompt=None, extra_pnginfo=None):
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
|
||||
# Validate the entire sidecar plan BEFORE writing anything, so a bad
|
||||
# chain (duplicate name, disallowed extension) writes no files at all.
|
||||
plan = sc.build_plan(sidecar)
|
||||
|
||||
h, w = int(images[0].shape[0]), int(images[0].shape[1])
|
||||
if os.path.isabs(output_folder):
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
output_dir = output_folder
|
||||
else:
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = \
|
||||
folder_paths.get_save_image_path(filename_prefix, output_dir, w, h)
|
||||
|
||||
results = []
|
||||
for batch_number, image in enumerate(images):
|
||||
arr = (255.0 * image.cpu().numpy()).clip(0, 255).astype(np.uint8)
|
||||
img = Image.fromarray(arr)
|
||||
|
||||
metadata = None
|
||||
if not args.disable_metadata:
|
||||
metadata = PngInfo()
|
||||
if prompt is not None:
|
||||
metadata.add_text("prompt", json.dumps(prompt))
|
||||
if extra_pnginfo is not None:
|
||||
for k in extra_pnginfo:
|
||||
metadata.add_text(k, json.dumps(extra_pnginfo[k]))
|
||||
|
||||
base = f"{filename.replace('%batch_num%', str(batch_number))}_{counter:05}_"
|
||||
img.save(self._safe_path(full_output_folder, base + ".png"),
|
||||
pnginfo=metadata, compress_level=4)
|
||||
for suffix, content in plan:
|
||||
with open(self._safe_path(full_output_folder, base + suffix),
|
||||
"w", encoding="utf-8") as f:
|
||||
f.write(content if content is not None else "")
|
||||
|
||||
results.append({"filename": base + ".png",
|
||||
"subfolder": subfolder, "type": "output"})
|
||||
counter += 1
|
||||
|
||||
return {"ui": {"images": results}}
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Sidecar": Sidecar,
|
||||
"SaveImageSidecars": SaveImageSidecars,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Sidecar": "Sidecar (text/json)",
|
||||
"SaveImageSidecars": "Save Image (Sidecars)",
|
||||
}
|
||||
+21
-8
@@ -15,33 +15,46 @@ ANY = AnyType("*")
|
||||
|
||||
|
||||
class TextGate:
|
||||
CATEGORY = "Datasete Gates"
|
||||
CATEGORY = "Dataset Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("STRING", ANY)
|
||||
RETURN_NAMES = ("text", "signal")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# `text` is optional so the node can run standalone in protected mode.
|
||||
# `protected` + `stored_text` are serializing widgets carrying the
|
||||
# authored text-node state (stored_text is hidden by the frontend).
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"forceInput": True}),
|
||||
},
|
||||
"optional": {
|
||||
"text": ("STRING", {"forceInput": True}),
|
||||
"signal": (ANY, {}),
|
||||
"protected": ("BOOLEAN", {"default": False}),
|
||||
# single-line so the frontend can fully hide it (the DOM editor
|
||||
# is the real text box); the value still holds arbitrary text.
|
||||
"stored_text": ("STRING", {"default": ""}),
|
||||
},
|
||||
"hidden": {"unique_id": "UNIQUE_ID"},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
def IS_CHANGED(cls, protected=False, stored_text="", **kwargs):
|
||||
# Protected = plain text node: cache on the authored text so downstream
|
||||
# only re-runs when it changes. Otherwise never cache (always pause).
|
||||
return stored_text if protected else float("nan")
|
||||
|
||||
def run(self, unique_id=None, text=None, signal=None,
|
||||
protected=False, stored_text=""):
|
||||
if protected:
|
||||
# Standalone text node: emit the authored text, ignore upstream, no
|
||||
# pause. Returns before importing comfy, so it stays import-safe.
|
||||
return (stored_text, signal)
|
||||
|
||||
def run(self, text, unique_id, signal=None):
|
||||
from . import gate_server
|
||||
import comfy.model_management as mm
|
||||
|
||||
gate_bus.GateBus.arm(unique_id)
|
||||
gate_server.send_text(unique_id, text)
|
||||
gate_server.send_text(unique_id, text or "")
|
||||
try:
|
||||
edited = gate_bus.GateBus.wait_payload(
|
||||
unique_id, should_cancel=mm.processing_interrupted)
|
||||
|
||||
+1
-1
@@ -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]
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
import torch
|
||||
from gates import bucket_node as bn
|
||||
|
||||
|
||||
def test_square_to_1280():
|
||||
out, m, w, h, label = bn.BucketResize().run(image=torch.rand((1, 1000, 1000, 3)))
|
||||
assert (w, h) == (1280, 1280)
|
||||
assert out.shape == (1, 1280, 1280, 3)
|
||||
assert m.shape == (1, 1280, 1280) and float(m.max()) == 0.0 # no mask -> zeros
|
||||
assert label == "1280x1280"
|
||||
|
||||
|
||||
def test_landscape_bucket_shapes():
|
||||
# tensor [B,H,W,3] with H=1000,W=2000 -> aspect 2.0 -> 1792x896
|
||||
out, m, w, h, label = bn.BucketResize().run(image=torch.rand((1, 1000, 2000, 3)))
|
||||
assert (w, h) == (1792, 896)
|
||||
assert out.shape == (1, 896, 1792, 3)
|
||||
assert label == "1792x896"
|
||||
|
||||
|
||||
def test_mask_resized_and_aligned():
|
||||
out, m, w, h, _ = bn.BucketResize().run(
|
||||
image=torch.rand((1, 1000, 1000, 3)), mask=torch.ones((1, 1000, 1000)))
|
||||
assert m.shape == (1, 1280, 1280) and float(m.min()) > 0.9
|
||||
|
||||
|
||||
def test_outputs_are_on_grid():
|
||||
out, m, w, h, _ = bn.BucketResize().run(
|
||||
image=torch.rand((1, 777, 1333, 3)), resolution=1280, divisible=64)
|
||||
assert w % 64 == 0 and h % 64 == 0
|
||||
assert out.shape[1] == h and out.shape[2] == w
|
||||
@@ -0,0 +1,47 @@
|
||||
from gates import buckets
|
||||
|
||||
# (iw, ih) -> expected (W, H) from KLEIN_BUCKET_SIZES.md, budget 1280, ÷64
|
||||
CASES = [
|
||||
(1000, 1000, 1280, 1280), # square
|
||||
(1000, 2000, 896, 1792), # a=0.50 portrait
|
||||
(1000, 1730, 960, 1664), # a≈0.58
|
||||
(1000, 1100, 1216, 1344), # a≈0.90 -> portrait-leaning
|
||||
(2000, 1000, 1792, 896), # a=2.00 landscape
|
||||
(1500, 1000, 1536, 1024), # a=1.50
|
||||
]
|
||||
|
||||
|
||||
def test_pick_bucket_matches_table():
|
||||
for iw, ih, W, H in CASES:
|
||||
assert buckets.pick_bucket(iw, ih, 1280, 64) == (W, H)
|
||||
|
||||
|
||||
def test_buckets_are_on_grid_and_within_budget():
|
||||
for iw, ih, *_ in CASES:
|
||||
W, H = buckets.pick_bucket(iw, ih, 1280, 64)
|
||||
assert W % 64 == 0 and H % 64 == 0
|
||||
assert W * H <= 1280 * 1280
|
||||
|
||||
|
||||
def test_square_is_exactly_1280():
|
||||
assert buckets.pick_bucket(512, 512, 1280, 64) == (1280, 1280)
|
||||
|
||||
|
||||
def test_cover_crop_exact_aspect_no_crop():
|
||||
# a=2.0 image onto 1792x896 bucket -> scale 0.896, no crop
|
||||
new_w, new_h, left, top, scale = buckets.cover_crop_params(2000, 1000, 1792, 896)
|
||||
assert (new_w, new_h) == (1792, 896)
|
||||
assert (left, top) == (0, 0)
|
||||
assert round(scale, 3) == 0.896
|
||||
|
||||
|
||||
def test_cover_crop_square_into_landscape_crops_height():
|
||||
new_w, new_h, left, top, scale = buckets.cover_crop_params(1000, 1000, 1792, 896)
|
||||
assert new_w == 1792 and new_h >= 896
|
||||
assert left == 0 and top == (new_h - 896) // 2 # centered vertical crop
|
||||
assert scale > 1.0 # upscaled to cover width
|
||||
|
||||
|
||||
def test_cover_crop_upscale_square():
|
||||
*_, scale = buckets.cover_crop_params(1000, 1000, 1280, 1280)
|
||||
assert round(scale, 2) == 1.28
|
||||
@@ -65,3 +65,10 @@ def test_wait_payload_should_cancel_raises():
|
||||
gb.GateBus.arm("p")
|
||||
with pytest.raises(gb.GateCancelled):
|
||||
gb.GateBus.wait_payload("p", should_cancel=lambda: True)
|
||||
|
||||
def test_wait_should_cancel_raises():
|
||||
# image gate: ComfyUI Interrupt (should_cancel) must abort the wait too
|
||||
gb.GateBus.arm("7")
|
||||
with pytest.raises(gb.GateCancelled):
|
||||
gb.GateBus.wait("7", should_cancel=lambda: True)
|
||||
assert gb.GateBus.cancelled is False
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
import pytest
|
||||
|
||||
from gates import sidecar
|
||||
|
||||
|
||||
def test_append_spec_builds_chain_without_mutating():
|
||||
c1 = sidecar.append_spec(None, "hello", "", ".txt")
|
||||
c2 = sidecar.append_spec(c1, "{}", "meta", ".json")
|
||||
assert c1 == [{"content": "hello", "name": "", "ext": ".txt"}]
|
||||
assert len(c2) == 2
|
||||
assert c2[1] == {"content": "{}", "name": "meta", "ext": ".json"}
|
||||
assert len(c1) == 1 # original chain untouched
|
||||
|
||||
|
||||
def test_normalize_ext_adds_dot_and_allowlists():
|
||||
assert sidecar.normalize_ext("txt") == ".txt"
|
||||
assert sidecar.normalize_ext(".json") == ".json"
|
||||
with pytest.raises(ValueError):
|
||||
sidecar.normalize_ext(".png")
|
||||
with pytest.raises(ValueError):
|
||||
sidecar.normalize_ext(".exe")
|
||||
|
||||
|
||||
def test_sanitize_name_strips_path_and_space():
|
||||
assert sidecar.sanitize_name(" variant_a ") == "variant_a"
|
||||
assert sidecar.sanitize_name("../evil") == "evil"
|
||||
assert sidecar.sanitize_name("a/b") == "b"
|
||||
assert sidecar.sanitize_name("") == ""
|
||||
|
||||
|
||||
def test_build_plan_resolves_suffixes():
|
||||
specs = [
|
||||
{"content": "cap", "name": "", "ext": ".txt"},
|
||||
{"content": "{}", "name": "", "ext": ".json"},
|
||||
{"content": "v", "name": "variant_a", "ext": ".txt"},
|
||||
]
|
||||
assert sidecar.build_plan(specs) == [
|
||||
(".txt", "cap"),
|
||||
(".json", "{}"),
|
||||
("variant_a.txt", "v"),
|
||||
]
|
||||
|
||||
|
||||
def test_build_plan_duplicate_empty_names_raises():
|
||||
specs = [
|
||||
{"content": "a", "name": "", "ext": ".txt"},
|
||||
{"content": "b", "name": "", "ext": ".txt"},
|
||||
]
|
||||
with pytest.raises(ValueError):
|
||||
sidecar.build_plan(specs)
|
||||
|
||||
|
||||
def test_build_plan_empty_txt_and_json_do_not_collide():
|
||||
specs = [
|
||||
{"content": "a", "name": "", "ext": ".txt"},
|
||||
{"content": "b", "name": "", "ext": ".json"},
|
||||
]
|
||||
assert len(sidecar.build_plan(specs)) == 2
|
||||
|
||||
|
||||
def test_build_plan_bad_extension_raises():
|
||||
with pytest.raises(ValueError):
|
||||
sidecar.build_plan([{"content": "x", "name": "", "ext": ".png"}])
|
||||
|
||||
|
||||
def test_build_plan_none_is_empty():
|
||||
assert sidecar.build_plan(None) == []
|
||||
@@ -0,0 +1,34 @@
|
||||
from gates import sidecar_node as sn
|
||||
|
||||
|
||||
def test_sidecar_run_builds_chain():
|
||||
node = sn.Sidecar()
|
||||
(chain,) = node.run(content="hello", name="", extension=".txt", sidecar=None)
|
||||
assert chain == [{"content": "hello", "name": "", "ext": ".txt"}]
|
||||
(chain2,) = node.run(content="{}", name="meta", extension=".json", sidecar=chain)
|
||||
assert len(chain2) == 2
|
||||
assert chain2[1] == {"content": "{}", "name": "meta", "ext": ".json"}
|
||||
|
||||
|
||||
def test_sidecar_io_shape():
|
||||
assert sn.Sidecar.RETURN_TYPES == ("SIDECAR",)
|
||||
it = sn.Sidecar.INPUT_TYPES()
|
||||
assert "content" in it["required"]
|
||||
assert "name" in it["required"]
|
||||
assert "extension" in it["required"]
|
||||
assert "sidecar" in it["optional"]
|
||||
|
||||
|
||||
def test_save_node_io_shape():
|
||||
assert sn.SaveImageSidecars.OUTPUT_NODE is True
|
||||
assert sn.SaveImageSidecars.RETURN_TYPES == ()
|
||||
it = sn.SaveImageSidecars.INPUT_TYPES()
|
||||
for k in ("images", "filename_prefix", "output_folder"):
|
||||
assert k in it["required"]
|
||||
assert "sidecar" in it["optional"]
|
||||
|
||||
|
||||
def test_mappings_present():
|
||||
assert "Sidecar" in sn.NODE_CLASS_MAPPINGS
|
||||
assert "SaveImageSidecars" in sn.NODE_CLASS_MAPPINGS
|
||||
assert sn.NODE_DISPLAY_NAME_MAPPINGS["SaveImageSidecars"]
|
||||
@@ -16,3 +16,31 @@ def test_textgate_io_shape():
|
||||
def test_textgate_is_changed_nan():
|
||||
v = textgate.TextGate.IS_CHANGED(text="hi", unique_id="1")
|
||||
assert math.isnan(v)
|
||||
|
||||
|
||||
def test_textgate_text_input_is_optional():
|
||||
it = textgate.TextGate.INPUT_TYPES()
|
||||
assert "text" in it["optional"]
|
||||
assert "protected" in it["optional"]
|
||||
assert "stored_text" in it["optional"]
|
||||
|
||||
|
||||
def test_textgate_protected_returns_stored_text_without_pause():
|
||||
# protected mode must return the stored text directly — no GateBus, no comfy
|
||||
out = textgate.TextGate().run(
|
||||
unique_id="1", text="from upstream", signal="sig",
|
||||
protected=True, stored_text="my authored text",
|
||||
)
|
||||
assert out == ("my authored text", "sig")
|
||||
|
||||
|
||||
def test_textgate_is_changed_protected_returns_stored_text():
|
||||
v = textgate.TextGate.IS_CHANGED(
|
||||
unique_id="1", protected=True, stored_text="frozen")
|
||||
assert v == "frozen"
|
||||
|
||||
|
||||
def test_textgate_is_changed_not_protected_is_nan():
|
||||
v = textgate.TextGate.IS_CHANGED(
|
||||
unique_id="1", protected=False, stored_text="ignored")
|
||||
assert math.isnan(v)
|
||||
|
||||
@@ -215,6 +215,16 @@ function showResolved(node, choiceLabel) {
|
||||
}
|
||||
|
||||
async function queueFromHere(node) {
|
||||
// Fire the same command the Run button / Ctrl+Enter use, so the prompt
|
||||
// actually EXECUTES. A bare app.queuePrompt(...) enqueues but skips the
|
||||
// command's run setup, so the 1.47 frontend doesn't kick off the run (you'd
|
||||
// have to press Run yourself). Fall back to app.queuePrompt on older
|
||||
// frontends without the command registry.
|
||||
const cmd = app.extensionManager?.command;
|
||||
if (cmd?.execute) {
|
||||
try { await cmd.execute("Comfy.QueuePrompt"); return; }
|
||||
catch (e) { /* fall through to the legacy path */ }
|
||||
}
|
||||
try {
|
||||
await app.queuePrompt(0, 1);
|
||||
} catch (e) {
|
||||
|
||||
+175
-10
@@ -7,6 +7,15 @@ import { api } from "../../scripts/api.js";
|
||||
// the server pushes via the "datasete-textgate-show" socket event and POSTs the
|
||||
// edited text back. Outputs are static (text, signal) — no dynamic slots.
|
||||
//
|
||||
// After Pass, a "▶ Run from here" button re-queues the prompt (Image Gate
|
||||
// parity): the gate re-arms every run and IS_CHANGED is NaN, so it re-pauses
|
||||
// each run. The edited text is sticky by INTENT: a Run-from-here re-queue keeps
|
||||
// YOUR edited text (even if a non-deterministic upstream regenerates it), while
|
||||
// a normal toolbar Queue shows whatever the upstream produced. Keying off which
|
||||
// button ran — not a text comparison — means a random/seeded upstream can't
|
||||
// clobber the edit on re-run. (Re-queuing still recomputes non-cacheable
|
||||
// upstream, as in any ComfyUI run; that regenerated text is simply ignored.)
|
||||
//
|
||||
// Sizing follows the Image Pool node: the editor is always present and FILLS the
|
||||
// node, with only a min-height floor (no max) so the node stays freely resizable
|
||||
// and the textarea grows with it.
|
||||
@@ -19,6 +28,48 @@ const MIN_EDITOR_H = 140; // textarea floor
|
||||
const BTN_ROW_H = 34; // Pass button row
|
||||
const MARGIN = 10; // ComfyUI DOM-widget inset, matches the other nodes
|
||||
|
||||
// ---- protected-mode widgets -------------------------------------------------
|
||||
// `protected` (BOOLEAN toggle) + `stored_text` (hidden STRING) are real backend
|
||||
// widgets. When protected, the node acts as a plain text node: it outputs
|
||||
// stored_text and ignores upstream (no pause). The DOM textarea is the visible
|
||||
// editor and mirrors its value into stored_text on EVERY change (typing, upstream
|
||||
// arrival, Pass) — so the editor content survives refresh / workflow reload in
|
||||
// BOTH modes (stored_text also reaches run() when protected).
|
||||
|
||||
function widgetByName(node, name) {
|
||||
return node.widgets?.find((w) => w.name === name);
|
||||
}
|
||||
|
||||
function isProtected(node) {
|
||||
return !!widgetByName(node, "protected")?.value;
|
||||
}
|
||||
|
||||
// mirror the editor text into the hidden stored_text widget (persist + backend)
|
||||
function syncStored(node) {
|
||||
const w = widgetByName(node, "stored_text");
|
||||
if (w) w.value = node._tg?.area?.value ?? "";
|
||||
}
|
||||
|
||||
// fully hide the auto-created stored_text widget (same as the pool node's
|
||||
// pool_id): getVisibleWidgets() filters on `hidden`, so it's dropped from both
|
||||
// draw and layout — computeSize alone (or type="hidden") does NOT hide it.
|
||||
// Serialization still iterates all widgets, so stored_text is saved/sent.
|
||||
function hideStoredWidget(node) {
|
||||
const w = widgetByName(node, "stored_text");
|
||||
if (!w) return;
|
||||
w.hidden = true;
|
||||
w.computeSize = () => [0, -4];
|
||||
}
|
||||
|
||||
// reflect the persisted stored_text + mode into the editor + UI. The editor text
|
||||
// is restored in BOTH modes so it survives a refresh / workflow reload; the mode
|
||||
// only selects the UI state (protected vs idle waiting-for-a-run).
|
||||
function applyPersistedMode(node) {
|
||||
if (!node._tg) return;
|
||||
node._tg.area.value = widgetByName(node, "stored_text")?.value ?? "";
|
||||
setState(node, isProtected(node) ? "protected" : "idle");
|
||||
}
|
||||
|
||||
// ---- server call ------------------------------------------------------------
|
||||
|
||||
async function postPass(node, text) {
|
||||
@@ -28,6 +79,47 @@ async function postPass(node, text) {
|
||||
await api.fetchApi(`${R}/pass`, { method: "POST", body: fd });
|
||||
}
|
||||
|
||||
// ---- run-from-here + state --------------------------------------------------
|
||||
// States: "idle" (pre-run), "paused" (waiting for Pass), "passed" (Run-from-here
|
||||
// shown). Re-queuing the whole prompt is enough to "resume" — cached upstream
|
||||
// re-pauses the gate, matching the Image Gate's queueFromHere.
|
||||
|
||||
async function queueFromHere(node) {
|
||||
// Fire the same command the Run button / Ctrl+Enter use, so the prompt
|
||||
// actually EXECUTES. A bare app.queuePrompt(...) enqueues but skips the
|
||||
// command's run setup, so the 1.47 frontend doesn't kick off the run (you'd
|
||||
// have to press Run yourself). Fall back to app.queuePrompt on older
|
||||
// frontends without the command registry.
|
||||
const cmd = app.extensionManager?.command;
|
||||
if (cmd?.execute) {
|
||||
try { await cmd.execute("Comfy.QueuePrompt"); return; }
|
||||
catch (e) { /* fall through to the legacy path */ }
|
||||
}
|
||||
try {
|
||||
await app.queuePrompt(0, 1);
|
||||
} catch (e) {
|
||||
try { await app.queuePrompt(0); } catch (e2) { console.error("[tgate] queue failed", e2); }
|
||||
}
|
||||
}
|
||||
|
||||
function setState(node, s) {
|
||||
node._tgState = s;
|
||||
const tg = node._tg;
|
||||
if (!tg) return;
|
||||
// Pass is hidden once passed AND in protected mode (no pause there);
|
||||
// Run-from-here only in the passed state.
|
||||
tg.pass.style.display = (s === "passed" || s === "protected") ? "none" : "";
|
||||
tg.runHere.style.display = s === "passed" ? "" : "none";
|
||||
if (s === "paused") tg.status.textContent = "edit, then Pass";
|
||||
else if (s === "passed") tg.status.textContent = "passed — Run from here to re-run";
|
||||
else if (s === "protected") tg.status.textContent = "🔒 protected — outputs this text (upstream ignored)";
|
||||
else tg.status.textContent = "";
|
||||
tg.area.placeholder = s === "protected"
|
||||
? "type text (used as a text node)…"
|
||||
: "waiting for a run…";
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
// ---- sizing (Image Pool pattern) --------------------------------------------
|
||||
|
||||
// Only a min-height FLOOR — no max — so the DOM widget fills the node and grows
|
||||
@@ -59,6 +151,8 @@ function injectStyles() {
|
||||
border:1px solid #555; color:#fff; }
|
||||
.tgate-pass { background:rgba(40,130,70,0.95); }
|
||||
.tgate-pass:hover { background:rgba(55,160,90,0.98); }
|
||||
.tgate-run { background:rgba(40,90,140,0.95); }
|
||||
.tgate-run:hover { background:rgba(60,120,180,0.98); }
|
||||
.tgate-status { font-size:11px; opacity:0.6; margin-left:auto; }
|
||||
`;
|
||||
const style = document.createElement("style");
|
||||
@@ -76,26 +170,55 @@ function setupTextGateNode(node) {
|
||||
const area = document.createElement("textarea");
|
||||
area.className = "tgate-area";
|
||||
area.placeholder = "waiting for a run…";
|
||||
// don't let typing/space toggle node selection or graph shortcuts
|
||||
area.onkeydown = (e) => e.stopPropagation();
|
||||
// Stop keys from reaching litegraph (so typing/space can't toggle node
|
||||
// selection or fire canvas shortcuts) — EXCEPT ComfyUI's prompt-weighting
|
||||
// shortcut (Ctrl/Cmd+↑/↓). That handler is a global `window` keydown listener
|
||||
// that wraps the selection in (token:weight); a blanket stopPropagation here
|
||||
// kept it from ever bubbling up, so weighting didn't work in this editor.
|
||||
// Its execCommand edit fires our oninput, so the weighted text still syncs.
|
||||
area.onkeydown = (e) => {
|
||||
const isWeight = (e.ctrlKey || e.metaKey) &&
|
||||
(e.key === "ArrowUp" || e.key === "ArrowDown");
|
||||
if (!isWeight) e.stopPropagation();
|
||||
};
|
||||
// keep the hidden stored_text widget mirrored so edits persist + reach run()
|
||||
area.oninput = () => syncStored(node);
|
||||
|
||||
const btns = document.createElement("div");
|
||||
btns.className = "tgate-btns";
|
||||
|
||||
const pass = document.createElement("button");
|
||||
pass.className = "tgate-pass";
|
||||
pass.textContent = "▶ Pass";
|
||||
pass.onclick = async () => {
|
||||
syncStored(node); // persist the passed text so a reload keeps it
|
||||
await postPass(node, area.value);
|
||||
setState(node, "passed");
|
||||
};
|
||||
|
||||
// Re-queue the prompt; cached upstream re-pauses the gate so you can run your
|
||||
// edited text downstream again without recomputing the graph above it.
|
||||
const runHere = document.createElement("button");
|
||||
runHere.className = "tgate-run";
|
||||
runHere.textContent = "▶ Run from here";
|
||||
runHere.style.display = "none";
|
||||
runHere.onclick = async () => {
|
||||
node._tgKeepEdit = true; // tell the next re-pause to preserve this edit
|
||||
node._tg.status.textContent = "re-running…";
|
||||
await queueFromHere(node);
|
||||
};
|
||||
|
||||
const status = document.createElement("span");
|
||||
status.className = "tgate-status";
|
||||
pass.onclick = async () => {
|
||||
await postPass(node, area.value);
|
||||
status.textContent = "passed";
|
||||
};
|
||||
|
||||
btns.appendChild(pass);
|
||||
btns.appendChild(runHere);
|
||||
btns.appendChild(status);
|
||||
|
||||
wrap.appendChild(area);
|
||||
wrap.appendChild(btns);
|
||||
node._tg = { wrap, area, status };
|
||||
node._tg = { wrap, area, status, pass, runHere };
|
||||
node._tgState = "idle";
|
||||
|
||||
// FILLS the node: floor-only min height, no max (Image Pool pattern).
|
||||
node._tgWidget = node.addDOMWidget("textgate_editor", "div", wrap, {
|
||||
@@ -111,24 +234,57 @@ function setupTextGateNode(node) {
|
||||
return r;
|
||||
};
|
||||
|
||||
// protected-mode wiring: hide the stored_text widget, label + react to the
|
||||
// toggle, and reflect the persisted mode/text into the editor.
|
||||
hideStoredWidget(node);
|
||||
const pw = widgetByName(node, "protected");
|
||||
if (pw) {
|
||||
pw.label = "🔒 Protected (text node)";
|
||||
const prev = pw.callback;
|
||||
pw.callback = function () {
|
||||
const r = prev?.apply(this, arguments);
|
||||
if (isProtected(node)) { syncStored(node); setState(node, "protected"); }
|
||||
else setState(node, "idle");
|
||||
return r;
|
||||
};
|
||||
}
|
||||
applyPersistedMode(node);
|
||||
|
||||
// sensible default size; the node stays freely resizable (no width floor lock)
|
||||
node.setSize([Math.max(node.size?.[0] || 0, MIN_W), node.computeSize()[1]]);
|
||||
syncWidgetWidth(node);
|
||||
}
|
||||
|
||||
// Build marker — lets you confirm the browser loaded THIS build (not a cached
|
||||
// old copy). If the editor comes back empty after reload but you don't see this
|
||||
// line in the devtools console, your tab is running stale JS: hard-refresh
|
||||
// (Ctrl/Cmd+Shift+R).
|
||||
const BUILD = "2026-07-03 persist+weight";
|
||||
|
||||
app.registerExtension({
|
||||
name: "datasete.gates.textgate",
|
||||
|
||||
// one global socket listener: route the server's pause event to the node
|
||||
setup() {
|
||||
console.info(`[datasete.textgate] loaded build ${BUILD}`);
|
||||
api.addEventListener("datasete-textgate-show", (e) => {
|
||||
const d = e.detail || {};
|
||||
const node = app.graph?.getNodeById?.(parseInt(d.id, 10));
|
||||
if (!node || node.type !== NODE || !node._tg) return;
|
||||
node._tg.area.value = d.text || "";
|
||||
node._tg.status.textContent = "edit, then Pass";
|
||||
if (isProtected(node)) return; // protected = no pause; ignore stray events
|
||||
// Sticky edit by intent: a Run-from-here re-queue (the _tgKeepEdit flag)
|
||||
// keeps YOUR edited text so the gate re-emits it downstream; a normal
|
||||
// Queue shows whatever the upstream produced. Keying off the button —
|
||||
// not a text comparison — means a non-deterministic upstream can't
|
||||
// clobber the edit on re-run.
|
||||
if (node._tgKeepEdit) {
|
||||
node._tgKeepEdit = false;
|
||||
} else {
|
||||
node._tg.area.value = d.text || "";
|
||||
}
|
||||
syncStored(node); // persist the shown text so a refresh/reload keeps it
|
||||
setState(node, "paused");
|
||||
try { node._tg.area.focus(); } catch (err) { /* ignore */ }
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
});
|
||||
},
|
||||
|
||||
@@ -141,5 +297,14 @@ app.registerExtension({
|
||||
setupTextGateNode(this);
|
||||
return r;
|
||||
};
|
||||
|
||||
// loaded workflows restore protected + stored_text after create — re-apply
|
||||
// the mode so the editor + UI match the saved state.
|
||||
const onConfigure = nodeType.prototype.onConfigure;
|
||||
nodeType.prototype.onConfigure = function () {
|
||||
const r = onConfigure?.apply(this, arguments);
|
||||
if (this._tg) applyPersistedMode(this);
|
||||
return r;
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user