3v324v23's picture
lfs
1e3b872
raw
history blame
1.5 kB
import { app } from "../../../scripts/app.js";
app.registerExtension({
name: "pysssss.KSamplerAdvDenoise",
async beforeRegisterNodeDef(nodeType) {
// Add menu options to conver to/from widgets
const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
nodeType.prototype.getExtraMenuOptions = function (_, options) {
const r = origGetExtraMenuOptions?.apply?.(this, arguments);
let stepsWidget = null;
let startAtWidget = null;
let endAtWidget = null;
for (const w of this.widgets || []) {
if (w.name === "steps") {
stepsWidget = w;
} else if (w.name === "start_at_step") {
startAtWidget = w;
} else if (w.name === "end_at_step") {
endAtWidget = w;
}
}
if (stepsWidget && startAtWidget && endAtWidget) {
options.push(
{
content: "Set Denoise",
callback: () => {
const steps = +prompt("How many steps do you want?", 15);
if (isNaN(steps)) {
return;
}
const denoise = +prompt("How much denoise? (0-1)", 0.5);
if (isNaN(denoise)) {
return;
}
stepsWidget.value = Math.floor(steps / Math.max(0, Math.min(1, denoise)));
stepsWidget.callback?.(stepsWidget.value);
startAtWidget.value = stepsWidget.value - steps;
startAtWidget.callback?.(startAtWidget.value);
endAtWidget.value = stepsWidget.value;
endAtWidget.callback?.(endAtWidget.value);
},
},
null
);
}
return r;
};
},
});