This commit is contained in:
justumen
2024-11-22 19:20:35 +01:00
parent 3d8d94597f
commit 0a968c1540
12 changed files with 270 additions and 170 deletions

View File

@@ -4,43 +4,64 @@ app.registerExtension({
name: "Bjornulf.LoopModelSelector",
async nodeCreated(node) {
if (node.comfyClass === "Bjornulf_LoopModelSelector") {
node.properties = node.properties || {};
const updateModelInputs = () => {
const initialWidth = node.size[0];
const numModelsWidget = node.widgets.find(w => w.name === "number_of_models");
if (!numModelsWidget) return;
const numModels = numModelsWidget.value;
const checkpointsList = node.widgets.find(w => w.name === "model_1").options.values;
const checkpointsList = node.widgets.find(w => w.name === "model_1")?.options?.values || [];
// Remove excess model widgets
node.widgets = node.widgets.filter(w => !w.name.startsWith("model_") || parseInt(w.name.split("_")[1]) <= numModels);
// Add new model widgets if needed
node.widgets = node.widgets.filter(w =>
!w.name.startsWith("model_") || parseInt(w.name.split("_")[1]) <= numModels
);
// Store current widget values in properties
node.widgets.forEach(w => {
if (w.name.startsWith("model_")) {
node.properties[w.name] = w.value;
}
});
// Remove all model-related widgets
node.widgets = node.widgets.filter(w =>
!w.name.startsWith("model_")
);
// Add new model widgets
for (let i = 1; i <= numModels; i++) {
const widgetName = `model_${i}`;
if (!node.widgets.find(w => w.name === widgetName)) {
const defaultIndex = Math.min(i - 1, checkpointsList.length - 1);
node.addWidget("combo", widgetName, checkpointsList[defaultIndex], () => {}, {
values: checkpointsList
});
}
const savedValue = node.properties[widgetName];
const defaultIndex = Math.min(i - 1, checkpointsList.length - 1);
const defaultValue = savedValue !== undefined ? savedValue : checkpointsList[defaultIndex];
const modelWidget = node.addWidget("combo", widgetName, defaultValue, (value) => {
node.properties[widgetName] = value;
}, {
values: checkpointsList
});
}
// Reorder widgets
node.widgets.sort((a, b) => {
if (a.name === "number_of_models") return -1;
if (b.name === "number_of_models") return 1;
if (a.name === "seed") return 1;
if (b.name === "seed") return -1;
if (a.name.startsWith("model_") && b.name.startsWith("model_")) {
return parseInt(a.name.split("_")[1]) - parseInt(b.name.split("_")[1]);
}
return a.name.localeCompare(b.name);
});
// Reorder widgets: number_of_models first, then the model widgets
const orderedWidgets = [node.widgets.find(w => w.name === "number_of_models")];
for (let i = 1; i <= numModels; i++) {
const modelWidgets = node.widgets.filter(w => w.name === `model_${i}`);
orderedWidgets.push(...modelWidgets);
}
// Add any remaining widgets
orderedWidgets.push(...node.widgets.filter(w => !orderedWidgets.includes(w)));
node.widgets = orderedWidgets;
// Adjust node size
node.setSize(node.computeSize());
node.size[0] = initialWidth; // Keep width fixed
};
// Set up number_of_models widget
// Set up number_of_models widget callback
const numModelsWidget = node.widgets.find(w => w.name === "number_of_models");
if (numModelsWidget) {
numModelsWidget.callback = () => {
@@ -49,48 +70,56 @@ app.registerExtension({
};
}
// Set seed widget to integer input
const seedWidget = node.widgets.find((w) => w.name === "seed");
if (seedWidget) {
seedWidget.type = "HIDDEN"; // Hide seed widget after restoring saved state
}
// Handle deserialization
const originalOnConfigure = node.onConfigure;
node.onConfigure = function(info) {
if (originalOnConfigure) {
originalOnConfigure.call(this, info);
}
// Restore model widgets based on saved properties
const savedProperties = info.properties;
if (savedProperties) {
Object.keys(savedProperties).forEach(key => {
if (key.startsWith("model_")) {
const widgetName = key;
const widgetValue = savedProperties[key];
const existingWidget = node.widgets.find(w => w.name === widgetName);
if (existingWidget) {
existingWidget.value = widgetValue;
} else {
node.addWidget("combo", widgetName, widgetValue, () => {}, {
values: node.widgets.find(w => w.name === "model_1").options.values
});
}
}
});
}
// const savedProperties = info.properties;
// if (savedProperties) {
// Object.keys(savedProperties).forEach(key => {
// if (key.startsWith("model_")) {
// const widgetName = key;
// const widgetValue = savedProperties[key];
// const existingWidget = node.widgets.find(w => w.name === widgetName);
// if (existingWidget) {
// existingWidget.value = widgetValue;
// } else {
// const checkpointsList = node.widgets.find(w => w.name === "model_1")?.options?.values || [];
// const defaultIndex = Math.min(parseInt(widgetName.split("_")[1]) - 1, checkpointsList.length - 1);
// node.addWidget("combo", widgetName, widgetValue || checkpointsList[defaultIndex], () => {}, {
// values: checkpointsList
// });
// }
// }
// });
// }
// Ensure seed is a valid integer
const seedWidget = node.widgets.find(w => w.name === "seed");
if (seedWidget && isNaN(parseInt(seedWidget.value))) {
seedWidget.value = 0; // Set a default value if invalid
// Save properties during serialization
const originalOnSerialize = node.onSerialize;
node.onSerialize = function(info) {
if (originalOnSerialize) {
originalOnSerialize.call(this, info);
}
info.properties = { ...this.properties };
};
// Update model inputs after restoring saved state
updateModelInputs();
};
// Serialize method to save properties
const originalSerialize = node.serialize;
node.serialize = function() {
const data = originalSerialize ? originalSerialize.call(this) : {};
data.properties = { ...this.properties };
return data;
};
// Initial update
updateModelInputs();
}