feat(graph): enhance node handling with graph identifiers and improve metadata updates, see #408, #538

This commit is contained in:
Will Miao
2025-10-07 23:22:38 +08:00
parent 9199950b74
commit 3118f3b43c
12 changed files with 574 additions and 103 deletions

View File

@@ -1,6 +1,7 @@
import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { addJsonDisplayWidget } from "./json_display_widget.js";
import { getNodeFromGraph } from "./utils.js";
app.registerExtension({
name: "LoraManager.DebugMetadata",
@@ -8,8 +9,8 @@ app.registerExtension({
setup() {
// Add message handler to listen for metadata updates from Python
api.addEventListener("metadata_update", (event) => {
const { id, metadata } = event.detail;
this.handleMetadataUpdate(id, metadata);
const { id, graph_id: graphId, metadata } = event.detail;
this.handleMetadataUpdate(id, graphId, metadata);
});
},
@@ -37,8 +38,8 @@ app.registerExtension({
},
// Handle metadata updates from Python
handleMetadataUpdate(id, metadata) {
const node = app.graph.getNodeById(+id);
handleMetadataUpdate(id, graphId, metadata) {
const node = getNodeFromGraph(graphId, id);
if (!node || node.comfyClass !== "Debug Metadata (LoraManager)") {
console.warn("Node not found or not a DebugMetadata node:", id);
return;

View File

@@ -7,6 +7,8 @@ import {
chainCallback,
mergeLoras,
setupInputWidgetWithAutocomplete,
getAllGraphNodes,
getNodeFromGraph,
} from "./utils.js";
import { addLorasWidget } from "./loras_widget.js";
@@ -16,23 +18,26 @@ app.registerExtension({
setup() {
// Add message handler to listen for messages from Python
api.addEventListener("lora_code_update", (event) => {
const { id, lora_code, mode } = event.detail;
this.handleLoraCodeUpdate(id, lora_code, mode);
this.handleLoraCodeUpdate(event.detail || {});
});
},
// Handle lora code updates from Python
handleLoraCodeUpdate(id, loraCode, mode) {
handleLoraCodeUpdate(message) {
const nodeId = message?.node_id ?? message?.id;
const graphId = message?.graph_id;
const loraCode = message?.lora_code ?? "";
const mode = message?.mode ?? "append";
const numericNodeId =
typeof nodeId === "string" ? Number(nodeId) : nodeId;
// Handle broadcast mode (for Desktop/non-browser support)
if (id === -1) {
if (numericNodeId === -1) {
// Find all Lora Loader nodes in the current graph
const loraLoaderNodes = [];
for (const nodeId in app.graph._nodes_by_id) {
const node = app.graph._nodes_by_id[nodeId];
if (node.comfyClass === "Lora Loader (LoraManager)") {
loraLoaderNodes.push(node);
}
}
const loraLoaderNodes = getAllGraphNodes(app.graph)
.map(({ node }) => node)
.filter((node) => node?.comfyClass === "Lora Loader (LoraManager)");
// Update each Lora Loader node found
if (loraLoaderNodes.length > 0) {
@@ -52,14 +57,18 @@ app.registerExtension({
}
// Standard mode - update a specific node
const node = app.graph.getNodeById(+id);
const node = getNodeFromGraph(graphId, numericNodeId);
if (
!node ||
(node.comfyClass !== "Lora Loader (LoraManager)" &&
node.comfyClass !== "Lora Stacker (LoraManager)" &&
node.comfyClass !== "WanVideo Lora Select (LoraManager)")
) {
console.warn("Node not found or not a LoraLoader:", id);
console.warn(
"Node not found or not a LoraLoader:",
graphId ?? "root",
nodeId
);
return;
}

View File

@@ -7,6 +7,8 @@ import {
chainCallback,
mergeLoras,
setupInputWidgetWithAutocomplete,
getLinkFromGraph,
getNodeKey,
} from "./utils.js";
import { addLorasWidget } from "./loras_widget.js";
@@ -124,17 +126,18 @@ app.registerExtension({
// Helper function to find and update downstream Lora Loader nodes
function updateDownstreamLoaders(startNode, visited = new Set()) {
if (visited.has(startNode.id)) return;
visited.add(startNode.id);
const nodeKey = getNodeKey(startNode);
if (!nodeKey || visited.has(nodeKey)) return;
visited.add(nodeKey);
// Check each output link
if (startNode.outputs) {
for (const output of startNode.outputs) {
if (output.links) {
for (const linkId of output.links) {
const link = app.graph.links[linkId];
const link = getLinkFromGraph(startNode.graph, linkId);
if (link) {
const targetNode = app.graph.getNodeById(link.target_id);
const targetNode = startNode.graph?.getNodeById?.(link.target_id);
// If target is a Lora Loader, collect all active loras in the chain and update
if (

View File

@@ -1,6 +1,6 @@
import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { CONVERTED_TYPE } from "./utils.js";
import { CONVERTED_TYPE, getNodeFromGraph } from "./utils.js";
import { addTagsWidget } from "./tags_widget.js";
// TriggerWordToggle extension for ComfyUI
@@ -10,8 +10,8 @@ app.registerExtension({
setup() {
// Add message handler to listen for messages from Python
api.addEventListener("trigger_word_update", (event) => {
const { id, message } = event.detail;
this.handleTriggerWordUpdate(id, message);
const { id, graph_id: graphId, message } = event.detail;
this.handleTriggerWordUpdate(id, graphId, message);
});
},
@@ -76,8 +76,8 @@ app.registerExtension({
},
// Handle trigger word updates from Python
handleTriggerWordUpdate(id, message) {
const node = app.graph.getNodeById(+id);
handleTriggerWordUpdate(id, graphId, message) {
const node = getNodeFromGraph(graphId, id);
if (!node || node.comfyClass !== "TriggerWord Toggle (LoraManager)") {
console.warn("Node not found or not a TriggerWordToggle:", id);
return;

View File

@@ -1,7 +1,7 @@
// ComfyUI extension to track model usage statistics
import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { showToast } from "./utils.js";
import { getAllGraphNodes, getNodeReference, showToast } from "./utils.js";
// Define target nodes and their widget configurations
const PATH_CORRECTION_TARGETS = [
@@ -56,25 +56,35 @@ app.registerExtension({
async refreshRegistry() {
try {
// Get current workflow nodes
const prompt = await app.graphToPrompt();
const workflow = prompt.workflow;
if (!workflow || !workflow.nodes) {
console.warn("No workflow nodes found for registry refresh");
return;
}
// Find all Lora nodes
const loraNodes = [];
for (const node of workflow.nodes.values()) {
if (node.type === "Lora Loader (LoraManager)" ||
node.type === "Lora Stacker (LoraManager)" ||
node.type === "WanVideo Lora Select (LoraManager)") {
const nodeEntries = getAllGraphNodes(app.graph);
for (const { graph, node } of nodeEntries) {
if (!node || !node.comfyClass) {
continue;
}
if (
node.comfyClass === "Lora Loader (LoraManager)" ||
node.comfyClass === "Lora Stacker (LoraManager)" ||
node.comfyClass === "WanVideo Lora Select (LoraManager)"
) {
const reference = getNodeReference(node);
if (!reference) {
continue;
}
const graphName = typeof graph?.name === "string" && graph.name.trim()
? graph.name
: null;
loraNodes.push({
node_id: node.id,
bgcolor: node.bgcolor || null,
title: node.title || node.type,
type: node.type
node_id: reference.node_id,
graph_id: reference.graph_id,
graph_name: graphName,
bgcolor: node.bgcolor ?? node.color ?? null,
title: node.title || node.comfyClass,
type: node.comfyClass,
});
}
}

View File

@@ -2,6 +2,120 @@ export const CONVERTED_TYPE = 'converted-widget';
import { app } from "../../scripts/app.js";
import { AutoComplete } from "./autocomplete.js";
const ROOT_GRAPH_ID = "root";
function isMapLike(collection) {
return collection && typeof collection.entries === "function" && typeof collection.values === "function";
}
function getChildGraphs(graph) {
if (!graph || !graph._subgraphs) {
return [];
}
const rawSubgraphs = isMapLike(graph._subgraphs)
? Array.from(graph._subgraphs.values())
: Object.values(graph._subgraphs);
return rawSubgraphs
.map((subgraph) => subgraph?.graph || subgraph?._graph || subgraph)
.filter((subgraph) => subgraph && subgraph !== graph);
}
function traverseGraphs(rootGraph, visitor, visited = new Set()) {
const graph = rootGraph || app.graph;
if (!graph) {
return;
}
const graphId = getGraphId(graph);
if (visited.has(graphId)) {
return;
}
visited.add(graphId);
visitor(graph);
for (const subgraph of getChildGraphs(graph)) {
traverseGraphs(subgraph, visitor, visited);
}
}
export function getGraphId(graph) {
return graph?.id ?? ROOT_GRAPH_ID;
}
export function getNodeGraphId(node) {
if (!node) {
return ROOT_GRAPH_ID;
}
return getGraphId(node.graph || app.graph);
}
export function getGraphById(graphId, rootGraph = app.graph) {
if (!graphId) {
return rootGraph;
}
let foundGraph = null;
traverseGraphs(rootGraph, (graph) => {
if (!foundGraph && getGraphId(graph) === graphId) {
foundGraph = graph;
}
});
return foundGraph;
}
export function getNodeFromGraph(graphId, nodeId) {
const graph = getGraphById(graphId) || app.graph;
if (!graph || typeof graph.getNodeById !== "function") {
return null;
}
const numericId = typeof nodeId === "string" ? Number(nodeId) : nodeId;
return graph.getNodeById(Number.isNaN(numericId) ? nodeId : numericId) || null;
}
export function getAllGraphNodes(rootGraph = app.graph) {
const nodes = [];
traverseGraphs(rootGraph, (graph) => {
if (Array.isArray(graph._nodes)) {
for (const node of graph._nodes) {
nodes.push({ graph, node });
}
}
});
return nodes;
}
export function getNodeReference(node) {
if (!node) {
return null;
}
return {
node_id: node.id,
graph_id: getNodeGraphId(node),
};
}
export function getNodeKey(node) {
if (!node) {
return null;
}
return `${getNodeGraphId(node)}:${node.id}`;
}
export function getLinkFromGraph(graph, linkId) {
if (!graph || graph.links == null) {
return null;
}
if (isMapLike(graph.links)) {
return graph.links.get(linkId) || null;
}
return graph.links[linkId] || null;
}
export function chainCallback(object, property, callback) {
if (object == undefined) {
//This should not happen.
@@ -103,42 +217,56 @@ export const LORA_PATTERN = /<lora:([^:]+):([-\d\.]+)(?::([-\d\.]+))?>/g;
// Get connected Lora Stacker nodes that feed into the current node
export function getConnectedInputStackers(node) {
const connectedStackers = [];
if (node.inputs) {
for (const input of node.inputs) {
if (input.name === "lora_stack" && input.link) {
const link = app.graph.links[input.link];
if (link) {
const sourceNode = app.graph.getNodeById(link.origin_id);
if (sourceNode && sourceNode.comfyClass === "Lora Stacker (LoraManager)") {
connectedStackers.push(sourceNode);
}
}
}
if (!node?.inputs) {
return connectedStackers;
}
for (const input of node.inputs) {
if (input.name !== "lora_stack" || !input.link) {
continue;
}
const link = getLinkFromGraph(node.graph, input.link);
if (!link) {
continue;
}
const sourceNode = node.graph?.getNodeById?.(link.origin_id);
if (sourceNode && sourceNode.comfyClass === "Lora Stacker (LoraManager)") {
connectedStackers.push(sourceNode);
}
}
return connectedStackers;
}
// Get connected TriggerWord Toggle nodes that receive output from the current node
export function getConnectedTriggerToggleNodes(node) {
const connectedNodes = [];
if (node.outputs && node.outputs.length > 0) {
for (const output of node.outputs) {
if (output.links && output.links.length > 0) {
for (const linkId of output.links) {
const link = app.graph.links[linkId];
if (link) {
const targetNode = app.graph.getNodeById(link.target_id);
if (targetNode && targetNode.comfyClass === "TriggerWord Toggle (LoraManager)") {
connectedNodes.push(targetNode.id);
}
}
}
if (!node?.outputs) {
return connectedNodes;
}
for (const output of node.outputs) {
if (!output?.links?.length) {
continue;
}
for (const linkId of output.links) {
const link = getLinkFromGraph(node.graph, linkId);
if (!link) {
continue;
}
const targetNode = node.graph?.getNodeById?.(link.target_id);
if (targetNode && targetNode.comfyClass === "TriggerWord Toggle (LoraManager)") {
connectedNodes.push(targetNode);
}
}
}
return connectedNodes;
}
@@ -161,11 +289,15 @@ export function getActiveLorasFromNode(node) {
// Recursively collect all active loras from a node and its input chain
export function collectActiveLorasFromChain(node, visited = new Set()) {
// Prevent infinite loops from circular references
if (visited.has(node.id)) {
const nodeKey = getNodeKey(node);
if (!nodeKey) {
return new Set();
}
visited.add(node.id);
if (visited.has(nodeKey)) {
return new Set();
}
visited.add(nodeKey);
// Get active loras from current node
const allActiveLoraNames = getActiveLorasFromNode(node);
@@ -181,14 +313,22 @@ export function collectActiveLorasFromChain(node, visited = new Set()) {
// Update trigger words for connected toggle nodes
export function updateConnectedTriggerWords(node, loraNames) {
const connectedNodeIds = getConnectedTriggerToggleNodes(node);
if (connectedNodeIds.length > 0) {
const connectedNodes = getConnectedTriggerToggleNodes(node);
if (connectedNodes.length > 0) {
const nodeIds = connectedNodes
.map((connectedNode) => getNodeReference(connectedNode))
.filter((reference) => reference !== null);
if (nodeIds.length === 0) {
return;
}
fetch("/api/lm/loras/get_trigger_words", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
lora_names: Array.from(loraNames),
node_ids: connectedNodeIds
node_ids: nodeIds
})
}).catch(err => console.error("Error fetching trigger words:", err));
}