From ceee482ecc15206276a5ffe0d931bc57b8de58c6 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Tue, 24 Jun 2025 16:36:15 +0800 Subject: [PATCH] feat: refactor Lora handling by introducing chainCallback for improved node initialization and widget management. Fixes #176 --- web/comfyui/lora_loader.js | 161 +++++++++++++++++++----------------- web/comfyui/lora_stacker.js | 54 ++++++------ web/comfyui/utils.js | 18 ++++ 3 files changed, 126 insertions(+), 107 deletions(-) diff --git a/web/comfyui/lora_loader.js b/web/comfyui/lora_loader.js index 431ce6c7..6c285c67 100644 --- a/web/comfyui/lora_loader.js +++ b/web/comfyui/lora_loader.js @@ -1,11 +1,12 @@ import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; import { - getLorasWidgetModule, LORA_PATTERN, collectActiveLorasFromChain, - updateConnectedTriggerWords + updateConnectedTriggerWords, + chainCallback } from "./utils.js"; -import { api } from "../../scripts/api.js"; +import { addLorasWidget } from "./loras_widget.js"; function mergeLoras(lorasText, lorasArr) { const result = []; @@ -82,7 +83,7 @@ app.registerExtension({ this.updateNodeLoraCode(node, loraCode, mode); }, - + // Helper method to update a single node's lora code updateNodeLoraCode(node, loraCode, mode) { // Update the input widget with new lora code @@ -107,90 +108,94 @@ app.registerExtension({ inputWidget.callback(inputWidget.value); } }, - - async nodeCreated(node) { - if (node.comfyClass === "Lora Loader (LoraManager)") { + + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeType.comfyClass == "Lora Loader (LoraManager)") { + chainCallback(nodeType.prototype, "onNodeCreated", function () { // Enable widget serialization - node.serialize_widgets = true; + this.serialize_widgets = true; - node.addInput('clip', 'CLIP', { - "shape": 7 + this.addInput("clip", "CLIP", { + shape: 7, }); - node.addInput("lora_stack", 'LORA_STACK', { - "shape": 7 // 7 is the shape of the optional input + this.addInput("lora_stack", "LORA_STACK", { + shape: 7, // 7 is the shape of the optional input }); - // Wait for node to be properly initialized - requestAnimationFrame(async () => { - // Restore saved value if exists - let existingLoras = []; - if (node.widgets_values && node.widgets_values.length > 0) { - // 0 for input widget, 1 for loras widget - const savedValue = node.widgets_values[1]; - existingLoras = savedValue || []; + // Restore saved value if exists + let existingLoras = []; + if (this.widgets_values && this.widgets_values.length > 0) { + // 0 for input widget, 1 for loras widget + const savedValue = this.widgets_values[1]; + existingLoras = savedValue || []; + } + // Merge the loras data + const mergedLoras = mergeLoras( + this.widgets[0].value, + existingLoras + ); + + // Add flag to prevent callback loops + let isUpdating = false; + + // Get the widget object directly from the returned object + this.lorasWidget = addLorasWidget( + this, + "loras", + { + defaultVal: mergedLoras, // Pass object directly + }, + (value) => { + // Collect all active loras from this node and its input chain + const allActiveLoraNames = collectActiveLorasFromChain(this); + + // Update trigger words for connected toggle nodes with the aggregated lora names + updateConnectedTriggerWords(this, allActiveLoraNames); + + // Prevent recursive calls + if (isUpdating) return; + isUpdating = true; + + try { + // Remove loras that are not in the value array + const inputWidget = this.widgets[0]; + const currentLoras = value.map((l) => l.name); + + // Use the constant pattern here as well + let newText = inputWidget.value.replace( + LORA_PATTERN, + (match, name, strength, clipStrength) => { + return currentLoras.includes(name) ? match : ""; + } + ); + + // Clean up multiple spaces and trim + newText = newText.replace(/\s+/g, " ").trim(); + + inputWidget.value = newText; + } finally { + isUpdating = false; } - // Merge the loras data - const mergedLoras = mergeLoras(node.widgets[0].value, existingLoras); - - // Add flag to prevent callback loops - let isUpdating = false; - - // Dynamically load the appropriate widget module - const lorasModule = await getLorasWidgetModule(); - const { addLorasWidget } = lorasModule; - - // Get the widget object directly from the returned object - const result = addLorasWidget(node, "loras", { - defaultVal: mergedLoras // Pass object directly - }, (value) => { - // Collect all active loras from this node and its input chain - const allActiveLoraNames = collectActiveLorasFromChain(node); - - // Update trigger words for connected toggle nodes with the aggregated lora names - updateConnectedTriggerWords(node, allActiveLoraNames); + } + ).widget; - // Prevent recursive calls - if (isUpdating) return; - isUpdating = true; + // Update input widget callback + const inputWidget = this.widgets[0]; + inputWidget.callback = (value) => { + if (isUpdating) return; + isUpdating = true; - try { - // Remove loras that are not in the value array - const inputWidget = node.widgets[0]; - const currentLoras = value.map(l => l.name); - - // Use the constant pattern here as well - let newText = inputWidget.value.replace(LORA_PATTERN, (match, name, strength, clipStrength) => { - return currentLoras.includes(name) ? match : ''; - }); - - // Clean up multiple spaces and trim - newText = newText.replace(/\s+/g, ' ').trim(); - - inputWidget.value = newText; - } finally { - isUpdating = false; - } - }); - - node.lorasWidget = result.widget; + try { + const currentLoras = this.lorasWidget.value || []; + const mergedLoras = mergeLoras(value, currentLoras); - // Update input widget callback - const inputWidget = node.widgets[0]; - inputWidget.callback = (value) => { - if (isUpdating) return; - isUpdating = true; - - try { - const currentLoras = node.lorasWidget.value || []; - const mergedLoras = mergeLoras(value, currentLoras); - - node.lorasWidget.value = mergedLoras; - } finally { - isUpdating = false; - } - }; - }); + this.lorasWidget.value = mergedLoras; + } finally { + isUpdating = false; + } + }; + }); } }, }); \ No newline at end of file diff --git a/web/comfyui/lora_stacker.js b/web/comfyui/lora_stacker.js index 9388483a..bd73b53e 100644 --- a/web/comfyui/lora_stacker.js +++ b/web/comfyui/lora_stacker.js @@ -1,11 +1,12 @@ import { app } from "../../scripts/app.js"; import { - getLorasWidgetModule, LORA_PATTERN, getActiveLorasFromNode, collectActiveLorasFromChain, - updateConnectedTriggerWords + updateConnectedTriggerWords, + chainCallback } from "./utils.js"; +import { addLorasWidget } from "./loras_widget.js"; function mergeLoras(lorasText, lorasArr) { const result = []; @@ -39,35 +40,30 @@ function mergeLoras(lorasText, lorasArr) { app.registerExtension({ name: "LoraManager.LoraStacker", - async nodeCreated(node) { - if (node.comfyClass === "Lora Stacker (LoraManager)") { - // Enable widget serialization - node.serialize_widgets = true; + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeType.comfyClass === "Lora Stacker (LoraManager)") { + chainCallback(nodeType.prototype, "onNodeCreated", async function() { + // Enable widget serialization + this.serialize_widgets = true; - node.addInput("lora_stack", 'LORA_STACK', { - "shape": 7 // 7 is the shape of the optional input - }); + this.addInput("lora_stack", 'LORA_STACK', { + "shape": 7 // 7 is the shape of the optional input + }); - // Wait for node to be properly initialized - requestAnimationFrame(async () => { // Restore saved value if exists let existingLoras = []; - if (node.widgets_values && node.widgets_values.length > 0) { + if (this.widgets_values && this.widgets_values.length > 0) { // 0 for input widget, 1 for loras widget - const savedValue = node.widgets_values[1]; + const savedValue = this.widgets_values[1]; existingLoras = savedValue || []; } // Merge the loras data - const mergedLoras = mergeLoras(node.widgets[0].value, existingLoras); + const mergedLoras = mergeLoras(this.widgets[0].value, existingLoras); // Add flag to prevent callback loops let isUpdating = false; - // Dynamically load the appropriate widget module - const lorasModule = await getLorasWidgetModule(); - const { addLorasWidget } = lorasModule; - - const result = addLorasWidget(node, "loras", { + const result = addLorasWidget(this, "loras", { defaultVal: mergedLoras // Pass object directly }, (value) => { // Prevent recursive calls @@ -76,7 +72,7 @@ app.registerExtension({ try { // Remove loras that are not in the value array - const inputWidget = node.widgets[0]; + const inputWidget = this.widgets[0]; const currentLoras = value.map(l => l.name); // Use the constant pattern here as well @@ -96,35 +92,35 @@ app.registerExtension({ activeLoraNames.add(lora.name); } }); - updateConnectedTriggerWords(node, activeLoraNames); + updateConnectedTriggerWords(this, activeLoraNames); // Find all Lora Loader nodes in the chain that might need updates - updateDownstreamLoaders(node); + updateDownstreamLoaders(this); } finally { isUpdating = false; } }); - node.lorasWidget = result.widget; + this.lorasWidget = result.widget; // Update input widget callback - const inputWidget = node.widgets[0]; + const inputWidget = this.widgets[0]; inputWidget.callback = (value) => { if (isUpdating) return; isUpdating = true; try { - const currentLoras = node.lorasWidget.value || []; + const currentLoras = this.lorasWidget.value || []; const mergedLoras = mergeLoras(value, currentLoras); - node.lorasWidget.value = mergedLoras; + this.lorasWidget.value = mergedLoras; // Update this stacker's direct trigger toggles with its own active loras - const activeLoraNames = getActiveLorasFromNode(node); - updateConnectedTriggerWords(node, activeLoraNames); + const activeLoraNames = getActiveLorasFromNode(this); + updateConnectedTriggerWords(this, activeLoraNames); // Find all Lora Loader nodes in the chain that might need updates - updateDownstreamLoaders(node); + updateDownstreamLoaders(this); } finally { isUpdating = false; } diff --git a/web/comfyui/utils.js b/web/comfyui/utils.js index 9fe00f9a..54924103 100644 --- a/web/comfyui/utils.js +++ b/web/comfyui/utils.js @@ -1,5 +1,23 @@ export const CONVERTED_TYPE = 'converted-widget'; +export function chainCallback(object, property, callback) { + if (object == undefined) { + //This should not happen. + console.error("Tried to add callback to non-existant object") + return; + } + if (property in object) { + const callback_orig = object[property] + object[property] = function () { + const r = callback_orig.apply(this, arguments); + callback.apply(this, arguments); + return r + }; + } else { + object[property] = callback; + } +} + export function getComfyUIFrontendVersion() { return window['__COMFYUI_FRONTEND_VERSION__'] || "0.0.0"; }