Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| aa909448d7 | |||
| 037cbf27db | |||
| 969463a4e9 | |||
| 7f90b6878f | |||
| 0413e25571 | |||
| cdd742c950 | |||
| 628a945514 | |||
| 95b3417ff6 |
+4
-2
@@ -18,14 +18,16 @@ if __package__:
|
|||||||
NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES
|
NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES
|
||||||
from .gates.profile_node import NODE_CLASS_MAPPINGS as _PROF_NODES, \
|
from .gates.profile_node import NODE_CLASS_MAPPINGS as _PROF_NODES, \
|
||||||
NODE_DISPLAY_NAME_MAPPINGS as _PROF_NAMES
|
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 import routes # noqa: F401 (registers aiohttp routes on import)
|
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 gate_server # noqa: F401 (registers /datasete_gate/* + text routes)
|
||||||
from .gates import profiles_routes # noqa: F401 (registers /grid_pool/profiles/*)
|
from .gates import profiles_routes # noqa: F401 (registers /grid_pool/profiles/*)
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES,
|
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES,
|
||||||
**_TEXT_NODES, **_PROF_NODES}
|
**_TEXT_NODES, **_PROF_NODES, **_BUCKET_NODES}
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES,
|
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES,
|
||||||
**_TEXT_NAMES, **_PROF_NAMES}
|
**_TEXT_NAMES, **_PROF_NAMES, **_BUCKET_NAMES}
|
||||||
else: # pragma: no cover - exercised only under pytest collection
|
else: # pragma: no cover - exercised only under pytest collection
|
||||||
NODE_CLASS_MAPPINGS = {}
|
NODE_CLASS_MAPPINGS = {}
|
||||||
NODE_DISPLAY_NAME_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,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 = "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)"}
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
Reference in New Issue
Block a user