Files
ethanfel-ComfyUI-LoRA-Optim…/js/lora_optimizer_bridge.js
T
EthanfelandClaude Opus 4.6 85f9e03f35 feat: add right-click swap between LoRA Optimizer and AutoTuner
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>
2026-03-08 18:04:08 +01:00

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);
};
},
});