新規作成
This commit is contained in:
るみなー先生
@@ -0,0 +1,351 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { api } from "../../scripts/api.js";
|
||||
|
||||
const EXTENSION_NAME = "ruminar.CheckpointWidgetRefresh";
|
||||
const REFRESHER_NODE_NAME = "CheckpointWidgetRefreshPanel";
|
||||
|
||||
let lastCheckpointList = null;
|
||||
|
||||
function getNodeTypeName(node) {
|
||||
return String(
|
||||
node?.type ??
|
||||
node?.comfyClass ??
|
||||
node?.constructor?.type ??
|
||||
""
|
||||
);
|
||||
}
|
||||
|
||||
function getNodeDataName(nodeData) {
|
||||
return String(nodeData?.name ?? "");
|
||||
}
|
||||
|
||||
function isRefresherNode(node) {
|
||||
return getNodeTypeName(node) === REFRESHER_NODE_NAME;
|
||||
}
|
||||
|
||||
function findCheckpointWidget(node) {
|
||||
const widgets = node?.widgets ?? [];
|
||||
return widgets.find((widget) =>
|
||||
widget?.name === "ckpt_name" &&
|
||||
Array.isArray(widget?.options?.values)
|
||||
);
|
||||
}
|
||||
|
||||
function outputNames(node) {
|
||||
return (node?.outputs ?? []).map((output) =>
|
||||
String(output?.name ?? "").toLowerCase()
|
||||
);
|
||||
}
|
||||
|
||||
function hasSelectorOutputs(node) {
|
||||
const names = outputNames(node);
|
||||
return names.includes("ckpt_name") && names.includes("ckpt_name_str");
|
||||
}
|
||||
|
||||
function hasLoaderOutputs(node) {
|
||||
const names = outputNames(node);
|
||||
return names.includes("model") && names.includes("clip") && names.includes("vae");
|
||||
}
|
||||
|
||||
function shouldTargetNode(node) {
|
||||
if (!node || isRefresherNode(node)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const typeName = getNodeTypeName(node);
|
||||
|
||||
if (typeName.includes("CheckpointNameSelector")) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (typeName === "CheckpointLoaderSimple") {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (hasSelectorOutputs(node) && Boolean(findCheckpointWidget(node))) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (hasLoaderOutputs(node) && Boolean(findCheckpointWidget(node))) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
function shouldPatchFutureNodeType(nodeData) {
|
||||
const name = getNodeDataName(nodeData);
|
||||
if (name === REFRESHER_NODE_NAME) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return name.includes("CheckpointNameSelector") || name === "CheckpointLoaderSimple";
|
||||
}
|
||||
|
||||
function chooseReplacement(oldValues, newValues, currentValue) {
|
||||
if (!Array.isArray(newValues) || newValues.length === 0) {
|
||||
return "";
|
||||
}
|
||||
|
||||
if (newValues.includes(currentValue)) {
|
||||
return currentValue;
|
||||
}
|
||||
|
||||
const oldIndex = Array.isArray(oldValues) ? oldValues.indexOf(currentValue) : -1;
|
||||
|
||||
if (oldIndex >= 0) {
|
||||
// Prefer the next surviving item from the old list.
|
||||
for (let index = oldIndex + 1; index < oldValues.length; index++) {
|
||||
const candidate = oldValues[index];
|
||||
if (newValues.includes(candidate)) {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
|
||||
// If the deleted item was at the end, choose the previous survivor.
|
||||
for (let index = oldIndex - 1; index >= 0; index--) {
|
||||
const candidate = oldValues[index];
|
||||
if (newValues.includes(candidate)) {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
|
||||
// Last fallback: same numeric position in the new list if possible.
|
||||
return newValues[Math.min(oldIndex, newValues.length - 1)] ?? newValues[0];
|
||||
}
|
||||
|
||||
return newValues[0];
|
||||
}
|
||||
|
||||
function arraysEqual(left, right) {
|
||||
if (!Array.isArray(left) || !Array.isArray(right)) {
|
||||
return false;
|
||||
}
|
||||
if (left.length !== right.length) {
|
||||
return false;
|
||||
}
|
||||
for (let index = 0; index < left.length; index++) {
|
||||
if (left[index] !== right[index]) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
function patchCheckpointSlotTypes(node, checkpoints) {
|
||||
const values = [...checkpoints];
|
||||
|
||||
for (const output of node.outputs ?? []) {
|
||||
if (String(output?.name ?? "") !== "ckpt_name") {
|
||||
continue;
|
||||
}
|
||||
|
||||
output.type = [...values];
|
||||
|
||||
const links = Array.isArray(output.links) ? output.links : [];
|
||||
for (const linkId of links) {
|
||||
const link = app.graph?.links?.[linkId];
|
||||
if (!link) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const targetNode = app.graph?.getNodeById?.(link.target_id);
|
||||
const targetInput = targetNode?.inputs?.[link.target_slot];
|
||||
if (targetInput && String(targetInput.name ?? "").toLowerCase().includes("ckpt")) {
|
||||
targetInput.type = [...values];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const input of node.inputs ?? []) {
|
||||
if (String(input?.name ?? "").toLowerCase().includes("ckpt")) {
|
||||
input.type = [...values];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function applyCheckpointListToNode(node, checkpoints) {
|
||||
if (!shouldTargetNode(node)) {
|
||||
return {
|
||||
matched: false,
|
||||
changed: false,
|
||||
valueChanged: false,
|
||||
};
|
||||
}
|
||||
|
||||
const widget = findCheckpointWidget(node);
|
||||
if (!widget) {
|
||||
return {
|
||||
matched: true,
|
||||
changed: false,
|
||||
valueChanged: false,
|
||||
reason: "checkpoint widget not found",
|
||||
};
|
||||
}
|
||||
|
||||
if (!widget.options) {
|
||||
widget.options = {};
|
||||
}
|
||||
|
||||
const oldValues = Array.isArray(widget.options.values)
|
||||
? [...widget.options.values]
|
||||
: [];
|
||||
|
||||
const oldValue = widget.value;
|
||||
const newValues = [...checkpoints];
|
||||
const newValue = chooseReplacement(oldValues, newValues, oldValue);
|
||||
|
||||
widget.options.values = newValues;
|
||||
widget.value = newValue;
|
||||
patchCheckpointSlotTypes(node, newValues);
|
||||
|
||||
const valuesChanged = !arraysEqual(oldValues, newValues);
|
||||
const valueChanged = oldValue !== newValue;
|
||||
|
||||
if (valueChanged && typeof widget.callback === "function") {
|
||||
try {
|
||||
widget.callback(widget.value);
|
||||
} catch (error) {
|
||||
console.warn("[CheckpointWidgetRefresh] widget callback failed", error);
|
||||
}
|
||||
}
|
||||
|
||||
if (valuesChanged || valueChanged) {
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
return {
|
||||
matched: true,
|
||||
changed: valuesChanged || valueChanged,
|
||||
valueChanged,
|
||||
oldValue,
|
||||
newValue,
|
||||
};
|
||||
}
|
||||
|
||||
async function fetchCheckpointList() {
|
||||
const response = await api.fetchApi("/checkpoint_widget_refresh/checkpoints", {
|
||||
cache: "no-store",
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}`);
|
||||
}
|
||||
|
||||
const payload = await response.json();
|
||||
return {
|
||||
checkpoints: Array.isArray(payload.checkpoints) ? payload.checkpoints : [],
|
||||
patchedClassCount: payload.patched_class_count ?? 0,
|
||||
patchedClasses: Array.isArray(payload.patched_classes) ? payload.patched_classes : [],
|
||||
};
|
||||
}
|
||||
|
||||
function updatePanelTitle(panelNode, text) {
|
||||
if (!panelNode) {
|
||||
return;
|
||||
}
|
||||
|
||||
panelNode.title = text;
|
||||
panelNode.setDirtyCanvas?.(true, true);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
async function refreshCheckpointWidgets(panelNode = null) {
|
||||
updatePanelTitle(panelNode, "Refreshing checkpoint widgets...");
|
||||
|
||||
const payload = await fetchCheckpointList();
|
||||
const checkpoints = payload.checkpoints;
|
||||
lastCheckpointList = checkpoints;
|
||||
|
||||
let matched = 0;
|
||||
let changed = 0;
|
||||
let valueChanged = 0;
|
||||
|
||||
for (const node of app.graph?._nodes ?? []) {
|
||||
const result = applyCheckpointListToNode(node, checkpoints);
|
||||
if (!result.matched) {
|
||||
continue;
|
||||
}
|
||||
|
||||
matched += 1;
|
||||
if (result.changed) {
|
||||
changed += 1;
|
||||
}
|
||||
if (result.valueChanged) {
|
||||
valueChanged += 1;
|
||||
console.log(
|
||||
`[CheckpointWidgetRefresh] ${getNodeTypeName(node)} #${node.id}: ` +
|
||||
`${result.oldValue} -> ${result.newValue}`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const title = `Updated ${changed}/${matched} ckpt widgets, ${checkpoints.length} ckpts`;
|
||||
updatePanelTitle(panelNode, title);
|
||||
|
||||
console.log(
|
||||
`[CheckpointWidgetRefresh] ${title}. ` +
|
||||
`valueChanged=${valueChanged}, backend patched=${payload.patchedClassCount}`,
|
||||
payload.patchedClasses
|
||||
);
|
||||
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function installRefreshPanelWidgets(node) {
|
||||
if (!node || node._checkpointWidgetRefreshWidgetsInstalled) {
|
||||
return;
|
||||
}
|
||||
|
||||
node.addWidget(
|
||||
"button",
|
||||
"Refresh Checkpoint Widgets",
|
||||
"refresh",
|
||||
async () => {
|
||||
try {
|
||||
await refreshCheckpointWidgets(node);
|
||||
} catch (error) {
|
||||
console.error("[CheckpointWidgetRefresh] refresh failed", error);
|
||||
updatePanelTitle(node, "Refresh failed. See console.");
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
node.title = "Checkpoint Widget Refresh";
|
||||
node.size = [320, 60];
|
||||
node._checkpointWidgetRefreshWidgetsInstalled = true;
|
||||
}
|
||||
|
||||
function installFuturePatchHook(nodeType, nodeData) {
|
||||
if (!shouldPatchFutureNodeType(nodeData)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
|
||||
nodeType.prototype.onNodeCreated = function (...args) {
|
||||
const result = originalOnNodeCreated?.apply(this, args);
|
||||
|
||||
if (getNodeDataName(nodeData) === REFRESHER_NODE_NAME) {
|
||||
installRefreshPanelWidgets(this);
|
||||
return result;
|
||||
}
|
||||
|
||||
if (lastCheckpointList) {
|
||||
setTimeout(() => {
|
||||
applyCheckpointListToNode(this, lastCheckpointList);
|
||||
app.graph?.setDirtyCanvas?.(true, true);
|
||||
}, 0);
|
||||
}
|
||||
|
||||
return result;
|
||||
};
|
||||
}
|
||||
|
||||
app.registerExtension({
|
||||
name: EXTENSION_NAME,
|
||||
|
||||
beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
installFuturePatchHook(nodeType, nodeData);
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user