317 lines
11 KiB
JavaScript
317 lines
11 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
import { ENDPOINTS } from "./constants.js";
|
|
|
|
const NODE_CLASS = "DistributedValue";
|
|
const WORKER_WIDGET_PREFIX = "_dv_worker_";
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Config helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
async function fetchWorkers() {
|
|
try {
|
|
const resp = await fetch(ENDPOINTS.CONFIG);
|
|
if (!resp.ok) return [];
|
|
const config = await resp.json();
|
|
return config.workers || [];
|
|
} catch {
|
|
return [];
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Hidden widget accessors (worker_values JSON storage)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function getWorkerValuesWidget(node) {
|
|
return node.widgets?.find((w) => w.name === "worker_values");
|
|
}
|
|
|
|
function getHiddenWidget(node) {
|
|
return node._dvHiddenWidget || getWorkerValuesWidget(node);
|
|
}
|
|
|
|
function getWorkerWidgets(node) {
|
|
return (node.widgets || []).filter((w) => w.name.startsWith(WORKER_WIDGET_PREFIX));
|
|
}
|
|
|
|
function readStore(node) {
|
|
const w = getHiddenWidget(node);
|
|
if (!w) return {};
|
|
try {
|
|
const parsed = JSON.parse(w.value || "{}");
|
|
return typeof parsed === "object" && parsed !== null ? parsed : {};
|
|
} catch {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
function writeStore(node, patch) {
|
|
const w = getHiddenWidget(node);
|
|
if (!w) return;
|
|
const store = readStore(node);
|
|
Object.assign(store, patch);
|
|
w.value = JSON.stringify(store);
|
|
}
|
|
|
|
function serializeWorkerValues(node) {
|
|
const store = readStore(node);
|
|
// Keep metadata keys (_type, _options) but overwrite worker values
|
|
const newStore = { _type: store._type || "STRING" };
|
|
if (store._options) newStore._options = store._options;
|
|
for (const w of getWorkerWidgets(node)) {
|
|
const key = w.name.slice(WORKER_WIDGET_PREFIX.length);
|
|
const val = w.value;
|
|
if (val !== undefined && val !== null && val !== "") {
|
|
newStore[key] = String(val);
|
|
}
|
|
}
|
|
const hidden = getHiddenWidget(node);
|
|
if (hidden) hidden.value = JSON.stringify(newStore);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Hide the raw worker_values widget
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function hideRawWidget(node) {
|
|
const w = getWorkerValuesWidget(node);
|
|
if (!w) return;
|
|
|
|
// Hide any DOM elements (multiline creates a textarea)
|
|
if (w.element) w.element.style.display = "none";
|
|
if (w.inputEl) w.inputEl.style.display = "none";
|
|
|
|
// Remove from visible widgets array
|
|
const idx = node.widgets.indexOf(w);
|
|
if (idx !== -1) node.widgets.splice(idx, 1);
|
|
|
|
// Keep reference
|
|
node._dvHiddenWidget = w;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Connection type detection
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function detectTargetType(node) {
|
|
// Find what our output (slot 0) is connected to
|
|
if (!node.outputs || !node.outputs[0] || !node.outputs[0].links) {
|
|
return { type: "STRING", options: null };
|
|
}
|
|
const linkIds = node.outputs[0].links;
|
|
if (!linkIds || linkIds.length === 0) {
|
|
return { type: "STRING", options: null };
|
|
}
|
|
|
|
const graph = node.graph || app.graph;
|
|
if (!graph) return { type: "STRING", options: null };
|
|
|
|
// Use the first connection to determine type
|
|
const linkId = linkIds[0];
|
|
const link = graph.links?.[linkId] || graph._links?.get?.(linkId);
|
|
if (!link) return { type: "STRING", options: null };
|
|
|
|
const targetNode = graph.getNodeById(link.target_id);
|
|
if (!targetNode) return { type: "STRING", options: null };
|
|
|
|
const targetInputName = targetNode.inputs?.[link.target_slot]?.name;
|
|
if (!targetInputName) return { type: "STRING", options: null };
|
|
|
|
// Find the corresponding widget on the target node to get type info
|
|
const targetWidget = targetNode.widgets?.find((w) => w.name === targetInputName);
|
|
if (!targetWidget) {
|
|
// Check node definition for input type
|
|
const nodeDef = targetNode.constructor?.nodeData;
|
|
if (nodeDef) {
|
|
const inputDef =
|
|
nodeDef.input?.required?.[targetInputName] ||
|
|
nodeDef.input?.optional?.[targetInputName];
|
|
if (inputDef) {
|
|
if (Array.isArray(inputDef[0])) {
|
|
return { type: "COMBO", options: inputDef[0] };
|
|
}
|
|
const typeName = inputDef[0];
|
|
if (typeName === "INT") return { type: "INT", options: null };
|
|
if (typeName === "FLOAT") return { type: "FLOAT", options: null };
|
|
}
|
|
}
|
|
return { type: "STRING", options: null };
|
|
}
|
|
|
|
// Detect type from widget
|
|
if (targetWidget.type === "combo") {
|
|
const options = targetWidget.options?.values || [];
|
|
return { type: "COMBO", options };
|
|
}
|
|
if (targetWidget.type === "number") {
|
|
// Distinguish INT vs FLOAT from the widget options
|
|
const step = targetWidget.options?.step;
|
|
if (step !== undefined && step % 1 === 0 && (!targetWidget.options?.precision || targetWidget.options.precision === 0)) {
|
|
return { type: "INT", options: null };
|
|
}
|
|
return { type: "FLOAT", options: null };
|
|
}
|
|
|
|
return { type: "STRING", options: null };
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Worker widget creation
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function removeWorkerWidgets(node) {
|
|
for (let i = node.widgets.length - 1; i >= 0; i--) {
|
|
if (node.widgets[i].name.startsWith(WORKER_WIDGET_PREFIX)) {
|
|
node.widgets.splice(i, 1);
|
|
}
|
|
}
|
|
}
|
|
|
|
function createWorkerWidgets(node, workers, inputType, comboOptions) {
|
|
removeWorkerWidgets(node);
|
|
|
|
const store = readStore(node);
|
|
|
|
for (let i = 0; i < workers.length; i++) {
|
|
const worker = workers[i];
|
|
const key = String(i + 1);
|
|
const widgetName = `${WORKER_WIDGET_PREFIX}${key}`;
|
|
const label = worker.name || worker.id || `Worker ${key}`;
|
|
const savedValue = store[key] || "";
|
|
|
|
let widget;
|
|
if (inputType === "COMBO" && Array.isArray(comboOptions) && comboOptions.length > 0) {
|
|
widget = node.addWidget(
|
|
"combo",
|
|
widgetName,
|
|
savedValue || comboOptions[0],
|
|
(value) => {
|
|
widget.value = value;
|
|
serializeWorkerValues(node);
|
|
},
|
|
{ values: comboOptions }
|
|
);
|
|
} else if (inputType === "INT") {
|
|
widget = node.addWidget(
|
|
"number",
|
|
widgetName,
|
|
savedValue ? Number(savedValue) : 0,
|
|
(value) => {
|
|
widget.value = value;
|
|
serializeWorkerValues(node);
|
|
},
|
|
{ min: -Infinity, max: Infinity, step: 1, precision: 0 }
|
|
);
|
|
} else if (inputType === "FLOAT") {
|
|
widget = node.addWidget(
|
|
"number",
|
|
widgetName,
|
|
savedValue ? Number(savedValue) : 0.0,
|
|
(value) => {
|
|
widget.value = value;
|
|
serializeWorkerValues(node);
|
|
},
|
|
{ min: -Infinity, max: Infinity, step: 0.1, precision: 3 }
|
|
);
|
|
} else {
|
|
// STRING (default)
|
|
widget = node.addWidget(
|
|
"string",
|
|
widgetName,
|
|
savedValue,
|
|
(value) => {
|
|
widget.value = value;
|
|
serializeWorkerValues(node);
|
|
},
|
|
{}
|
|
);
|
|
}
|
|
widget.label = label;
|
|
}
|
|
|
|
// Persist type metadata
|
|
const patch = { _type: inputType };
|
|
if (comboOptions) patch._options = comboOptions;
|
|
writeStore(node, patch);
|
|
|
|
// Re-serialize worker values after creating widgets
|
|
serializeWorkerValues(node);
|
|
|
|
// Resize node to fit
|
|
const size = node.computeSize();
|
|
size[0] = Math.max(size[0], node.size?.[0] || 0);
|
|
node.setSize(size);
|
|
if (node.setDirtyCanvas) node.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Rebuild widgets based on current connection
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function rebuildWidgets(node, workers) {
|
|
const { type, options } = detectTargetType(node);
|
|
createWorkerWidgets(node, workers, type, options);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Extension registration
|
|
// ---------------------------------------------------------------------------
|
|
|
|
app.registerExtension({
|
|
name: "Distributed.DistributedValue",
|
|
async nodeCreated(node) {
|
|
if (node.comfyClass !== NODE_CLASS) return;
|
|
|
|
try {
|
|
const workers = await fetchWorkers();
|
|
node._dvWorkers = workers;
|
|
|
|
// Hide the raw worker_values widget
|
|
hideRawWidget(node);
|
|
|
|
// Read stored type or detect from connection
|
|
const store = readStore(node);
|
|
const storedType = store._type || "STRING";
|
|
const storedOptions = store._options || null;
|
|
|
|
if (workers.length > 0) {
|
|
createWorkerWidgets(node, workers, storedType, storedOptions);
|
|
}
|
|
|
|
// Listen for connection changes to adapt widget type
|
|
const originalOnConnectionsChange = node.onConnectionsChange;
|
|
node.onConnectionsChange = function (side, slotIndex, connected, linkInfo, ioSlot) {
|
|
if (originalOnConnectionsChange) {
|
|
originalOnConnectionsChange.call(this, side, slotIndex, connected, linkInfo, ioSlot);
|
|
}
|
|
// side 2 = output; slot 0 = our "value" output
|
|
if (slotIndex === 0 && this._dvWorkers?.length > 0) {
|
|
// Small delay to let litegraph finish link setup
|
|
setTimeout(() => rebuildWidgets(this, this._dvWorkers), 50);
|
|
}
|
|
};
|
|
|
|
// Override configure for workflow load
|
|
const originalConfigure = node.configure;
|
|
node.configure = function (data) {
|
|
const result = originalConfigure
|
|
? originalConfigure.call(this, data)
|
|
: undefined;
|
|
|
|
setTimeout(() => {
|
|
hideRawWidget(this);
|
|
if (this._dvWorkers?.length > 0) {
|
|
const s = readStore(this);
|
|
createWorkerWidgets(this, this._dvWorkers, s._type || "STRING", s._options || null);
|
|
}
|
|
}, 100);
|
|
|
|
return result;
|
|
};
|
|
} catch (error) {
|
|
console.error("Error in DistributedValue extension:", error);
|
|
}
|
|
},
|
|
});
|