Adds "Swap to LoRA AutoTuner" / "Swap to LoRA Optimizer" to the right-click context menu. Transfers shared widget values, reconnects all matching inputs/outputs by name, and removes the old node. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
373 lines
12 KiB
JavaScript
373 lines
12 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
|
|
/**
|
|
* AutoTuner ↔ Optimizer Bridge
|
|
*
|
|
* 1. Widget sync: when the Optimizer runs with settings_source="from_autotuner",
|
|
* applied_settings in the UI message updates the Optimizer's knobs.
|
|
*
|
|
* 2. Opposing switches: when tuner_data is connected between an AutoTuner and
|
|
* an Optimizer, their mode switches stay in sync as opposites:
|
|
* Optimizer "from_autotuner" ↔ AutoTuner "merge"
|
|
* Optimizer "manual" ↔ AutoTuner "tuning_only"
|
|
*
|
|
* 3. Conditional visibility: settings_source (Optimizer) and output_mode
|
|
* (AutoTuner) are hidden until tuner_data is connected.
|
|
*/
|
|
|
|
const HIDDEN_TAG = "loraopt_bridge_hidden";
|
|
const _origProps = {};
|
|
|
|
function toggleWidget(node, widget, show) {
|
|
if (!widget) return;
|
|
const key = node.id + ":" + widget.name;
|
|
if (!_origProps[key]) {
|
|
_origProps[key] = {
|
|
origType: widget.type,
|
|
origComputeSize: widget.computeSize,
|
|
};
|
|
}
|
|
widget.hidden = !show;
|
|
widget.type = show ? _origProps[key].origType : HIDDEN_TAG;
|
|
widget.computeSize = show ? _origProps[key].origComputeSize : () => [0, -4];
|
|
}
|
|
|
|
const BRIDGE_WIDGETS = [
|
|
"optimization_mode",
|
|
"merge_quality",
|
|
"sparsification",
|
|
"sparsification_density",
|
|
"dare_dampening",
|
|
"auto_strength",
|
|
];
|
|
|
|
// Mapping between the two switches
|
|
const OPTIMIZER_TO_AUTOTUNER = {
|
|
"from_autotuner": "merge",
|
|
"manual": "tuning_only",
|
|
};
|
|
const AUTOTUNER_TO_OPTIMIZER = {
|
|
"merge": "from_autotuner",
|
|
"tuning_only": "manual",
|
|
};
|
|
|
|
function findWidget(node, name) {
|
|
return node.widgets ? node.widgets.find((w) => w.name === name) : null;
|
|
}
|
|
|
|
/**
|
|
* Find the AutoTuner node connected to this Optimizer's tuner_data input,
|
|
* or the Optimizer node connected to this AutoTuner's tuner_data output.
|
|
*/
|
|
function findConnectedAutoTuner(optimizerNode) {
|
|
// Find the tuner_data input slot index
|
|
if (!optimizerNode.inputs) return null;
|
|
for (let i = 0; i < optimizerNode.inputs.length; i++) {
|
|
const input = optimizerNode.inputs[i];
|
|
if (input.name === "tuner_data" && input.link != null) {
|
|
const linkInfo = app.graph.links[input.link];
|
|
if (linkInfo) {
|
|
const sourceNode = app.graph.getNodeById(linkInfo.origin_id);
|
|
if (sourceNode && sourceNode.comfyClass === "LoRAAutoTuner") {
|
|
return sourceNode;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
function findConnectedOptimizer(autotunerNode) {
|
|
// Find tuner_data output slot
|
|
if (!autotunerNode.outputs) return null;
|
|
for (let i = 0; i < autotunerNode.outputs.length; i++) {
|
|
const output = autotunerNode.outputs[i];
|
|
if (output.name === "tuner_data" && output.links) {
|
|
for (const linkId of output.links) {
|
|
const linkInfo = app.graph.links[linkId];
|
|
if (linkInfo) {
|
|
const targetNode = app.graph.getNodeById(linkInfo.target_id);
|
|
if (targetNode && targetNode.comfyClass === "LoRAOptimizer") {
|
|
return targetNode;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
/**
|
|
* Check if the Optimizer's tuner_data input slot is connected.
|
|
*/
|
|
function isTunerDataInputConnected(optimizerNode) {
|
|
if (!optimizerNode.inputs) return false;
|
|
for (const input of optimizerNode.inputs) {
|
|
if (input.name === "tuner_data" && input.link != null) return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
/**
|
|
* Check if the AutoTuner's tuner_data output slot has any connections.
|
|
*/
|
|
function isTunerDataOutputConnected(autotunerNode) {
|
|
if (!autotunerNode.outputs) return false;
|
|
for (const output of autotunerNode.outputs) {
|
|
if (output.name === "tuner_data" && output.links && output.links.length > 0) return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
function resizeNode(node) {
|
|
const computed = node.computeSize();
|
|
node.setSize([Math.max(node.size[0], computed[0]), computed[1]]);
|
|
app.canvas?.setDirty?.(true, true);
|
|
}
|
|
|
|
function updateOptimizerVisibility(node) {
|
|
const w = findWidget(node, "settings_source");
|
|
if (!w) return;
|
|
const connected = isTunerDataInputConnected(node);
|
|
toggleWidget(node, w, connected);
|
|
if (!connected) w.value = "manual";
|
|
resizeNode(node);
|
|
}
|
|
|
|
function updateAutoTunerVisibility(node) {
|
|
const w = findWidget(node, "output_mode");
|
|
if (!w) return;
|
|
const connected = isTunerDataOutputConnected(node);
|
|
toggleWidget(node, w, connected);
|
|
if (!connected) w.value = "merge";
|
|
resizeNode(node);
|
|
}
|
|
|
|
// Guard against infinite recursion during sync
|
|
let _syncing = false;
|
|
|
|
function syncOppositeSwitch(sourceWidget, sourceNode, targetFinder, mapping) {
|
|
if (_syncing) return;
|
|
const targetNode = targetFinder(sourceNode);
|
|
if (!targetNode) return;
|
|
|
|
const targetWidgetName = sourceWidget.name === "settings_source" ? "output_mode" : "settings_source";
|
|
const targetWidget = findWidget(targetNode, targetWidgetName);
|
|
if (!targetWidget) return;
|
|
|
|
const expectedValue = mapping[sourceWidget.value];
|
|
if (expectedValue && targetWidget.value !== expectedValue) {
|
|
_syncing = true;
|
|
try {
|
|
targetWidget.value = expectedValue;
|
|
} finally {
|
|
_syncing = false;
|
|
}
|
|
app.canvas?.setDirty?.(true, true);
|
|
}
|
|
}
|
|
|
|
// --- Node Swap: right-click "Swap to AutoTuner / Optimizer" ---
|
|
|
|
// Widgets that exist on both nodes (values transfer directly on swap)
|
|
const SHARED_WIDGETS = [
|
|
"output_strength", "clip_strength_multiplier", "normalize_keys",
|
|
"architecture_preset", "auto_strength_floor", "cache_patches", "vram_budget",
|
|
];
|
|
|
|
const SWAP_PAIRS = {
|
|
"LoRAOptimizer": "LoRAAutoTuner",
|
|
"LoRAAutoTuner": "LoRAOptimizer",
|
|
};
|
|
|
|
const SWAP_LABELS = {
|
|
"LoRAOptimizer": "Swap to LoRA AutoTuner",
|
|
"LoRAAutoTuner": "Swap to LoRA Optimizer",
|
|
};
|
|
|
|
// Output slot name mapping: src output name → dst output name
|
|
// Optimizer: model, clip, analysis_report, tuner_data, lora_data
|
|
// AutoTuner: model, clip, report, analysis_report, tuner_data, lora_data
|
|
// We match by name; "report" (AutoTuner-only) has no Optimizer equivalent.
|
|
// Input slots match by name (model, lora_stack, clip, etc.).
|
|
|
|
function swapNode(oldNode) {
|
|
const targetClass = SWAP_PAIRS[oldNode.comfyClass];
|
|
if (!targetClass) return;
|
|
|
|
const newNode = LiteGraph.createNode(targetClass);
|
|
if (!newNode) return;
|
|
|
|
app.graph.add(newNode);
|
|
|
|
// Copy position and size
|
|
newNode.pos[0] = oldNode.pos[0];
|
|
newNode.pos[1] = oldNode.pos[1];
|
|
newNode.size[0] = Math.max(oldNode.size[0], newNode.size[0]);
|
|
|
|
// Transfer shared widget values
|
|
for (const name of SHARED_WIDGETS) {
|
|
const src = findWidget(oldNode, name);
|
|
const dst = findWidget(newNode, name);
|
|
if (src && dst) {
|
|
dst.value = src.value;
|
|
}
|
|
}
|
|
|
|
// Reconnect inputs by matching slot name
|
|
if (oldNode.inputs) {
|
|
for (let i = 0; i < oldNode.inputs.length; i++) {
|
|
const oldInput = oldNode.inputs[i];
|
|
if (oldInput.link == null) continue;
|
|
const linkInfo = app.graph.links[oldInput.link];
|
|
if (!linkInfo) continue;
|
|
|
|
// Find matching input slot on new node by name
|
|
const newSlotIdx = newNode.inputs
|
|
? newNode.inputs.findIndex((inp) => inp.name === oldInput.name)
|
|
: -1;
|
|
if (newSlotIdx === -1) continue;
|
|
|
|
const sourceNode = app.graph.getNodeById(linkInfo.origin_id);
|
|
if (!sourceNode) continue;
|
|
|
|
sourceNode.connect(linkInfo.origin_slot, newNode, newSlotIdx);
|
|
}
|
|
}
|
|
|
|
// Reconnect outputs by matching slot name
|
|
if (oldNode.outputs) {
|
|
for (let i = 0; i < oldNode.outputs.length; i++) {
|
|
const oldOutput = oldNode.outputs[i];
|
|
if (!oldOutput.links || oldOutput.links.length === 0) continue;
|
|
|
|
// Find matching output slot on new node by name
|
|
const newSlotIdx = newNode.outputs
|
|
? newNode.outputs.findIndex((out) => out.name === oldOutput.name)
|
|
: -1;
|
|
if (newSlotIdx === -1) continue;
|
|
|
|
// Iterate a copy — link list mutates as we disconnect/reconnect
|
|
for (const linkId of [...oldOutput.links]) {
|
|
const linkInfo = app.graph.links[linkId];
|
|
if (!linkInfo) continue;
|
|
const targetNode = app.graph.getNodeById(linkInfo.target_id);
|
|
if (!targetNode) continue;
|
|
|
|
newNode.connect(newSlotIdx, targetNode, linkInfo.target_slot);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Remove old node
|
|
app.graph.remove(oldNode);
|
|
|
|
// Resize new node to fit widgets
|
|
setTimeout(() => {
|
|
const computed = newNode.computeSize();
|
|
newNode.setSize([Math.max(newNode.size[0], computed[0]), computed[1]]);
|
|
app.canvas?.setDirty?.(true, true);
|
|
}, 0);
|
|
}
|
|
|
|
function addSwapMenuOption(node) {
|
|
const label = SWAP_LABELS[node.comfyClass];
|
|
if (!label) return;
|
|
|
|
const origGetExtra = node.getExtraMenuOptions;
|
|
node.getExtraMenuOptions = function (_, options) {
|
|
if (origGetExtra) origGetExtra.apply(this, arguments);
|
|
|
|
// Insert before the last item (which is usually "Properties")
|
|
options.splice(-1, 0, null, {
|
|
content: label,
|
|
callback: () => swapNode(this),
|
|
});
|
|
};
|
|
}
|
|
|
|
// --- LoRAOptimizer extension ---
|
|
app.registerExtension({
|
|
name: "LoRAOptimizer.AutoTunerBridge",
|
|
nodeCreated(node) {
|
|
if (node.comfyClass !== "LoRAOptimizer") return;
|
|
|
|
addSwapMenuOption(node);
|
|
|
|
// Widget sync on execution (applied_settings from Python)
|
|
const origOnExecuted = node.onExecuted;
|
|
node.onExecuted = function (message) {
|
|
if (origOnExecuted) {
|
|
origOnExecuted.call(this, message);
|
|
}
|
|
|
|
if (!message || !message.applied_settings || !message.applied_settings[0]) {
|
|
return;
|
|
}
|
|
|
|
let config;
|
|
try {
|
|
config = JSON.parse(message.applied_settings[0]);
|
|
} catch {
|
|
return;
|
|
}
|
|
|
|
for (const name of BRIDGE_WIDGETS) {
|
|
const widget = findWidget(this, name);
|
|
if (widget && config[name] !== undefined) {
|
|
widget.value = config[name];
|
|
}
|
|
}
|
|
|
|
app.canvas?.setDirty?.(true, true);
|
|
};
|
|
|
|
// Opposing switch sync: settings_source → output_mode
|
|
const settingsWidget = findWidget(node, "settings_source");
|
|
if (settingsWidget) {
|
|
const origCallback = settingsWidget.callback;
|
|
settingsWidget.callback = function (value) {
|
|
if (origCallback) origCallback.call(this, value);
|
|
syncOppositeSwitch(settingsWidget, node, findConnectedAutoTuner, OPTIMIZER_TO_AUTOTUNER);
|
|
};
|
|
}
|
|
|
|
// Hide settings_source until tuner_data is connected
|
|
// Defer initial hide so node layout is finalized first
|
|
setTimeout(() => updateOptimizerVisibility(node), 0);
|
|
const origOnConnChange = node.onConnectionsChange;
|
|
node.onConnectionsChange = function (side, slot, connected, linkInfo, ioSlot) {
|
|
if (origOnConnChange) origOnConnChange.apply(this, arguments);
|
|
if (side === 1) updateOptimizerVisibility(this);
|
|
};
|
|
},
|
|
});
|
|
|
|
// --- LoRAAutoTuner extension ---
|
|
app.registerExtension({
|
|
name: "LoRAOptimizer.AutoTunerOutputMode",
|
|
nodeCreated(node) {
|
|
if (node.comfyClass !== "LoRAAutoTuner") return;
|
|
|
|
addSwapMenuOption(node);
|
|
|
|
// Opposing switch sync: output_mode → settings_source
|
|
const outputModeWidget = findWidget(node, "output_mode");
|
|
if (outputModeWidget) {
|
|
const origCallback = outputModeWidget.callback;
|
|
outputModeWidget.callback = function (value) {
|
|
if (origCallback) origCallback.call(this, value);
|
|
syncOppositeSwitch(outputModeWidget, node, findConnectedOptimizer, AUTOTUNER_TO_OPTIMIZER);
|
|
};
|
|
}
|
|
|
|
// Hide output_mode until tuner_data output is connected
|
|
setTimeout(() => updateAutoTunerVisibility(node), 0);
|
|
const origOnConnChange = node.onConnectionsChange;
|
|
node.onConnectionsChange = function (side, slot, connected, linkInfo, ioSlot) {
|
|
if (origOnConnChange) origOnConnChange.apply(this, arguments);
|
|
if (side === 2) updateAutoTunerVisibility(this);
|
|
};
|
|
},
|
|
});
|