69 lines
1.8 KiB
JavaScript
69 lines
1.8 KiB
JavaScript
const { app } = window.comfyAPI.app;
|
|
|
|
const HIDDEN_LOADER_WIDGETS = new Set([
|
|
"backend",
|
|
"dtype",
|
|
"quantization",
|
|
"local_files_only",
|
|
"unload_policy",
|
|
"max_memory_gb",
|
|
]);
|
|
|
|
function widget(node, name) {
|
|
return node.widgets?.find((w) => w.name === name);
|
|
}
|
|
|
|
function migrateValues(node) {
|
|
const temp = widget(node, "temperature");
|
|
if (!temp) return;
|
|
|
|
const tempNumber = Number(temp.value);
|
|
if (!Number.isFinite(tempNumber)) {
|
|
temp.value = 0.45;
|
|
}
|
|
}
|
|
|
|
function positiveNumber(value) {
|
|
const number = Number(value);
|
|
return Number.isFinite(number) && number > 0 ? number : null;
|
|
}
|
|
|
|
function hideWidgets(node, names) {
|
|
const widthBeforeHide = positiveNumber(node.size?.[0]);
|
|
for (const w of node.widgets || []) {
|
|
if (!names.has(w.name) || w.__dgHidden) continue;
|
|
w.__dgHidden = true;
|
|
w.hidden = true;
|
|
w.computeSize = () => [0, -4];
|
|
}
|
|
const computedSize = node.computeSize?.() || node.size;
|
|
const width = widthBeforeHide || positiveNumber(computedSize?.[0]);
|
|
const height = positiveNumber(computedSize?.[1]) || positiveNumber(node.size?.[1]);
|
|
if (width && height) node.setSize?.([width, height]);
|
|
}
|
|
|
|
function simplifyLoader(node) {
|
|
migrateValues(node);
|
|
hideWidgets(node, HIDDEN_LOADER_WIDGETS);
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "DiffusionGemmaPromptBuilder.ModelLoaderUI",
|
|
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
if (nodeData?.name !== "DiffusionGemmaModelLoader") return;
|
|
|
|
const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function (...args) {
|
|
originalOnNodeCreated?.apply(this, args);
|
|
simplifyLoader(this);
|
|
};
|
|
|
|
const originalOnConfigure = nodeType.prototype.onConfigure;
|
|
nodeType.prototype.onConfigure = function (...args) {
|
|
originalOnConfigure?.apply(this, args);
|
|
simplifyLoader(this);
|
|
};
|
|
},
|
|
});
|