feat: refactor Lora handling by introducing chainCallback for improved node initialization and widget management. Fixes #176

This commit is contained in:
Will Miao
2025-06-24 16:36:15 +08:00
parent d0ed1213d8
commit ceee482ecc
3 changed files with 126 additions and 107 deletions

View File

@@ -1,11 +1,12 @@
import { app } from "../../scripts/app.js"; import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { import {
getLorasWidgetModule,
LORA_PATTERN, LORA_PATTERN,
collectActiveLorasFromChain, collectActiveLorasFromChain,
updateConnectedTriggerWords updateConnectedTriggerWords,
chainCallback
} from "./utils.js"; } from "./utils.js";
import { api } from "../../scripts/api.js"; import { addLorasWidget } from "./loras_widget.js";
function mergeLoras(lorasText, lorasArr) { function mergeLoras(lorasText, lorasArr) {
const result = []; const result = [];
@@ -82,7 +83,7 @@ app.registerExtension({
this.updateNodeLoraCode(node, loraCode, mode); this.updateNodeLoraCode(node, loraCode, mode);
}, },
// Helper method to update a single node's lora code // Helper method to update a single node's lora code
updateNodeLoraCode(node, loraCode, mode) { updateNodeLoraCode(node, loraCode, mode) {
// Update the input widget with new lora code // Update the input widget with new lora code
@@ -107,90 +108,94 @@ app.registerExtension({
inputWidget.callback(inputWidget.value); inputWidget.callback(inputWidget.value);
} }
}, },
async nodeCreated(node) { async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (node.comfyClass === "Lora Loader (LoraManager)") { if (nodeType.comfyClass == "Lora Loader (LoraManager)") {
chainCallback(nodeType.prototype, "onNodeCreated", function () {
// Enable widget serialization // Enable widget serialization
node.serialize_widgets = true; this.serialize_widgets = true;
node.addInput('clip', 'CLIP', { this.addInput("clip", "CLIP", {
"shape": 7 shape: 7,
}); });
node.addInput("lora_stack", 'LORA_STACK', { this.addInput("lora_stack", "LORA_STACK", {
"shape": 7 // 7 is the shape of the optional input shape: 7, // 7 is the shape of the optional input
}); });
// Wait for node to be properly initialized // Restore saved value if exists
requestAnimationFrame(async () => { let existingLoras = [];
// Restore saved value if exists if (this.widgets_values && this.widgets_values.length > 0) {
let existingLoras = []; // 0 for input widget, 1 for loras widget
if (node.widgets_values && node.widgets_values.length > 0) { const savedValue = this.widgets_values[1];
// 0 for input widget, 1 for loras widget existingLoras = savedValue || [];
const savedValue = node.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); ).widget;
// 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);
// Prevent recursive calls // Update input widget callback
if (isUpdating) return; const inputWidget = this.widgets[0];
isUpdating = true; inputWidget.callback = (value) => {
if (isUpdating) return;
isUpdating = true;
try { try {
// Remove loras that are not in the value array const currentLoras = this.lorasWidget.value || [];
const inputWidget = node.widgets[0]; const mergedLoras = mergeLoras(value, currentLoras);
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;
// Update input widget callback this.lorasWidget.value = mergedLoras;
const inputWidget = node.widgets[0]; } finally {
inputWidget.callback = (value) => { isUpdating = false;
if (isUpdating) return; }
isUpdating = true; };
});
try {
const currentLoras = node.lorasWidget.value || [];
const mergedLoras = mergeLoras(value, currentLoras);
node.lorasWidget.value = mergedLoras;
} finally {
isUpdating = false;
}
};
});
} }
}, },
}); });

View File

@@ -1,11 +1,12 @@
import { app } from "../../scripts/app.js"; import { app } from "../../scripts/app.js";
import { import {
getLorasWidgetModule,
LORA_PATTERN, LORA_PATTERN,
getActiveLorasFromNode, getActiveLorasFromNode,
collectActiveLorasFromChain, collectActiveLorasFromChain,
updateConnectedTriggerWords updateConnectedTriggerWords,
chainCallback
} from "./utils.js"; } from "./utils.js";
import { addLorasWidget } from "./loras_widget.js";
function mergeLoras(lorasText, lorasArr) { function mergeLoras(lorasText, lorasArr) {
const result = []; const result = [];
@@ -39,35 +40,30 @@ function mergeLoras(lorasText, lorasArr) {
app.registerExtension({ app.registerExtension({
name: "LoraManager.LoraStacker", name: "LoraManager.LoraStacker",
async nodeCreated(node) { async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (node.comfyClass === "Lora Stacker (LoraManager)") { if (nodeType.comfyClass === "Lora Stacker (LoraManager)") {
// Enable widget serialization chainCallback(nodeType.prototype, "onNodeCreated", async function() {
node.serialize_widgets = true; // Enable widget serialization
this.serialize_widgets = true;
node.addInput("lora_stack", 'LORA_STACK', { this.addInput("lora_stack", 'LORA_STACK', {
"shape": 7 // 7 is the shape of the optional input "shape": 7 // 7 is the shape of the optional input
}); });
// Wait for node to be properly initialized
requestAnimationFrame(async () => {
// Restore saved value if exists // Restore saved value if exists
let existingLoras = []; 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 // 0 for input widget, 1 for loras widget
const savedValue = node.widgets_values[1]; const savedValue = this.widgets_values[1];
existingLoras = savedValue || []; existingLoras = savedValue || [];
} }
// Merge the loras data // 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 // Add flag to prevent callback loops
let isUpdating = false; let isUpdating = false;
// Dynamically load the appropriate widget module const result = addLorasWidget(this, "loras", {
const lorasModule = await getLorasWidgetModule();
const { addLorasWidget } = lorasModule;
const result = addLorasWidget(node, "loras", {
defaultVal: mergedLoras // Pass object directly defaultVal: mergedLoras // Pass object directly
}, (value) => { }, (value) => {
// Prevent recursive calls // Prevent recursive calls
@@ -76,7 +72,7 @@ app.registerExtension({
try { try {
// Remove loras that are not in the value array // 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); const currentLoras = value.map(l => l.name);
// Use the constant pattern here as well // Use the constant pattern here as well
@@ -96,35 +92,35 @@ app.registerExtension({
activeLoraNames.add(lora.name); activeLoraNames.add(lora.name);
} }
}); });
updateConnectedTriggerWords(node, activeLoraNames); updateConnectedTriggerWords(this, activeLoraNames);
// Find all Lora Loader nodes in the chain that might need updates // Find all Lora Loader nodes in the chain that might need updates
updateDownstreamLoaders(node); updateDownstreamLoaders(this);
} finally { } finally {
isUpdating = false; isUpdating = false;
} }
}); });
node.lorasWidget = result.widget; this.lorasWidget = result.widget;
// Update input widget callback // Update input widget callback
const inputWidget = node.widgets[0]; const inputWidget = this.widgets[0];
inputWidget.callback = (value) => { inputWidget.callback = (value) => {
if (isUpdating) return; if (isUpdating) return;
isUpdating = true; isUpdating = true;
try { try {
const currentLoras = node.lorasWidget.value || []; const currentLoras = this.lorasWidget.value || [];
const mergedLoras = mergeLoras(value, currentLoras); 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 // Update this stacker's direct trigger toggles with its own active loras
const activeLoraNames = getActiveLorasFromNode(node); const activeLoraNames = getActiveLorasFromNode(this);
updateConnectedTriggerWords(node, activeLoraNames); updateConnectedTriggerWords(this, activeLoraNames);
// Find all Lora Loader nodes in the chain that might need updates // Find all Lora Loader nodes in the chain that might need updates
updateDownstreamLoaders(node); updateDownstreamLoaders(this);
} finally { } finally {
isUpdating = false; isUpdating = false;
} }

View File

@@ -1,5 +1,23 @@
export const CONVERTED_TYPE = 'converted-widget'; 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() { export function getComfyUIFrontendVersion() {
return window['__COMFYUI_FRONTEND_VERSION__'] || "0.0.0"; return window['__COMFYUI_FRONTEND_VERSION__'] || "0.0.0";
} }