From 85da7175bcdbf4ecafbcb5204cbdf6c89fd1e850 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Tue, 23 Jun 2026 16:24:04 +0800 Subject: [PATCH] feat: add Node Marker system with right-click marking --- py/routes/handlers/misc_handlers.py | 16 +++- static/js/utils/uiHelpers.js | 7 +- web/comfyui/node_marker.js | 126 ++++++++++++++++++++++++++++ web/comfyui/workflow_registry.js | 80 ++++++++++++++---- 4 files changed, 207 insertions(+), 22 deletions(-) create mode 100644 web/comfyui/node_marker.js diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index 3eca6366..c3b86981 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -535,6 +535,7 @@ class NodeRegistry: "capabilities": capabilities, "widget_names": widget_names, "mode": node.get("mode"), + "marker_role": node.get("marker_role"), } logger.debug("Registered %s nodes in registry", len(nodes)) self._registry_updated.set() @@ -3104,13 +3105,17 @@ class NodeRegistryHandler: try: data = await request.json() widget_name = data.get("widget_name") + action = data.get("action") value = data.get("value") mode = data.get("mode", "replace") node_ids = data.get("node_ids") - if not isinstance(widget_name, str) or not widget_name: + if not action and (not isinstance(widget_name, str) or not widget_name): return web.json_response( - {"success": False, "error": "Missing widget_name parameter"}, + { + "success": False, + "error": "Missing parameter: provide either 'action' or 'widget_name'", + }, status=400, ) @@ -3149,12 +3154,15 @@ class NodeRegistryHandler: except (TypeError, ValueError): parsed_node_id = node_identifier - payload = { + payload: dict = { "id": parsed_node_id, - "widget_name": widget_name, "value": value, "mode": mode, } + if action: + payload["action"] = action + if widget_name: + payload["widget_name"] = widget_name if graph_identifier is not None: payload["graph_id"] = str(graph_identifier) diff --git a/static/js/utils/uiHelpers.js b/static/js/utils/uiHelpers.js index d53b43fd..afe5b9fe 100644 --- a/static/js/utils/uiHelpers.js +++ b/static/js/utils/uiHelpers.js @@ -915,7 +915,7 @@ async function sendTextToNodes(nodeIds, nodesMap, text, mode, messages = {}) { 'Content-Type': 'application/json', }, body: JSON.stringify({ - widget_name: 'text', + action: 'inject_text', value: text, mode: mode || 'append', node_ids: references, @@ -948,7 +948,10 @@ export async function sendEmbeddingToWorkflow(embeddingCode) { if (!isNodeEnabled(node)) { return false; } - return node.capabilities?.has_text_widget === true; + return ( + node.capabilities?.has_text_widget === true || + node.marker_role === "send_prompt_target" + ); }); const nodeKeys = Object.keys(textNodes); diff --git a/web/comfyui/node_marker.js b/web/comfyui/node_marker.js new file mode 100644 index 00000000..9d344c22 --- /dev/null +++ b/web/comfyui/node_marker.js @@ -0,0 +1,126 @@ +import { app } from "../../scripts/app.js"; + +// ============================================================================= +// Node Marker – right-click node marking (no dedicated node required) +// +// Adds a "Mark as →" submenu with role options to any node's context menu. +// Roles are stored in ``node.properties.lm_marker_role`` and automatically +// persist with the workflow JSON. +// +// The workflow registry reads these markers and makes them available to the +// standalone UI (e.g. ``sendEmbeddingToWorkflow`` also considers nodes marked +// as ``send_prompt_target``). +// ============================================================================= + +const ROLES = { + send_prompt_target: { + label: "Send Prompt Target", + emoji: "\uD83D\uDCDD", + }, + send_gen_params: { + label: "Send Gen Params Target", + emoji: "\uD83D\uDD27", + }, +}; + +// ---- Helpers ---------------------------------------------------------------- + +function getMarker(node) { + return node?.properties?.lm_marker_role ?? null; +} + +function setMarker(node, roleKey) { + if (!node || !ROLES[roleKey]) return; + node.properties = node.properties || {}; + node.properties.lm_marker_role = roleKey; + + // Save original title if not already saved, then prefix with emoji + if (!node.properties.lm_marker_original_title) { + node.properties.lm_marker_original_title = node.title || ""; + } + const def = ROLES[roleKey]; + node.title = `${def.emoji} ${node.properties.lm_marker_original_title}`; + + if (typeof node.setDirtyCanvas === "function") { + node.setDirtyCanvas(true, true); + } + triggerRegistryRefresh(); +} + +function clearMarker(node) { + if (!node) return; + delete node.properties.lm_marker_role; + + // Restore original title: prefer stripping emoji from current title + // (captures user renames after marking), fall back to saved original. + const cleaned = node.title?.replace( + /^(\u2709\uFE0F?|\u2699\uFE0F?|\uD83D\uDCDD|\uD83C\uDF9B\uFE0F?|\uD83D\uDD27)\s*/, + '' + ); + if (cleaned && cleaned !== node.title) { + node.title = cleaned; + } else { + const orig = node.properties.lm_marker_original_title; + if (orig !== undefined) { + node.title = orig; + } + } + delete node.properties.lm_marker_original_title; + + if (typeof node.setDirtyCanvas === "function") { + node.setDirtyCanvas(true, true); + } + triggerRegistryRefresh(); +} + +function triggerRegistryRefresh() { + // workflow_registry.js listens for this event to re-scan the graph. + window.dispatchEvent(new CustomEvent("lm_marker_changed")); +} + +// ---- Submenu builder -------------------------------------------------------- + +function buildSubmenuOptions(node) { + const currentRole = getMarker(node); + const options = []; + + for (const [key, def] of Object.entries(ROLES)) { + const isActive = currentRole === key; + options.push({ + content: `${isActive ? "\u2713 " : ""}${def.label}`, + disabled: isActive, + callback: () => setMarker(node, key), + }); + } + + if (currentRole) { + options.push({ + content: "Clear marker", + callback: () => clearMarker(node), + }); + } + + return options; +} + +function buildMenuItems(node) { + return [ + null, + { + content: "Mark as", + has_submenu: true, + submenu: { + options: buildSubmenuOptions(node), + }, + }, + ]; +} + +// ---- Extension -------------------------------------------------------------- + +app.registerExtension({ + name: "LoraManager.NodeMarker", + getNodeMenuItems(node) { + return buildMenuItems(node); + }, +}); diff --git a/web/comfyui/workflow_registry.js b/web/comfyui/workflow_registry.js index fffcbaeb..69f31c6d 100644 --- a/web/comfyui/workflow_registry.js +++ b/web/comfyui/workflow_registry.js @@ -28,6 +28,11 @@ app.registerExtension({ api.addEventListener("lm_widget_update", (event) => { this.applyWidgetUpdate(event?.detail ?? {}); }); + + // React to marker changes from the Node Marker extension + window.addEventListener("lm_marker_changed", () => { + this.refreshRegistry(); + }); }, async refreshRegistry() { @@ -49,8 +54,10 @@ app.registerExtension({ const supportsLora = LORA_NODE_CLASSES.has(node.comfyClass); const hasTargetWidget = widgetNames.some((name) => TARGET_WIDGET_NAMES.has(name)); const hasTextWidget = TEXT_CAPABLE_CLASSES.has(node.comfyClass); + const markerRole = node.properties?.lm_marker_role ?? null; - if (!supportsLora && !hasTargetWidget && !hasTextWidget) { + // Skip nodes with no relevant capability UNLESS they are marked + if (!supportsLora && !hasTargetWidget && !hasTextWidget && !markerRole) { continue; } @@ -71,6 +78,7 @@ app.registerExtension({ type: node.comfyClass, comfy_class: node.comfyClass, mode: node.mode, + marker_role: markerRole, capabilities: { supports_lora: supportsLora, has_text_widget: hasTextWidget, @@ -102,11 +110,12 @@ app.registerExtension({ applyWidgetUpdate(message) { const nodeId = message?.node_id ?? message?.id; const graphId = message?.graph_id; + const action = message?.action; const widgetName = message?.widget_name; const value = message?.value; const mode = message?.mode ?? "replace"; - if (nodeId == null || !widgetName) { + if (nodeId == null || (!action && !widgetName)) { console.warn("LoRA Manager: invalid widget update payload", message); return; } @@ -126,33 +135,72 @@ app.registerExtension({ return; } - const widgetIndex = node.widgets.findIndex((widget) => widget?.name === widgetName); - if (widgetIndex === -1) { - console.warn( - "LoRA Manager: target widget not found on node", - widgetName, - node - ); + // ---- Resolve target widget ---- + let targetWidget = null; + + if (action === "inject_text") { + // Find the first text-capable widget by type. + // Normalise to lowercase for case-insensitive matching. + const TEXT_TYPES = new Set(["string", "customtext"]); + targetWidget = node.widgets.find((w) => { + const t = typeof w?.type === "string" ? w.type.toLowerCase() : ""; + if (TEXT_TYPES.has(t)) return true; + // Broad fallback for unknown composite types. + if (t.includes("string")) { + return true; + } + return false; + }); + if (!targetWidget) { + // Last resort: pick the first widget that is not a hidden/internal type + targetWidget = node.widgets.find((w) => w?.name && !w.name.startsWith("_")); + if (!targetWidget) { + console.warn( + "LoRA Manager: no suitable widget for inject_text on node", + node.id + ); + return; + } + } + } else if (widgetName) { + // Legacy: find widget by name + targetWidget = node.widgets.find((w) => w?.name === widgetName); + if (!targetWidget) { + console.warn( + "LoRA Manager: target widget not found on node", + widgetName, + node + ); + return; + } + } else { + console.warn("LoRA Manager: no action or widget_name in payload", message); return; } - const widget = node.widgets[widgetIndex]; + // ---- Update widget value ---- + const widgetIndex = node.widgets.indexOf(targetWidget); let newValue = value; if (mode === "append") { - const separator = widget.value && widget.value.length > 0 ? " " : ""; - newValue = widget.value + separator + value; + const separator = + targetWidget.value && targetWidget.value.length > 0 ? " " : ""; + newValue = targetWidget.value + separator + value; } - widget.value = newValue; + targetWidget.value = newValue; - if (Array.isArray(node.widgets_values) && node.widgets_values.length > widgetIndex) { + if ( + Array.isArray(node.widgets_values) && + widgetIndex >= 0 && + node.widgets_values.length > widgetIndex + ) { node.widgets_values[widgetIndex] = newValue; } - if (typeof widget.callback === "function") { + if (typeof targetWidget.callback === "function") { try { - widget.callback(newValue); + targetWidget.callback(newValue); } catch (callbackError) { console.error("LoRA Manager: widget callback failed", callbackError); }