Merge feat/text-gate: Text Gate (Manual Pass) node
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+5
-3
@@ -14,11 +14,13 @@ if __package__:
|
||||
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.textgate import NODE_CLASS_MAPPINGS as _TEXT_NODES, \
|
||||
NODE_DISPLAY_NAME_MAPPINGS as _TEXT_NAMES
|
||||
from .gates import routes # noqa: F401 (registers aiohttp routes on import)
|
||||
from .gates import gate_server # noqa: F401 (registers /datasete_gate/* routes)
|
||||
from .gates import gate_server # noqa: F401 (registers /datasete_gate/* + text routes)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES}
|
||||
NODE_CLASS_MAPPINGS = {**_POOL_NODES, **_LOADER_NODES, **_GATE_NODES, **_TEXT_NODES}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**_POOL_NAMES, **_LOADER_NAMES, **_GATE_NAMES, **_TEXT_NAMES}
|
||||
else: # pragma: no cover - exercised only under pytest collection
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
@@ -9,12 +9,14 @@ class GateCancelled(Exception):
|
||||
class GateBus:
|
||||
messages = {} # node_id(str) -> chosen int (1-based)
|
||||
masks = {} # node_id(str) -> PNG bytes
|
||||
payloads = {} # node_id(str) -> arbitrary payload (e.g., edited text)
|
||||
cancelled = False
|
||||
|
||||
@classmethod
|
||||
def arm(cls, node_id):
|
||||
cls.messages.pop(str(node_id), None)
|
||||
cls.masks.pop(str(node_id), None)
|
||||
cls.payloads.pop(str(node_id), None)
|
||||
cls.cancelled = False
|
||||
|
||||
@classmethod
|
||||
@@ -41,3 +43,17 @@ class GateBus:
|
||||
@classmethod
|
||||
def pop_mask(cls, node_id):
|
||||
return cls.masks.pop(str(node_id), None)
|
||||
|
||||
@classmethod
|
||||
def put_payload(cls, node_id, value):
|
||||
cls.payloads[str(node_id)] = value
|
||||
|
||||
@classmethod
|
||||
def wait_payload(cls, node_id, period=0.1, should_cancel=None):
|
||||
sid = str(node_id)
|
||||
while sid not in cls.payloads:
|
||||
if cls.cancelled or (should_cancel is not None and should_cancel()):
|
||||
cls.cancelled = False
|
||||
raise GateCancelled()
|
||||
time.sleep(period)
|
||||
return cls.payloads.pop(sid)
|
||||
|
||||
@@ -42,3 +42,16 @@ async def _mask(request):
|
||||
if node_id is not None:
|
||||
GateBus.put_mask(node_id, data)
|
||||
return web.json_response({})
|
||||
|
||||
|
||||
def send_text(node_id, text):
|
||||
PromptServer.instance.send_sync(
|
||||
"datasete-textgate-show", {"id": str(node_id), "text": text or ""}
|
||||
)
|
||||
|
||||
|
||||
@routes.post("/datasete_text_gate/pass")
|
||||
async def _text_pass(request):
|
||||
post = await request.post()
|
||||
GateBus.put_payload(post.get("id"), post.get("text", ""))
|
||||
return web.json_response({})
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# gates/textgate.py
|
||||
from . import gate_bus
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
|
||||
class AnyType(str):
|
||||
"""Type that compares equal to any other type (ComfyUI wildcard convention)."""
|
||||
def __ne__(self, other):
|
||||
return False
|
||||
|
||||
|
||||
ANY = AnyType("*")
|
||||
|
||||
|
||||
class TextGate:
|
||||
CATEGORY = "Datasete Gates"
|
||||
FUNCTION = "run"
|
||||
RETURN_TYPES = ("STRING", ANY)
|
||||
RETURN_NAMES = ("text", "signal")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"forceInput": True}),
|
||||
},
|
||||
"optional": {
|
||||
"signal": (ANY, {}),
|
||||
},
|
||||
"hidden": {"unique_id": "UNIQUE_ID"},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
|
||||
def run(self, text, unique_id, signal=None):
|
||||
from . import gate_server
|
||||
import comfy.model_management as mm
|
||||
|
||||
gate_bus.GateBus.arm(unique_id)
|
||||
gate_server.send_text(unique_id, text)
|
||||
try:
|
||||
edited = gate_bus.GateBus.wait_payload(
|
||||
unique_id, should_cancel=mm.processing_interrupted)
|
||||
except gate_bus.GateCancelled:
|
||||
raise mm.InterruptProcessingException()
|
||||
return (edited, signal)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"TextGate": TextGate}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"TextGate": "Text Gate (Manual Pass)"}
|
||||
@@ -36,3 +36,32 @@ def test_arm_clears_mask():
|
||||
gb.GateBus.put_mask("9", b"x")
|
||||
gb.GateBus.arm("9")
|
||||
assert gb.GateBus.pop_mask("9") is None
|
||||
|
||||
def test_payload_roundtrip():
|
||||
gb.GateBus.arm("p")
|
||||
gb.GateBus.put_payload("p", "hello edited")
|
||||
assert gb.GateBus.wait_payload("p") == "hello edited"
|
||||
|
||||
def test_payload_consumed():
|
||||
gb.GateBus.arm("p")
|
||||
gb.GateBus.put_payload("p", "x")
|
||||
gb.GateBus.wait_payload("p")
|
||||
assert "p" not in gb.GateBus.payloads
|
||||
|
||||
def test_arm_clears_payload():
|
||||
gb.GateBus.put_payload("p", "stale")
|
||||
gb.GateBus.arm("p")
|
||||
assert "p" not in gb.GateBus.payloads
|
||||
|
||||
def test_wait_payload_cancel_flag_raises():
|
||||
import pytest
|
||||
gb.GateBus.arm("p")
|
||||
gb.GateBus.cancelled = True
|
||||
with pytest.raises(gb.GateCancelled):
|
||||
gb.GateBus.wait_payload("p")
|
||||
|
||||
def test_wait_payload_should_cancel_raises():
|
||||
import pytest
|
||||
gb.GateBus.arm("p")
|
||||
with pytest.raises(gb.GateCancelled):
|
||||
gb.GateBus.wait_payload("p", should_cancel=lambda: True)
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
# tests/test_textgate.py
|
||||
import math
|
||||
|
||||
from gates import textgate
|
||||
|
||||
def test_anytype_is_compatible_with_everything():
|
||||
assert (textgate.ANY != "IMAGE") is False
|
||||
assert (textgate.ANY != "LATENT") is False
|
||||
assert isinstance(textgate.ANY, str)
|
||||
|
||||
def test_textgate_io_shape():
|
||||
assert textgate.TextGate.RETURN_NAMES == ("text", "signal")
|
||||
assert textgate.TextGate.RETURN_TYPES[0] == "STRING"
|
||||
assert textgate.TextGate.RETURN_TYPES[1] == textgate.ANY
|
||||
|
||||
def test_textgate_is_changed_nan():
|
||||
v = textgate.TextGate.IS_CHANGED(text="hi", unique_id="1")
|
||||
assert math.isnan(v)
|
||||
+22
-1
@@ -104,13 +104,26 @@ function previewHeight(node) {
|
||||
return 2 * MARGIN + computeImgH(node) + BTN_ROW_H;
|
||||
}
|
||||
|
||||
// DomWidgets sizes the preview container from the widget width, which can lag
|
||||
// node.size[0] on this frontend — pin it so the image/buttons reflow to fill.
|
||||
function syncWidgetWidth(node) {
|
||||
if (node._previewWidget) node._previewWidget.width = node.size?.[0] || 220;
|
||||
}
|
||||
|
||||
function resizePreview(node) {
|
||||
// Fully remove the preview element from layout when idle — collapsing the
|
||||
// widget height to 0 isn't enough: the <img> 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]]);
|
||||
// Image Pool pattern: grow to fit the content floor but preserve a larger
|
||||
// user-set size (so the node stays freely resizable); collapse exactly when
|
||||
// idle. Forcing the height on every call would lock the node.
|
||||
const target = shown
|
||||
? Math.max(node.size?.[1] || 0, node.computeSize()[1])
|
||||
: node.computeSize()[1];
|
||||
node.setSize([w, target]);
|
||||
syncWidgetWidth(node);
|
||||
node.setDirtyCanvas(true, true);
|
||||
}
|
||||
|
||||
@@ -494,6 +507,14 @@ function setupGateNode(node) {
|
||||
getMinHeight: () => previewHeight(node),
|
||||
});
|
||||
|
||||
// keep the preview width synced on manual resize so the image/buttons reflow
|
||||
const onResize = node.onResize;
|
||||
node.onResize = function () {
|
||||
const r = onResize?.apply(this, arguments);
|
||||
syncWidgetWidth(node);
|
||||
return r;
|
||||
};
|
||||
|
||||
// sync visible route outputs to the routes widget, now and on change
|
||||
applyRouteCount(node, getRouteCount(node));
|
||||
const rw = routesWidget(node);
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { api } from "../../scripts/api.js";
|
||||
|
||||
// Text Gate (Manual Pass) — pauses a running prompt, shows the incoming text in
|
||||
// an editable textarea, and on Pass emits the (edited) text. The Python node
|
||||
// blocks in run() on GateBus.wait_payload(); this extension renders the editor
|
||||
// the server pushes via the "datasete-textgate-show" socket event and POSTs the
|
||||
// edited text back. Outputs are static (text, signal) — no dynamic slots.
|
||||
//
|
||||
// Sizing follows the Image Pool node: the editor is always present and FILLS the
|
||||
// node, with only a min-height floor (no max) so the node stays freely resizable
|
||||
// and the textarea grows with it.
|
||||
|
||||
const NODE = "TextGate";
|
||||
const R = "/datasete_text_gate";
|
||||
|
||||
const MIN_W = 320; // default node width (freely resizable)
|
||||
const MIN_EDITOR_H = 140; // textarea floor
|
||||
const BTN_ROW_H = 34; // Pass button row
|
||||
const MARGIN = 10; // ComfyUI DOM-widget inset, matches the other nodes
|
||||
|
||||
// ---- server call ------------------------------------------------------------
|
||||
|
||||
async function postPass(node, text) {
|
||||
const fd = new FormData();
|
||||
fd.append("id", String(node.id));
|
||||
fd.append("text", text);
|
||||
await api.fetchApi(`${R}/pass`, { method: "POST", body: fd });
|
||||
}
|
||||
|
||||
// ---- sizing (Image Pool pattern) --------------------------------------------
|
||||
|
||||
// Only a min-height FLOOR — no max — so the DOM widget fills the node and grows
|
||||
// when the user resizes it. (A fixed height, or forcing node height on every
|
||||
// interaction, would lock the node and leave dead grey space below the editor.)
|
||||
function widgetFloor() {
|
||||
return 2 * MARGIN + MIN_EDITOR_H + BTN_ROW_H;
|
||||
}
|
||||
|
||||
// DomWidgets sizes the editor container from the widget width, which can lag
|
||||
// node.size[0] on this frontend — pin it so the textarea reflows to fill.
|
||||
function syncWidgetWidth(node) {
|
||||
if (node._tgWidget) node._tgWidget.width = node.size?.[0] || MIN_W;
|
||||
}
|
||||
|
||||
// ---- styles + node setup ----------------------------------------------------
|
||||
|
||||
function injectStyles() {
|
||||
if (document.getElementById("tgate-styles")) return;
|
||||
const css = `
|
||||
.tgate-wrap { display:flex; flex-direction:column; gap:6px; box-sizing:border-box;
|
||||
height:100%; min-height:0; }
|
||||
.tgate-area { flex:1 1 auto; min-height:0; width:100%; box-sizing:border-box; resize:none;
|
||||
font-size:12px; line-height:1.4; padding:6px; border-radius:4px;
|
||||
border:1px solid #555; background:rgba(0,0,0,0.25); color:#fff;
|
||||
font-family:ui-monospace, monospace; overflow:auto; }
|
||||
.tgate-btns { display:flex; gap:6px; align-items:center; flex:0 0 auto; }
|
||||
.tgate-btns button { font-size:12px; padding:3px 14px; cursor:pointer; border-radius:3px;
|
||||
border:1px solid #555; color:#fff; }
|
||||
.tgate-pass { background:rgba(40,130,70,0.95); }
|
||||
.tgate-pass:hover { background:rgba(55,160,90,0.98); }
|
||||
.tgate-status { font-size:11px; opacity:0.6; margin-left:auto; }
|
||||
`;
|
||||
const style = document.createElement("style");
|
||||
style.id = "tgate-styles";
|
||||
style.textContent = css;
|
||||
document.head.appendChild(style);
|
||||
}
|
||||
|
||||
function setupTextGateNode(node) {
|
||||
injectStyles();
|
||||
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "tgate-wrap";
|
||||
|
||||
const area = document.createElement("textarea");
|
||||
area.className = "tgate-area";
|
||||
area.placeholder = "waiting for a run…";
|
||||
// don't let typing/space toggle node selection or graph shortcuts
|
||||
area.onkeydown = (e) => e.stopPropagation();
|
||||
|
||||
const btns = document.createElement("div");
|
||||
btns.className = "tgate-btns";
|
||||
const pass = document.createElement("button");
|
||||
pass.className = "tgate-pass";
|
||||
pass.textContent = "▶ Pass";
|
||||
const status = document.createElement("span");
|
||||
status.className = "tgate-status";
|
||||
pass.onclick = async () => {
|
||||
await postPass(node, area.value);
|
||||
status.textContent = "passed";
|
||||
};
|
||||
btns.appendChild(pass);
|
||||
btns.appendChild(status);
|
||||
|
||||
wrap.appendChild(area);
|
||||
wrap.appendChild(btns);
|
||||
node._tg = { wrap, area, status };
|
||||
|
||||
// FILLS the node: floor-only min height, no max (Image Pool pattern).
|
||||
node._tgWidget = node.addDOMWidget("textgate_editor", "div", wrap, {
|
||||
serialize: false,
|
||||
getMinHeight: () => widgetFloor(),
|
||||
});
|
||||
|
||||
// keep the editor width synced on manual resize so the textarea reflows
|
||||
const onResize = node.onResize;
|
||||
node.onResize = function () {
|
||||
const r = onResize?.apply(this, arguments);
|
||||
syncWidgetWidth(node);
|
||||
return r;
|
||||
};
|
||||
|
||||
// sensible default size; the node stays freely resizable (no width floor lock)
|
||||
node.setSize([Math.max(node.size?.[0] || 0, MIN_W), node.computeSize()[1]]);
|
||||
syncWidgetWidth(node);
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: "datasete.gates.textgate",
|
||||
|
||||
// one global socket listener: route the server's pause event to the node
|
||||
setup() {
|
||||
api.addEventListener("datasete-textgate-show", (e) => {
|
||||
const d = e.detail || {};
|
||||
const node = app.graph?.getNodeById?.(parseInt(d.id, 10));
|
||||
if (!node || node.type !== NODE || !node._tg) return;
|
||||
node._tg.area.value = d.text || "";
|
||||
node._tg.status.textContent = "edit, then Pass";
|
||||
try { node._tg.area.focus(); } catch (err) { /* ignore */ }
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
});
|
||||
},
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== NODE) return;
|
||||
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function () {
|
||||
const r = onNodeCreated?.apply(this, arguments);
|
||||
setupTextGateNode(this);
|
||||
return r;
|
||||
};
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user