Files
robertvoy-ComfyUI-Distributed/web/main.js
T

479 lines
16 KiB
JavaScript

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { DistributedUI } from './ui.js';
import { createStateManager } from './stateManager.js';
import { createApiClient } from './apiClient.js';
import { renderSidebarContent, updateWorkerCard } from './sidebarRenderer.js';
import { handleInterruptWorkers, handleClearMemory } from './workerUtils.js';
import { setupInterceptor } from './executionUtils.js';
import { PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js';
import { updateTunnelUIElements, refreshTunnelStatus, handleTunnelToggle } from './tunnelManager.js';
import { checkAllWorkerStatuses, checkWorkerStatus, loadManagedWorkers } from './workerLifecycle.js';
import { detectMasterIP } from './masterDetection.js';
import { parseHostInput, getMasterUrl as buildMasterUrl } from './urlUtils.js';
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
class DistributedExtension {
constructor() {
this.config = null;
this.originalQueuePrompt = api.queuePrompt.bind(api);
this.logAutoRefreshInterval = null;
this.masterSettingsExpanded = false;
this.app = app; // Store app reference for toast notifications
this.tunnelStatus = { status: "unknown" };
this.tunnelElements = {};
// Initialize centralized state
this.state = createStateManager();
// Initialize UI component factory
this.ui = new DistributedUI();
// Initialize API client
this.api = createApiClient(window.location.origin);
// Initialize status check timeout reference
this.statusCheckTimeout = null;
// Initialize abort controller for status checks
this.statusCheckAbortController = null;
this.themeMutationObserver = null;
// Inject CSS for pulsing animation
this.injectStyles();
this.loadConfig().then(async () => {
this.registerSidebarTab();
this.setupInterceptor();
// Don't start polling until panel opens
// this.startStatusChecking();
loadManagedWorkers(this);
// Detect master IP after everything is set up
this.detectMasterIP();
// Listen for Nodes 2.0 setting changes (once, for the lifetime of the extension)
this._setupNodes2Listener();
});
}
// Debug logging helpers
log(message, level = "info") {
if (level === "debug" && !this.config?.settings?.debug) return;
if (level === "error") {
console.error(`[Distributed] ${message}`);
} else {
console.log(`[Distributed] ${message}`);
}
}
injectStyles() {
const styleId = 'distributed-styles';
if (!document.getElementById(styleId)) {
const style = document.createElement('style');
style.id = styleId;
style.textContent = PULSE_ANIMATION_CSS;
document.head.appendChild(style);
}
const fileStyleId = 'distributed-file-styles';
if (!document.getElementById(fileStyleId)) {
const style = document.createElement('style');
style.id = fileStyleId;
fetch(new URL('./distributed.css', import.meta.url))
.then((response) => response.text())
.then((cssText) => {
style.textContent = cssText;
})
.catch((error) => {
this.log(`Failed to load distributed.css: ${error.message}`, "error");
});
document.head.appendChild(style);
}
}
// --- State & Config Management (Single Source of Truth) ---
get enabledWorkers() {
return this.config?.workers?.filter(w => w.enabled) || [];
}
get isEnabled() {
return this.enabledWorkers.length > 0;
}
isMasterParticipationEnabled() {
return !Boolean(this.config?.settings?.master_delegate_only);
}
isMasterFallbackActive() {
return Boolean(this.config?.settings?.master_delegate_only) && this.enabledWorkers.length === 0;
}
isMasterParticipating() {
return this.isMasterParticipationEnabled() || this.isMasterFallbackActive();
}
async updateMasterParticipation(enabled) {
if (!this.config?.settings) {
this.config.settings = {};
}
const delegateOnly = !enabled;
if (this.config.settings.master_delegate_only === delegateOnly) {
return;
}
await this._updateSetting('master_delegate_only', delegateOnly);
if (this.panelElement) {
renderSidebarContent(this, this.panelElement);
}
}
async loadConfig() {
try {
this.config = await this.api.getConfig();
this.log("Loaded config: " + JSON.stringify(this.config), "debug");
// Ensure default flag values
if (!this.config.settings) {
this.config.settings = {};
}
if (this.config.settings.has_auto_populated_workers === undefined) {
this.config.settings.has_auto_populated_workers = false;
}
// Load stored master CUDA device
this.masterCudaDevice = this.config?.master?.cuda_device ?? undefined;
// Sync to state
if (this.config.workers) {
this.config.workers.forEach(w => {
this.state.updateWorker(w.id, { enabled: w.enabled });
});
}
this._emitWorkersChanged();
} catch (error) {
this.log("Failed to load config: " + error.message, "error");
this.config = { workers: [], settings: { has_auto_populated_workers: false } };
}
}
_emitWorkersChanged() {
if (typeof window === "undefined" || typeof window.dispatchEvent !== "function") {
return;
}
window.dispatchEvent(new CustomEvent(WORKERS_CHANGED_EVENT, {
detail: { workers: this.config?.workers || [] },
}));
}
_applyMasterHost(host) {
if (!host || !this.config) return;
if (!this.config.master) this.config.master = {};
this.config.master.host = host;
const hostInput = document.getElementById('master-host');
if (hostInput) {
hostInput.value = host;
}
}
_parseHostInput(value) {
return parseHostInput(value);
}
updateTunnelUIElements(isRunning, isStarting) {
return updateTunnelUIElements(this, isRunning, isStarting);
}
async refreshTunnelStatus() {
return refreshTunnelStatus(this);
}
async handleTunnelToggle(button) {
return handleTunnelToggle(this, button);
}
async updateWorkerEnabled(workerId, enabled) {
const worker = this.config.workers.find(w => w.id === workerId);
if (worker) {
worker.enabled = enabled;
this.state.updateWorker(workerId, { enabled });
this._emitWorkersChanged();
// Immediately update status dot based on enabled state
const statusDot = document.getElementById(`status-${workerId}`);
if (statusDot) {
if (enabled) {
// Enabled: Start with checking state and trigger check
this.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Checking status...", false);
setTimeout(() => checkWorkerStatus(this, worker), TIMEOUTS.STATUS_CHECK_DELAY);
} else {
// Disabled: Set to gray
this.ui.updateStatusDot(workerId, STATUS_COLORS.DISABLED_GRAY, "Disabled", false);
}
}
}
try {
await this.api.updateWorker(workerId, { enabled });
} catch (error) {
this.log("Error updating worker: " + error.message, "error");
}
if (this.panelElement) {
await renderSidebarContent(this, this.panelElement);
}
}
async _updateSetting(key, value) {
// Update local config
if (!this.config.settings) {
this.config.settings = {};
}
this.config.settings[key] = value;
try {
await this.api.updateSetting(key, value);
const prettyKey = key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase());
let detail;
if (key === 'worker_timeout_seconds') {
const secs = parseInt(value, 10);
detail = `Worker Timeout set to ${Number.isFinite(secs) ? secs : value}s`;
} else if (typeof value === 'boolean') {
detail = `${prettyKey} ${value ? 'enabled' : 'disabled'}`;
} else {
detail = `${prettyKey} set to ${value}`;
}
app.extensionManager.toast.add({
severity: "success",
summary: "Setting Updated",
detail,
life: 2000
});
} catch (error) {
this.log(`Error updating setting '${key}': ${error.message}`, "error");
app.extensionManager.toast.add({
severity: "error",
summary: "Setting Update Failed",
detail: error.message,
life: 3000
});
}
}
// --- UI Rendering ---
registerSidebarTab() {
app.extensionManager.registerSidebarTab({
id: "distributed",
icon: "pi pi-server",
title: "Distributed",
tooltip: "Distributed Control Panel",
type: "custom",
render: (el) => {
this.panelElement = el;
this.onPanelOpen();
renderSidebarContent(this, el);
this._applyNodes2Style();
this._applyThemeToneClass();
},
destroy: () => {
this.onPanelClose();
}
});
}
onPanelOpen() {
this.log("Panel opened - starting status polling", "debug");
if (!this.statusCheckTimeout) {
checkAllWorkerStatuses(this);
}
this._startThemeObserver();
this._applyThemeToneClass();
}
onPanelClose() {
this.log("Panel closed - stopping status polling", "debug");
// Cancel any pending status checks
if (this.statusCheckAbortController) {
this.statusCheckAbortController.abort();
this.statusCheckAbortController = null;
}
// Clear the timeout
if (this.statusCheckTimeout) {
clearTimeout(this.statusCheckTimeout);
this.statusCheckTimeout = null;
}
this._stopThemeObserver();
this.panelElement = null;
}
_applyNodes2Style() {
if (!this.panelElement) return;
const enabled = app.ui.settings.getSettingValue("Comfy.VueNodes.Enabled") ?? false;
this.panelElement.classList.toggle('distributed-panel--nodes2', Boolean(enabled));
}
_parseColorToRgba(colorValue) {
if (!colorValue || typeof colorValue !== "string") {
return null;
}
const color = colorValue.trim().toLowerCase();
if (!color || color === "transparent") {
return null;
}
const rgbMatch = color.match(/^rgba?\(([^)]+)\)$/);
if (rgbMatch) {
const parts = rgbMatch[1].split(",").map((part) => Number(part.trim()));
if (parts.length < 3 || parts.slice(0, 3).some((part) => Number.isNaN(part))) {
return null;
}
const alpha = parts.length >= 4 && Number.isFinite(parts[3]) ? parts[3] : 1;
return {
r: Math.max(0, Math.min(255, parts[0])),
g: Math.max(0, Math.min(255, parts[1])),
b: Math.max(0, Math.min(255, parts[2])),
a: Math.max(0, Math.min(1, alpha)),
};
}
const hexMatch = color.match(/^#([0-9a-f]{3}|[0-9a-f]{6})$/i);
if (hexMatch) {
const value = hexMatch[1];
const expanded = value.length === 3
? value.split("").map((c) => `${c}${c}`).join("")
: value;
const r = parseInt(expanded.slice(0, 2), 16);
const g = parseInt(expanded.slice(2, 4), 16);
const b = parseInt(expanded.slice(4, 6), 16);
return { r, g, b, a: 1 };
}
return null;
}
_isPanelLightTheme() {
const fallbackLight = window.matchMedia?.("(prefers-color-scheme: light)")?.matches || false;
if (!this.panelElement) {
return fallbackLight;
}
let current = this.panelElement;
while (current) {
const bg = getComputedStyle(current).backgroundColor;
const rgba = this._parseColorToRgba(bg);
if (rgba && rgba.a > 0.02) {
// Relative luminance approximation (0..1)
const luminance = (0.2126 * rgba.r + 0.7152 * rgba.g + 0.0722 * rgba.b) / 255;
return luminance > 0.58;
}
current = current.parentElement;
}
return fallbackLight;
}
_applyThemeToneClass() {
if (!this.panelElement) {
return;
}
this.panelElement.classList.toggle("distributed-panel--light", this._isPanelLightTheme());
}
_startThemeObserver() {
if (this.themeMutationObserver) {
return;
}
this.themeMutationObserver = new MutationObserver(() => {
this._applyThemeToneClass();
});
this.themeMutationObserver.observe(document.documentElement, {
attributes: true,
attributeFilter: ["class", "style"],
});
if (document.body) {
this.themeMutationObserver.observe(document.body, {
attributes: true,
attributeFilter: ["class", "style"],
});
}
}
_stopThemeObserver() {
if (!this.themeMutationObserver) {
return;
}
this.themeMutationObserver.disconnect();
this.themeMutationObserver = null;
}
_setupNodes2Listener() {
app.ui.settings.addEventListener("Comfy.VueNodes.Enabled.change", (e) => {
const enabled = e.detail?.value ?? false;
if (this.panelElement) {
this.panelElement.classList.toggle('distributed-panel--nodes2', Boolean(enabled));
this._applyThemeToneClass();
}
});
}
// --- Core Logic & Execution ---
setupInterceptor() {
setupInterceptor(this);
}
updateWorkerCard(workerId, newStatus) {
return updateWorkerCard(this, workerId, newStatus);
}
/**
* Cleanup method to stop intervals and listeners
*/
cleanup() {
if (this.logAutoRefreshInterval) {
clearInterval(this.logAutoRefreshInterval);
this.logAutoRefreshInterval = null;
}
if (this.statusCheckTimeout) {
clearTimeout(this.statusCheckTimeout);
this.statusCheckTimeout = null;
}
this.log("Cleaned up intervals", "debug");
}
getMasterUrl() {
return buildMasterUrl(this.config, window.location, (message, level) => this.log(message, level));
}
async detectMasterIP() {
return detectMasterIP(this);
}
_handleInterruptWorkers(button) {
return handleInterruptWorkers(this, button);
}
_handleClearMemory(button) {
return handleClearMemory(this, button);
}
}
app.registerExtension({
name: "Distributed.Panel",
async setup() {
new DistributedExtension();
}
});