475 lines
16 KiB
JavaScript
475 lines
16 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
import { ENDPOINTS } from "./constants.js";
|
|
|
|
const NODE_CLASS = "DistributedValue";
|
|
const CONVERTED_WIDGET = "converted-widget";
|
|
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
|
|
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
|
|
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
|
|
|
|
const trackedNodes = new Set();
|
|
let workersChangedListenerAttached = false;
|
|
|
|
function filterEnabledWorkers(workers) {
|
|
if (!Array.isArray(workers)) return [];
|
|
return workers.filter((worker) => Boolean(worker?.enabled));
|
|
}
|
|
|
|
async function fetchWorkers() {
|
|
try {
|
|
const resp = await fetch(ENDPOINTS.CONFIG);
|
|
if (!resp.ok) return [];
|
|
const config = await resp.json();
|
|
return filterEnabledWorkers(config.workers);
|
|
} catch {
|
|
return [];
|
|
}
|
|
}
|
|
|
|
function getRawDefaultWidget(node) {
|
|
return node.widgets?.find((w) => w.name === "default_value");
|
|
}
|
|
|
|
function getRawWorkerValuesWidget(node) {
|
|
return node.widgets?.find((w) => w.name === "worker_values");
|
|
}
|
|
|
|
function getDynamicDefaultWidget(node) {
|
|
return node.widgets?.find((w) => w.name === DYNAMIC_DEFAULT_WIDGET);
|
|
}
|
|
|
|
function getDynamicWorkerWidgets(node) {
|
|
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
|
|
}
|
|
|
|
function hideWidgetForGood(node, widget, suffix = "") {
|
|
if (!widget) return;
|
|
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
|
|
|
|
widget.origType = widget.type;
|
|
widget.origComputeSize = widget.computeSize;
|
|
widget.origSerializeValue = widget.serializeValue;
|
|
widget.computeSize = () => [0, -4];
|
|
widget.type = `${CONVERTED_WIDGET}${suffix}`;
|
|
|
|
// Hide any attached DOM element (multiline widgets).
|
|
if (widget.element) widget.element.style.display = "none";
|
|
if (widget.inputEl) widget.inputEl.style.display = "none";
|
|
|
|
if (widget.linkedWidgets) {
|
|
for (const linked of widget.linkedWidgets) {
|
|
hideWidgetForGood(node, linked, `:${widget.name}`);
|
|
}
|
|
}
|
|
}
|
|
|
|
function hideRawWidgets(node) {
|
|
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
|
|
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
|
|
}
|
|
|
|
function removeDynamicDefaultWidget(node) {
|
|
const idx = node.widgets?.findIndex((w) => w.name === DYNAMIC_DEFAULT_WIDGET);
|
|
if (idx != null && idx >= 0) {
|
|
node.widgets.splice(idx, 1);
|
|
}
|
|
}
|
|
|
|
function removeDynamicWorkerWidgets(node) {
|
|
if (!node.widgets) return;
|
|
for (let i = node.widgets.length - 1; i >= 0; i--) {
|
|
if (node.widgets[i].name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX)) {
|
|
node.widgets.splice(i, 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
function readWorkerStore(node) {
|
|
const raw = getRawWorkerValuesWidget(node);
|
|
if (!raw) return {};
|
|
try {
|
|
const parsed = JSON.parse(raw.value || "{}");
|
|
return typeof parsed === "object" && parsed !== null ? parsed : {};
|
|
} catch {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
function writeWorkerStore(node, store) {
|
|
const raw = getRawWorkerValuesWidget(node);
|
|
if (!raw) return;
|
|
raw.value = JSON.stringify(store);
|
|
}
|
|
|
|
function normalizeComboOptions(options) {
|
|
if (!options) return null;
|
|
if (Array.isArray(options)) return options;
|
|
if (Array.isArray(options.values)) return options.values;
|
|
return null;
|
|
}
|
|
|
|
function resolveGraphLink(graph, linkId) {
|
|
const links = graph.links || graph._links;
|
|
if (!links) return null;
|
|
const link = links[linkId] ?? (typeof links.get === "function" ? links.get(linkId) : null);
|
|
if (!link) return null;
|
|
if (Array.isArray(link)) {
|
|
return {
|
|
target_id: link[2],
|
|
target_slot: link[3],
|
|
};
|
|
}
|
|
return link;
|
|
}
|
|
|
|
function detectTargetType(node) {
|
|
const out = node.outputs?.[0];
|
|
const linkIds = out?.links || [];
|
|
if (!linkIds.length) {
|
|
return { connected: false, type: "STRING", options: null };
|
|
}
|
|
|
|
const graph = node.graph || app.graph;
|
|
if (!graph) {
|
|
return { connected: false, type: "STRING", options: null };
|
|
}
|
|
|
|
const link = resolveGraphLink(graph, linkIds[0]);
|
|
if (!link) {
|
|
return { connected: false, type: "STRING", options: null };
|
|
}
|
|
|
|
const targetNode = graph.getNodeById(link.target_id);
|
|
if (!targetNode) {
|
|
return { connected: false, type: "STRING", options: null };
|
|
}
|
|
|
|
const targetInputName = targetNode.inputs?.[link.target_slot]?.name;
|
|
if (!targetInputName) {
|
|
return { connected: false, type: "STRING", options: null };
|
|
}
|
|
|
|
const targetWidget = targetNode.widgets?.find((w) => w.name === targetInputName);
|
|
if (targetWidget) {
|
|
if (targetWidget.type === "combo") {
|
|
const comboOptions = normalizeComboOptions(targetWidget.options);
|
|
return { connected: true, type: "COMBO", options: comboOptions };
|
|
}
|
|
if (targetWidget.type === "number") {
|
|
const step = targetWidget.options?.step;
|
|
const precision = targetWidget.options?.precision;
|
|
const isInt = Number.isInteger(step) && (precision === 0 || precision == null);
|
|
return { connected: true, type: isInt ? "INT" : "FLOAT", options: null };
|
|
}
|
|
}
|
|
|
|
const nodeDef = targetNode.constructor?.nodeData;
|
|
const inputDef = nodeDef?.input?.required?.[targetInputName] || nodeDef?.input?.optional?.[targetInputName];
|
|
if (inputDef) {
|
|
const defType = inputDef[0];
|
|
if (Array.isArray(defType)) {
|
|
return { connected: true, type: "COMBO", options: defType };
|
|
}
|
|
if (defType === "INT") return { connected: true, type: "INT", options: null };
|
|
if (defType === "FLOAT") return { connected: true, type: "FLOAT", options: null };
|
|
}
|
|
|
|
return { connected: true, type: "STRING", options: null };
|
|
}
|
|
|
|
function normalizeNumber(value, fallback) {
|
|
const parsed = Number(value);
|
|
return Number.isFinite(parsed) ? parsed : fallback;
|
|
}
|
|
|
|
function getDefaultInitialValue(node, inputType, comboOptions) {
|
|
const rawDefault = getRawDefaultWidget(node);
|
|
const current = rawDefault?.value;
|
|
|
|
if (inputType === "INT") {
|
|
return Math.trunc(normalizeNumber(current, 0));
|
|
}
|
|
if (inputType === "FLOAT") {
|
|
return normalizeNumber(current, 0);
|
|
}
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
|
|
const currentText = current == null ? "" : String(current);
|
|
return comboOptions.includes(currentText) ? currentText : comboOptions[0];
|
|
}
|
|
return current == null ? "" : String(current);
|
|
}
|
|
|
|
function setRawDefaultValue(node, value) {
|
|
const rawDefault = getRawDefaultWidget(node);
|
|
if (!rawDefault) return;
|
|
rawDefault.value = value;
|
|
}
|
|
|
|
function serializeWorkerStoreFromWidgets(node, inputType, comboOptions) {
|
|
const nextStore = { _type: inputType };
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions)) {
|
|
nextStore._options = comboOptions;
|
|
}
|
|
const valuesByWorkerId = {};
|
|
|
|
for (const widget of getDynamicWorkerWidgets(node)) {
|
|
const key = widget.name.slice(DYNAMIC_WORKER_WIDGET_PREFIX.length);
|
|
if (widget.value !== "" && widget.value !== null && widget.value !== undefined) {
|
|
const value = String(widget.value);
|
|
nextStore[key] = value;
|
|
if (widget._dvWorkerId) {
|
|
valuesByWorkerId[widget._dvWorkerId] = value;
|
|
}
|
|
}
|
|
}
|
|
if (Object.keys(valuesByWorkerId).length) {
|
|
nextStore._by_worker_id = valuesByWorkerId;
|
|
}
|
|
|
|
writeWorkerStore(node, nextStore);
|
|
}
|
|
|
|
function updateWorkerStoreTypeMetadata(node, inputType, comboOptions) {
|
|
const store = readWorkerStore(node);
|
|
store._type = inputType;
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions)) {
|
|
store._options = comboOptions;
|
|
} else {
|
|
delete store._options;
|
|
}
|
|
writeWorkerStore(node, store);
|
|
}
|
|
|
|
function createDynamicDefaultWidget(node, inputType, comboOptions) {
|
|
removeDynamicDefaultWidget(node);
|
|
const initial = getDefaultInitialValue(node, inputType, comboOptions);
|
|
let widget;
|
|
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
|
|
widget = node.addWidget(
|
|
"combo",
|
|
DYNAMIC_DEFAULT_WIDGET,
|
|
initial,
|
|
(value) => {
|
|
widget.value = value;
|
|
setRawDefaultValue(node, String(value));
|
|
},
|
|
{ values: comboOptions }
|
|
);
|
|
} else if (inputType === "INT") {
|
|
widget = node.addWidget(
|
|
"number",
|
|
DYNAMIC_DEFAULT_WIDGET,
|
|
initial,
|
|
(value) => {
|
|
widget.value = Math.trunc(normalizeNumber(value, 0));
|
|
setRawDefaultValue(node, widget.value);
|
|
},
|
|
{ min: -Infinity, max: Infinity, step: 1, precision: 0 }
|
|
);
|
|
} else if (inputType === "FLOAT") {
|
|
widget = node.addWidget(
|
|
"number",
|
|
DYNAMIC_DEFAULT_WIDGET,
|
|
initial,
|
|
(value) => {
|
|
widget.value = normalizeNumber(value, 0);
|
|
setRawDefaultValue(node, widget.value);
|
|
},
|
|
{ min: -Infinity, max: Infinity, step: 0.1, precision: 3 }
|
|
);
|
|
} else {
|
|
widget = node.addWidget(
|
|
"string",
|
|
DYNAMIC_DEFAULT_WIDGET,
|
|
initial,
|
|
(value) => {
|
|
widget.value = value ?? "";
|
|
setRawDefaultValue(node, widget.value);
|
|
},
|
|
{}
|
|
);
|
|
}
|
|
|
|
widget.label = "default_value";
|
|
}
|
|
|
|
function getWorkerInitialValue(store, key, workerId, inputType, comboOptions) {
|
|
const byWorkerId = store?._by_worker_id;
|
|
const saved = (byWorkerId && workerId && byWorkerId[workerId] != null)
|
|
? byWorkerId[workerId]
|
|
: store[key];
|
|
if (saved == null) {
|
|
if (inputType === "INT" || inputType === "FLOAT") return 0;
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
|
|
return comboOptions[0];
|
|
}
|
|
return "";
|
|
}
|
|
|
|
if (inputType === "INT") return Math.trunc(normalizeNumber(saved, 0));
|
|
if (inputType === "FLOAT") return normalizeNumber(saved, 0);
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
|
|
const savedText = String(saved);
|
|
return comboOptions.includes(savedText) ? savedText : comboOptions[0];
|
|
}
|
|
return String(saved);
|
|
}
|
|
|
|
function createWorkerWidgets(node, workers, inputType, comboOptions) {
|
|
removeDynamicWorkerWidgets(node);
|
|
const store = readWorkerStore(node);
|
|
|
|
for (let i = 0; i < workers.length; i++) {
|
|
const key = String(i + 1);
|
|
const worker = workers[i];
|
|
const label = worker.name || worker.id || `Worker ${key}`;
|
|
const widgetName = `${DYNAMIC_WORKER_WIDGET_PREFIX}${key}`;
|
|
const initial = getWorkerInitialValue(store, key, worker.id, inputType, comboOptions);
|
|
let widget;
|
|
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length) {
|
|
widget = node.addWidget(
|
|
"combo",
|
|
widgetName,
|
|
initial,
|
|
(value) => {
|
|
widget.value = value;
|
|
serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
|
|
},
|
|
{ values: comboOptions }
|
|
);
|
|
} else if (inputType === "INT") {
|
|
widget = node.addWidget(
|
|
"number",
|
|
widgetName,
|
|
initial,
|
|
(value) => {
|
|
widget.value = Math.trunc(normalizeNumber(value, 0));
|
|
serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
|
|
},
|
|
{ min: -Infinity, max: Infinity, step: 1, precision: 0 }
|
|
);
|
|
} else if (inputType === "FLOAT") {
|
|
widget = node.addWidget(
|
|
"number",
|
|
widgetName,
|
|
initial,
|
|
(value) => {
|
|
widget.value = normalizeNumber(value, 0);
|
|
serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
|
|
},
|
|
{ min: -Infinity, max: Infinity, step: 0.1, precision: 3 }
|
|
);
|
|
} else {
|
|
widget = node.addWidget(
|
|
"string",
|
|
widgetName,
|
|
initial,
|
|
(value) => {
|
|
widget.value = value ?? "";
|
|
serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
|
|
},
|
|
{}
|
|
);
|
|
}
|
|
|
|
widget.label = label;
|
|
widget._dvWorkerId = worker.id;
|
|
}
|
|
|
|
serializeWorkerStoreFromWidgets(node, inputType, comboOptions);
|
|
}
|
|
|
|
function rebuildWidgets(node) {
|
|
hideRawWidgets(node);
|
|
const workers = node._dvWorkers || [];
|
|
const store = readWorkerStore(node);
|
|
const detected = detectTargetType(node);
|
|
const inputType = detected.connected ? detected.type : (store._type || "STRING");
|
|
const comboOptions = detected.connected ? detected.options : normalizeComboOptions(store._options);
|
|
|
|
createDynamicDefaultWidget(node, inputType, comboOptions);
|
|
if (workers.length > 0) {
|
|
createWorkerWidgets(node, workers, inputType, comboOptions);
|
|
} else {
|
|
removeDynamicWorkerWidgets(node);
|
|
updateWorkerStoreTypeMetadata(node, inputType, comboOptions);
|
|
}
|
|
|
|
const size = node.computeSize();
|
|
size[0] = Math.max(size[0], node.size?.[0] || 0);
|
|
node.setSize(size);
|
|
if (node.setDirtyCanvas) node.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
function refreshNodeWorkers(node, workers) {
|
|
if (!node || !node.graph) return;
|
|
node._dvWorkers = workers;
|
|
rebuildWidgets(node);
|
|
}
|
|
|
|
async function refreshTrackedNodes(workers = null) {
|
|
const nextWorkers = workers || (await fetchWorkers());
|
|
for (const node of trackedNodes) {
|
|
refreshNodeWorkers(node, nextWorkers);
|
|
}
|
|
}
|
|
|
|
function attachWorkersChangedListener() {
|
|
if (workersChangedListenerAttached) return;
|
|
if (typeof window === "undefined" || typeof window.addEventListener !== "function") return;
|
|
|
|
window.addEventListener(WORKERS_CHANGED_EVENT, (event) => {
|
|
const changedWorkers = filterEnabledWorkers(event?.detail?.workers);
|
|
if (changedWorkers.length > 0 || Array.isArray(event?.detail?.workers)) {
|
|
void refreshTrackedNodes(changedWorkers);
|
|
return;
|
|
}
|
|
void refreshTrackedNodes();
|
|
});
|
|
|
|
workersChangedListenerAttached = true;
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "Distributed.DistributedValue",
|
|
async nodeCreated(node) {
|
|
if (node.comfyClass !== NODE_CLASS) return;
|
|
|
|
try {
|
|
attachWorkersChangedListener();
|
|
trackedNodes.add(node);
|
|
node._dvWorkers = await fetchWorkers();
|
|
rebuildWidgets(node);
|
|
|
|
const originalOnConnectionsChange = node.onConnectionsChange;
|
|
node.onConnectionsChange = function (type, index, connected, linkInfo, ioSlot) {
|
|
if (originalOnConnectionsChange) {
|
|
originalOnConnectionsChange.call(this, type, index, connected, linkInfo, ioSlot);
|
|
}
|
|
if (type === 2 && index === 0) {
|
|
setTimeout(() => rebuildWidgets(this), 20);
|
|
}
|
|
};
|
|
|
|
const originalConfigure = node.configure;
|
|
node.configure = function (data) {
|
|
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
|
|
setTimeout(() => rebuildWidgets(this), 20);
|
|
return result;
|
|
};
|
|
|
|
const originalOnRemoved = node.onRemoved;
|
|
node.onRemoved = function () {
|
|
trackedNodes.delete(this);
|
|
if (originalOnRemoved) {
|
|
return originalOnRemoved.call(this);
|
|
}
|
|
};
|
|
} catch (error) {
|
|
console.error("Error in DistributedValue extension:", error);
|
|
}
|
|
},
|
|
});
|