Files

85 lines
4.3 KiB
JavaScript

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
const NODE = "PROMPTFORGE_Sampler";
const PRESET_FIELDS = ["steps", "cfg", "sampler_name", "scheduler", "denoise", "runs", "step_cache", "cache_threshold"];
const NO_PRESET = "(no preset)";
app.registerExtension({
name: "PromptForge.Sampler",
async beforeRegisterNodeDef(nodeType, nodeData) {
if (nodeData.name !== NODE) return;
const onCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
onCreated?.apply(this, arguments);
const node = this;
const widget = (name) => node.widgets.find((item) => item.name === name);
const apply = (settings) => {
for (const [name, value] of Object.entries(settings)) {
const target = widget(name);
if (!target || value === undefined) continue;
if (target.options?.values && !target.options.values.includes(value)) continue;
target.value = value;
}
node.setDirtyCanvas(true, true);
};
const loaderNode = () => {
const link = app.graph.links[node.inputs?.find((input) => input.name === "pipe")?.link];
return app.graph.getNodeById(link?.origin_id);
};
const model = () => loaderNode()?.promptforgeModelInfo?.();
let showAll = false;
const presetWidget = node.addWidget("combo", "preset", NO_PRESET, (name) => name !== NO_PRESET && apply(presets[name] ?? {}), { values: [NO_PRESET] });
presetWidget.serialize = false;
let presets = {};
const fits = (entry) => {
const current = model();
if (!current || !entry.family) return true;
return entry.family === current.family && (!entry.model || entry.model === current.model);
};
const reload = (data) => {
presets = data;
const names = Object.keys(presets).filter((name) => showAll || fits(presets[name]));
presetWidget.options.values = [NO_PRESET, ...names];
if (presetWidget.value !== NO_PRESET && !names.includes(presetWidget.value)) presetWidget.value = NO_PRESET;
};
api.fetchApi("/promptforge/sampler_presets").then((response) => response.json()).then(reload);
const all = node.addWidget("toggle", "show all presets", false, (value) => { showAll = value; reload(presets); }, {});
node.promptforgeOnPipe = () => { showAll = false; all.value = false; api.fetchApi("/promptforge/sampler_presets").then((response) => response.json()).then(reload); };
all.serialize = false;
const save = node.addWidget("button", "Save as preset", null, async () => {
const name = window.prompt("Preset name", presetWidget.value !== NO_PRESET ? presetWidget.value : "");
if (!name) return;
const current = model() || {};
await send({ name, family: current.family || "", model: current.model || "",
settings: Object.fromEntries(PRESET_FIELDS.map((field) => [field, widget(field).value])) });
presetWidget.value = name;
});
const remove = node.addWidget("button", "Delete preset", null, async () => {
if (presetWidget.value === NO_PRESET) return;
await send({ name: presetWidget.value, delete: true });
presetWidget.value = NO_PRESET;
});
// only ever runs when the button is pressed; the loader is whatever is wired to the pipe input right now
const loader = loaderNode;
const say = (summary, detail) => {
const toast = app.extensionManager?.toast;
if (toast) toast.add({ severity: "info", summary, detail, life: 4000 });
else window.alert(`${summary}\n${detail}`);
};
const recommend = async () => {
const info = loader()?.promptforgeModelInfo?.();
if (!info) return say("PromptForge Sampler", "Connect a PromptForge Loader to the pipe input first.");
const { settings, source } = await (await api.fetchApi("/promptforge/recommend?" + new URLSearchParams(info))).json();
apply(settings);
say("Recommended settings applied", Object.entries(settings).map(([key, value]) => `${key}: ${value} (${source[key]})`).join("\n"));
};
const recommendButton = node.addWidget("button", "Use recommended settings", null, recommend);
[save, remove, recommendButton].forEach((button) => { button.serialize = false; });
};
},
});