From 9148dfec25bac7114f936115d86e962bbbc17ae5 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:35:04 +0200 Subject: [PATCH 01/10] feat: gate_bus blocking choice waiter Co-Authored-By: Claude Opus 4.8 --- gates/gate_bus.py | 35 +++++++++++++++++++++++++++++++++++ tests/test_gate_bus.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 gates/gate_bus.py create mode 100644 tests/test_gate_bus.py diff --git a/gates/gate_bus.py b/gates/gate_bus.py new file mode 100644 index 0000000..0afff1c --- /dev/null +++ b/gates/gate_bus.py @@ -0,0 +1,35 @@ +"""Blocking choice bus for the Image Gate node. Stdlib only — no comfy/torch.""" +import time + + +class GateCancelled(Exception): + pass + + +class GateBus: + messages = {} # node_id(str) -> chosen int (1-based) + masks = {} # node_id(str) -> PNG bytes + cancelled = False + + @classmethod + def arm(cls, node_id): + cls.messages.pop(str(node_id), None) + cls.masks.pop(str(node_id), None) + cls.cancelled = False + + @classmethod + def put(cls, node_id, message): + if message == "__cancel__": + cls.cancelled = True + else: + cls.messages[str(node_id)] = int(message) + + @classmethod + def wait(cls, node_id, period=0.1): + sid = str(node_id) + while sid not in cls.messages: + if cls.cancelled: + cls.cancelled = False + raise GateCancelled() + time.sleep(period) + return cls.messages.pop(sid) diff --git a/tests/test_gate_bus.py b/tests/test_gate_bus.py new file mode 100644 index 0000000..a111e41 --- /dev/null +++ b/tests/test_gate_bus.py @@ -0,0 +1,28 @@ +# tests/test_gate_bus.py +import pytest +from gates import gate_bus as gb + +def test_put_and_wait_returns_choice(): + gb.GateBus.arm("7") + gb.GateBus.put("7", "3") + assert gb.GateBus.wait("7") == 3 + +def test_wait_consumes_message(): + gb.GateBus.arm("7") + gb.GateBus.put("7", "2") + gb.GateBus.wait("7") + assert "7" not in gb.GateBus.messages + +def test_cancel_raises_and_resets(): + gb.GateBus.arm("7") + gb.GateBus.put("7", "__cancel__") + with pytest.raises(gb.GateCancelled): + gb.GateBus.wait("7") + assert gb.GateBus.cancelled is False # reset after raising + +def test_arm_clears_stale_state(): + gb.GateBus.put("1", "5") + gb.GateBus.cancelled = True + gb.GateBus.arm("1") + assert "1" not in gb.GateBus.messages + assert gb.GateBus.cancelled is False From 11772bc29d1d6eb445dd2de91a3a53eb9ef7f640 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:35:28 +0200 Subject: [PATCH 02/10] feat: gate_bus mask stash Co-Authored-By: Claude Opus 4.8 --- gates/gate_bus.py | 8 ++++++++ tests/test_gate_bus.py | 10 ++++++++++ 2 files changed, 18 insertions(+) diff --git a/gates/gate_bus.py b/gates/gate_bus.py index 0afff1c..4013ed1 100644 --- a/gates/gate_bus.py +++ b/gates/gate_bus.py @@ -33,3 +33,11 @@ class GateBus: raise GateCancelled() time.sleep(period) return cls.messages.pop(sid) + + @classmethod + def put_mask(cls, node_id, data): + cls.masks[str(node_id)] = data + + @classmethod + def pop_mask(cls, node_id): + return cls.masks.pop(str(node_id), None) diff --git a/tests/test_gate_bus.py b/tests/test_gate_bus.py index a111e41..231b29b 100644 --- a/tests/test_gate_bus.py +++ b/tests/test_gate_bus.py @@ -26,3 +26,13 @@ def test_arm_clears_stale_state(): gb.GateBus.arm("1") assert "1" not in gb.GateBus.messages assert gb.GateBus.cancelled is False + +def test_mask_stash_roundtrip(): + gb.GateBus.put_mask("9", b"PNGDATA") + assert gb.GateBus.pop_mask("9") == b"PNGDATA" + assert gb.GateBus.pop_mask("9") is None # popped + +def test_arm_clears_mask(): + gb.GateBus.put_mask("9", b"x") + gb.GateBus.arm("9") + assert gb.GateBus.pop_mask("9") is None From f0f8676eaaed6324c6c7c91d1176edb0dfb164ca Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:35:51 +0200 Subject: [PATCH 03/10] feat: gate route_tuple helper Co-Authored-By: Claude Opus 4.8 --- gates/gate.py | 15 +++++++++++++++ tests/test_gate.py | 11 +++++++++++ 2 files changed, 26 insertions(+) create mode 100644 gates/gate.py create mode 100644 tests/test_gate.py diff --git a/gates/gate.py b/gates/gate.py new file mode 100644 index 0000000..d7bcfc7 --- /dev/null +++ b/gates/gate.py @@ -0,0 +1,15 @@ +# gates/gate.py +import io +import math + +import numpy as np +import torch +from PIL import Image + +from . import gate_bus + +MAX_ROUTES = 10 + + +def route_tuple(chosen, image, blocker, max_routes=MAX_ROUTES): + return tuple(image if i == chosen else blocker for i in range(max_routes)) diff --git a/tests/test_gate.py b/tests/test_gate.py new file mode 100644 index 0000000..2af7a61 --- /dev/null +++ b/tests/test_gate.py @@ -0,0 +1,11 @@ +# tests/test_gate.py +from gates import gate + +def test_route_tuple_places_image_at_chosen(): + B = object() + t = gate.route_tuple(2, "IMG", B, max_routes=5) + assert t == (B, B, "IMG", B, B) + +def test_route_tuple_length_is_max(): + B = object() + assert len(gate.route_tuple(0, "IMG", B, max_routes=10)) == 10 From ea3438567acf29f84fc6e5c03a70e8e253e7a20a Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:42:04 +0200 Subject: [PATCH 04/10] feat: gate mask_from_stash (paint or zeros) Co-Authored-By: Claude Opus 4.8 --- gates/gate.py | 9 +++++++++ tests/test_gate.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/gates/gate.py b/gates/gate.py index d7bcfc7..6ac5e29 100644 --- a/gates/gate.py +++ b/gates/gate.py @@ -13,3 +13,12 @@ MAX_ROUTES = 10 def route_tuple(chosen, image, blocker, max_routes=MAX_ROUTES): return tuple(image if i == chosen else blocker for i in range(max_routes)) + + +def mask_from_stash(data, image): + b, h, w = image.shape[0], image.shape[1], image.shape[2] + if not data: + return torch.zeros((b, h, w), dtype=torch.float32) + m = Image.open(io.BytesIO(data)).convert("L") + arr = np.array(m, dtype=np.float32) / 255.0 + return torch.from_numpy(arr).unsqueeze(0) diff --git a/tests/test_gate.py b/tests/test_gate.py index 2af7a61..b0c3bdc 100644 --- a/tests/test_gate.py +++ b/tests/test_gate.py @@ -1,4 +1,9 @@ # tests/test_gate.py +import io + +import torch +from PIL import Image + from gates import gate def test_route_tuple_places_image_at_chosen(): @@ -9,3 +14,14 @@ def test_route_tuple_places_image_at_chosen(): def test_route_tuple_length_is_max(): B = object() assert len(gate.route_tuple(0, "IMG", B, max_routes=10)) == 10 + +def test_mask_from_stash_none_is_zeros(): + img = torch.zeros((1, 6, 4, 3)) + m = gate.mask_from_stash(None, img) + assert m.shape == (1, 6, 4) and float(m.max()) == 0.0 + +def test_mask_from_stash_decodes_png(): + buf = io.BytesIO(); Image.new("L", (4, 6), 255).save(buf, "PNG") + img = torch.zeros((1, 6, 4, 3)) + m = gate.mask_from_stash(buf.getvalue(), img) + assert m.shape == (1, 6, 4) and float(m.min()) > 0.99 From d8dbc4fb4b66266202f1cede8d0a16773a45cee1 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:42:54 +0200 Subject: [PATCH 05/10] =?UTF-8?q?feat:=20ImageGate=20node=20=E2=80=94=20pa?= =?UTF-8?q?use,=20route=20via=20ExecutionBlocker,=20mask=20out?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.8 --- gates/gate.py | 42 ++++++++++++++++++++++++++++++++++++++++++ tests/test_gate.py | 10 ++++++++++ 2 files changed, 52 insertions(+) diff --git a/gates/gate.py b/gates/gate.py index 6ac5e29..ae62b52 100644 --- a/gates/gate.py +++ b/gates/gate.py @@ -22,3 +22,45 @@ def mask_from_stash(data, image): m = Image.open(io.BytesIO(data)).convert("L") arr = np.array(m, dtype=np.float32) / 255.0 return torch.from_numpy(arr).unsqueeze(0) + + +class ImageGate: + CATEGORY = "Datasete Gates" + FUNCTION = "run" + RETURN_TYPES = ("MASK",) + ("IMAGE",) * MAX_ROUTES + RETURN_NAMES = ("mask",) + tuple(f"route_{i + 1}" for i in range(MAX_ROUTES)) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "routes": ("INT", {"default": 2, "min": 1, "max": MAX_ROUTES}), + }, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("nan") # always pause; never cached + + def run(self, image, routes, unique_id): + from comfy_execution.graph_utils import ExecutionBlocker + from . import gate_server + + 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() + + mask = mask_from_stash(gate_bus.GateBus.pop_mask(unique_id), image) + chosen = max(0, min(chosen_1 - 1, routes - 1)) + blocker = ExecutionBlocker(None) + return (mask,) + route_tuple(chosen, image, blocker, MAX_ROUTES) + + +NODE_CLASS_MAPPINGS = {"ImageGate": ImageGate} +NODE_DISPLAY_NAME_MAPPINGS = {"ImageGate": "Image Gate (Manual Router)"} diff --git a/tests/test_gate.py b/tests/test_gate.py index b0c3bdc..273427a 100644 --- a/tests/test_gate.py +++ b/tests/test_gate.py @@ -1,5 +1,6 @@ # tests/test_gate.py import io +import math import torch from PIL import Image @@ -25,3 +26,12 @@ def test_mask_from_stash_decodes_png(): img = torch.zeros((1, 6, 4, 3)) m = gate.mask_from_stash(buf.getvalue(), img) assert m.shape == (1, 6, 4) and float(m.min()) > 0.99 + +def test_is_changed_always_nan(): + v = gate.ImageGate.IS_CHANGED(image=None, routes=2, unique_id="1") + assert math.isnan(v) + +def test_return_types_shape(): + assert gate.ImageGate.RETURN_TYPES[0] == "MASK" + assert len(gate.ImageGate.RETURN_TYPES) == gate.MAX_ROUTES + 1 + assert all(t == "IMAGE" for t in gate.ImageGate.RETURN_TYPES[1:]) From 8e8eb317f77435ea6640ac2f7c3e6d0d45c719d8 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:43:53 +0200 Subject: [PATCH 06/10] feat: gate server routes + preview + register ImageGate Co-Authored-By: Claude Opus 4.8 --- __init__.py | 7 +++++-- gates/gate_server.py | 44 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 gates/gate_server.py diff --git a/__init__.py b/__init__.py index 5c5c5da..be141f7 100644 --- a/__init__.py +++ b/__init__.py @@ -12,10 +12,13 @@ if __package__: NODE_DISPLAY_NAME_MAPPINGS as _POOL_NAMES from .gates.loader import NODE_CLASS_MAPPINGS as _LOADER_NODES, \ NODE_DISPLAY_NAME_MAPPINGS as _LOADER_NAMES + from .gates.gate import NODE_CLASS_MAPPINGS as _GATE_NODES, \ + NODE_DISPLAY_NAME_MAPPINGS as _GATE_NAMES from .gates import routes # noqa: F401 (registers aiohttp routes on import) + from .gates import gate_server # noqa: F401 (registers /datasete_gate/* routes) - NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES} - NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES} + NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES} + NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES} else: # pragma: no cover - exercised only under pytest collection NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/gates/gate_server.py b/gates/gate_server.py new file mode 100644 index 0000000..43b020a --- /dev/null +++ b/gates/gate_server.py @@ -0,0 +1,44 @@ +# gates/gate_server.py +import base64 +import io + +import numpy as np +from aiohttp import web +from PIL import Image +from server import PromptServer + +from .gate_bus import GateBus + +routes = PromptServer.instance.routes + + +def send_preview(node_id, image, n_routes): + arr = (image[0].cpu().numpy() * 255.0).clip(0, 255).astype("uint8") + buf = io.BytesIO() + Image.fromarray(arr).save(buf, "PNG") + b64 = base64.b64encode(buf.getvalue()).decode() + PromptServer.instance.send_sync( + "datasete-gate-show", + {"id": str(node_id), "image": b64, "routes": int(n_routes)}, + ) + + +@routes.post("/datasete_gate/choice") +async def _choice(request): + post = await request.post() + GateBus.put(post.get("id"), post.get("message")) + return web.json_response({}) + + +@routes.post("/datasete_gate/mask") +async def _mask(request): + reader = await request.multipart() + node_id, data = None, None + async for part in reader: + if part.name == "id": + node_id = await part.text() + elif part.name == "mask": + data = await part.read(decode=False) + if node_id is not None: + GateBus.put_mask(node_id, data) + return web.json_response({}) From 63647d24882975414025ec040b8ec1e5b5ff0a3f Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:48:55 +0200 Subject: [PATCH 07/10] =?UTF-8?q?feat:=20image=20gate=20frontend=20?= =?UTF-8?q?=E2=80=94=20preview,=20dynamic=20outputs,=20route/stop/mask?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.8 --- web/image_gate.js | 397 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 397 insertions(+) create mode 100644 web/image_gate.js diff --git a/web/image_gate.js b/web/image_gate.js new file mode 100644 index 0000000..b30ea8c --- /dev/null +++ b/web/image_gate.js @@ -0,0 +1,397 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +// Image Gate (Manual Router) — pauses a running prompt, shows the image with N +// labeled route buttons + an Edit-mask + a Stop button, and routes the image down +// the clicked output (others ExecutionBlocker-ed server-side). The Python node +// blocks in run() on GateBus.wait(); this extension renders the preview that the +// server pushes via the "datasete-gate-show" socket event and POSTs the choice. + +const NODE = "ImageGate"; +const MAX_ROUTES = 10; +const R = "/datasete_gate"; + +const PREVIEW_IMG_H = 240; // fixed image area (object-fit:contain) +const BTN_ROW_H = 64; // buttons area (route buttons wrap + mask/stop) +const MARGIN = 10; // ComfyUI DOM-widget inset, matches the pool node + +// ---- routes widget + label store ------------------------------------------- + +function routesWidget(node) { + return node.widgets?.find((w) => w.name === "routes"); +} + +function getRouteCount(node) { + let n = parseInt(routesWidget(node)?.value ?? 2, 10); + if (isNaN(n)) n = 2; + return Math.max(1, Math.min(MAX_ROUTES, n)); +} + +// Labels live in node.properties (litegraph serializes properties for free, so +// they survive reload without a fake serializing widget — route_labels is not a +// backend input, so we must NOT push it into widgets_values). +function labelStore(node) { + if (!Array.isArray(node.properties.routeLabels)) node.properties.routeLabels = []; + return node.properties.routeLabels; +} + +function labelFor(node, route) { // route is 1-based + const v = labelStore(node)[route - 1]; + return (v != null && String(v).trim()) || String(route); +} + +function setRouteLabel(node, route, text) { + labelStore(node)[route - 1] = text; + applyOutputLabels(node); + if (node._gateActive) renderButtons(node); // live-update visible buttons + node.setDirtyCanvas?.(true, true); +} + +// ---- dynamic route outputs -------------------------------------------------- +// Slot 0 is the always-visible `mask` output; slots 1..N are route_1..route_N. +// We only ever add/remove from the TAIL so existing slot indices (and the +// backend's index→RETURN_TYPES mapping) stay stable and connections are kept. + +function applyOutputLabels(node) { + for (let i = 1; i < node.outputs.length; i++) { + node.outputs[i].label = labelFor(node, i); + } +} + +function applyRouteCount(node, n) { + if (!node.outputs || node.outputs.length === 0) return; + let cur = node.outputs.length - 1; // current route outputs + while (cur < n) { node.addOutput(`route_${cur + 1}`, "IMAGE"); cur++; } + while (cur > n) { node.removeOutput(node.outputs.length - 1); cur--; } + applyOutputLabels(node); + node.setDirtyCanvas?.(true, true); +} + +// ---- server calls ----------------------------------------------------------- + +async function postChoice(node, message) { + const fd = new FormData(); + fd.append("id", String(node.id)); + fd.append("message", String(message)); + await api.fetchApi(`${R}/choice`, { method: "POST", body: fd }); +} + +async function postMask(node, blob) { + const fd = new FormData(); + fd.append("id", String(node.id)); + fd.append("mask", blob, "mask.png"); + await api.fetchApi(`${R}/mask`, { method: "POST", body: fd }); +} + +// ---- preview DOM widget ----------------------------------------------------- + +function previewHeight(node) { + return node._gateActive ? 2 * MARGIN + PREVIEW_IMG_H + BTN_ROW_H : 0; +} + +function resizePreview(node) { + if (node._previewWidget) node._previewWidget.computedHeight = previewHeight(node); + const w = node.size?.[0] || 220; + node.setSize([w, node.computeSize()[1]]); + node.setDirtyCanvas(true, true); +} + +function renderButtons(node) { + const { btns } = node._gate; + btns.innerHTML = ""; + const routes = node._gateRoutes || getRouteCount(node); + for (let i = 1; i <= routes; i++) { + const b = document.createElement("button"); + b.className = "dgate-route"; + b.textContent = labelFor(node, i); + b.onclick = async () => { await postChoice(node, i); hidePreview(node); }; + btns.appendChild(b); + } + const edit = document.createElement("button"); + edit.className = "dgate-edit"; + edit.textContent = "🖌 Edit mask"; + edit.onclick = () => openMaskEditor(node); + btns.appendChild(edit); + + const stop = document.createElement("button"); + stop.className = "dgate-stop"; + stop.textContent = "■ Stop"; + stop.onclick = async () => { await postChoice(node, "__cancel__"); hidePreview(node); }; + btns.appendChild(stop); +} + +function showPreview(node, b64, routes) { + node._gateActive = true; + node._gateRoutes = Math.max(1, Math.min(MAX_ROUTES, parseInt(routes, 10) || getRouteCount(node))); + node._previewB64 = b64; + node._gate.img.src = `data:image/png;base64,${b64}`; + renderButtons(node); + resizePreview(node); +} + +function hidePreview(node) { + node._gateActive = false; + node._previewB64 = null; + if (node._gate) { + node._gate.img.removeAttribute("src"); + node._gate.btns.innerHTML = ""; + } + resizePreview(node); +} + +// ---- mask editor (reuses ComfyUI MaskEditor, like the pool node) ------------ +// The preview arrives as base64 (no server file), so upload it to input/ first, +// point the MaskEditor at it, then poll node.images for the saved clipspace ref. + +function b64ToBlob(b64, type) { + const bin = atob(b64); + const arr = new Uint8Array(bin.length); + for (let i = 0; i < bin.length; i++) arr[i] = bin.charCodeAt(i); + return new Blob([arr], { type }); +} + +function blobToImage(blob) { + return new Promise((resolve, reject) => { + const img = new Image(); + img.onload = () => resolve(img); + img.onerror = reject; + img.src = URL.createObjectURL(blob); + }); +} + +function comfyAppClass() { + try { return app.constructor; } catch (e) { return null; } +} + +// MaskEditor registers the painted image as this node's output; clear those +// stores so nothing repopulates node.imgs (we draw our own preview). +function clearNodeOutputs(node) { + try { + for (const map of [app.nodeOutputs, app.nodePreviewImages]) { + if (!map) continue; + for (const k of Object.keys(map)) { + if (k === String(node.id) || k.endsWith(`:${node.id}`)) delete map[k]; + } + } + } catch (e) { /* best effort */ } +} + +function cleanupMaskState(node) { + if (node._maskPoll) { clearInterval(node._maskPoll); node._maskPoll = null; } + node._maskActive = false; + try { + node.images = undefined; + node.previewMediaType = undefined; + } catch (e) { /* best effort */ } + clearNodeOutputs(node); + node.setDirtyCanvas?.(true, true); +} + +async function uploadPreview(node) { + const blob = b64ToBlob(node._previewB64, "image/png"); + const fd = new FormData(); + fd.append("image", blob, `gate_${node.id}.png`); + fd.append("subfolder", "datasete_gate"); + fd.append("type", "input"); + fd.append("overwrite", "true"); + const res = await api.fetchApi("/upload/image", { method: "POST", body: fd }); + const j = await res.json(); + return { filename: j.name, subfolder: j.subfolder || "datasete_gate", type: j.type || "input" }; +} + +async function captureMask(node, ref) { + try { + const sub = ref.subfolder ?? "clipspace"; + const type = ref.type ?? "input"; + const url = `/view?filename=${encodeURIComponent(ref.filename)}&subfolder=${encodeURIComponent(sub)}&type=${encodeURIComponent(type)}&r=${Date.now()}`; + const resp = await api.fetchApi(url); + const blob = await resp.blob(); + const img = await blobToImage(blob); + const c = document.createElement("canvas"); + c.width = img.naturalWidth || img.width; + c.height = img.naturalHeight || img.height; + const ctx = c.getContext("2d"); + ctx.drawImage(img, 0, 0); + const d = ctx.getImageData(0, 0, c.width, c.height); + const px = d.data; + // MaskEditor stores the mask in the ALPHA channel; painted areas come through + // as alpha 0, so invert (255 - a) into grayscale -> white = painted (MASK). + for (let i = 0; i < px.length; i += 4) { + const a = px[i + 3]; + px[i] = px[i + 1] = px[i + 2] = 255 - a; + px[i + 3] = 255; + } + ctx.putImageData(d, 0, 0); + const maskBlob = await new Promise((res) => c.toBlob(res, "image/png")); + await postMask(node, maskBlob); + } catch (e) { + console.error("[dgate] mask capture failed", e); + } finally { + cleanupMaskState(node); + } +} + +async function openMaskEditor(node) { + if (!node._previewB64) return; + cleanupMaskState(node); + let ref; + try { + ref = await uploadPreview(node); + } catch (e) { + console.error("[dgate] preview upload failed", e); + return; + } + + node.images = [ref]; + node.previewMediaType = "image"; + node.imageIndex = 0; + node._maskActive = true; + + const Comfy = comfyAppClass(); + try { if (Comfy) Comfy.clipspace_return_node = node; } catch (e) { /* ignore */ } + + // No save callback in frontend 1.45 — poll for the editor writing clipspace. + let waited = 0; + node._maskPoll = setInterval(() => { + waited += 300; + const r = node.images && node.images[0]; + if (node._maskActive && r && r.subfolder === "clipspace") { + clearInterval(node._maskPoll); node._maskPoll = null; + captureMask(node, r); + } else if (waited > 10 * 60 * 1000) { + cleanupMaskState(node); + } + }, 300); + + try { app.canvas?.selectNode?.(node); } catch (e) { /* ignore */ } + const cmd = app.extensionManager?.command; + if (cmd?.execute) { + cmd.execute("Comfy.MaskEditor.OpenMaskEditor"); + } else if (Comfy?.open_maskeditor) { + Comfy.open_maskeditor(); + } else { + console.error("[dgate] no MaskEditor entry point found"); + cleanupMaskState(node); + } +} + +// ---- styles + node setup ---------------------------------------------------- + +function injectStyles() { + if (document.getElementById("dgate-styles")) return; + const css = ` + .dgate-wrap { display:flex; flex-direction:column; gap:6px; box-sizing:border-box; + height:100%; min-height:0; } + .dgate-img { width:100%; height:${PREVIEW_IMG_H}px; object-fit:contain; display:block; + background:rgba(0,0,0,0.25); border-radius:4px; } + .dgate-btns { display:flex; flex-wrap:wrap; gap:6px; align-items:center; } + .dgate-btns button { font-size:12px; padding:3px 10px; cursor:pointer; border-radius:3px; + border:1px solid #555; color:#fff; } + .dgate-route { background:rgba(40,90,140,0.9); } + .dgate-route:hover { background:rgba(60,120,180,0.95); } + .dgate-edit { background:rgba(40,40,40,0.9); margin-left:auto; } + .dgate-stop { background:rgba(160,40,40,0.9); } + `; + const style = document.createElement("style"); + style.id = "dgate-styles"; + style.textContent = css; + document.head.appendChild(style); +} + +function setupGateNode(node) { + injectStyles(); + + // Never let the MaskEditor's source image render as an output preview on us — + // we draw the preview ourselves in the DOM widget below. + try { + Object.defineProperty(node, "imgs", { + configurable: true, + get() { return undefined; }, + set() { /* suppress */ }, + }); + } catch (e) { /* ignore */ } + + const wrap = document.createElement("div"); + wrap.className = "dgate-wrap"; + const img = document.createElement("img"); + img.className = "dgate-img"; + const btns = document.createElement("div"); + btns.className = "dgate-btns"; + wrap.appendChild(img); + wrap.appendChild(btns); + node._gate = { wrap, img, btns }; + + node._previewWidget = node.addDOMWidget("gate_preview", "div", wrap, { + serialize: false, + getMinHeight: () => previewHeight(node), + }); + + // sync visible route outputs to the routes widget, now and on change + applyRouteCount(node, getRouteCount(node)); + const rw = routesWidget(node); + if (rw) { + const prev = rw.callback; + rw.callback = function () { + const r = prev?.apply(this, arguments); + applyRouteCount(node, getRouteCount(node)); + return r; + }; + } + + node._gateActive = false; + resizePreview(node); +} + +app.registerExtension({ + name: "datasete.gates.imagegate", + + // one global socket listener: route the server's pause event to the node + setup() { + api.addEventListener("datasete-gate-show", (e) => { + const d = e.detail || {}; + const node = app.graph?.getNodeById?.(parseInt(d.id, 10)); + if (!node || node.type !== NODE) return; + showPreview(node, d.image, d.routes); + }); + }, + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE) return; + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated?.apply(this, arguments); + setupGateNode(this); + return r; + }; + + // loaded workflows restore the routes widget + properties after create — + // re-sync output count/labels to match. + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + const r = onConfigure?.apply(this, arguments); + if (this.outputs) { + applyRouteCount(this, getRouteCount(this)); + } + return r; + }; + + // per-route "Rename…" entries (editable labels, persisted in properties) + const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (canvas, options) { + const r = getExtraMenuOptions?.apply(this, arguments); + const node = this; + const routes = getRouteCount(node); + for (let i = 1; i <= routes; i++) { + options.push({ + content: `Rename route ${i} (“${labelFor(node, i)}”)…`, + callback: () => { + const text = prompt(`Label for route ${i}:`, labelFor(node, i)); + if (text != null) setRouteLabel(node, i, text); + }, + }); + } + return r; + }; + }, +}); From 45e16e113412ff19500810ea88e028c1212ad434 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 17:59:04 +0200 Subject: [PATCH 08/10] fix: hide gate preview element when idle (no stray black box) Co-Authored-By: Claude Opus 4.8 --- web/image_gate.js | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/web/image_gate.js b/web/image_gate.js index b30ea8c..69cf85d 100644 --- a/web/image_gate.js +++ b/web/image_gate.js @@ -90,7 +90,10 @@ function previewHeight(node) { } function resizePreview(node) { - if (node._previewWidget) node._previewWidget.computedHeight = previewHeight(node); + // Fully remove the preview element from layout when idle — collapsing the + // widget height to 0 isn't enough: the fixed-height would still paint as + // a black box hanging below the node frame. + if (node._gate) node._gate.wrap.style.display = node._gateActive ? "flex" : "none"; const w = node.size?.[0] || 220; node.setSize([w, node.computeSize()[1]]); node.setDirtyCanvas(true, true); From f9f924942e6ade260fd7012563bfa4fb3c9782af Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 18:17:37 +0200 Subject: [PATCH 09/10] feat: sticky mask + keep preview after routing + width-scaled preview After a route choice the node now keeps the image and shows a 'Run from here' re-queue button instead of blanking. The last painted mask is remembered and auto-re-stashed on each new pause (with a Clear control) so it is not lost between runs. The preview image area now scales with the node width. Co-Authored-By: Claude Opus 4.8 --- web/image_gate.js | 180 +++++++++++++++++++++++++++++++++++----------- 1 file changed, 138 insertions(+), 42 deletions(-) diff --git a/web/image_gate.js b/web/image_gate.js index 69cf85d..12967d5 100644 --- a/web/image_gate.js +++ b/web/image_gate.js @@ -11,8 +11,9 @@ const NODE = "ImageGate"; const MAX_ROUTES = 10; const R = "/datasete_gate"; -const PREVIEW_IMG_H = 240; // fixed image area (object-fit:contain) -const BTN_ROW_H = 64; // buttons area (route buttons wrap + mask/stop) +const MIN_IMG_H = 140; // preview image area clamps (scales with node width) +const MAX_IMG_H = 600; +const BTN_ROW_H = 78; // buttons area (route buttons wrap + actions) const MARGIN = 10; // ComfyUI DOM-widget inset, matches the pool node // ---- routes widget + label store ------------------------------------------- @@ -43,7 +44,7 @@ function labelFor(node, route) { // route is 1-based function setRouteLabel(node, route, text) { labelStore(node)[route - 1] = text; applyOutputLabels(node); - if (node._gateActive) renderButtons(node); // live-update visible buttons + if (node._gateState && node._gateState !== "idle") render(node); // live-update node.setDirtyCanvas?.(true, true); } @@ -83,65 +84,138 @@ async function postMask(node, blob) { await api.fetchApi(`${R}/mask`, { method: "POST", body: fd }); } -// ---- preview DOM widget ----------------------------------------------------- +// ---- preview DOM widget + state machine ------------------------------------- +// States: "idle" (collapsed, before the first run), "paused" (waiting for a +// route choice — route buttons shown), "resolved" (a route was picked — image + +// mask kept, a "Run from here" re-queue button shown). The node never blanks +// once a run has happened, so the previewed image and the sticky mask stay for +// context and the painted mask is reused on the next run until cleared. + +function computeImgH(node) { + // image area scales with node WIDTH and the image's aspect ratio, so a wider + // node shows a bigger preview (getMinHeight is polled each layout frame). + const w = Math.max(120, (node.size?.[0] || 220) - 2 * MARGIN); + const h = Math.round(w * (node._imgAspect || 1)); + return Math.max(MIN_IMG_H, Math.min(h, MAX_IMG_H)); +} function previewHeight(node) { - return node._gateActive ? 2 * MARGIN + PREVIEW_IMG_H + BTN_ROW_H : 0; + if (!node._gateState || node._gateState === "idle") return 0; + return 2 * MARGIN + computeImgH(node) + BTN_ROW_H; } function resizePreview(node) { // Fully remove the preview element from layout when idle — collapsing the - // widget height to 0 isn't enough: the fixed-height would still paint as - // a black box hanging below the node frame. - if (node._gate) node._gate.wrap.style.display = node._gateActive ? "flex" : "none"; + // widget height to 0 isn't enough: the would still paint below the node. + const shown = node._gateState && node._gateState !== "idle"; + if (node._gate) node._gate.wrap.style.display = shown ? "flex" : "none"; const w = node.size?.[0] || 220; node.setSize([w, node.computeSize()[1]]); node.setDirtyCanvas(true, true); } -function renderButtons(node) { - const { btns } = node._gate; - btns.innerHTML = ""; - const routes = node._gateRoutes || getRouteCount(node); - for (let i = 1; i <= routes; i++) { - const b = document.createElement("button"); - b.className = "dgate-route"; - b.textContent = labelFor(node, i); - b.onclick = async () => { await postChoice(node, i); hidePreview(node); }; - btns.appendChild(b); - } +function hasMask(node) { return !!node._stickyMask; } + +function maskControls(node) { + // Edit / Clear buttons + a small "mask retained" badge, shared by both states. + const els = []; const edit = document.createElement("button"); edit.className = "dgate-edit"; edit.textContent = "🖌 Edit mask"; edit.onclick = () => openMaskEditor(node); - btns.appendChild(edit); - - const stop = document.createElement("button"); - stop.className = "dgate-stop"; - stop.textContent = "■ Stop"; - stop.onclick = async () => { await postChoice(node, "__cancel__"); hidePreview(node); }; - btns.appendChild(stop); + els.push(edit); + if (hasMask(node)) { + const clr = document.createElement("button"); + clr.className = "dgate-clear"; + clr.textContent = "✕ Clear mask"; + clr.onclick = () => clearMask(node); + els.push(clr); + } + const badge = document.createElement("span"); + badge.className = "dgate-status"; + badge.textContent = hasMask(node) ? "🎭 mask retained" : "no mask"; + badge.style.opacity = hasMask(node) ? "0.9" : "0.45"; + els.push(badge); + return els; } -function showPreview(node, b64, routes) { - node._gateActive = true; +function render(node) { + const { btns } = node._gate; + btns.innerHTML = ""; + const routes = node._gateRoutes || getRouteCount(node); + + if (node._gateState === "paused") { + for (let i = 1; i <= routes; i++) { + const b = document.createElement("button"); + b.className = "dgate-route"; + b.textContent = labelFor(node, i); + b.onclick = async () => { + await postChoice(node, i); + showResolved(node, labelFor(node, i)); + }; + btns.appendChild(b); + } + maskControls(node).forEach((el) => btns.appendChild(el)); + const stop = document.createElement("button"); + stop.className = "dgate-stop"; + stop.textContent = "■ Stop"; + stop.onclick = async () => { + await postChoice(node, "__cancel__"); + showResolved(node, "stopped"); + }; + btns.appendChild(stop); + } else if (node._gateState === "resolved") { + const status = document.createElement("span"); + status.className = "dgate-status"; + status.textContent = `✓ routed to ${node._gateChoice ?? "?"}`; + btns.appendChild(status); + const run = document.createElement("button"); + run.className = "dgate-run"; + run.textContent = "▶ Run from here"; + run.onclick = () => queueFromHere(node); + btns.appendChild(run); + maskControls(node).forEach((el) => btns.appendChild(el)); + } +} + +function showPaused(node, b64, routes) { + node._gateState = "paused"; node._gateRoutes = Math.max(1, Math.min(MAX_ROUTES, parseInt(routes, 10) || getRouteCount(node))); node._previewB64 = b64; node._gate.img.src = `data:image/png;base64,${b64}`; - renderButtons(node); + // sticky mask: re-stash the last painted mask for THIS run before the user + // picks a route. run() does arm()→clear, then send_preview→this event, then + // blocks in wait(), so this POST always lands before the choice is made. + if (node._stickyMask) { + postMask(node, b64ToBlob(node._stickyMask, "image/png")).catch(() => {}); + } + render(node); resizePreview(node); } -function hidePreview(node) { - node._gateActive = false; - node._previewB64 = null; - if (node._gate) { - node._gate.img.removeAttribute("src"); - node._gate.btns.innerHTML = ""; - } +function showResolved(node, choiceLabel) { + node._gateState = "resolved"; + node._gateChoice = choiceLabel; + render(node); resizePreview(node); } +async function queueFromHere(node) { + try { + await app.queuePrompt(0, 1); + } catch (e) { + try { await app.queuePrompt(0); } catch (e2) { console.error("[dgate] queue failed", e2); } + } +} + +async function clearMask(node) { + node._stickyMask = null; + // zero the current run's stash: an empty mask part -> server stores b"" -> + // mask_from_stash() treats it as falsy -> zeros. + try { await postMask(node, new Blob([], { type: "image/png" })); } catch (e) { /* ignore */ } + render(node); +} + // ---- mask editor (reuses ComfyUI MaskEditor, like the pool node) ------------ // The preview arrives as base64 (no server file), so upload it to input/ first, // point the MaskEditor at it, then poll node.images for the saved clipspace ref. @@ -162,6 +236,15 @@ function blobToImage(blob) { }); } +function blobToB64(blob) { + return new Promise((resolve, reject) => { + const fr = new FileReader(); + fr.onload = () => resolve(String(fr.result).split(",")[1] || ""); + fr.onerror = reject; + fr.readAsDataURL(blob); + }); +} + function comfyAppClass() { try { return app.constructor; } catch (e) { return null; } } @@ -227,10 +310,13 @@ async function captureMask(node, ref) { ctx.putImageData(d, 0, 0); const maskBlob = await new Promise((res) => c.toBlob(res, "image/png")); await postMask(node, maskBlob); + // remember it so it auto-applies on the next run until the user clears it + try { node._stickyMask = await blobToB64(maskBlob); } catch (e) { /* ignore */ } } catch (e) { console.error("[dgate] mask capture failed", e); } finally { cleanupMaskState(node); + if (node._gateState && node._gateState !== "idle") render(node); // show badge } } @@ -285,15 +371,18 @@ function injectStyles() { const css = ` .dgate-wrap { display:flex; flex-direction:column; gap:6px; box-sizing:border-box; height:100%; min-height:0; } - .dgate-img { width:100%; height:${PREVIEW_IMG_H}px; object-fit:contain; display:block; + .dgate-img { width:100%; flex:1 1 auto; min-height:0; object-fit:contain; display:block; background:rgba(0,0,0,0.25); border-radius:4px; } - .dgate-btns { display:flex; flex-wrap:wrap; gap:6px; align-items:center; } + .dgate-btns { display:flex; flex-wrap:wrap; gap:6px; align-items:center; flex:0 0 auto; } .dgate-btns button { font-size:12px; padding:3px 10px; cursor:pointer; border-radius:3px; border:1px solid #555; color:#fff; } .dgate-route { background:rgba(40,90,140,0.9); } .dgate-route:hover { background:rgba(60,120,180,0.95); } - .dgate-edit { background:rgba(40,40,40,0.9); margin-left:auto; } - .dgate-stop { background:rgba(160,40,40,0.9); } + .dgate-edit { background:rgba(40,40,40,0.9); } + .dgate-clear { background:rgba(90,60,30,0.9); } + .dgate-run { background:rgba(40,130,70,0.95); } + .dgate-stop { background:rgba(160,40,40,0.9); margin-left:auto; } + .dgate-status { font-size:11px; opacity:0.8; padding:0 4px; align-self:center; } `; const style = document.createElement("style"); style.id = "dgate-styles"; @@ -318,6 +407,13 @@ function setupGateNode(node) { wrap.className = "dgate-wrap"; const img = document.createElement("img"); img.className = "dgate-img"; + // capture the image aspect so the preview area scales with the node width + img.onload = () => { + const w = img.naturalWidth || 1; + const h = img.naturalHeight || 1; + node._imgAspect = h / w; + resizePreview(node); + }; const btns = document.createElement("div"); btns.className = "dgate-btns"; wrap.appendChild(img); @@ -341,7 +437,7 @@ function setupGateNode(node) { }; } - node._gateActive = false; + node._gateState = "idle"; resizePreview(node); } @@ -354,7 +450,7 @@ app.registerExtension({ const d = e.detail || {}; const node = app.graph?.getNodeById?.(parseInt(d.id, 10)); if (!node || node.type !== NODE) return; - showPreview(node, d.image, d.routes); + showPaused(node, d.image, d.routes); }); }, From 6e27da0dce9ad37c684ebbeb97a8979d76df5846 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 21 Jun 2026 18:26:33 +0200 Subject: [PATCH 10/10] feat: show painted mask as a translucent red overlay on the gate preview Co-Authored-By: Claude Opus 4.8 --- web/image_gate.js | 81 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 75 insertions(+), 6 deletions(-) diff --git a/web/image_gate.js b/web/image_gate.js index 12967d5..6efc345 100644 --- a/web/image_gate.js +++ b/web/image_gate.js @@ -176,6 +176,7 @@ function render(node) { btns.appendChild(run); maskControls(node).forEach((el) => btns.appendChild(el)); } + updateMaskOverlay(node); } function showPaused(node, b64, routes) { @@ -210,12 +211,64 @@ async function queueFromHere(node) { async function clearMask(node) { node._stickyMask = null; + node._stickyMaskOverlay = null; // zero the current run's stash: an empty mask part -> server stores b"" -> // mask_from_stash() treats it as falsy -> zeros. try { await postMask(node, new Blob([], { type: "image/png" })); } catch (e) { /* ignore */ } render(node); } +// ---- mask overlay (show the painted region over the preview, semi-transparent) +// The sticky mask is grayscale (white = painted). Recolor it into an RGBA layer +// where alpha = paint intensity and RGB = a highlight color, so unpainted areas +// are fully transparent and only the painted region tints the image. + +function maskToOverlay(b64) { + return new Promise((resolve, reject) => { + const im = new Image(); + im.onload = () => { + const c = document.createElement("canvas"); + c.width = im.naturalWidth || im.width; + c.height = im.naturalHeight || im.height; + const ctx = c.getContext("2d"); + ctx.drawImage(im, 0, 0); + const d = ctx.getImageData(0, 0, c.width, c.height); + const px = d.data; + for (let i = 0; i < px.length; i += 4) { + const v = px[i]; // grayscale luminance (R=G=B) + px[i] = 255; px[i + 1] = 64; px[i + 2] = 64; // highlight = red + px[i + 3] = v; // alpha = paint intensity + } + ctx.putImageData(d, 0, 0); + resolve(c.toDataURL("image/png")); + }; + im.onerror = reject; + im.src = `data:image/png;base64,${b64}`; + }); +} + +async function setStickyMask(node, b64) { + node._stickyMask = b64; + try { + node._stickyMaskOverlay = b64 ? await maskToOverlay(b64) : null; + } catch (e) { + node._stickyMaskOverlay = null; + } + updateMaskOverlay(node); +} + +function updateMaskOverlay(node) { + const mi = node._gate?.maskImg; + if (!mi) return; + if (node._gateState && node._gateState !== "idle" && node._stickyMaskOverlay) { + mi.src = node._stickyMaskOverlay; + mi.style.display = "block"; + } else { + mi.removeAttribute("src"); + mi.style.display = "none"; + } +} + // ---- mask editor (reuses ComfyUI MaskEditor, like the pool node) ------------ // The preview arrives as base64 (no server file), so upload it to input/ first, // point the MaskEditor at it, then poll node.images for the saved clipspace ref. @@ -310,8 +363,9 @@ async function captureMask(node, ref) { ctx.putImageData(d, 0, 0); const maskBlob = await new Promise((res) => c.toBlob(res, "image/png")); await postMask(node, maskBlob); - // remember it so it auto-applies on the next run until the user clears it - try { node._stickyMask = await blobToB64(maskBlob); } catch (e) { /* ignore */ } + // remember it so it auto-applies on the next run until the user clears it, + // and build the colored overlay shown over the preview. + try { await setStickyMask(node, await blobToB64(maskBlob)); } catch (e) { /* ignore */ } } catch (e) { console.error("[dgate] mask capture failed", e); } finally { @@ -371,8 +425,12 @@ function injectStyles() { const css = ` .dgate-wrap { display:flex; flex-direction:column; gap:6px; box-sizing:border-box; height:100%; min-height:0; } - .dgate-img { width:100%; flex:1 1 auto; min-height:0; object-fit:contain; display:block; - background:rgba(0,0,0,0.25); border-radius:4px; } + .dgate-imgbox { position:relative; flex:1 1 auto; min-height:0; width:100%; + background:rgba(0,0,0,0.25); border-radius:4px; overflow:hidden; } + .dgate-img { position:absolute; inset:0; width:100%; height:100%; object-fit:contain; + display:block; } + .dgate-mask { position:absolute; inset:0; width:100%; height:100%; object-fit:contain; + opacity:0.5; pointer-events:none; } .dgate-btns { display:flex; flex-wrap:wrap; gap:6px; align-items:center; flex:0 0 auto; } .dgate-btns button { font-size:12px; padding:3px 10px; cursor:pointer; border-radius:3px; border:1px solid #555; color:#fff; } @@ -405,6 +463,11 @@ function setupGateNode(node) { const wrap = document.createElement("div"); wrap.className = "dgate-wrap"; + + // image + mask overlay share a container so both letterbox identically and + // stay pixel-aligned (object-fit:contain on same-size, same-aspect layers). + const imgbox = document.createElement("div"); + imgbox.className = "dgate-imgbox"; const img = document.createElement("img"); img.className = "dgate-img"; // capture the image aspect so the preview area scales with the node width @@ -414,11 +477,17 @@ function setupGateNode(node) { node._imgAspect = h / w; resizePreview(node); }; + const maskImg = document.createElement("img"); + maskImg.className = "dgate-mask"; + maskImg.style.display = "none"; + imgbox.appendChild(img); + imgbox.appendChild(maskImg); + const btns = document.createElement("div"); btns.className = "dgate-btns"; - wrap.appendChild(img); + wrap.appendChild(imgbox); wrap.appendChild(btns); - node._gate = { wrap, img, btns }; + node._gate = { wrap, imgbox, img, maskImg, btns }; node._previewWidget = node.addDOMWidget("gate_preview", "div", wrap, { serialize: false,