diff --git a/README.md b/README.md index e62d29b..02507e8 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ 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. -All nodes appear under the **“Datasete Gates”** category. +All nodes appear under the **“Dataset Gates”** category. ## Nodes at a glance @@ -14,6 +14,7 @@ All nodes appear under the **“Datasete Gates”** category. | **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. | +| **Image Chooser Gate (Batch)** | `ImageChooserGate` | Pauses the run, displays every image in an incoming batch, and passes the **selected subset** onward as a batch. | | **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 @@ -172,6 +173,33 @@ silently skipped. Built for manual dataset sorting. --- +## Image Chooser Gate (Batch) + +Pauses the running prompt and shows the complete incoming image batch as a +scrollable thumbnail grid. Click one or more thumbnails, then click **Pass +selected** to emit just those images as a new batch. + +### Inputs / Outputs + +| Port | Type | Description | +|------|------|-------------| +| `images` (input) | IMAGE | The batch to review. It must contain at least one image. | +| `images` (output) | IMAGE | The selected images as a batch, in their original input order. | + +### How it works + +- Click a thumbnail to select or deselect it. **Select all** and **Clear** help + with larger batches; at least one image is required before passing. +- Each new run starts with an empty selection, so a stale choice cannot be + applied to a different batch. The gate always pauses and is never cached. +- The browser receives only small JPEG thumbnails. The output comes directly + from the original tensor, so choosing images does not resize or recompress + them. +- **Stop** cancels the run. After passing, **Run from here** queues the workflow + again and presents a fresh choice. + +--- + ## Text Gate (Manual Pass) Pauses the run, shows the incoming text in an **editable** box, and emits it @@ -200,10 +228,11 @@ Pauses every run; ComfyUI's global **Cancel** unblocks it cleanly (no deadlock). ### 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. +Image Gate, Image Chooser Gate, and Text Gate **block the executor thread** +during a run and wait for a click (a small server-side waiter plus an HTTP 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 @@ -242,9 +271,11 @@ Layout: - `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/gate_bus.py` — pure blocking choice/text/mask/selection waiter for the gates. - `gates/imaging.py` — torch/PIL tensor loaders. -- `gates/node.py` · `loader.py` · `gate.py` · `textgate.py` · `profile_node.py` — the nodes. +- `gates/node.py` · `loader.py` · `gate.py` · `image_chooser.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/*`). + (`/grid_pool/*`, `/datasete_gate/*`, `/datasete_image_chooser/*`, + `/grid_pool/profiles/*`). - `web/*.js` — the in-node UIs (grid + MaskEditor, gate previews, profile dropdown). diff --git a/__init__.py b/__init__.py index 8cda9e3..4de77dd 100644 --- a/__init__.py +++ b/__init__.py @@ -16,6 +16,8 @@ if __package__: NODE_DISPLAY_NAME_MAPPINGS as _GATE_NAMES from .gates.textgate import NODE_CLASS_MAPPINGS as _TEXT_NODES, \ NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES + from .gates.image_chooser import NODE_CLASS_MAPPINGS as _CHOOSER_NODES, \ + NODE_DISPLAY_NAME_MAPPINGS as _CHOOSER_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, \ @@ -27,11 +29,11 @@ if __package__: from .gates import profiles_routes # noqa: F401 (registers /grid_pool/profiles/*) NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES, - **_TEXT_NODES, **_PROF_NODES, **_BUCKET_NODES, - **_SC_NODES} + **_TEXT_NODES, **_CHOOSER_NODES, **_PROF_NODES, + **_BUCKET_NODES, **_SC_NODES} NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES, - **_TEXT_NAMES, **_PROF_NAMES, **_BUCKET_NAMES, - **_SC_NAMES} + **_TEXT_NAMES, **_CHOOSER_NAMES, **_PROF_NAMES, + **_BUCKET_NAMES, **_SC_NAMES} else: # pragma: no cover - exercised only under pytest collection NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/gates/gate_bus.py b/gates/gate_bus.py index b2499f8..a31e96a 100644 --- a/gates/gate_bus.py +++ b/gates/gate_bus.py @@ -1,5 +1,7 @@ -"""Blocking choice bus for the Image Gate node. Stdlib only — no comfy/torch.""" +"""Blocking coordination for manual gate nodes. Stdlib only — no comfy/torch.""" +import threading import time +import uuid class GateCancelled(Exception): @@ -11,6 +13,11 @@ class GateBus: masks = {} # node_id(str) -> PNG bytes payloads = {} # node_id(str) -> arbitrary payload (e.g., edited text) cancelled = False + active_tokens = {} # node_id(str) -> per-run token for scoped waiters + token_payloads = {} # (node_id, token) -> arbitrary payload + token_cancelled = set() + token_contexts = {} # (node_id, token) -> waiter-specific validation data + token_lock = threading.Lock() @classmethod def arm(cls, node_id): @@ -57,3 +64,78 @@ class GateBus: raise GateCancelled() time.sleep(period) return cls.payloads.pop(sid) + + @classmethod + def arm_token(cls, node_id, context=None): + """Open a run-scoped channel and invalidate any older run for the node.""" + sid = str(node_id) + with cls.token_lock: + old_token = cls.active_tokens.get(sid) + if old_token is not None: + old_key = (sid, old_token) + cls.token_payloads.pop(old_key, None) + cls.token_cancelled.discard(old_key) + cls.token_contexts.pop(old_key, None) + + token = uuid.uuid4().hex + key = (sid, token) + cls.active_tokens[sid] = token + cls.token_contexts[key] = context + return token + + @classmethod + def token_context(cls, node_id, token): + sid = str(node_id) + with cls.token_lock: + if cls.active_tokens.get(sid) != token: + return None + return cls.token_contexts.get((sid, token)) + + @classmethod + def put_token_payload(cls, node_id, token, value): + sid = str(node_id) + with cls.token_lock: + if cls.active_tokens.get(sid) != token: + return False + cls.token_payloads[(sid, token)] = value + return True + + @classmethod + def cancel_token(cls, node_id, token): + sid = str(node_id) + with cls.token_lock: + if cls.active_tokens.get(sid) != token: + return False + cls.token_cancelled.add((sid, token)) + return True + + @classmethod + def wait_token_payload(cls, node_id, token, period=0.1, should_cancel=None): + sid = str(node_id) + key = (sid, token) + while True: + with cls.token_lock: + superseded = cls.active_tokens.get(sid) != token + cancelled = key in cls.token_cancelled + if superseded or cancelled: + cls.token_cancelled.discard(key) + raise GateCancelled() + if key in cls.token_payloads: + return cls.token_payloads.pop(key) + + if should_cancel is not None and should_cancel(): + with cls.token_lock: + cls.token_cancelled.discard(key) + raise GateCancelled() + time.sleep(period) + + @classmethod + def disarm_token(cls, node_id, token): + sid = str(node_id) + key = (sid, token) + with cls.token_lock: + if cls.active_tokens.get(sid) == token: + cls.active_tokens.pop(sid, None) + cls.token_payloads.pop(key, None) + cls.token_cancelled.discard(key) + cls.token_contexts.pop(key, None) diff --git a/gates/gate_server.py b/gates/gate_server.py index 6dcce61..79e1b5a 100644 --- a/gates/gate_server.py +++ b/gates/gate_server.py @@ -8,6 +8,7 @@ from PIL import Image from server import PromptServer from .gate_bus import GateBus +from .image_chooser import encode_previews, normalize_selection routes = PromptServer.instance.routes @@ -23,6 +24,22 @@ def send_preview(node_id, image, n_routes): ) +def send_image_choices(node_id, token, images): + """Show a lightweight preview of every image to the queuing client.""" + server = PromptServer.instance + server.send_sync( + "datasete-image-chooser-show", + { + "id": str(node_id), + "display_id": str(getattr(server, "last_node_id", None) or node_id), + "token": token, + "images": encode_previews(images), + "count": int(images.shape[0]), + }, + getattr(server, "client_id", None), + ) + + @routes.post("/datasete_gate/choice") async def _choice(request): post = await request.post() @@ -44,6 +61,35 @@ async def _mask(request): return web.json_response({}) +@routes.post("/datasete_image_chooser/select") +async def _image_chooser_select(request): + post = await request.post() + node_id = post.get("id") + token = post.get("token") + if node_id is None or token is None: + return web.json_response({"error": "missing node id or token"}, status=400) + + batch_size = GateBus.token_context(node_id, token) + if batch_size is None: + return web.json_response({"error": "chooser run is no longer active"}, status=409) + + if post.get("action") == "cancel": + accepted = GateBus.cancel_token(node_id, token) + else: + selection = post.get("selection") + if selection is None: + return web.json_response({"error": "missing selection"}, status=400) + try: + selection = normalize_selection(selection, batch_size) + except ValueError as exc: + return web.json_response({"error": str(exc)}, status=400) + accepted = GateBus.put_token_payload(node_id, token, selection) + + if not accepted: + return web.json_response({"error": "chooser run is no longer active"}, status=409) + return web.json_response({}) + + def send_text(node_id, text): PromptServer.instance.send_sync( "datasete-textgate-show", {"id": str(node_id), "text": text or ""} diff --git a/gates/image_chooser.py b/gates/image_chooser.py new file mode 100644 index 0000000..b513eef --- /dev/null +++ b/gates/image_chooser.py @@ -0,0 +1,121 @@ +"""Manual gate that selects a subset from an IMAGE batch. + +The selection helpers and thumbnail encoder intentionally avoid importing +ComfyUI so this module remains straightforward to unit-test. +""" +import base64 +import io +import json + +import numpy as np +from PIL import Image + +from . import gate_bus + + +PREVIEW_MAX_SIDE = 256 +PREVIEW_JPEG_QUALITY = 82 + + +def normalize_selection(selection, batch_size): + """Return validated, unique indices in their original batch order.""" + if isinstance(selection, str): + try: + selection = json.loads(selection) + except json.JSONDecodeError as exc: + raise ValueError("Selection must be a JSON array of image indices") from exc + + if not isinstance(selection, (list, tuple)): + raise ValueError("Selection must be a list of image indices") + if not selection: + raise ValueError("Select at least one image") + + unique = set() + for index in selection: + if isinstance(index, bool) or not isinstance(index, int): + raise ValueError("Every selected image index must be an integer") + if index < 0 or index >= batch_size: + raise ValueError( + f"Selected image index {index} is outside batch size {batch_size}" + ) + unique.add(index) + + # Batch order is deterministic and does not depend on click order. + return tuple(sorted(unique)) + + +def select_batch(images, selection): + """Select one or more images while preserving the IMAGE batch dimension.""" + batch_size = int(images.shape[0]) + indices = normalize_selection(selection, batch_size) + return images[list(indices)] + + +def encode_previews(images, max_side=PREVIEW_MAX_SIDE, + jpeg_quality=PREVIEW_JPEG_QUALITY): + """Encode small JPEG previews without modifying the original tensor.""" + previews = [] + for index, image in enumerate(images): + array = (image.detach().cpu().float().numpy() * 255.0).clip(0, 255).astype( + np.uint8 + ) + pil = Image.fromarray(array) + if pil.mode != "RGB": + pil = pil.convert("RGB") + source_width, source_height = pil.size + pil.thumbnail((max_side, max_side), Image.Resampling.LANCZOS) + + buffer = io.BytesIO() + pil.save(buffer, "JPEG", quality=jpeg_quality) + previews.append({ + "index": index, + "image": base64.b64encode(buffer.getvalue()).decode("ascii"), + "width": source_width, + "height": source_height, + }) + return previews + + +class ImageChooserGate: + CATEGORY = "Dataset Gates" + FUNCTION = "run" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {"images": ("IMAGE",)}, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("nan") + + def run(self, images, unique_id): + batch_size = int(images.shape[0]) + if batch_size < 1: + raise ValueError("Image Chooser Gate requires a non-empty image batch") + + from . import gate_server + import comfy.model_management as mm + + token = gate_bus.GateBus.arm_token(unique_id, context=batch_size) + try: + gate_server.send_image_choices(unique_id, token, images) + selection = gate_bus.GateBus.wait_token_payload( + unique_id, token, should_cancel=mm.processing_interrupted + ) + except gate_bus.GateCancelled: + raise mm.InterruptProcessingException() + finally: + gate_bus.GateBus.disarm_token(unique_id, token) + + return (select_batch(images, selection),) + + +NODE_CLASS_MAPPINGS = {"ImageChooserGate": ImageChooserGate} +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageChooserGate": "Image Chooser Gate (Batch)", +} diff --git a/tests/test_gate_bus.py b/tests/test_gate_bus.py index a7fe16f..bc8950f 100644 --- a/tests/test_gate_bus.py +++ b/tests/test_gate_bus.py @@ -72,3 +72,32 @@ def test_wait_should_cancel_raises(): with pytest.raises(gb.GateCancelled): gb.GateBus.wait("7", should_cancel=lambda: True) assert gb.GateBus.cancelled is False + + +def test_token_payload_roundtrip_and_context(): + token = gb.GateBus.arm_token("chooser", context=4) + assert gb.GateBus.token_context("chooser", token) == 4 + assert gb.GateBus.put_token_payload("chooser", token, (0, 3)) is True + assert gb.GateBus.wait_token_payload("chooser", token) == (0, 3) + gb.GateBus.disarm_token("chooser", token) + assert gb.GateBus.token_context("chooser", token) is None + + +def test_stale_token_cannot_answer_or_cancel_new_run(): + old_token = gb.GateBus.arm_token("chooser", context=2) + new_token = gb.GateBus.arm_token("chooser", context=5) + + assert gb.GateBus.put_token_payload("chooser", old_token, [0]) is False + assert gb.GateBus.cancel_token("chooser", old_token) is False + assert gb.GateBus.token_context("chooser", new_token) == 5 + with pytest.raises(gb.GateCancelled): + gb.GateBus.wait_token_payload("chooser", old_token) + gb.GateBus.disarm_token("chooser", new_token) + + +def test_token_cancel_only_cancels_matching_waiter(): + token = gb.GateBus.arm_token("chooser", context=1) + assert gb.GateBus.cancel_token("chooser", token) is True + with pytest.raises(gb.GateCancelled): + gb.GateBus.wait_token_payload("chooser", token) + gb.GateBus.disarm_token("chooser", token) diff --git a/tests/test_image_chooser.py b/tests/test_image_chooser.py new file mode 100644 index 0000000..56c4b4e --- /dev/null +++ b/tests/test_image_chooser.py @@ -0,0 +1,117 @@ +import base64 +import io +import math +import sys +import types + +import pytest +import torch +from PIL import Image + +import gates +from gates import image_chooser + + +def test_normalize_selection_deduplicates_and_preserves_batch_order(): + assert image_chooser.normalize_selection("[3, 1, 3]", 4) == (1, 3) + + +@pytest.mark.parametrize( + "selection", + [ + [], + "[]", + "not json", + "1", + [True], + [1.0], + ["1"], + [-1], + [3], + ], +) +def test_normalize_selection_rejects_invalid_choices(selection): + with pytest.raises(ValueError): + image_chooser.normalize_selection(selection, 3) + + +def test_select_batch_keeps_batch_dimension_and_original_pixels(): + images = torch.arange(4 * 2 * 3 * 3, dtype=torch.float32).reshape(4, 2, 3, 3) + + one = image_chooser.select_batch(images, [2]) + many = image_chooser.select_batch(images, [3, 0, 3]) + + assert one.shape == (1, 2, 3, 3) + assert torch.equal(one[0], images[2]) + assert many.shape == (2, 2, 3, 3) + assert torch.equal(many, images[[0, 3]]) + assert many.dtype == images.dtype + + +def test_encode_previews_returns_small_jpegs_and_source_dimensions(): + images = torch.zeros((2, 80, 40, 3), dtype=torch.float32) + images[1] = 1.0 + + previews = image_chooser.encode_previews(images, max_side=32) + + assert [preview["index"] for preview in previews] == [0, 1] + assert [(preview["width"], preview["height"]) for preview in previews] == [ + (40, 80), + (40, 80), + ] + decoded = [ + Image.open(io.BytesIO(base64.b64decode(preview["image"]))) + for preview in previews + ] + assert all(image.format == "JPEG" and image.mode == "RGB" for image in decoded) + assert [image.size for image in decoded] == [(16, 32), (16, 32)] + + +def test_encode_previews_accepts_bfloat16_images(): + images = torch.zeros((1, 8, 8, 3), dtype=torch.bfloat16) + assert len(image_chooser.encode_previews(images)) == 1 + + +def test_run_waits_for_token_scoped_selection(monkeypatch): + fake_server = types.ModuleType("gates.gate_server") + + def send_image_choices(node_id, token, images): + assert images.shape[0] == 3 + assert image_chooser.gate_bus.GateBus.put_token_payload( + node_id, token, [2, 0] + ) + + fake_server.send_image_choices = send_image_choices + monkeypatch.setitem(sys.modules, "gates.gate_server", fake_server) + monkeypatch.setattr(gates, "gate_server", fake_server, raising=False) + + class InterruptProcessingException(Exception): + pass + + fake_mm = types.ModuleType("comfy.model_management") + fake_mm.processing_interrupted = lambda: False + fake_mm.InterruptProcessingException = InterruptProcessingException + fake_comfy = types.ModuleType("comfy") + fake_comfy.model_management = fake_mm + monkeypatch.setitem(sys.modules, "comfy", fake_comfy) + monkeypatch.setitem(sys.modules, "comfy.model_management", fake_mm) + + images = torch.arange(3, dtype=torch.float32).reshape(3, 1, 1, 1) + selected, = image_chooser.ImageChooserGate().run(images, unique_id="12") + + assert torch.equal(selected.flatten(), torch.tensor([0.0, 2.0])) + assert "12" not in image_chooser.gate_bus.GateBus.active_tokens + + +def test_image_chooser_node_contract(): + inputs = image_chooser.ImageChooserGate.INPUT_TYPES() + + assert inputs["required"] == {"images": ("IMAGE",)} + assert inputs["hidden"] == {"unique_id": "UNIQUE_ID"} + assert image_chooser.ImageChooserGate.RETURN_TYPES == ("IMAGE",) + assert image_chooser.ImageChooserGate.RETURN_NAMES == ("images",) + assert image_chooser.ImageChooserGate.FUNCTION == "run" + assert image_chooser.ImageChooserGate.CATEGORY == "Dataset Gates" + assert math.isnan(image_chooser.ImageChooserGate.IS_CHANGED(images=None)) + assert image_chooser.NODE_CLASS_MAPPINGS["ImageChooserGate"] \ + is image_chooser.ImageChooserGate diff --git a/web/image_chooser.js b/web/image_chooser.js new file mode 100644 index 0000000..8870a6b --- /dev/null +++ b/web/image_chooser.js @@ -0,0 +1,479 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +// Image Chooser Gate (Batch) — pauses a prompt, displays every image in the +// incoming batch, and resumes with the selected subset as a new IMAGE batch. + +const NODE = "ImageChooserGate"; +const R = "/datasete_image_chooser"; + +const MIN_W = 420; +const MIN_GRID_H = 190; +const TOOLBAR_H = 72; +const MARGIN = 10; +const chooserNodes = new Set(); + +function graphNodeById(graph, id) { + if (!graph?.getNodeById) return null; + return graph.getNodeById(id) ?? graph.getNodeById(parseInt(id, 10)); +} + +function chooserByLocator(locator) { + const direct = graphNodeById(app.graph, locator); + if (direct?.type === NODE) return direct; + + // Native subgraphs identify executing nodes as colon-separated locators: + // root subgraph-node id -> nested node id -> ... -> chooser id. + let graph = app.rootGraph ?? app.graph?.rootGraph ?? app.graph; + let node = null; + for (const id of String(locator).split(":")) { + node = graphNodeById(graph, id); + if (!node) return null; + graph = node.subgraph; + } + return node?.type === NODE ? node : null; +} + +function widgetFloor(node) { + return node._icgState === "idle" ? 0 : 2 * MARGIN + MIN_GRID_H + TOOLBAR_H; +} + +function syncWidgetWidth(node) { + if (node._icgWidget) node._icgWidget.width = node.size?.[0] || MIN_W; +} + +function resizeChooser(node) { + const shown = node._icgState !== "idle"; + if (node._icg?.wrap) node._icg.wrap.style.display = shown ? "flex" : "none"; + const width = node.size?.[0] || MIN_W; + const target = shown + ? Math.max(node.size?.[1] || 0, node.computeSize()[1]) + : node.computeSize()[1]; + node.setSize([width, target]); + syncWidgetWidth(node); + node.setDirtyCanvas?.(true, true); +} + +async function postChooser(node, fields) { + const form = new FormData(); + form.append("id", String(node._icgWaitId ?? node.id)); + form.append("token", String(node._icgToken ?? "")); + for (const [key, value] of Object.entries(fields)) form.append(key, value); + const response = await api.fetchApi(`${R}/select`, { method: "POST", body: form }); + if (!response.ok) throw new Error(`chooser request failed (${response.status})`); +} + +async function queueFromHere() { + const command = app.extensionManager?.command; + if (command?.execute) { + try { + await command.execute("Comfy.QueuePrompt"); + return true; + } catch (e) { /* use the legacy path below */ } + } + try { + await app.queuePrompt(0, 1); + return true; + } catch (e) { + try { + await app.queuePrompt(0); + return true; + } catch (fallbackError) { + console.error("[image-chooser] queue failed", fallbackError); + return false; + } + } +} + +function selectedIndices(node) { + return [...node._icgSelected].sort((a, b) => a - b); +} + +function updateCells(node) { + for (const [index, cell] of node._icgCells) { + const selected = node._icgSelected.has(index); + cell.classList.toggle("icg-selected", selected); + cell.setAttribute("aria-pressed", selected ? "true" : "false"); + const check = cell.querySelector(".icg-check"); + if (check) check.textContent = selected ? "✓" : ""; + } +} + +function select(node, indices) { + if (node._icgState !== "paused" || node._icgBusy) return; + node._icgSelected = new Set(indices); + node._icgError = ""; + updateCells(node); + renderToolbar(node); +} + +function button(text, className, onclick) { + const element = document.createElement("button"); + element.type = "button"; + element.textContent = text; + element.className = className; + element.onclick = onclick; + return element; +} + +function renderToolbar(node) { + const { toolbar } = node._icg; + toolbar.innerHTML = ""; + const count = node._icgItems.length; + const chosen = selectedIndices(node); + + if (node._icgState === "paused") { + const all = button("Select all", "icg-secondary", () => { + select(node, node._icgItems.map((item) => item.index)); + }); + all.disabled = node._icgBusy || chosen.length === count; + + const clear = button("Clear", "icg-secondary", () => select(node, [])); + clear.disabled = node._icgBusy || chosen.length === 0; + + const pass = button( + chosen.length ? `▶ Pass selected (${chosen.length})` : "▶ Pass selected", + "icg-pass", + async () => { + if (!chosen.length || node._icgBusy) return; + const revision = node._icgRevision; + node._icgBusy = true; + node._icgError = ""; + renderGridState(node); + renderToolbar(node); + try { + await postChooser(node, { selection: JSON.stringify(chosen) }); + if (node._icgRevision === revision) { + node._icgState = "resolved"; + node._icgResolvedCount = chosen.length; + } + } catch (error) { + console.error("[image-chooser] pass failed", error); + if (node._icgRevision === revision) { + node._icgError = "Could not pass the selection"; + } + } finally { + if (node._icgRevision !== revision) return; + node._icgBusy = false; + renderGridState(node); + renderToolbar(node); + resizeChooser(node); + } + }, + ); + pass.disabled = node._icgBusy || chosen.length === 0; + + const stop = button("■ Stop", "icg-stop", async () => { + if (node._icgBusy) return; + const revision = node._icgRevision; + node._icgBusy = true; + node._icgError = ""; + renderGridState(node); + renderToolbar(node); + try { + await postChooser(node, { action: "cancel" }); + if (node._icgRevision === revision) { + node._icgState = "resolved"; + node._icgStopped = true; + } + } catch (error) { + console.error("[image-chooser] stop failed", error); + if (node._icgRevision === revision) { + node._icgError = "Could not stop the run"; + } + } finally { + if (node._icgRevision !== revision) return; + node._icgBusy = false; + renderGridState(node); + renderToolbar(node); + } + }); + stop.disabled = node._icgBusy; + + const status = document.createElement("span"); + status.className = `icg-status${node._icgError ? " icg-error" : ""}`; + status.textContent = node._icgError || (node._icgBusy + ? "sending…" + : `${chosen.length} of ${count} selected`); + + toolbar.appendChild(all); + toolbar.appendChild(clear); + toolbar.appendChild(pass); + toolbar.appendChild(status); + toolbar.appendChild(stop); + } else if (node._icgState === "resolved") { + const status = document.createElement("span"); + status.className = `icg-status icg-resolved${node._icgError ? " icg-error" : ""}`; + status.textContent = node._icgError || (node._icgQueueBusy + ? "queuing…" + : node._icgStopped + ? "■ run stopped" + : `✓ passed ${node._icgResolvedCount} of ${count}`); + const run = button("▶ Run from here", "icg-run", async () => { + if (node._icgQueueBusy) return; + node._icgQueueBusy = true; + node._icgError = ""; + renderToolbar(node); + if (!await queueFromHere()) { + node._icgQueueBusy = false; + node._icgError = "Could not queue the workflow"; + renderToolbar(node); + } + }); + run.disabled = node._icgQueueBusy; + toolbar.appendChild(status); + toolbar.appendChild(run); + } +} + +function renderGridState(node) { + const editable = node._icgState === "paused" && !node._icgBusy; + for (const cell of node._icgCells.values()) cell.disabled = !editable; + updateCells(node); +} + +function toggleImage(node, index) { + if (node._icgState !== "paused" || node._icgBusy) return; + const next = new Set(node._icgSelected); + if (next.has(index)) next.delete(index); + else next.add(index); + select(node, next); +} + +function renderGrid(node) { + const { grid } = node._icg; + grid.innerHTML = ""; + node._icgCells = new Map(); + + for (const item of node._icgItems) { + const cell = document.createElement("button"); + cell.type = "button"; + cell.className = "icg-cell"; + cell.title = `Image ${item.index + 1} — ${item.width}×${item.height}`; + cell.setAttribute("aria-label", `Select image ${item.index + 1}`); + cell.onclick = () => toggleImage(node, item.index); + + const image = document.createElement("img"); + image.className = "icg-thumb"; + image.alt = `Batch image ${item.index + 1}`; + image.draggable = false; + image.src = `data:image/jpeg;base64,${item.image}`; + + const badge = document.createElement("span"); + badge.className = "icg-badge"; + badge.textContent = String(item.index + 1); + + const dimensions = document.createElement("span"); + dimensions.className = "icg-dimensions"; + dimensions.textContent = `${item.width}×${item.height}`; + + const check = document.createElement("span"); + check.className = "icg-check"; + + cell.appendChild(image); + cell.appendChild(badge); + cell.appendChild(dimensions); + cell.appendChild(check); + grid.appendChild(cell); + node._icgCells.set(item.index, cell); + } + renderGridState(node); +} + +function showBatch(node, waitId, token, items) { + node._icgRevision += 1; + node._icgWaitId = waitId; + node._icgToken = token; + node._icgItems = Array.isArray(items) ? items : []; + node._icgSelected = new Set(); + node._icgBusy = false; + node._icgError = ""; + node._icgStopped = false; + node._icgResolvedCount = 0; + node._icgQueueBusy = false; + node._icgState = "paused"; + renderGrid(node); + renderToolbar(node); + resizeChooser(node); +} + +function injectStyles() { + if (document.getElementById("icg-styles")) return; + const css = ` + .icg-wrap { display:flex; flex-direction:column; gap:7px; box-sizing:border-box; + height:100%; min-height:0; } + .icg-grid { flex:1 1 auto; min-height:0; overflow:auto; display:grid; + grid-template-columns:repeat(auto-fill,minmax(108px,1fr)); + grid-auto-rows:max-content; gap:7px; align-content:start; padding:2px; } + .icg-cell { position:relative; aspect-ratio:1; min-width:0; overflow:hidden; padding:0; + border:2px solid transparent; border-radius:5px; background:#1d1d1d; + cursor:pointer; transition:border-color .1s, box-shadow .1s; } + .icg-cell:hover:not(:disabled) { border-color:#777; } + .icg-cell:disabled { cursor:default; opacity:1; } + .icg-cell.icg-selected { border-color:#55d78a; box-shadow:0 0 0 1px #55d78a; } + .icg-thumb { position:absolute; inset:0; width:100%; height:100%; object-fit:contain; + background:#181818; pointer-events:none; } + .icg-badge, .icg-dimensions, .icg-check { position:absolute; color:#fff; + background:rgba(0,0,0,.72); border-radius:3px; pointer-events:none; } + .icg-badge { top:4px; left:4px; min-width:18px; padding:1px 4px; font-size:11px; + text-align:center; } + .icg-dimensions { left:4px; bottom:4px; padding:1px 4px; font-size:9px; opacity:.78; } + .icg-check { top:4px; right:4px; width:20px; height:20px; line-height:20px; + text-align:center; font-size:14px; font-weight:bold; } + .icg-cell.icg-selected .icg-check { background:rgba(35,145,80,.92); } + .icg-toolbar { flex:0 0 auto; display:flex; flex-wrap:wrap; gap:6px; align-items:center; } + .icg-toolbar button { font-size:12px; padding:4px 10px; cursor:pointer; border-radius:3px; + border:1px solid #555; color:#fff; } + .icg-toolbar button:disabled { opacity:.42; cursor:default; } + .icg-secondary { background:rgba(45,45,45,.95); } + .icg-pass { background:rgba(40,130,70,.95); } + .icg-pass:hover:not(:disabled) { background:rgba(55,160,90,.98); } + .icg-run { background:rgba(40,90,140,.95); } + .icg-stop { background:rgba(160,40,40,.9); margin-left:auto; } + .icg-status { font-size:11px; opacity:.72; padding:0 3px; } + .icg-error { color:#ff8a8a; opacity:1; } + .icg-resolved { color:#83e5aa; opacity:.95; } + `; + const style = document.createElement("style"); + style.id = "icg-styles"; + style.textContent = css; + document.head.appendChild(style); +} + +function setupChooserNode(node) { + injectStyles(); + chooserNodes.add(node); + + // This node owns its preview UI; suppress ComfyUI's output-image preview. + try { + Object.defineProperty(node, "imgs", { + configurable: true, + get() { return undefined; }, + set() { /* suppress */ }, + }); + } catch (e) { /* best effort */ } + + const wrap = document.createElement("div"); + wrap.className = "icg-wrap"; + const grid = document.createElement("div"); + grid.className = "icg-grid"; + const toolbar = document.createElement("div"); + toolbar.className = "icg-toolbar"; + wrap.appendChild(grid); + wrap.appendChild(toolbar); + + node._icg = { wrap, grid, toolbar }; + node._icgItems = []; + node._icgCells = new Map(); + node._icgSelected = new Set(); + node._icgState = "idle"; + node._icgBusy = false; + node._icgRevision = 0; + node._icgWaitId = null; + node._icgToken = null; + node._icgQueueBusy = false; + + node._icgWidget = node.addDOMWidget("image_chooser", "div", wrap, { + serialize: false, + getMinHeight: () => widgetFloor(node), + }); + + const onResize = node.onResize; + node.onResize = function () { + const result = onResize?.apply(this, arguments); + syncWidgetWidth(node); + return result; + }; + + // Removing an actively waiting chooser must not strand the executor. + const onRemoved = node.onRemoved; + node.onRemoved = function () { + if (node._icgState === "paused" && node._icgToken) { + postChooser(node, { action: "cancel" }).catch(() => {}); + } + chooserNodes.delete(node); + return onRemoved?.apply(this, arguments); + }; + + node.setSize([Math.max(node.size?.[0] || 0, MIN_W), node.computeSize()[1]]); + resizeChooser(node); +} + +app.registerExtension({ + name: "datasete.gates.imagechooser", + + setup() { + api.addEventListener("datasete-image-chooser-show", (event) => { + const data = event.detail || {}; + const displayId = data.display_id ?? data.id; + const node = chooserByLocator(displayId); + if (!node || node.type !== NODE || !node._icg) return; + showBatch(node, data.id, data.token, data.images); + }); + + // Collapse the previous review as a new execution starts; a fresh chooser + // event will reopen it if this run reaches the node. An interrupt leaves an + // honest stopped state. + api.addEventListener("execution_start", () => { + for (const node of chooserNodes) { + if (node.type !== NODE || node._icgState === "idle") continue; + node._icgRevision += 1; + node._icgState = "idle"; + node._icgBusy = false; + node._icgQueueBusy = false; + node._icgWaitId = null; + node._icgToken = null; + resizeChooser(node); + } + }); + api.addEventListener("execution_interrupted", () => { + for (const node of chooserNodes) { + if (node.type !== NODE + || (node._icgState !== "paused" && !node._icgQueueBusy)) continue; + node._icgRevision += 1; + node._icgState = "resolved"; + node._icgBusy = false; + node._icgQueueBusy = false; + node._icgStopped = true; + node._icgWaitId = null; + node._icgToken = null; + renderGridState(node); + renderToolbar(node); + } + }); + api.addEventListener("execution_error", () => { + for (const node of chooserNodes) { + if (node.type !== NODE + || (node._icgState !== "paused" && !node._icgQueueBusy)) continue; + node._icgRevision += 1; + node._icgState = "resolved"; + node._icgBusy = false; + node._icgQueueBusy = false; + node._icgStopped = true; + node._icgError = "Execution stopped with an error"; + node._icgWaitId = null; + node._icgToken = null; + renderGridState(node); + renderToolbar(node); + } + }); + }, + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE) return; + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated?.apply(this, arguments); + setupChooserNode(this); + return result; + }; + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + const result = onConfigure?.apply(this, arguments); + if (this._icg) syncWidgetWidth(this); + return result; + }; + }, +});