From aa5af8ab1473b924511a4209237e12d27d3d4d72 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 11 Jul 2025 09:59:21 +1000 Subject: [PATCH 01/93] initial support for cloud workers --- distributed.py | 12 +++++++- web/gpupanel.js | 78 +++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 73 insertions(+), 17 deletions(-) diff --git a/distributed.py b/distributed.py index 3189eab..cea68e0 100644 --- a/distributed.py +++ b/distributed.py @@ -1338,10 +1338,20 @@ class DistributedCollectorNode: try: session = await get_client_session() - async with session.post(f"{master_url}/distributed/job_complete", data=data) as response: + url = f"{master_url}/distributed/job_complete" + debug_log(f"Worker - Sending image to URL: {url}") + # Disable SSL verification for HTTPS URLs (temporary fix for tunnels) + ssl_context = None + if url.startswith('https://'): + import ssl + ssl_context = ssl.create_default_context() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + async with session.post(url, data=data, ssl=ssl_context) as response: response.raise_for_status() except Exception as e: log(f"Worker - Failed to send image {image_index+1} to master: {e}") + debug_log(f"Worker - Full error details: URL={master_url}/distributed/job_complete") async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): if is_worker: diff --git a/web/gpupanel.js b/web/gpupanel.js index a98fc3e..7a3f85d 100644 --- a/web/gpupanel.js +++ b/web/gpupanel.js @@ -1326,16 +1326,28 @@ class DistributedExtension { async checkWorkerStatus(worker) { const host = worker.host || window.location.hostname; - const url = `http://${host}:${worker.port}/prompt`; - const statusDot = document.getElementById(`status-${worker.id}`); - // Status dot will be updated when response comes back + // Detect if we should use HTTPS based on hostname or port + const useHttps = host.includes('trycloudflare.com') || + host.includes('ngrok') || + host.includes('.ts.net') || + worker.port === 443; + + const protocol = useHttps ? 'https' : 'http'; + + // Don't add port for standard HTTPS/HTTP ports + const defaultPort = useHttps ? 443 : 80; + const port = worker.port === defaultPort ? '' : `:${worker.port}`; + + const url = `${protocol}://${host}${port}/prompt`; + + const statusDot = document.getElementById(`status-${worker.id}`); try { const response = await fetch(url, { method: 'GET', mode: 'cors', - signal: AbortSignal.timeout(300) // 300ms for faster response + signal: AbortSignal.timeout(900) // Use your increased timeout }); if (response.ok) { @@ -1360,8 +1372,6 @@ class DistributedExtension { // Clear launching state since worker is now online if (this.launchingWorkers.has(worker.id)) { this.launchingWorkers.delete(worker.id); - - // Also clear launching flag in backend this.clearLaunchingFlag(worker.id); } } else { @@ -1382,6 +1392,8 @@ class DistributedExtension { // Only update to red if not currently launching this.updateStatusDot(worker.id, "#c04c4c", "Offline - Cannot connect", false); } + + this.debugLog(`Worker ${worker.id} status check failed: ${error.message}`); } // Update control buttons based on new status @@ -1389,8 +1401,21 @@ class DistributedExtension { } async _dispatchToWorker(worker, prompt, workflow, imageReferences) { - const host = worker.host || window.location.hostname; - const workerUrl = `http://${host}:${worker.port}`; + const host = worker.host || "localhost"; + + // Detect if we should use HTTPS based on hostname or port + const useHttps = host.includes('trycloudflare.com') || + host.includes('ngrok') || + host.includes('.ts.net') || + worker.port === 443; + + const protocol = useHttps ? 'https' : 'http'; + + // Don't add port for standard HTTPS/HTTP ports + const defaultPort = useHttps ? 443 : 80; + const port = worker.port === defaultPort ? '' : `:${worker.port}`; + + const workerUrl = `${protocol}://${host}${port}`; // Debug logging - always log to console for debugging console.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`); @@ -1421,7 +1446,7 @@ class DistributedExtension { const promptToSend = { prompt, - extra_data: { }, // Removed workflow from extra_pnginfo to reduce payload size + extra_data: { extra_pnginfo: { workflow } }, client_id: api.clientId }; @@ -1552,13 +1577,29 @@ class DistributedExtension { const startTime = Date.now(); const checkPromises = workers.map(async (worker) => { - const host = worker.host || window.location.hostname; - const url = `http://${host}:${worker.port}/prompt`; + const host = worker.host || "localhost"; + + // Detect if we should use HTTPS based on hostname or port + const useHttps = host.includes('trycloudflare.com') || + host.includes('ngrok') || + host.includes('.ts.net') || + worker.port === 443; + + const protocol = useHttps ? 'https' : 'http'; + + // Don't add port for standard HTTPS/HTTP ports + const defaultPort = useHttps ? 443 : 80; + const port = worker.port === defaultPort ? '' : `:${worker.port}`; + + const url = `${protocol}://${host}${port}/prompt`; + + this.debugLog(`Pre-flight checking ${worker.name} at: ${url}`); + try { const response = await fetch(url, { method: 'GET', mode: 'cors', - signal: AbortSignal.timeout(300) // 300ms timeout for pre-flight + signal: AbortSignal.timeout(900) // Use your increased timeout }); if (response.ok) { @@ -1569,7 +1610,7 @@ class DistributedExtension { return { worker, active: false }; } } catch (error) { - this.debugLog(`Worker ${worker.name} is offline or unreachable`); + this.debugLog(`Worker ${worker.name} is offline or unreachable: ${error.message}`); return { worker, active: false }; } }); @@ -1593,7 +1634,6 @@ class DistributedExtension { return activeWorkers; } - async launchWorker(workerId) { const worker = this.config.workers.find(w => w.id === workerId); const launchBtn = document.querySelector(`#controls-${workerId} button`); @@ -2175,7 +2215,13 @@ class DistributedExtension { getMasterUrl() { // Always use the detected/configured master IP for consistency if (this.config?.master?.host) { - return `http://${this.config.master.host}:${window.location.port}`; + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + // For standard ports, don't include them in the URL + if ((window.location.protocol === 'https:' && port === '443') || + (window.location.protocol === 'http:' && port === '80')) { + return `${window.location.protocol}//${this.config.master.host}`; + } + return `${window.location.protocol}//${this.config.master.host}:${port}`; } // If no master IP is set but we're on a network address, use it @@ -2469,7 +2515,7 @@ class DistributedExtension { return; } - if (isNaN(port) || port < 1024 || port > 65535) { + if (isNaN(port) || port < 1 || port > 65535) { app.extensionManager.toast.add({ severity: "error", summary: "Validation Error", From c1b82179db844f2b147ee08cf9bc1dbbc6404b95 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 11 Jul 2025 13:50:36 +1000 Subject: [PATCH 02/93] include master in UI panel, auto detect local workers on first launch --- distributed.py | 60 +++++- web/gpupanel.js | 544 ++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 514 insertions(+), 90 deletions(-) diff --git a/distributed.py b/distributed.py index cea68e0..e9a7c0f 100644 --- a/distributed.py +++ b/distributed.py @@ -95,6 +95,53 @@ async def get_network_info_endpoint(request): """Get network interfaces and recommend best IP for master.""" import socket + # Get CUDA device if available + cuda_device = None + cuda_device_count = 0 + physical_device_count = 0 + + if torch.cuda.is_available(): + try: + import os + import subprocess + + # Get visible device count (what PyTorch sees) + cuda_device_count = torch.cuda.device_count() + + # Try to get actual physical device info + cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '') + + # Method 1: Parse CUDA_VISIBLE_DEVICES + if cuda_visible and cuda_visible.strip(): + visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()] + if visible_devices: + # Get the first visible device as the actual physical device + cuda_device = visible_devices[0] + + # Try to get total physical device count using nvidia-smi + try: + result = subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], + capture_output=True, text=True, timeout=5) + if result.returncode == 0: + physical_device_count = len(result.stdout.strip().split('\n')) + else: + physical_device_count = max(visible_devices) + 1 # Best guess + except: + physical_device_count = max(visible_devices) + 1 # Best guess + else: + cuda_device = 0 + physical_device_count = cuda_device_count + else: + # No CUDA_VISIBLE_DEVICES set, current device is actual device + cuda_device = torch.cuda.current_device() + physical_device_count = cuda_device_count + + except Exception as e: + debug_log(f"CUDA detection error: {e}") + cuda_device = None + cuda_device_count = 0 + physical_device_count = 0 + def get_network_ips(): """Get all network IPs, trying multiple methods.""" ips = [] @@ -195,6 +242,8 @@ async def get_network_info_endpoint(request): "hostname": hostname, "all_ips": all_ips, "recommended_ip": recommended_ip, + "cuda_device": cuda_device, + "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count, "message": "Auto-detected network configuration" }) except Exception as e: @@ -347,6 +396,8 @@ async def update_master_endpoint(request): config['master'] = {} # Update all provided fields + if "name" in data: + config['master']['name'] = data['name'] if "host" in data: config['master']['host'] = data['host'] if "port" in data: @@ -1340,14 +1391,7 @@ class DistributedCollectorNode: session = await get_client_session() url = f"{master_url}/distributed/job_complete" debug_log(f"Worker - Sending image to URL: {url}") - # Disable SSL verification for HTTPS URLs (temporary fix for tunnels) - ssl_context = None - if url.startswith('https://'): - import ssl - ssl_context = ssl.create_default_context() - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE - async with session.post(url, data=data, ssl=ssl_context) as response: + async with session.post(url, data=data) as response: response.raise_for_status() except Exception as e: log(f"Worker - Failed to send image {image_index+1} to master: {e}") diff --git a/web/gpupanel.js b/web/gpupanel.js index 7a3f85d..74c818b 100644 --- a/web/gpupanel.js +++ b/web/gpupanel.js @@ -75,12 +75,13 @@ class DistributedExtension { // Inject CSS for pulsing animation this.injectStyles(); - this.loadConfig().then(() => { + this.loadConfig().then(async () => { this.registerSidebarTab(); this.setupInterceptor(); this.startStatusChecking(); this.loadManagedWorkers(); - this.detectMasterIP(); // Auto-detect master IP on startup + // Detect master IP after everything is set up + this.detectMasterIP(); }); } @@ -120,7 +121,7 @@ class DistributedExtension { const response = await fetch(`${window.location.origin}/distributed/config`); if (response.ok) { this.config = await response.json(); - this.debugLog("Loaded config: " + JSON.stringify(this.config)); + this.log("Loaded config: " + JSON.stringify(this.config)); } else { this.log("Failed to load config"); this.config = { workers: [], settings: {} }; @@ -143,6 +144,7 @@ class DistributedExtension { headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ worker_id: workerId, enabled: enabled }) }); + } catch (error) { this.log("Error updating worker: " + error.message); } @@ -211,7 +213,7 @@ class DistributedExtension { try { const targetsToInterrupt = this.enabledWorkers.map(w => ({ name: w.name, - url: `http://${w.host || window.location.hostname}:${w.port}` + url: this.getWorkerUrl(w) })); if (targetsToInterrupt.length === 0) { @@ -266,7 +268,7 @@ class DistributedExtension { try { const urlsToClear = this.enabledWorkers.map(w => ({ name: w.name, - url: `http://${w.host || window.location.hostname}:${w.port}` + url: this.getWorkerUrl(w) })); if (urlsToClear.length === 0) { @@ -324,7 +326,7 @@ class DistributedExtension { } - renderSidebarContent(el) { + async renderSidebarContent(el) { // Panel is being opened/rendered this.debugLog("Panel opened"); @@ -333,8 +335,19 @@ class DistributedExtension { return; } - // Store reference to the panel element - this.panelElement = el; + // Prevent infinite recursion + if (this._isRendering) { + this.debugLog("Already rendering, skipping"); + return; + } + this._isRendering = true; + + try { + // Store reference to the panel element + this.panelElement = el; + + // Reload config when panel opens to pick up any external changes + await this.loadConfig(); // Reload managed workers when panel opens this.loadManagedWorkers().then(() => { @@ -349,20 +362,275 @@ class DistributedExtension { setTimeout(() => this.checkAllWorkerStatuses(), 0); el.innerHTML = ''; + + // Create toolbar header to match ComfyUI style + const toolbar = document.createElement("div"); + toolbar.className = "p-toolbar p-component border-x-0 border-t-0 rounded-none px-2 py-1 min-h-8"; + toolbar.style.cssText = "border-bottom: 1px solid #444; background: transparent;"; + + const toolbarStart = document.createElement("div"); + toolbarStart.className = "p-toolbar-start"; + + const titleSpan = document.createElement("span"); + titleSpan.className = "text-xs 2xl:text-sm truncate"; + titleSpan.textContent = "COMFYUI DISTRIBUTED"; + titleSpan.title = "ComfyUI Distributed"; + + toolbarStart.appendChild(titleSpan); + toolbar.appendChild(toolbarStart); + + const toolbarCenter = document.createElement("div"); + toolbarCenter.className = "p-toolbar-center"; + toolbar.appendChild(toolbarCenter); + + const toolbarEnd = document.createElement("div"); + toolbarEnd.className = "p-toolbar-end"; + toolbar.appendChild(toolbarEnd); + + el.appendChild(toolbar); + + // Main container with adjusted padding const container = document.createElement("div"); - container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: 100%;"; + container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: calc(100% - 32px);"; - const title = document.createElement("h3"); - title.textContent = "ComfyUI Distributed"; - title.style.cssText = "margin: 0 0 15px 0; padding-bottom: 10px; border-bottom: 1px solid #444;"; - container.appendChild(title); + // Detect master info on panel open (in case CUDA info wasn't available at startup) + this.log(`Panel opened. CUDA device count: ${this.cudaDeviceCount}, Workers: ${this.config?.workers?.length || 0}`); + if (!this.cudaDeviceCount) { + await this.detectMasterIP(); + } + // Check if we should auto-populate workers after detection + if (this.cudaDeviceCount > 0 && (!this.config.workers || this.config.workers.length === 0)) { + this.log(`Auto-populating workers based on ${this.cudaDeviceCount} CUDA devices (excluding master on CUDA ${this.masterCudaDevice})`); + + const newWorkers = []; + let workerNum = 1; + let portOffset = 0; + + for (let i = 0; i < this.cudaDeviceCount; i++) { + // Skip the CUDA device used by master + if (i === this.masterCudaDevice) { + this.log(`Skipping CUDA ${i} (used by master)`); + continue; + } + + const worker = { + id: Date.now() + i, + name: `Worker ${workerNum}`, + host: "localhost", + port: 8189 + portOffset, + cuda_device: i, + enabled: true + }; + newWorkers.push(worker); + workerNum++; + portOffset++; + } + + // Only proceed if we have workers to add + if (newWorkers.length > 0) { + this.log(`Auto-populating ${newWorkers.length} workers`); + + // Add workers to config + this.config.workers = newWorkers; + + // Save each worker using the update endpoint + for (const worker of newWorkers) { + try { + const response = await fetch(`${window.location.origin}/distributed/config/update_worker`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + worker_id: worker.id, + ...worker + }) + }); + + if (!response.ok) { + this.log(`Failed to save worker ${worker.name}`); + } + } catch (error) { + this.log(`Error saving worker ${worker.name}: ${error.message}`); + } + } + + this.log(`Auto-populated ${newWorkers.length} workers and saved config`); + + // Reload the config to include the new workers + await this.loadConfig(); + + // Continue rendering with the updated config + } else { + this.log("No additional CUDA devices available for workers (all used by master)"); + } + } + + // Master Node Section - styled exactly like a worker card + const masterDiv = document.createElement("div"); + masterDiv.style.cssText = "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;"; + + // Left column - checkbox area (always checked and disabled) + const checkboxColumn = document.createElement("div"); + checkboxColumn.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444; cursor: default;"; + checkboxColumn.title = "Master node is always enabled"; + + // Greyed out checkbox that's always checked + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.checked = true; + checkbox.disabled = true; + checkbox.style.cssText = "cursor: default; width: 16px; height: 16px; opacity: 0.6;"; + + checkboxColumn.appendChild(checkbox); + + // Right column - content area + const contentColumn = document.createElement("div"); + contentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; + + // First row: master info + const infoRow = document.createElement("div"); + infoRow.style.cssText = "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;"; + infoRow.title = "Click to expand settings"; + + // Worker content + const workerContent = document.createElement("div"); + workerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + + const statusDot = document.createElement("span"); + statusDot.id = "master-status"; + statusDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; background-color: #4CAF50;"; + statusDot.title = "Online"; + + const gpuInfo = document.createElement("span"); + const cudaInfo = this.masterCudaDevice !== undefined ? `CUDA ${this.masterCudaDevice} • ` : ''; + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + gpuInfo.innerHTML = `${this.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; + + workerContent.appendChild(statusDot); + workerContent.appendChild(gpuInfo); + + // Settings arrow + const settingsArrow = document.createElement("span"); + settingsArrow.className = "settings-arrow"; + settingsArrow.innerHTML = "▶"; + settingsArrow.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;"; + + infoRow.appendChild(workerContent); + infoRow.appendChild(settingsArrow); + + // Click handler for info row + let masterSettingsExpanded = false; + infoRow.onclick = () => { + masterSettingsExpanded = !masterSettingsExpanded; + if (masterSettingsExpanded) { + masterSettingsDiv.classList.add("expanded"); + masterSettingsDiv.style.padding = "12px"; + masterSettingsDiv.style.marginTop = "8px"; + masterSettingsDiv.style.marginBottom = "8px"; + settingsArrow.style.transform = "rotate(90deg)"; + } else { + masterSettingsDiv.classList.remove("expanded"); + masterSettingsDiv.style.padding = "0 12px"; + masterSettingsDiv.style.marginTop = "0"; + masterSettingsDiv.style.marginBottom = "0"; + settingsArrow.style.transform = "rotate(0deg)"; + } + }; + + // Hover effect for entire content column + contentColumn.onmouseover = () => { + contentColumn.style.backgroundColor = "#333"; + settingsArrow.style.color = "#fff"; + }; + contentColumn.onmouseout = () => { + contentColumn.style.backgroundColor = "transparent"; + settingsArrow.style.color = "#888"; + }; + + contentColumn.appendChild(infoRow); + + // Second row: master controls (styled like remote worker) + const controlsDiv = document.createElement("div"); + controlsDiv.style.cssText = "padding: 0 8px 6px 8px; display: flex; gap: 4px;"; + + // Master info box styled like remote worker + const masterInfo = document.createElement("div"); + masterInfo.style.cssText = "background-color: #333; color: #888; padding: 4px 12px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1;"; + masterInfo.textContent = "Master"; + controlsDiv.appendChild(masterInfo); + + // Third row: expandable settings panel + const masterSettingsDiv = document.createElement("div"); + masterSettingsDiv.id = "master-settings"; + masterSettingsDiv.className = "worker-settings"; + masterSettingsDiv.style.cssText = "margin: 0 8px; padding: 0 12px; background: #1e1e1e; border-radius: 4px;"; + + // Create settings form using reusable components + const settingsForm = document.createElement("div"); + settingsForm.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; + + // Name field - reuse createFormGroup + const nameGroup = this.createFormGroup("Name:", this.config?.master?.name || "Master", "master-name"); + settingsForm.appendChild(nameGroup); + + // Host field + const hostGroup = this.createFormGroup("Host:", this.config?.master?.host || "", "master-host", "text", "Auto-detect if empty"); + settingsForm.appendChild(hostGroup); + + // Buttons row - using same style as workers + const buttonGroup = document.createElement("div"); + buttonGroup.style.cssText = "display: flex; gap: 8px; margin-top: 8px;"; + + const saveBtn = this._createButton("Save", + async () => { + const nameInput = document.getElementById('master-name'); + const hostInput = document.getElementById('master-host'); + + // Update config + if (!this.config.master) this.config.master = {}; + this.config.master.name = nameInput.value.trim() || "Master"; + + // Save both host and name + await this.saveMasterConfig({ + host: hostInput.value.trim(), + name: this.config.master.name + }); + + // Update display + document.getElementById('master-name-display').textContent = this.config.master.name; + this.updateMasterDisplay(); + + saveBtn.textContent = "Saved!"; + setTimeout(() => { saveBtn.textContent = "Save"; }, 2000); + }, + "background-color: #4a7c4a;"); + saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; + + const cancelBtn = this._createButton("Cancel", + () => { + // Reset inputs to original values + document.getElementById('master-name').value = this.config?.master?.name || "Master"; + document.getElementById('master-host').value = this.config?.master?.host || ""; + }, + "background-color: #555;"); + cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; + + buttonGroup.appendChild(saveBtn); + buttonGroup.appendChild(cancelBtn); + + settingsForm.appendChild(buttonGroup); + masterSettingsDiv.appendChild(settingsForm); + + contentColumn.appendChild(controlsDiv); + contentColumn.appendChild(masterSettingsDiv); + + masterDiv.appendChild(checkboxColumn); + masterDiv.appendChild(contentColumn); + container.appendChild(masterDiv); + + // Workers Section (no heading) const gpuSection = document.createElement("div"); gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; - const gpuTitle = document.createElement("h4"); - gpuTitle.textContent = "Available Workers"; - gpuTitle.style.cssText = "margin: 0 0 10px 0; font-size: 14px;"; - gpuSection.appendChild(gpuTitle); const gpuList = document.createElement("div"); const workers = this.config?.workers || []; @@ -687,23 +955,23 @@ class DistributedExtension { const settingsHeader = document.createElement("div"); settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;"; - const settingsTitle = document.createElement("h4"); - settingsTitle.textContent = "Settings"; - settingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + const workerSettingsTitle = document.createElement("h4"); + workerSettingsTitle.textContent = "Settings"; + workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;"; - const settingsToggle = document.createElement("span"); - settingsToggle.textContent = "▶"; // Right arrow when collapsed - settingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + const workerSettingsToggle = document.createElement("span"); + workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed + workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; - settingsHeader.appendChild(settingsTitle); - settingsHeader.appendChild(settingsToggle); + settingsHeader.appendChild(workerSettingsTitle); + settingsHeader.appendChild(workerSettingsToggle); // Hover effect for header settingsHeader.onmouseover = () => { - settingsToggle.style.color = "#fff"; + workerSettingsToggle.style.color = "#fff"; }; settingsHeader.onmouseout = () => { - settingsToggle.style.color = "#888"; + workerSettingsToggle.style.color = "#888"; }; // Collapsible settings content @@ -719,10 +987,10 @@ class DistributedExtension { settingsExpanded = !settingsExpanded; if (settingsExpanded) { settingsContent.style.maxHeight = "200px"; - settingsToggle.style.transform = "rotate(90deg)"; + workerSettingsToggle.style.transform = "rotate(90deg)"; } else { settingsContent.style.maxHeight = "0"; - settingsToggle.style.transform = "rotate(0deg)"; + workerSettingsToggle.style.transform = "rotate(0deg)"; } }; @@ -798,6 +1066,29 @@ class DistributedExtension { container.appendChild(summarySection); el.appendChild(container); this.updateSummary(); + } finally { + // Always reset the rendering flag + this._isRendering = false; + } + } + + updateMasterDisplay() { + // Update CUDA info if element exists + const cudaInfo = document.getElementById('master-cuda-info'); + if (cudaInfo) { + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + if (this.masterCudaDevice !== undefined && this.masterCudaDevice !== null) { + cudaInfo.textContent = `CUDA ${this.masterCudaDevice} • Port ${port}`; + } else { + cudaInfo.textContent = `Port ${port}`; + } + } + + // Update name if changed + const nameDisplay = document.getElementById('master-name-display'); + if (nameDisplay && this.config?.master?.name) { + nameDisplay.textContent = this.config.master.name; + } } updateSummary() { @@ -1314,6 +1605,9 @@ class DistributedExtension { } async checkAllWorkerStatuses() { + // Check master status + this.checkMasterStatus(); + if (!this.config || !this.config.workers) return; for (const worker of this.config.workers) { @@ -1324,23 +1618,57 @@ class DistributedExtension { } } - async checkWorkerStatus(worker) { + async checkMasterStatus() { + try { + const response = await fetch(`${window.location.origin}/prompt`, { + method: 'GET', + signal: AbortSignal.timeout(900) + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + const isProcessing = queueRemaining > 0; + + // Update master status dot + const statusDot = document.getElementById('master-status'); + if (statusDot) { + if (isProcessing) { + statusDot.style.backgroundColor = "#f0ad4e"; + statusDot.title = `Processing (${queueRemaining} in queue)`; + } else { + statusDot.style.backgroundColor = "#4CAF50"; + statusDot.title = "Online"; + } + } + } + } catch (error) { + // Master is always online (we're running on it), so keep it green + const statusDot = document.getElementById('master-status'); + if (statusDot) { + statusDot.style.backgroundColor = "#4CAF50"; + statusDot.title = "Online"; + } + } + } + + // Helper to build worker URL + getWorkerUrl(worker, endpoint = '') { const host = worker.host || window.location.hostname; - // Detect if we should use HTTPS based on hostname or port - const useHttps = host.includes('trycloudflare.com') || - host.includes('ngrok') || - host.includes('.ts.net') || - worker.port === 443; - + // Simple rule: port 443 = HTTPS, anything else = HTTP + const useHttps = worker.port === 443; const protocol = useHttps ? 'https' : 'http'; // Don't add port for standard HTTPS/HTTP ports const defaultPort = useHttps ? 443 : 80; const port = worker.port === defaultPort ? '' : `:${worker.port}`; - const url = `${protocol}://${host}${port}/prompt`; - + return `${protocol}://${host}${port}${endpoint}`; + } + + async checkWorkerStatus(worker) { + const url = this.getWorkerUrl(worker, '/prompt'); const statusDot = document.getElementById(`status-${worker.id}`); try { @@ -1401,21 +1729,7 @@ class DistributedExtension { } async _dispatchToWorker(worker, prompt, workflow, imageReferences) { - const host = worker.host || "localhost"; - - // Detect if we should use HTTPS based on hostname or port - const useHttps = host.includes('trycloudflare.com') || - host.includes('ngrok') || - host.includes('.ts.net') || - worker.port === 443; - - const protocol = useHttps ? 'https' : 'http'; - - // Don't add port for standard HTTPS/HTTP ports - const defaultPort = useHttps ? 443 : 80; - const port = worker.port === defaultPort ? '' : `:${worker.port}`; - - const workerUrl = `${protocol}://${host}${port}`; + const workerUrl = this.getWorkerUrl(worker); // Debug logging - always log to console for debugging console.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`); @@ -1577,21 +1891,7 @@ class DistributedExtension { const startTime = Date.now(); const checkPromises = workers.map(async (worker) => { - const host = worker.host || "localhost"; - - // Detect if we should use HTTPS based on hostname or port - const useHttps = host.includes('trycloudflare.com') || - host.includes('ngrok') || - host.includes('.ts.net') || - worker.port === 443; - - const protocol = useHttps ? 'https' : 'http'; - - // Don't add port for standard HTTPS/HTTP ports - const defaultPort = useHttps ? 443 : 80; - const port = worker.port === defaultPort ? '' : `:${worker.port}`; - - const url = `${protocol}://${host}${port}/prompt`; + const url = this.getWorkerUrl(worker, '/prompt'); this.debugLog(`Pre-flight checking ${worker.name} at: ${url}`); @@ -2215,13 +2515,29 @@ class DistributedExtension { getMasterUrl() { // Always use the detected/configured master IP for consistency if (this.config?.master?.host) { - const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - // For standard ports, don't include them in the URL - if ((window.location.protocol === 'https:' && port === '443') || - (window.location.protocol === 'http:' && port === '80')) { - return `${window.location.protocol}//${this.config.master.host}`; + const configuredHost = this.config.master.host; + + // If the configured host already includes protocol, use as-is + if (configuredHost.startsWith('http://') || configuredHost.startsWith('https://')) { + return configuredHost; + } + + // For domain names (not IPs), default to HTTPS + const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost); + const isLocalhost = configuredHost === 'localhost' || configuredHost === '127.0.0.1'; + + if (!isIP && !isLocalhost && configuredHost.includes('.')) { + // It's a domain name, use HTTPS + return `https://${configuredHost}`; + } else { + // For IPs and localhost, use current access method + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + if ((window.location.protocol === 'https:' && port === '443') || + (window.location.protocol === 'http:' && port === '80')) { + return `${window.location.protocol}//${configuredHost}`; + } + return `${window.location.protocol}//${configuredHost}:${port}`; } - return `${window.location.protocol}//${this.config.master.host}:${port}`; } // If no master IP is set but we're on a network address, use it @@ -2245,7 +2561,49 @@ class DistributedExtension { } const data = await response.json(); - this.debugLog("Network info: " + JSON.stringify(data)); + this.log("Network info: " + JSON.stringify(data)); + + // Store CUDA device info + if (data.cuda_device !== null && data.cuda_device !== undefined) { + this.masterCudaDevice = data.cuda_device; + // Update the master display with CUDA info + this.updateMasterDisplay(); + } + + // Store CUDA device count for auto-population + if (data.cuda_device_count > 0) { + this.cudaDeviceCount = data.cuda_device_count; + this.log(`Detected ${this.cudaDeviceCount} CUDA devices`); + + // Auto-populate workers if none exist + this.log(`Current workers: ${this.config.workers ? this.config.workers.length : 'null'}`); + if (!this.config.workers || this.config.workers.length === 0) { + this.log(`Auto-populating ${this.cudaDeviceCount} workers based on CUDA devices`); + + const newWorkers = []; + for (let i = 0; i < this.cudaDeviceCount; i++) { + const worker = { + id: Date.now() + i, + name: `Worker ${i + 1}`, + host: "localhost", + port: 8189 + i, + cuda_device: i, + enabled: true + }; + newWorkers.push(worker); + } + + // Update config and save + this.config.workers = newWorkers; + await this.saveConfig(); + + // Refresh the panel if it's already open + const panel = document.getElementById("distributed-panel"); + if (panel) { + this.renderSidebarContent(panel); + } + } + } // Check if we already have a master host configured if (this.config?.master?.host) { @@ -2258,7 +2616,7 @@ class DistributedExtension { this.log(`Auto-detected master IP: ${data.recommended_ip}`); // Save the detected IP (pass true to suppress notification) - await this.saveMasterIp(data.recommended_ip, true); + await this.saveMasterConfig({ host: data.recommended_ip }, true); // Show a single combined notification if (app.extensionManager?.toast) { @@ -2275,14 +2633,12 @@ class DistributedExtension { } } - async saveMasterIp(masterIp, suppressNotification = false) { + async saveMasterConfig(updates, suppressNotification = false) { try { const response = await fetch(`${window.location.origin}/distributed/config/update_master`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - host: masterIp.trim() - }) + body: JSON.stringify(updates) }); if (response.ok) { @@ -2290,7 +2646,8 @@ class DistributedExtension { if (!this.config.master) { this.config.master = {}; } - this.config.master.host = masterIp.trim(); + // Update all provided fields + Object.assign(this.config.master, updates); // Only show notification if not suppressed if (!suppressNotification) { @@ -2459,6 +2816,24 @@ class DistributedExtension { return form; } + createSettingsToggle() { + const settingsRow = document.createElement("div"); + settingsRow.style.cssText = "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;"; + + const settingsTitle = document.createElement("h4"); + settingsTitle.textContent = "Settings"; + settingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + + const settingsToggle = document.createElement("span"); + settingsToggle.textContent = "▶"; // Right arrow when collapsed + settingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + + settingsRow.appendChild(settingsToggle); + settingsRow.appendChild(settingsTitle); + + return { settingsRow, settingsToggle }; + } + createFormGroup(label, value, id, type = "text", placeholder = "") { const group = document.createElement("div"); group.style.cssText = "display: flex; align-items: center; gap: 8px;"; @@ -2519,7 +2894,7 @@ class DistributedExtension { app.extensionManager.toast.add({ severity: "error", summary: "Validation Error", - detail: "Port must be between 1024 and 65535", + detail: "Port must be between 1 and 65535", life: 3000 }); return; @@ -2697,6 +3072,11 @@ class DistributedExtension { } } + // Backward compatibility wrapper + async saveMasterIp(masterIp, suppressNotification = false) { + return this.saveMasterConfig({ host: masterIp.trim() }, suppressNotification); + } + async addNewWorker() { // Generate new worker ID const newId = `worker_${Date.now()}`; From 8dfe63a12d55e829f59aa10036812c4cae089090 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 11 Jul 2025 17:30:33 +1000 Subject: [PATCH 03/93] refactored the js --- web/gpupanel.js | 1083 +++++++++++++++++++++++++++++++---------------- 1 file changed, 725 insertions(+), 358 deletions(-) diff --git a/web/gpupanel.js b/web/gpupanel.js index 74c818b..619b647 100644 --- a/web/gpupanel.js +++ b/web/gpupanel.js @@ -65,12 +65,17 @@ class DistributedExtension { constructor() { this.config = null; this.originalQueuePrompt = api.queuePrompt.bind(api); - this.workerStatuses = new Map(); // Track worker statuses - this.managedWorkers = new Map(); // Track UI-managed workers this.statusCheckInterval = null; - this.logAutoRefreshInterval = null; // Track log auto-refresh - this.launchingWorkers = new Set(); // Track workers currently launching - this.expandedWorkers = new Set(); // Track which workers have expanded settings + this.logAutoRefreshInterval = null; + + // Initialize centralized state + this.state = this._createStateManager(); + + // Initialize UI component factory + this.ui = this._createUIFactory(); + + // Initialize API client + this.api = this._createApiClient(); // Inject CSS for pulsing animation this.injectStyles(); @@ -95,6 +100,563 @@ class DistributedExtension { log(message) { console.log(`[Distributed] ${message}`); } + + // --- UI Component Factory --- + _createUIFactory() { + const styles = { + statusDot: "display: inline-block; width: 8px; height: 8px; border-radius: 50%;", + controlsDiv: "padding: 0 8px 6px 8px; display: flex; gap: 4px;", + checkboxColumn: "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444;", + workerCard: "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;", + contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;", + infoRow: "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;", + workerContent: "display: flex; align-items: center; gap: 8px; flex: 1;", + settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;", + formGroup: "display: flex; align-items: center; gap: 8px; margin-bottom: 8px;", + formLabel: "flex: 0 0 100px; font-size: 12px; color: #ccc;", + formInput: "flex: 1; padding: 4px 8px; background: #1a1a1a; border: 1px solid #444; border-radius: 4px; color: #fff; font-size: 12px;", + infoBox: "background-color: #333; color: #888; padding: 4px 12px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1;" + }; + + return { + createStatusDot(id, color = "#666", title = "Status") { + const dot = document.createElement("span"); + if (id) dot.id = id; + dot.style.cssText = styles.statusDot + ` background-color: ${color};`; + dot.title = title; + return dot; + }, + + createButton(text, onClick, customStyle = "") { + const button = document.createElement("button"); + button.textContent = text; + button.className = "distributed-button"; + button.style.cssText = DistributedExtension.BUTTON_STYLES.base + customStyle; + if (onClick) button.onclick = onClick; + return button; + }, + + createButtonGroup(buttons, style = "") { + const group = document.createElement("div"); + group.style.cssText = "display: flex; gap: 4px; margin-top: 10px;" + style; + buttons.forEach(button => group.appendChild(button)); + return group; + }, + + createWorkerControls(workerId, handlers = {}) { + const controlsDiv = document.createElement("div"); + controlsDiv.id = `controls-${workerId}`; + controlsDiv.style.cssText = styles.controlsDiv; + + const buttons = []; + + if (handlers.launch) { + const launchBtn = this.createButton('Launch', handlers.launch); + launchBtn.id = `launch-${workerId}`; + launchBtn.title = "Launch this worker instance"; + buttons.push(launchBtn); + } + + if (handlers.stop) { + const stopBtn = this.createButton('Stop', handlers.stop); + stopBtn.id = `stop-${workerId}`; + stopBtn.title = "Stop this worker instance"; + buttons.push(stopBtn); + } + + if (handlers.viewLog) { + const logBtn = this.createButton('View Log', handlers.viewLog); + logBtn.id = `log-${workerId}`; + logBtn.title = "View worker log file"; + buttons.push(logBtn); + } + + buttons.forEach(btn => controlsDiv.appendChild(btn)); + return controlsDiv; + }, + + createFormGroup(label, value, id, type = "text", placeholder = "") { + const group = document.createElement("div"); + group.style.cssText = styles.formGroup; + + const labelEl = document.createElement("label"); + labelEl.textContent = label; + labelEl.htmlFor = id; + labelEl.style.cssText = styles.formLabel; + + const input = document.createElement("input"); + input.type = type; + input.id = id; + input.value = value; + input.placeholder = placeholder; + input.style.cssText = styles.formInput; + + group.appendChild(labelEl); + group.appendChild(input); + return { group, input }; + }, + + createWorkerCard() { + const card = document.createElement("div"); + card.style.cssText = styles.workerCard; + + const checkboxColumn = document.createElement("div"); + checkboxColumn.style.cssText = styles.checkboxColumn; + + const contentColumn = document.createElement("div"); + contentColumn.style.cssText = styles.contentColumn; + + const infoRow = document.createElement("div"); + infoRow.style.cssText = styles.infoRow; + + const workerContent = document.createElement("div"); + workerContent.style.cssText = styles.workerContent; + + const settingsArrow = document.createElement("span"); + settingsArrow.className = "settings-arrow"; + settingsArrow.innerHTML = "▶"; + settingsArrow.style.cssText = styles.settingsArrow; + + const controlsDiv = document.createElement("div"); + controlsDiv.style.cssText = styles.controlsDiv; + + return { + card, + checkboxColumn, + contentColumn, + infoRow, + workerContent, + settingsArrow, + controlsDiv + }; + }, + + createInfoBox(text) { + const box = document.createElement("div"); + box.style.cssText = styles.infoBox; + box.textContent = text; + return box; + }, + + addHoverEffect(element, onHover, onLeave) { + element.onmouseover = onHover; + element.onmouseout = onLeave; + }, + + createCard(type = 'worker', options = {}) { + const card = document.createElement("div"); + const baseStyle = "margin-bottom: 12px; border-radius: 4px; overflow: hidden; display: flex;"; + + switch(type) { + case 'master': + case 'worker': + card.style.cssText = baseStyle + "background: #2a2a2a;"; + break; + case 'blueprint': + card.style.cssText = baseStyle + "border: 2px solid #666; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);"; + if (options.onClick) card.onclick = options.onClick; + if (options.title) card.title = options.title; + break; + case 'add': + card.style.cssText = "margin-top: 10px; " + baseStyle + "border: 1px solid #444; cursor: pointer; transition: all 0.2s ease;"; + if (options.onClick) card.onclick = options.onClick; + if (options.title) card.title = options.title; + break; + } + + if (options.onMouseEnter) { + card.addEventListener('mouseenter', options.onMouseEnter); + } + if (options.onMouseLeave) { + card.addEventListener('mouseleave', options.onMouseLeave); + } + + return card; + }, + + createCardColumn(type = 'checkbox', options = {}) { + const column = document.createElement("div"); + const baseStyle = "display: flex; align-items: center; justify-content: center;"; + + switch(type) { + case 'checkbox': + column.style.cssText = baseStyle + "flex: 0 0 40px; border-right: 1px solid #444; cursor: default;"; + if (options.title) column.title = options.title; + break; + case 'icon': + column.style.cssText = baseStyle + "width: 50px; flex-shrink: 0; font-size: 20px; color: #666;"; + break; + case 'content': + column.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; + break; + } + + return column; + }, + + createInfoRow(options = {}) { + const row = document.createElement("div"); + row.style.cssText = "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;"; + if (options.onClick) row.onclick = options.onClick; + return row; + }, + + createWorkerContent() { + const content = document.createElement("div"); + content.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + return content; + }, + + createSettingsForm(fields = [], options = {}) { + const form = document.createElement("div"); + form.style.cssText = "display: flex; flex-direction: column; gap: 10px;"; + + fields.forEach(field => { + if (field.type === 'checkbox') { + const group = document.createElement("div"); + group.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; + + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = field.id; + checkbox.checked = field.checked || false; + if (field.onChange) checkbox.onchange = field.onChange; + + const label = document.createElement("label"); + label.htmlFor = field.id; + label.textContent = field.label; + label.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + group.appendChild(checkbox); + group.appendChild(label); + form.appendChild(group); + } else { + const result = this.createFormGroup(field.label, field.value, field.id, field.type, field.placeholder); + if (field.groupId) result.group.id = field.groupId; + if (field.display) result.group.style.display = field.display; + form.appendChild(result.group); + } + }); + + if (options.buttons) { + const buttonGroup = this.createButtonGroup(options.buttons, options.buttonStyle || " margin-top: 8px;"); + form.appendChild(buttonGroup); + } + + return form; + }, + + createCompleteCard(type, options = {}) { + // Create the main card + const card = this.createCard(type, options); + + // Create left column + let leftColumn; + if (type === 'master' || type === 'worker') { + leftColumn = this.createCardColumn('checkbox', options.checkboxOptions); + + // Add checkbox + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + if (options.checkboxId) checkbox.id = options.checkboxId; + checkbox.checked = options.checked || false; + checkbox.disabled = options.disabled || false; + checkbox.style.cssText = `cursor: ${options.disabled ? 'default' : 'pointer'}; width: 16px; height: 16px;${options.checkboxStyle || ''}`; + if (options.checkboxTitle) checkbox.title = options.checkboxTitle; + + leftColumn.appendChild(checkbox); + + if (options.onCheckboxClick && !options.disabled) { + checkbox.style.pointerEvents = "none"; + leftColumn.style.cursor = "pointer"; + leftColumn.onclick = options.onCheckboxClick; + } + } else { + // For blueprint and add cards, create icon column + leftColumn = document.createElement("div"); + const borderColor = type === 'blueprint' ? '#666' : '#444'; + leftColumn.style.cssText = `flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: ${type === 'blueprint' ? '2px' : '1px'} solid ${borderColor}; color: ${type === 'blueprint' ? '#888' : '#555'}; font-size: ${type === 'blueprint' ? '24px' : '18px'}; font-weight: ${type === 'blueprint' ? 'bold' : 'normal'};`; + leftColumn.innerHTML = "+"; + + if (options.onHover) { + card.addEventListener('mouseenter', () => { + leftColumn.style.color = type === 'blueprint' ? '#aaa' : '#888'; + leftColumn.style.borderColor = type === 'blueprint' ? '#888' : '#666'; + }); + card.addEventListener('mouseleave', () => { + leftColumn.style.color = type === 'blueprint' ? '#888' : '#555'; + leftColumn.style.borderColor = borderColor; + }); + } + } + + // Create right column + const rightColumn = this.createCardColumn('content'); + + // Add hover effect for master/worker cards + if ((type === 'master' || type === 'worker') && options.onHover !== false) { + rightColumn.onmouseover = () => rightColumn.style.backgroundColor = "#333"; + rightColumn.onmouseout = () => rightColumn.style.backgroundColor = "transparent"; + } + + // Create info row + const infoRow = this.createInfoRow(options.infoRowOptions); + + // Create worker content + const workerContent = this.createWorkerContent(); + + // Build the card + infoRow.appendChild(workerContent); + rightColumn.appendChild(infoRow); + + // Add controls div if specified + if (options.includeControls) { + const controlsDiv = document.createElement("div"); + controlsDiv.id = options.controlsId; + controlsDiv.style.cssText = styles.controlsDiv; + rightColumn.appendChild(controlsDiv); + } + + // Add settings div if specified + if (options.includeSettings) { + const settingsDiv = document.createElement("div"); + settingsDiv.id = options.settingsId; + settingsDiv.className = "worker-settings"; + settingsDiv.style.cssText = styles.workerSettings; + rightColumn.appendChild(settingsDiv); + } + + card.appendChild(leftColumn); + card.appendChild(rightColumn); + + return { + card, + leftColumn, + rightColumn, + infoRow, + workerContent, + checkbox: leftColumn.querySelector('input[type="checkbox"]') + }; + }, + + styles // Expose styles for custom usage + }; + } + + // --- State Management --- + _createStateManager() { + const state = { + workers: new Map(), // Unified worker state: { status, managed, launching, expanded, ... } + masterStatus: 'online', + domCache: new Map() // Cache DOM elements + }; + + return { + // Worker state management + getWorker(workerId) { + return state.workers.get(String(workerId)) || {}; + }, + + updateWorker(workerId, updates) { + const id = String(workerId); + const current = state.workers.get(id) || {}; + state.workers.set(id, { ...current, ...updates }); + return state.workers.get(id); + }, + + setWorkerStatus(workerId, status) { + return this.updateWorker(workerId, { status }); + }, + + setWorkerManaged(workerId, info) { + return this.updateWorker(workerId, { managed: info }); + }, + + setWorkerLaunching(workerId, launching) { + return this.updateWorker(workerId, { launching }); + }, + + setWorkerExpanded(workerId, expanded) { + return this.updateWorker(workerId, { expanded }); + }, + + isWorkerLaunching(workerId) { + return this.getWorker(workerId).launching || false; + }, + + isWorkerExpanded(workerId) { + return this.getWorker(workerId).expanded || false; + }, + + isWorkerManaged(workerId) { + return !!this.getWorker(workerId).managed; + }, + + getWorkerStatus(workerId) { + return this.getWorker(workerId).status || {}; + }, + + // Master state + setMasterStatus(status) { + state.masterStatus = status; + }, + + getMasterStatus() { + return state.masterStatus; + }, + + // DOM cache + cacheElement(key, element) { + state.domCache.set(key, element); + }, + + getCachedElement(key) { + return state.domCache.get(key); + }, + + clearCache() { + state.domCache.clear(); + } + }; + } + + // --- API Client --- + _createApiClient() { + const baseUrl = window.location.origin; + + const request = async (endpoint, options = {}) => { + try { + const response = await fetch(`${baseUrl}${endpoint}`, { + headers: { 'Content-Type': 'application/json' }, + ...options + }); + + if (!response.ok) { + const error = await response.json().catch(() => ({ message: 'Request failed' })); + throw new Error(error.message || `HTTP ${response.status}`); + } + + return await response.json(); + } catch (error) { + this.log(`API Error: ${endpoint} - ${error.message}`); + throw error; + } + }; + + return { + // Config endpoints + async getConfig() { + return request('/distributed/config'); + }, + + async updateWorker(workerId, data) { + return request('/distributed/config/update_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId, ...data }) + }); + }, + + async deleteWorker(workerId) { + return request('/distributed/config/delete_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async updateSetting(key, value) { + return request('/distributed/config/update_setting', { + method: 'POST', + body: JSON.stringify({ key, value }) + }); + }, + + async updateMaster(data) { + return request('/distributed/config/update_master', { + method: 'POST', + body: JSON.stringify(data) + }); + }, + + // Worker management endpoints + async launchWorker(workerId) { + return request('/distributed/launch_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async stopWorker(workerId) { + return request('/distributed/stop_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async getManagedWorkers() { + return request('/distributed/managed_workers'); + }, + + async getWorkerLog(workerId, lines = 1000) { + return request(`/distributed/worker_log/${workerId}?lines=${lines}`); + }, + + async clearLaunchingFlag(workerId) { + return request('/distributed/worker/clear_launching', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + // Job preparation + async prepareJob(multiJobId) { + return request('/distributed/prepare_job', { + method: 'POST', + body: JSON.stringify({ multi_job_id: multiJobId }) + }); + }, + + // Image loading + async loadImage(imagePath) { + return request('/distributed/load_image', { + method: 'POST', + body: JSON.stringify({ image_path: imagePath }) + }); + }, + + // Network info + async getNetworkInfo() { + return request('/distributed/network_info'); + }, + + // Status checking (with timeout) + async checkStatus(url, timeout = 900) { + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), timeout); + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: controller.signal + }); + clearTimeout(timeoutId); + + if (!response.ok) throw new Error(`HTTP ${response.status}`); + return await response.json(); + } catch (error) { + clearTimeout(timeoutId); + throw error; + } + }, + + // Batch status checking + async checkMultipleStatuses(urls) { + return Promise.allSettled( + urls.map(url => this.checkStatus(url)) + ); + } + }; + } injectStyles() { const styleId = 'distributed-styles'; @@ -118,16 +680,10 @@ class DistributedExtension { async loadConfig() { try { - const response = await fetch(`${window.location.origin}/distributed/config`); - if (response.ok) { - this.config = await response.json(); - this.log("Loaded config: " + JSON.stringify(this.config)); - } else { - this.log("Failed to load config"); - this.config = { workers: [], settings: {} }; - } + this.config = await this.api.getConfig(); + this.log("Loaded config: " + JSON.stringify(this.config)); } catch (error) { - this.log("Error loading config: " + error.message); + this.log("Failed to load config: " + error.message); this.config = { workers: [], settings: {} }; } } @@ -139,12 +695,7 @@ class DistributedExtension { } try { - await fetch(`${window.location.origin}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId, enabled: enabled }) - }); - + await this.api.updateWorker(workerId, { enabled }); } catch (error) { this.log("Error updating worker: " + error.message); } @@ -158,11 +709,7 @@ class DistributedExtension { this.config.settings[key] = value; try { - await fetch(`${window.location.origin}/distributed/config/update_setting`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ key, value }) - }); + await this.api.updateSetting(key, value); app.extensionManager.toast.add({ severity: "success", @@ -195,39 +742,36 @@ class DistributedExtension { } _createButton(text, onClick, style) { - const button = document.createElement("button"); - button.textContent = text; - button.style.cssText = DistributedExtension.BUTTON_STYLES.base + style; - button.onclick = onClick; - button.className = "distributed-button"; - return button; + return this.ui.createButton(text, onClick, style); } - - async _handleInterruptWorkers(button) { + + // Generic handler for parallel worker operations + async _handleWorkerOperation(button, operation, successText, errorText) { const originalText = button.textContent; const originalStyle = button.style.cssText; - button.textContent = "Interrupting..."; + + button.textContent = operation.loadingText; button.disabled = true; - button.style.backgroundColor = "#c04c4c"; try { - const targetsToInterrupt = this.enabledWorkers.map(w => ({ + const urlsToProcess = this.enabledWorkers.map(w => ({ name: w.name, url: this.getWorkerUrl(w) })); - if (targetsToInterrupt.length === 0) { + if (urlsToProcess.length === 0) { button.textContent = "No Workers"; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + button.style.backgroundColor = "#c04c4c"; setTimeout(() => { button.textContent = originalText; button.style.cssText = originalStyle; + button.disabled = false; }, 3000); return; } - const promises = targetsToInterrupt.map(target => - fetch(`${target.url}/interrupt`, { + const promises = urlsToProcess.map(target => + fetch(`${target.url}${operation.endpoint}`, { method: 'POST', mode: 'cors' }) @@ -239,14 +783,13 @@ class DistributedExtension { const failures = results.filter(r => !r.ok); if (failures.length === 0) { - button.textContent = "Interrupted!"; + button.textContent = successText; button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); - // Update status after interrupt - setTimeout(() => this.checkAllWorkerStatuses(), 500); + if (operation.onSuccess) operation.onSuccess(); } else { - button.textContent = "Error! See Console"; + button.textContent = errorText; button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); - this.log("Failed to interrupt: " + failures.map(f => f.name).join(", ")); + this.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`); } setTimeout(() => { @@ -258,71 +801,21 @@ class DistributedExtension { } } + async _handleInterruptWorkers(button) { + return this._handleWorkerOperation(button, { + name: "Interrupt", + endpoint: "/interrupt", + loadingText: "Interrupting...", + onSuccess: () => setTimeout(() => this.checkAllWorkerStatuses(), 500) + }, "Interrupted!", "Error! See Console"); + } + async _handleClearMemory(button) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; // Store complete original styling - - button.textContent = "Clearing..."; - button.disabled = true; - - try { - const urlsToClear = this.enabledWorkers.map(w => ({ - name: w.name, - url: this.getWorkerUrl(w) - })); - - if (urlsToClear.length === 0) { - button.textContent = "No Workers"; - // Preserve original cssText but update background color - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #c04c4c;'; - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; // Restore complete original styling - button.disabled = false; - }, 3000); - return; - } - - const promises = urlsToClear.map(target => - fetch(`${target.url}/distributed/clear_memory`, { - method: 'POST', - mode: 'cors' - }) - .then(response => ({ ok: response.ok, name: target.name })) - .catch(() => ({ ok: false, name: target.name })) - ); - - const results = await Promise.all(promises); - const failures = results.filter(r => !r.ok); - - if (failures.length === 0) { - button.textContent = "Success!"; - // Preserve original cssText but update background color - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #3ca03c;'; - } else { - button.textContent = "Error! See Console"; - // Preserve original cssText but update background color - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #c04c4c;'; - this.log("Failed to clear memory on: " + failures.map(f => f.name).join(", ")); - } - - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; // Restore complete original styling - button.disabled = false; - }, 3000); - } catch (error) { - // Handle any unexpected errors - button.textContent = "Error! See Console"; - button.style.cssText = originalStyle.replace(/background-color:[^;]*;?/, '') + ' background-color: #c04c4c;'; - this.log("Unexpected error in clear memory: " + error.message); - - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - button.disabled = false; - }, 3000); - } + return this._handleWorkerOperation(button, { + name: "Clear memory", + endpoint: "/distributed/clear_memory", + loadingText: "Clearing..." + }, "Success!", "Error! See Console"); } @@ -437,18 +930,7 @@ class DistributedExtension { // Save each worker using the update endpoint for (const worker of newWorkers) { try { - const response = await fetch(`${window.location.origin}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - worker_id: worker.id, - ...worker - }) - }); - - if (!response.ok) { - this.log(`Failed to save worker ${worker.name}`); - } + await this.api.updateWorker(worker.id, worker); } catch (error) { this.log(`Error saving worker ${worker.name}: ${error.message}`); } @@ -466,13 +948,10 @@ class DistributedExtension { } // Master Node Section - styled exactly like a worker card - const masterDiv = document.createElement("div"); - masterDiv.style.cssText = "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;"; + const masterDiv = this.ui.createCard('master'); // Left column - checkbox area (always checked and disabled) - const checkboxColumn = document.createElement("div"); - checkboxColumn.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444; cursor: default;"; - checkboxColumn.title = "Master node is always enabled"; + const checkboxColumn = this.ui.createCardColumn('checkbox', { title: "Master node is always enabled" }); // Greyed out checkbox that's always checked const checkbox = document.createElement("input"); @@ -484,22 +963,16 @@ class DistributedExtension { checkboxColumn.appendChild(checkbox); // Right column - content area - const contentColumn = document.createElement("div"); - contentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; + const contentColumn = this.ui.createCardColumn('content'); // First row: master info - const infoRow = document.createElement("div"); - infoRow.style.cssText = "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;"; + const infoRow = this.ui.createInfoRow({ onClick: null }); infoRow.title = "Click to expand settings"; // Worker content - const workerContent = document.createElement("div"); - workerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + const workerContent = this.ui.createWorkerContent(); - const statusDot = document.createElement("span"); - statusDot.id = "master-status"; - statusDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; background-color: #4CAF50;"; - statusDot.title = "Online"; + const statusDot = this.ui.createStatusDot("master-status", "#4CAF50", "Online"); const gpuInfo = document.createElement("span"); const cudaInfo = this.masterCudaDevice !== undefined ? `CUDA ${this.masterCudaDevice} • ` : ''; @@ -569,18 +1042,15 @@ class DistributedExtension { const settingsForm = document.createElement("div"); settingsForm.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; - // Name field - reuse createFormGroup - const nameGroup = this.createFormGroup("Name:", this.config?.master?.name || "Master", "master-name"); - settingsForm.appendChild(nameGroup); + // Name field + const nameResult = this.ui.createFormGroup("Name:", this.config?.master?.name || "Master", "master-name"); + settingsForm.appendChild(nameResult.group); // Host field - const hostGroup = this.createFormGroup("Host:", this.config?.master?.host || "", "master-host", "text", "Auto-detect if empty"); - settingsForm.appendChild(hostGroup); + const hostResult = this.ui.createFormGroup("Host:", this.config?.master?.host || "", "master-host", "text", "Auto-detect if empty"); + settingsForm.appendChild(hostResult.group); // Buttons row - using same style as workers - const buttonGroup = document.createElement("div"); - buttonGroup.style.cssText = "display: flex; gap: 8px; margin-top: 8px;"; - const saveBtn = this._createButton("Save", async () => { const nameInput = document.getElementById('master-name'); @@ -615,9 +1085,7 @@ class DistributedExtension { "background-color: #555;"); cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - buttonGroup.appendChild(saveBtn); - buttonGroup.appendChild(cancelBtn); - + const buttonGroup = this.ui.createButtonGroup([saveBtn, cancelBtn], " margin-top: 8px;"); settingsForm.appendChild(buttonGroup); masterSettingsDiv.appendChild(settingsForm); @@ -637,9 +1105,10 @@ class DistributedExtension { // If no workers exist, show a full blueprint placeholder first if (workers.length === 0) { - const blueprintDiv = document.createElement("div"); - blueprintDiv.style.cssText = "margin-bottom: 12px; border: 2px solid #666; border-radius: 4px; overflow: hidden; display: flex; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);"; - blueprintDiv.title = "Click to add your first worker"; + const blueprintDiv = this.ui.createCard('blueprint', { + title: "Click to add your first worker", + onClick: () => this.addNewWorker() + }); // Left column - plus icon const blueprintCheckbox = document.createElement("div"); @@ -647,18 +1116,15 @@ class DistributedExtension { blueprintCheckbox.innerHTML = "+"; // Right column - const blueprintContent = document.createElement("div"); - blueprintContent.style.cssText = "flex: 1; display: flex; flex-direction: column;"; + const blueprintContent = this.ui.createCardColumn('content'); // Info row with placeholder content - const blueprintInfo = document.createElement("div"); - blueprintInfo.style.cssText = "display: flex; align-items: center; padding: 8px; min-height: 60px;"; + const blueprintInfo = this.ui.createInfoRow(); - const blueprintWorkerContent = document.createElement("div"); - blueprintWorkerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + const blueprintWorkerContent = this.ui.createWorkerContent(); - const blueprintDot = document.createElement("span"); - blueprintDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; border: 1px solid #555;"; + const blueprintDot = this.ui.createStatusDot(null, "transparent", ""); + blueprintDot.style.border = "1px solid #555"; const blueprintText = document.createElement("span"); blueprintText.innerHTML = `Add New Worker
[CUDA] • [Port]`; @@ -683,9 +1149,6 @@ class DistributedExtension { blueprintDiv.appendChild(blueprintCheckbox); blueprintDiv.appendChild(blueprintContent); - // Click handler - blueprintDiv.onclick = () => this.addNewWorker(); - // Hover effect blueprintDiv.onmouseover = () => { blueprintDiv.style.borderColor = "#888"; @@ -705,12 +1168,11 @@ class DistributedExtension { // Show existing workers workers.forEach(worker => { - const gpuDiv = document.createElement("div"); - gpuDiv.style.cssText = "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;"; + const gpuDiv = this.ui.createCard('worker'); // Left column - checkbox area (extends full height) - const checkboxColumn = document.createElement("div"); - checkboxColumn.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444; cursor: pointer;"; + const checkboxColumn = this.ui.createCardColumn('checkbox'); + checkboxColumn.style.cursor = "pointer"; checkboxColumn.title = "Enable/disable this worker"; // Standard checkbox @@ -730,22 +1192,16 @@ class DistributedExtension { }; // Right column - content area - const contentColumn = document.createElement("div"); - contentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; + const contentColumn = this.ui.createCardColumn('content'); // First row: worker info - const infoRow = document.createElement("div"); - infoRow.style.cssText = "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;"; + const infoRow = this.ui.createInfoRow(); infoRow.title = "Click to expand settings"; // Worker content - const workerContent = document.createElement("div"); - workerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + const workerContent = this.ui.createWorkerContent(); - const statusDot = document.createElement("span"); - statusDot.id = `status-${worker.id}`; - statusDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; background-color: #666;"; - statusDot.title = "Checking status..."; + const statusDot = this.ui.createStatusDot(`status-${worker.id}`, "#666", "Checking status..."); const gpuInfo = document.createElement("span"); const isRemote = this.isRemoteWorker(worker); @@ -767,7 +1223,7 @@ class DistributedExtension { settingsArrow.className = "settings-arrow"; settingsArrow.innerHTML = "▶"; settingsArrow.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;"; - if (this.expandedWorkers.has(worker.id)) { + if (this.state.isWorkerExpanded(worker.id)) { settingsArrow.style.transform = "rotate(90deg)"; } @@ -809,39 +1265,40 @@ class DistributedExtension { controlsDiv.appendChild(remoteInfo); } else { // Local workers get launch/stop/log buttons - const launchBtn = this._createButton('Launch', - () => this.launchWorker(worker.id), - 'background-color: #555;'); // Neutral color initially - launchBtn.style.cssText += ' padding: 3px 8px; font-size: 11px;'; + const controls = this.ui.createWorkerControls(worker.id, { + launch: () => this.launchWorker(worker.id), + stop: () => this.stopWorker(worker.id), + viewLog: () => this.viewWorkerLog(worker.id) + }); + + // Apply custom button styles + const launchBtn = controls.querySelector(`#launch-${worker.id}`); + const stopBtn = controls.querySelector(`#stop-${worker.id}`); + const logBtn = controls.querySelector(`#log-${worker.id}`); + + launchBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; background-color: #555;'; launchBtn.title = "Launch worker (runs in background with logging)"; - const stopBtn = this._createButton('Stop', - () => this.stopWorker(worker.id), - 'background-color: #555;'); // Neutral color initially - stopBtn.style.cssText += ' padding: 3px 8px; font-size: 11px;'; + stopBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; background-color: #555;'; stopBtn.title = "Stop worker"; - const logBtn = this._createButton('View Log', - () => this.viewWorkerLog(worker.id), - 'background-color: #685434;'); // ComfyUI darker yellow - logBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; display: none;'; - logBtn.id = `log-${worker.id}`; - logBtn.title = "View worker log file"; + logBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; display: none; background-color: #685434;'; - controlsDiv.appendChild(launchBtn); - controlsDiv.appendChild(stopBtn); - controlsDiv.appendChild(logBtn); + // Move controls into the controlsDiv + while (controls.firstChild) { + controlsDiv.appendChild(controls.firstChild); + } } // Third row: expandable settings panel const settingsDiv = document.createElement("div"); settingsDiv.id = `settings-${worker.id}`; settingsDiv.className = "worker-settings"; - if (this.expandedWorkers.has(worker.id)) { + if (this.state.isWorkerExpanded(worker.id)) { settingsDiv.classList.add("expanded"); } settingsDiv.style.cssText = "margin: 0 8px; padding: 0 12px; background: #1e1e1e; border-radius: 4px;"; - if (this.expandedWorkers.has(worker.id)) { + if (this.state.isWorkerExpanded(worker.id)) { settingsDiv.style.padding = "12px"; settingsDiv.style.marginTop = "8px"; settingsDiv.style.marginBottom = "8px"; @@ -866,9 +1323,10 @@ class DistributedExtension { // Only show the minimal "Add Worker" box if there are existing workers if (workers.length > 0) { // Add Worker placeholder box - const addWorkerDiv = document.createElement("div"); - addWorkerDiv.style.cssText = "margin-top: 10px; margin-bottom: 12px; border: 1px solid #444; border-radius: 4px; overflow: hidden; display: flex; cursor: pointer; transition: all 0.2s ease;"; - addWorkerDiv.title = "Click to add a new worker"; + const addWorkerDiv = this.ui.createCard('add', { + title: "Click to add a new worker", + onClick: () => this.addNewWorker() + }); // Left column - plus icon area const plusColumn = document.createElement("div"); @@ -876,20 +1334,18 @@ class DistributedExtension { plusColumn.innerHTML = "+"; // Right column - content area - const addContentColumn = document.createElement("div"); - addContentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column;"; + const addContentColumn = this.ui.createCardColumn('content'); // Info row - const addInfoRow = document.createElement("div"); - addInfoRow.style.cssText = "display: flex; align-items: center; padding: 8px; min-height: 42px;"; // Reduced from 60px to 42px (30% reduction) + const addInfoRow = this.ui.createInfoRow(); + addInfoRow.style.minHeight = "42px"; // Reduced from 60px to 42px (30% reduction) // Placeholder content - const addContent = document.createElement("div"); - addContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + const addContent = this.ui.createWorkerContent(); // Placeholder dot - const placeholderDot = document.createElement("span"); - placeholderDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; border: 1px solid #555;"; + const placeholderDot = this.ui.createStatusDot(null, "transparent", ""); + placeholderDot.style.border = "1px solid #555"; // Placeholder text const placeholderInfo = document.createElement("span"); @@ -918,9 +1374,6 @@ class DistributedExtension { plusColumn.style.borderColor = "#444"; }; - // Click handler - addWorkerDiv.onclick = () => this.addNewWorker(); - gpuSection.appendChild(addWorkerDiv); } @@ -1544,11 +1997,7 @@ class DistributedExtension { async _prepareDistributedJob(multi_job_id) { try { - await fetch(`${window.location.origin}/distributed/prepare_job`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ multi_job_id }) - }); + await this.api.prepareJob(multi_job_id); } catch (error) { this.log("Error preparing job: " + error.message); throw error; @@ -1611,8 +2060,8 @@ class DistributedExtension { if (!this.config || !this.config.workers) return; for (const worker of this.config.workers) { - // Only check status for enabled workers - if (worker.enabled) { + // Check status for enabled workers OR workers that are launching + if (worker.enabled || this.state.isWorkerLaunching(worker.id)) { this.checkWorkerStatus(worker); } } @@ -1684,7 +2133,7 @@ class DistributedExtension { const isProcessing = queueRemaining > 0; // Update status - this.workerStatuses.set(worker.id, { + this.state.setWorkerStatus(worker.id, { online: true, processing: isProcessing, queueCount: queueRemaining @@ -1698,8 +2147,8 @@ class DistributedExtension { } // Clear launching state since worker is now online - if (this.launchingWorkers.has(worker.id)) { - this.launchingWorkers.delete(worker.id); + if (this.state.isWorkerLaunching(worker.id)) { + this.state.setWorkerLaunching(worker.id, false); this.clearLaunchingFlag(worker.id); } } else { @@ -1707,14 +2156,14 @@ class DistributedExtension { } } catch (error) { // Worker is offline or unreachable - this.workerStatuses.set(worker.id, { + this.state.setWorkerStatus(worker.id, { online: false, processing: false, queueCount: 0 }); // Check if worker is launching - if (this.launchingWorkers.has(worker.id)) { + if (this.state.isWorkerLaunching(worker.id)) { this.updateStatusDot(worker.id, "#f0ad4e", "Launching...", true); } else { // Only update to red if not currently launching @@ -1796,14 +2245,8 @@ class DistributedExtension { } // Load image from master's filesystem via API - const response = await fetch(`${window.location.origin}/distributed/load_image`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ image_path: filename }) - }); - - if (response.ok) { - const data = await response.json(); + try { + const data = await this.api.loadImage(filename); const imageData = { name: filename, image: data.image_data @@ -1813,8 +2256,9 @@ class DistributedExtension { // Cache the image for future use this._imageCache.set(filename, imageData); this.debugLog(`Loaded and cached image: ${filename}`); - } else { - this.log(`Failed to load image ${filename}: ${response.statusText}`); + } catch (loadError) { + this.log(`Failed to load image ${filename}: ${loadError.message}`); + throw loadError; } } catch (error) { this.log(`Error loading image ${filename}: ${error.message}`); @@ -1939,11 +2383,11 @@ class DistributedExtension { const launchBtn = document.querySelector(`#controls-${workerId} button`); this.updateStatusDot(workerId, "#f0ad4e", "Launching...", true); - this.launchingWorkers.add(workerId); + this.state.setWorkerLaunching(workerId, true); // Allow 90 seconds for worker to launch (model loading can take time) setTimeout(() => { - this.launchingWorkers.delete(workerId); + this.state.setWorkerLaunching(workerId, false); }, 90000); if (!launchBtn) return; @@ -1952,20 +2396,14 @@ class DistributedExtension { // Disable button immediately launchBtn.disabled = true; - const response = await fetch(`${window.location.origin}/distributed/launch_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - - const result = await response.json(); - if (response.ok) { + const result = await this.api.launchWorker(workerId); + if (result) { this.log(`Launched ${worker.name} (PID: ${result.pid})`); if (result.log_file) { this.debugLog(`Log file: ${result.log_file}`); } - this.managedWorkers.set(String(workerId), { + this.state.setWorkerManaged(workerId, { pid: result.pid, log_file: result.log_file, started_at: Date.now() @@ -1974,31 +2412,20 @@ class DistributedExtension { // Update controls immediately to hide launch button and show stop/log buttons this.updateWorkerControls(workerId); setTimeout(() => this.checkWorkerStatus(worker), 2000); - } else { - throw new Error(result.message || 'Failed to launch worker'); - } - - // Check if worker was already running - if (!response.ok && result && result.message && result.message.includes("already running")) { - if (result.log_file) { - const existing = this.managedWorkers.get(String(workerId)); - if (!existing) { - this.managedWorkers.set(String(workerId), { - log_file: result.log_file - }); - } else { - existing.log_file = result.log_file; - } - } - this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 100); } } catch (error) { - this.log(`Error launching worker: ${error}`); - - // Re-enable button on error - if (launchBtn) { - launchBtn.disabled = false; + // Check if worker was already running + if (error.message && error.message.includes("already running")) { + this.log(`Worker ${worker.name} is already running`); + this.updateWorkerControls(workerId); + setTimeout(() => this.checkWorkerStatus(worker), 100); + } else { + this.log(`Error launching worker: ${error.message || error}`); + + // Re-enable button on error + if (launchBtn) { + launchBtn.disabled = false; + } } } } @@ -2015,16 +2442,10 @@ class DistributedExtension { } try { - const response = await fetch(`${window.location.origin}/distributed/stop_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - - const result = await response.json(); - if (response.ok) { + const result = await this.api.stopWorker(workerId); + if (result) { this.log(`Stopped worker: ${result.message}`); - this.managedWorkers.delete(String(workerId)); + this.state.setWorkerManaged(workerId, null); // Flash success feedback if (stopBtn) { @@ -2065,14 +2486,10 @@ class DistributedExtension { async clearLaunchingFlag(workerId) { try { - await fetch(`${window.location.origin}/distributed/worker/clear_launching`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); + await this.api.clearLaunchingFlag(workerId); this.debugLog(`Cleared launching flag for worker ${workerId}`); } catch (error) { - this.debugLog(`Error clearing launching flag: ${error}`); + this.debugLog(`Error clearing launching flag: ${error.message || error}`); } } @@ -2146,27 +2563,25 @@ class DistributedExtension { async loadManagedWorkers() { try { - const response = await fetch(`${window.location.origin}/distributed/managed_workers`); - if (response.ok) { - const result = await response.json(); - this.managedWorkers.clear(); + const result = await this.api.getManagedWorkers(); + // Clear all managed worker info + this.state.clearCache(); + + // Check for launching workers + for (const [workerId, info] of Object.entries(result.managed_workers)) { + this.state.setWorkerManaged(workerId, info); - // Check for launching workers - for (const [workerId, info] of Object.entries(result.managed_workers)) { - this.managedWorkers.set(workerId, info); - - // If worker is marked as launching, add to launchingWorkers set - if (info.launching) { - this.launchingWorkers.add(parseInt(workerId)); - this.debugLog(`Worker ${workerId} is in launching state`); - } - } - - // Update UI for all workers - if (this.config?.workers) { - this.config.workers.forEach(w => this.updateWorkerControls(w.id)); + // If worker is marked as launching, add to launchingWorkers set + if (info.launching) { + this.state.setWorkerLaunching(workerId, true); + this.debugLog(`Worker ${workerId} is in launching state`); } } + + // Update UI for all workers + if (this.config?.workers) { + this.config.workers.forEach(w => this.updateWorkerControls(w.id)); + } } catch (error) { this.debugLog(`Error loading managed workers: ${error}`); } @@ -2186,8 +2601,8 @@ class DistributedExtension { } // Ensure we check for string ID - const managedInfo = this.managedWorkers.get(String(workerId)); - const status = this.workerStatuses.get(workerId); + const managedInfo = this.state.getWorker(workerId).managed; + const status = this.state.getWorkerStatus(workerId); // Update button states const buttons = controlsDiv.querySelectorAll('button'); @@ -2228,7 +2643,7 @@ class DistributedExtension { } async viewWorkerLog(workerId) { - const managedInfo = this.managedWorkers.get(String(workerId)); + const managedInfo = this.state.getWorker(workerId).managed; if (!managedInfo?.log_file) return; const logBtn = document.getElementById(`log-${workerId}`); @@ -2242,12 +2657,7 @@ class DistributedExtension { try { // Fetch log content - const response = await fetch(`${window.location.origin}/distributed/worker_log/${workerId}?lines=1000`); - if (!response.ok) { - throw new Error(`Failed to fetch log: ${response.statusText}`); - } - - const data = await response.json(); + const data = await this.api.getWorkerLog(workerId, 1000); // Create modal dialog this.showLogModal(workerId, data); @@ -2449,10 +2859,7 @@ class DistributedExtension { if (!logContent) return; try { - const response = await fetch(`${window.location.origin}/distributed/worker_log/${workerId}?lines=1000`); - if (!response.ok) throw new Error('Failed to fetch log'); - - const data = await response.json(); + const data = await this.api.getWorkerLog(workerId, 1000); // Update content const shouldAutoScroll = logContent.scrollTop + logContent.clientHeight >= logContent.scrollHeight - 50; @@ -2486,21 +2893,6 @@ class DistributedExtension { } } - copyLogPath(logPath) { - if (navigator.clipboard) { - navigator.clipboard.writeText(logPath).then(() => { - app.extensionManager.toast.add({ - severity: "success", - summary: "Path Copied", - detail: "Log file path copied to clipboard", - life: 2000 - }); - }).catch(err => { - this.debugLog('Failed to copy: ' + err.message); - }); - } - } - formatFileSize(bytes) { if (bytes < 1024) return bytes + ' B'; if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; @@ -2554,13 +2946,7 @@ class DistributedExtension { async detectMasterIP() { try { - const response = await fetch(`${window.location.origin}/distributed/network_info`); - if (!response.ok) { - this.debugLog("Failed to get network info"); - return; - } - - const data = await response.json(); + const data = await this.api.getNetworkInfo(); this.log("Network info: " + JSON.stringify(data)); // Store CUDA device info @@ -2595,7 +2981,6 @@ class DistributedExtension { // Update config and save this.config.workers = newWorkers; - await this.saveConfig(); // Refresh the panel if it's already open const panel = document.getElementById("distributed-panel"); @@ -2695,8 +3080,8 @@ class DistributedExtension { if (!settingsDiv) return; - if (this.expandedWorkers.has(workerId)) { - this.expandedWorkers.delete(workerId); + if (this.state.isWorkerExpanded(workerId)) { + this.state.setWorkerExpanded(workerId, false); settingsDiv.classList.remove("expanded"); if (settingsArrow) { settingsArrow.style.transform = "rotate(0deg)"; @@ -2706,7 +3091,7 @@ class DistributedExtension { settingsDiv.style.marginTop = "0"; settingsDiv.style.marginBottom = "0"; } else { - this.expandedWorkers.add(workerId); + this.state.setWorkerExpanded(workerId, true); settingsDiv.classList.add("expanded"); if (settingsArrow) { settingsArrow.style.transform = "rotate(90deg)"; @@ -2835,26 +3220,8 @@ class DistributedExtension { } createFormGroup(label, value, id, type = "text", placeholder = "") { - const group = document.createElement("div"); - group.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const labelEl = document.createElement("label"); - labelEl.textContent = label; - labelEl.style.cssText = "font-size: 12px; color: #ccc; width: 80px;"; - - const input = document.createElement("input"); - input.type = type; - input.id = id; - input.value = value; - if (placeholder) { - input.placeholder = placeholder; - } - input.style.cssText = "flex: 1; padding: 4px 8px; background: #2a2a2a; border: 1px solid #444; border-radius: 3px; color: #fff; font-size: 12px;"; - - group.appendChild(labelEl); - group.appendChild(input); - - return group; + const result = this.ui.createFormGroup(label, value, id, type, placeholder); + return result.group; } async saveWorkerSettings(workerId) { @@ -3130,7 +3497,7 @@ class DistributedExtension { }); // Refresh UI and expand the new worker - this.expandedWorkers.add(newId); + this.state.setWorkerExpanded(newId, true); if (this.panelElement) { this.renderSidebarContent(this.panelElement); } From f5aa01b48b66416f813cdf2d6125270cd9dbae61 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sun, 13 Jul 2025 09:26:50 +1000 Subject: [PATCH 04/93] major refactoring of js --- web/apiClient.js | 144 ++ web/constants.js | 92 ++ web/executionUtils.js | 447 +++++ web/gpupanel.js | 3521 ---------------------------------------- web/main.js | 1178 ++++++++++++++ web/sidebarRenderer.js | 273 ++++ web/stateManager.js | 61 + web/ui.js | 997 ++++++++++++ web/workerUtils.js | 266 +++ 9 files changed, 3458 insertions(+), 3521 deletions(-) create mode 100644 web/apiClient.js create mode 100644 web/constants.js create mode 100644 web/executionUtils.js delete mode 100644 web/gpupanel.js create mode 100644 web/main.js create mode 100644 web/sidebarRenderer.js create mode 100644 web/stateManager.js create mode 100644 web/ui.js create mode 100644 web/workerUtils.js diff --git a/web/apiClient.js b/web/apiClient.js new file mode 100644 index 0000000..061428d --- /dev/null +++ b/web/apiClient.js @@ -0,0 +1,144 @@ +export function createApiClient(baseUrl) { + const request = async (endpoint, options = {}, retries = 3) => { + let lastError; + let delay = 1000; // Initial delay for exponential backoff + + for (let attempt = 0; attempt < retries; attempt++) { + try { + const response = await fetch(`${baseUrl}${endpoint}`, { + headers: { 'Content-Type': 'application/json' }, + ...options + }); + + if (!response.ok) { + const error = await response.json().catch(() => ({ message: 'Request failed' })); + throw new Error(error.message || `HTTP ${response.status}`); + } + + return await response.json(); + } catch (error) { + lastError = error; + console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`); + if (attempt < retries - 1) { + await new Promise(resolve => setTimeout(resolve, delay)); + delay *= 2; // Exponential backoff + } + } + } + throw lastError; + }; + + return { + // Config endpoints + async getConfig() { + return request('/distributed/config'); + }, + + async updateWorker(workerId, data) { + return request('/distributed/config/update_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId, ...data }) + }); + }, + + async deleteWorker(workerId) { + return request('/distributed/config/delete_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async updateSetting(key, value) { + return request('/distributed/config/update_setting', { + method: 'POST', + body: JSON.stringify({ key, value }) + }); + }, + + async updateMaster(data) { + return request('/distributed/config/update_master', { + method: 'POST', + body: JSON.stringify(data) + }); + }, + + // Worker management endpoints + async launchWorker(workerId) { + return request('/distributed/launch_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async stopWorker(workerId) { + return request('/distributed/stop_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + async getManagedWorkers() { + return request('/distributed/managed_workers'); + }, + + async getWorkerLog(workerId, lines = 1000) { + return request(`/distributed/worker_log/${workerId}?lines=${lines}`); + }, + + async clearLaunchingFlag(workerId) { + return request('/distributed/worker/clear_launching', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }) + }); + }, + + // Job preparation + async prepareJob(multiJobId) { + return request('/distributed/prepare_job', { + method: 'POST', + body: JSON.stringify({ multi_job_id: multiJobId }) + }); + }, + + // Image loading + async loadImage(imagePath) { + return request('/distributed/load_image', { + method: 'POST', + body: JSON.stringify({ image_path: imagePath }) + }); + }, + + // Network info + async getNetworkInfo() { + return request('/distributed/network_info'); + }, + + // Status checking (with timeout) + async checkStatus(url, timeout = 5000) { + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), timeout); + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: controller.signal + }); + clearTimeout(timeoutId); + + if (!response.ok) throw new Error(`HTTP ${response.status}`); + return await response.json(); + } catch (error) { + clearTimeout(timeoutId); + throw error; + } + }, + + // Batch status checking + async checkMultipleStatuses(urls) { + return Promise.allSettled( + urls.map(url => this.checkStatus(url)) + ); + } + }; +} \ No newline at end of file diff --git a/web/constants.js b/web/constants.js new file mode 100644 index 0000000..223807d --- /dev/null +++ b/web/constants.js @@ -0,0 +1,92 @@ +export const BUTTON_STYLES = { + base: "width: 100%; padding: 8px 12px; font-size: 12px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: all 0.2s; font-weight: 500;", + clearMemory: "background-color: #555;", + success: "background-color: #3ca03c;", + error: "background-color: #c04c4c;", + launch: "background-color: #4a7c4a;", + stop: "background-color: #7c4a4a;", + grey: "background-color: #555;", +}; + +export const STATUS_COLORS = { + DISABLED_GRAY: "#666", + OFFLINE_RED: "#c04c4c", + ONLINE_GREEN: "#3ca03c", + PROCESSING_YELLOW: "#f0ad4e" +}; + +export const PULSE_ANIMATION_CSS = ` + @keyframes pulse { + 0% { + opacity: 1; + transform: scale(1); + box-shadow: 0 0 0 0 rgba(240, 173, 78, 0.7); + } + 50% { + opacity: 0.3; + transform: scale(1.3); + box-shadow: 0 0 0 6px rgba(240, 173, 78, 0); + } + 100% { + opacity: 1; + transform: scale(1); + box-shadow: 0 0 0 0 rgba(240, 173, 78, 0); + } + } + .status-pulsing { + animation: pulse 1.2s ease-in-out infinite; + } + + /* Button hover effects */ + .distributed-button:hover:not(:disabled) { + filter: brightness(1.2); + transition: filter 0.2s ease; + } + .distributed-button:disabled { + opacity: 0.6; + cursor: not-allowed; + } + + /* Settings button animation */ + .settings-btn { + transition: transform 0.2s ease; + } + + + /* Expanded settings panel */ + .worker-settings { + max-height: 0; + overflow: hidden; + opacity: 0; + transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease; + } + .worker-settings.expanded { + max-height: 500px; + opacity: 1; + padding: 12px 0; + } +`; + +export const UI_STYLES = { + statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;", + controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;", + formGroup: "display: flex; flex-direction: column; gap: 5px;", + formLabel: "font-size: 12px; color: #ccc; font-weight: 500;", + formInput: "padding: 6px 10px; background: #2a2a2a; border: 1px solid #444; color: white; font-size: 12px; border-radius: 4px; transition: border-color 0.2s;", + workerCard: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex; background: #2a2a2a;", + checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);", + contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;", + infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer;", + workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;", + settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto; padding: 4px;", + infoBox: "background-color: #333; color: #999; padding: 5px 14px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;", + workerSettings: "margin: 0 12px; padding: 0 12px; background: #1e1e1e; border-radius: 4px; border: 1px solid #2a2a2a;" +}; + +export const TIMEOUTS = { + DEFAULT_FETCH: 5000, // ms for general API calls + STATUS_CHECK: 2000, // ms for status checks + LAUNCH: 90000, // ms for worker launch (longer for model loading) + RETRY_DELAY: 1000, // initial delay for exponential backoff + MAX_RETRIES: 3 // max retry attempts +}; \ No newline at end of file diff --git a/web/executionUtils.js b/web/executionUtils.js new file mode 100644 index 0000000..d26fcba --- /dev/null +++ b/web/executionUtils.js @@ -0,0 +1,447 @@ +import { api } from "../../scripts/api.js"; +import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker } from './workerUtils.js'; +import { TIMEOUTS } from './constants.js'; + +export function setupInterceptor(extension) { + api.queuePrompt = async (number, prompt) => { + if (extension.isEnabled) { + const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0; + const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0; + + if (hasCollector || hasDistUpscale) { + const result = await executeParallelDistributed(extension, prompt); + // Check status after dispatching jobs + setTimeout(() => extension.checkAllWorkerStatuses(), 500); + return result; + } + } + return extension.originalQueuePrompt(number, prompt); + }; +} + +export async function executeParallelDistributed(extension, promptWrapper) { + try { + const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution + const enabledWorkers = extension.enabledWorkers; + + // Pre-flight health check on all enabled workers + const activeWorkers = await performPreflightCheck(extension, enabledWorkers); + + if (activeWorkers.length === 0 && enabledWorkers.length > 0) { + extension.log("No active workers found. Running on master only."); + // Fall back to master-only execution + return extension.originalQueuePrompt(0, promptWrapper); + } + + extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug"); + + // Find all distributed nodes in the workflow + const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector"); + const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); + const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + + // Map original node IDs to truly unique job IDs for this specific run + const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`])); + + // Prepare a separate job queue on the backend for each unique job ID + const preparePromises = Array.from(job_id_map.values()).map(uniqueId => prepareDistributedJob(extension, uniqueId)); + await Promise.all(preparePromises); + + const jobs = []; + // Use only active workers + const participants = ['master', ...activeWorkers.map(w => w.id)]; + + for (const participantId of participants) { + const options = { + enabled_worker_ids: activeWorkers.map(w => w.id), + workflow: promptWrapper.workflow, + job_id_map: job_id_map // Pass the map of unique IDs + }; + + const jobApiPrompt = prepareApiPromptForParticipant( + extension, promptWrapper.output, participantId, options + ); + + if (participantId === 'master') { + jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); + } else { + const worker = activeWorkers.find(w => w.id === participantId); + if (worker) { + const job = { + type: 'worker', + worker, + prompt: jobApiPrompt, + workflow: promptWrapper.workflow + }; + + // Add image references if found for remote workers + if (options._imageReferences) { + job.imageReferences = options._imageReferences; + } + + jobs.push(job); + } + } + } + + const result = await executeJobs(extension, jobs); + return result; + } catch (error) { + extension.log("Parallel execution failed: " + error.message, "error"); + throw error; + } +} + +export function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) { + let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); + const isMaster = participantId === 'master'; + + // Find all distributed nodes once (before pruning) + const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector"); + const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed"); + const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; + + // For workers, prune the workflow to only include distributed node dependencies + if (!isMaster && allDistributedNodes.length > 0) { + jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes); + } + + // Handle image references for remote workers + if (!isMaster && options.enabled_worker_ids) { + // Check if this is a remote worker + const workerId = participantId; + const workerInfo = extension.config.workers.find(w => w.id === workerId); + const isRemote = workerInfo && workerInfo.host; + + if (isRemote) { + // Find all image references in the pruned workflow + const imageReferences = findImageReferences(extension, jobApiPrompt); + if (imageReferences.size > 0) { + extension.log(`Found ${imageReferences.size} image references for remote worker ${workerId}`, "debug"); + // Store image references for later processing + options._imageReferences = imageReferences; + } + } + } + + // Handle Distributed seed nodes + const distributorNodes = findNodesByClass(jobApiPrompt, "DistributedSeed"); + if (distributorNodes.length > 0) { + extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug"); + } + + for (const seedNode of distributorNodes) { + const { inputs } = jobApiPrompt[seedNode.id]; + inputs.is_worker = !isMaster; + if (!isMaster) { + const workerIndex = options.enabled_worker_ids.indexOf(participantId); + inputs.worker_id = `worker_${workerIndex}`; + extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug"); + } + } + + // Handle Distributed collector nodes (already found above) + for (const collector of collectorNodes) { + const { inputs } = jobApiPrompt[collector.id]; + + // Check if this collector is downstream from a distributed upscaler + const hasUpstreamDistributedUpscaler = hasUpstreamNode( + jobApiPrompt, + collector.id, + 'UltimateSDUpscaleDistributed' + ); + + if (hasUpstreamDistributedUpscaler) { + // Set pass_through mode for this collector + inputs.pass_through = true; + extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug"); + } else { + // Normal collector behavior + // Get the unique job ID from the map created for this execution + const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id; + + // Use the truly unique ID for this execution + inputs.multi_job_id = uniqueJobId; + inputs.is_worker = !isMaster; + if (isMaster) { + inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + } else { + inputs.master_url = extension.getMasterUrl(); + // Also make the worker_job_id unique to prevent potential caching issues + inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; + inputs.worker_id = participantId; + } + } + } + + // Handle Ultimate SD Upscale Distributed nodes + for (const upscaleNode of upscaleNodes) { + const { inputs } = jobApiPrompt[upscaleNode.id]; + + // Get the unique job ID from the map + const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id; + + inputs.multi_job_id = uniqueJobId; + inputs.is_worker = !isMaster; + + if (isMaster) { + inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + } else { + inputs.master_url = extension.getMasterUrl(); + inputs.worker_id = participantId; + // Workers also need the enabled_worker_ids to calculate tile distribution + inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + } + } + + return jobApiPrompt; +} + +export async function prepareDistributedJob(extension, multi_job_id) { + try { + await extension.api.prepareJob(multi_job_id); + } catch (error) { + extension.log("Error preparing job: " + error.message, "error"); + throw error; + } +} + +export async function executeJobs(extension, jobs) { + let masterPromptId = null; + + // Pre-load all unique images before dispatching to workers + const allImageReferences = new Map(); + for (const job of jobs) { + if (job.type === 'worker' && job.imageReferences) { + for (const [filename, info] of job.imageReferences) { + allImageReferences.set(filename, info); + } + } + } + + if (allImageReferences.size > 0) { + extension.log(`Pre-loading ${allImageReferences.size} unique image(s) for all workers`, "debug"); + await loadImagesForWorker(extension, allImageReferences); + } + + // Now dispatch jobs in parallel + const promises = jobs.map(job => { + if (job.type === 'master') { + return extension.originalQueuePrompt(0, job.promptWrapper).then(result => { + masterPromptId = result; + return result; + }); + } else { + return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences); + } + }); + await Promise.all(promises); + return masterPromptId || { "prompt_id": "distributed-job-dispatched" }; +} + +async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) { + const workerUrl = extension.getWorkerUrl(worker); + + // Debug logging - always log to console for debugging + extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug"); + extension.log('[Distributed] Worker URL: ' + workerUrl, "debug"); + + // Handle image uploads for remote workers + if (imageReferences && imageReferences.size > 0) { + extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug"); + + try { + // Load images from master + const images = await loadImagesForWorker(extension, imageReferences); + + // Upload images to worker + if (images.length > 0) { + await uploadImagesToWorker(extension, workerUrl, images); + extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug"); + } + } catch (error) { + extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error"); + // Continue with workflow execution even if image upload fails + } + } + + const promptToSend = { + prompt, + extra_data: { extra_pnginfo: { workflow } }, + client_id: api.clientId + }; + + extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug"); + + try { + await fetch(`${workerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify(promptToSend) + }); + } catch (e) { + extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error"); + } +} + +export async function loadImagesForWorker(extension, imageReferences) { + const images = []; + + // Use a cache to avoid loading the same image multiple times + if (!extension._imageCache) { + extension._imageCache = new Map(); + } + + for (const [filename, info] of imageReferences) { + try { + // Check cache first + if (extension._imageCache.has(filename)) { + images.push(extension._imageCache.get(filename)); + extension.log(`Using cached image: ${filename}`, "debug"); + continue; + } + + // Limit cache size + if (extension._imageCache.size >= 10) { + const oldestKey = extension._imageCache.keys().next().value; + extension._imageCache.delete(oldestKey); + extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug"); + } + + // Load image from master's filesystem via API + try { + const data = await extension.api.loadImage(filename); + const imageData = { + name: filename, + image: data.image_data + }; + images.push(imageData); + + // Cache the image for future use + extension._imageCache.set(filename, imageData); + extension.log(`Loaded and cached image: ${filename}`, "debug"); + } catch (loadError) { + extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error"); + throw loadError; + } + } catch (error) { + extension.log(`Error loading image ${filename}: ${error.message}`, "error"); + } + } + + // Clear cache after a reasonable time to avoid memory issues + setTimeout(() => { + if (extension._imageCache && extension._imageCache.size > 0) { + extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug"); + extension._imageCache.clear(); + } + }, 30000); // Clear after 30 seconds + + return images; +} + +export async function uploadImagesToWorker(extension, workerUrl, images) { + // Upload images to worker's ComfyUI instance + for (const imageData of images) { + const formData = new FormData(); + + // Detect MIME type from base64 header + const mimeMatch = imageData.image.match(/^data:(image\/\w+);base64,/); + const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected + + // Convert base64 to blob + const base64Data = imageData.image.replace(/^data:image\/\w+;base64,/, ''); + const byteCharacters = atob(base64Data); + const byteNumbers = new Array(byteCharacters.length); + for (let i = 0; i < byteCharacters.length; i++) { + byteNumbers[i] = byteCharacters.charCodeAt(i); + } + const byteArray = new Uint8Array(byteNumbers); + const blob = new Blob([byteArray], { type: mimeType }); + + // Use original filename without heavy cleaning + let cleanName = imageData.name; + let subfolder = ''; + + // Extract subfolder if present + if (cleanName.includes('/')) { + const parts = cleanName.split('/'); + subfolder = parts.slice(0, -1).join('/'); + cleanName = parts[parts.length - 1]; + } + + formData.append('image', blob, cleanName); + formData.append('type', 'input'); + formData.append('subfolder', subfolder); + formData.append('overwrite', 'true'); + + try { + const response = await fetch(`${workerUrl}/upload/image`, { + method: 'POST', + mode: 'cors', + body: formData + }); + + if (!response.ok) { + throw new Error(`Upload failed: ${response.statusText}`); + } + + extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug"); + } catch (error) { + extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error"); + // Continue with other images + } + } +} + +export async function performPreflightCheck(extension, workers) { + if (workers.length === 0) return []; + + extension.log(`Performing pre-flight health check on ${workers.length} workers...`, "debug"); + const startTime = Date.now(); + + const checkPromises = workers.map(async (worker) => { + const url = extension.getWorkerUrl(worker, '/prompt'); + + extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug"); + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + }); + + if (response.ok) { + extension.log(`Worker ${worker.name} is active`, "debug"); + return { worker, active: true }; + } else { + extension.log(`Worker ${worker.name} returned ${response.status}`, "debug"); + return { worker, active: false }; + } + } catch (error) { + extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug"); + return { worker, active: false }; + } + }); + + const results = await Promise.all(checkPromises); + const activeWorkers = results.filter(r => r.active).map(r => r.worker); + + const elapsed = Date.now() - startTime; + extension.log(`Pre-flight check completed in ${elapsed}ms. Active workers: ${activeWorkers.length}/${workers.length}`, "debug"); + + // Update UI status indicators for inactive workers + results.filter(r => !r.active).forEach(r => { + const statusDot = document.getElementById(`status-${r.worker.id}`); + if (statusDot) { + // Remove pulsing animation once status is determined + statusDot.classList.remove('status-pulsing'); + statusDot.style.backgroundColor = "#c04c4c"; // Red for offline + statusDot.title = "Offline - Cannot connect"; + } + }); + + return activeWorkers; +} \ No newline at end of file diff --git a/web/gpupanel.js b/web/gpupanel.js deleted file mode 100644 index 619b647..0000000 --- a/web/gpupanel.js +++ /dev/null @@ -1,3521 +0,0 @@ -import { app } from "../../scripts/app.js"; -import { api } from "../../scripts/api.js"; - -class DistributedExtension { - // Button styling constants - static BUTTON_STYLES = { - base: "width: 100%; padding: 8px; font-size: 12px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: background-color 0.2s;", - clearMemory: "background-color: #555;", - success: "background-color: #3ca03c;", - error: "background-color: #c04c4c;", - }; - - // CSS for pulsing animation and button hover effects - static PULSE_ANIMATION_CSS = ` - @keyframes pulse { - 0% { - opacity: 1; - transform: scale(1); - box-shadow: 0 0 0 0 rgba(240, 173, 78, 0.7); - } - 50% { - opacity: 0.3; - transform: scale(1.3); - box-shadow: 0 0 0 6px rgba(240, 173, 78, 0); - } - 100% { - opacity: 1; - transform: scale(1); - box-shadow: 0 0 0 0 rgba(240, 173, 78, 0); - } - } - .status-pulsing { - animation: pulse 1.2s ease-in-out infinite; - } - - /* Button hover effects */ - .distributed-button:hover:not(:disabled) { - filter: brightness(1.2); - transition: filter 0.2s ease; - } - .distributed-button:disabled { - opacity: 0.6; - cursor: not-allowed; - } - - /* Settings button animation */ - .settings-btn { - transition: transform 0.2s ease; - } - - - /* Expanded settings panel */ - .worker-settings { - max-height: 0; - overflow: hidden; - opacity: 0; - transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease; - } - .worker-settings.expanded { - max-height: 300px; - opacity: 1; - } - `; - - constructor() { - this.config = null; - this.originalQueuePrompt = api.queuePrompt.bind(api); - this.statusCheckInterval = null; - this.logAutoRefreshInterval = null; - - // Initialize centralized state - this.state = this._createStateManager(); - - // Initialize UI component factory - this.ui = this._createUIFactory(); - - // Initialize API client - this.api = this._createApiClient(); - - // Inject CSS for pulsing animation - this.injectStyles(); - - this.loadConfig().then(async () => { - this.registerSidebarTab(); - this.setupInterceptor(); - this.startStatusChecking(); - this.loadManagedWorkers(); - // Detect master IP after everything is set up - this.detectMasterIP(); - }); - } - - // Debug logging helpers - debugLog(message) { - if (this.config?.settings?.debug) { - console.log(`[Distributed] ${message}`); - } - } - - log(message) { - console.log(`[Distributed] ${message}`); - } - - // --- UI Component Factory --- - _createUIFactory() { - const styles = { - statusDot: "display: inline-block; width: 8px; height: 8px; border-radius: 50%;", - controlsDiv: "padding: 0 8px 6px 8px; display: flex; gap: 4px;", - checkboxColumn: "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444;", - workerCard: "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;", - contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;", - infoRow: "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;", - workerContent: "display: flex; align-items: center; gap: 8px; flex: 1;", - settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;", - formGroup: "display: flex; align-items: center; gap: 8px; margin-bottom: 8px;", - formLabel: "flex: 0 0 100px; font-size: 12px; color: #ccc;", - formInput: "flex: 1; padding: 4px 8px; background: #1a1a1a; border: 1px solid #444; border-radius: 4px; color: #fff; font-size: 12px;", - infoBox: "background-color: #333; color: #888; padding: 4px 12px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1;" - }; - - return { - createStatusDot(id, color = "#666", title = "Status") { - const dot = document.createElement("span"); - if (id) dot.id = id; - dot.style.cssText = styles.statusDot + ` background-color: ${color};`; - dot.title = title; - return dot; - }, - - createButton(text, onClick, customStyle = "") { - const button = document.createElement("button"); - button.textContent = text; - button.className = "distributed-button"; - button.style.cssText = DistributedExtension.BUTTON_STYLES.base + customStyle; - if (onClick) button.onclick = onClick; - return button; - }, - - createButtonGroup(buttons, style = "") { - const group = document.createElement("div"); - group.style.cssText = "display: flex; gap: 4px; margin-top: 10px;" + style; - buttons.forEach(button => group.appendChild(button)); - return group; - }, - - createWorkerControls(workerId, handlers = {}) { - const controlsDiv = document.createElement("div"); - controlsDiv.id = `controls-${workerId}`; - controlsDiv.style.cssText = styles.controlsDiv; - - const buttons = []; - - if (handlers.launch) { - const launchBtn = this.createButton('Launch', handlers.launch); - launchBtn.id = `launch-${workerId}`; - launchBtn.title = "Launch this worker instance"; - buttons.push(launchBtn); - } - - if (handlers.stop) { - const stopBtn = this.createButton('Stop', handlers.stop); - stopBtn.id = `stop-${workerId}`; - stopBtn.title = "Stop this worker instance"; - buttons.push(stopBtn); - } - - if (handlers.viewLog) { - const logBtn = this.createButton('View Log', handlers.viewLog); - logBtn.id = `log-${workerId}`; - logBtn.title = "View worker log file"; - buttons.push(logBtn); - } - - buttons.forEach(btn => controlsDiv.appendChild(btn)); - return controlsDiv; - }, - - createFormGroup(label, value, id, type = "text", placeholder = "") { - const group = document.createElement("div"); - group.style.cssText = styles.formGroup; - - const labelEl = document.createElement("label"); - labelEl.textContent = label; - labelEl.htmlFor = id; - labelEl.style.cssText = styles.formLabel; - - const input = document.createElement("input"); - input.type = type; - input.id = id; - input.value = value; - input.placeholder = placeholder; - input.style.cssText = styles.formInput; - - group.appendChild(labelEl); - group.appendChild(input); - return { group, input }; - }, - - createWorkerCard() { - const card = document.createElement("div"); - card.style.cssText = styles.workerCard; - - const checkboxColumn = document.createElement("div"); - checkboxColumn.style.cssText = styles.checkboxColumn; - - const contentColumn = document.createElement("div"); - contentColumn.style.cssText = styles.contentColumn; - - const infoRow = document.createElement("div"); - infoRow.style.cssText = styles.infoRow; - - const workerContent = document.createElement("div"); - workerContent.style.cssText = styles.workerContent; - - const settingsArrow = document.createElement("span"); - settingsArrow.className = "settings-arrow"; - settingsArrow.innerHTML = "▶"; - settingsArrow.style.cssText = styles.settingsArrow; - - const controlsDiv = document.createElement("div"); - controlsDiv.style.cssText = styles.controlsDiv; - - return { - card, - checkboxColumn, - contentColumn, - infoRow, - workerContent, - settingsArrow, - controlsDiv - }; - }, - - createInfoBox(text) { - const box = document.createElement("div"); - box.style.cssText = styles.infoBox; - box.textContent = text; - return box; - }, - - addHoverEffect(element, onHover, onLeave) { - element.onmouseover = onHover; - element.onmouseout = onLeave; - }, - - createCard(type = 'worker', options = {}) { - const card = document.createElement("div"); - const baseStyle = "margin-bottom: 12px; border-radius: 4px; overflow: hidden; display: flex;"; - - switch(type) { - case 'master': - case 'worker': - card.style.cssText = baseStyle + "background: #2a2a2a;"; - break; - case 'blueprint': - card.style.cssText = baseStyle + "border: 2px solid #666; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);"; - if (options.onClick) card.onclick = options.onClick; - if (options.title) card.title = options.title; - break; - case 'add': - card.style.cssText = "margin-top: 10px; " + baseStyle + "border: 1px solid #444; cursor: pointer; transition: all 0.2s ease;"; - if (options.onClick) card.onclick = options.onClick; - if (options.title) card.title = options.title; - break; - } - - if (options.onMouseEnter) { - card.addEventListener('mouseenter', options.onMouseEnter); - } - if (options.onMouseLeave) { - card.addEventListener('mouseleave', options.onMouseLeave); - } - - return card; - }, - - createCardColumn(type = 'checkbox', options = {}) { - const column = document.createElement("div"); - const baseStyle = "display: flex; align-items: center; justify-content: center;"; - - switch(type) { - case 'checkbox': - column.style.cssText = baseStyle + "flex: 0 0 40px; border-right: 1px solid #444; cursor: default;"; - if (options.title) column.title = options.title; - break; - case 'icon': - column.style.cssText = baseStyle + "width: 50px; flex-shrink: 0; font-size: 20px; color: #666;"; - break; - case 'content': - column.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; - break; - } - - return column; - }, - - createInfoRow(options = {}) { - const row = document.createElement("div"); - row.style.cssText = "display: flex; align-items: center; padding: 8px; cursor: pointer; min-height: 60px;"; - if (options.onClick) row.onclick = options.onClick; - return row; - }, - - createWorkerContent() { - const content = document.createElement("div"); - content.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; - return content; - }, - - createSettingsForm(fields = [], options = {}) { - const form = document.createElement("div"); - form.style.cssText = "display: flex; flex-direction: column; gap: 10px;"; - - fields.forEach(field => { - if (field.type === 'checkbox') { - const group = document.createElement("div"); - group.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; - - const checkbox = document.createElement("input"); - checkbox.type = "checkbox"; - checkbox.id = field.id; - checkbox.checked = field.checked || false; - if (field.onChange) checkbox.onchange = field.onChange; - - const label = document.createElement("label"); - label.htmlFor = field.id; - label.textContent = field.label; - label.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - group.appendChild(checkbox); - group.appendChild(label); - form.appendChild(group); - } else { - const result = this.createFormGroup(field.label, field.value, field.id, field.type, field.placeholder); - if (field.groupId) result.group.id = field.groupId; - if (field.display) result.group.style.display = field.display; - form.appendChild(result.group); - } - }); - - if (options.buttons) { - const buttonGroup = this.createButtonGroup(options.buttons, options.buttonStyle || " margin-top: 8px;"); - form.appendChild(buttonGroup); - } - - return form; - }, - - createCompleteCard(type, options = {}) { - // Create the main card - const card = this.createCard(type, options); - - // Create left column - let leftColumn; - if (type === 'master' || type === 'worker') { - leftColumn = this.createCardColumn('checkbox', options.checkboxOptions); - - // Add checkbox - const checkbox = document.createElement("input"); - checkbox.type = "checkbox"; - if (options.checkboxId) checkbox.id = options.checkboxId; - checkbox.checked = options.checked || false; - checkbox.disabled = options.disabled || false; - checkbox.style.cssText = `cursor: ${options.disabled ? 'default' : 'pointer'}; width: 16px; height: 16px;${options.checkboxStyle || ''}`; - if (options.checkboxTitle) checkbox.title = options.checkboxTitle; - - leftColumn.appendChild(checkbox); - - if (options.onCheckboxClick && !options.disabled) { - checkbox.style.pointerEvents = "none"; - leftColumn.style.cursor = "pointer"; - leftColumn.onclick = options.onCheckboxClick; - } - } else { - // For blueprint and add cards, create icon column - leftColumn = document.createElement("div"); - const borderColor = type === 'blueprint' ? '#666' : '#444'; - leftColumn.style.cssText = `flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: ${type === 'blueprint' ? '2px' : '1px'} solid ${borderColor}; color: ${type === 'blueprint' ? '#888' : '#555'}; font-size: ${type === 'blueprint' ? '24px' : '18px'}; font-weight: ${type === 'blueprint' ? 'bold' : 'normal'};`; - leftColumn.innerHTML = "+"; - - if (options.onHover) { - card.addEventListener('mouseenter', () => { - leftColumn.style.color = type === 'blueprint' ? '#aaa' : '#888'; - leftColumn.style.borderColor = type === 'blueprint' ? '#888' : '#666'; - }); - card.addEventListener('mouseleave', () => { - leftColumn.style.color = type === 'blueprint' ? '#888' : '#555'; - leftColumn.style.borderColor = borderColor; - }); - } - } - - // Create right column - const rightColumn = this.createCardColumn('content'); - - // Add hover effect for master/worker cards - if ((type === 'master' || type === 'worker') && options.onHover !== false) { - rightColumn.onmouseover = () => rightColumn.style.backgroundColor = "#333"; - rightColumn.onmouseout = () => rightColumn.style.backgroundColor = "transparent"; - } - - // Create info row - const infoRow = this.createInfoRow(options.infoRowOptions); - - // Create worker content - const workerContent = this.createWorkerContent(); - - // Build the card - infoRow.appendChild(workerContent); - rightColumn.appendChild(infoRow); - - // Add controls div if specified - if (options.includeControls) { - const controlsDiv = document.createElement("div"); - controlsDiv.id = options.controlsId; - controlsDiv.style.cssText = styles.controlsDiv; - rightColumn.appendChild(controlsDiv); - } - - // Add settings div if specified - if (options.includeSettings) { - const settingsDiv = document.createElement("div"); - settingsDiv.id = options.settingsId; - settingsDiv.className = "worker-settings"; - settingsDiv.style.cssText = styles.workerSettings; - rightColumn.appendChild(settingsDiv); - } - - card.appendChild(leftColumn); - card.appendChild(rightColumn); - - return { - card, - leftColumn, - rightColumn, - infoRow, - workerContent, - checkbox: leftColumn.querySelector('input[type="checkbox"]') - }; - }, - - styles // Expose styles for custom usage - }; - } - - // --- State Management --- - _createStateManager() { - const state = { - workers: new Map(), // Unified worker state: { status, managed, launching, expanded, ... } - masterStatus: 'online', - domCache: new Map() // Cache DOM elements - }; - - return { - // Worker state management - getWorker(workerId) { - return state.workers.get(String(workerId)) || {}; - }, - - updateWorker(workerId, updates) { - const id = String(workerId); - const current = state.workers.get(id) || {}; - state.workers.set(id, { ...current, ...updates }); - return state.workers.get(id); - }, - - setWorkerStatus(workerId, status) { - return this.updateWorker(workerId, { status }); - }, - - setWorkerManaged(workerId, info) { - return this.updateWorker(workerId, { managed: info }); - }, - - setWorkerLaunching(workerId, launching) { - return this.updateWorker(workerId, { launching }); - }, - - setWorkerExpanded(workerId, expanded) { - return this.updateWorker(workerId, { expanded }); - }, - - isWorkerLaunching(workerId) { - return this.getWorker(workerId).launching || false; - }, - - isWorkerExpanded(workerId) { - return this.getWorker(workerId).expanded || false; - }, - - isWorkerManaged(workerId) { - return !!this.getWorker(workerId).managed; - }, - - getWorkerStatus(workerId) { - return this.getWorker(workerId).status || {}; - }, - - // Master state - setMasterStatus(status) { - state.masterStatus = status; - }, - - getMasterStatus() { - return state.masterStatus; - }, - - // DOM cache - cacheElement(key, element) { - state.domCache.set(key, element); - }, - - getCachedElement(key) { - return state.domCache.get(key); - }, - - clearCache() { - state.domCache.clear(); - } - }; - } - - // --- API Client --- - _createApiClient() { - const baseUrl = window.location.origin; - - const request = async (endpoint, options = {}) => { - try { - const response = await fetch(`${baseUrl}${endpoint}`, { - headers: { 'Content-Type': 'application/json' }, - ...options - }); - - if (!response.ok) { - const error = await response.json().catch(() => ({ message: 'Request failed' })); - throw new Error(error.message || `HTTP ${response.status}`); - } - - return await response.json(); - } catch (error) { - this.log(`API Error: ${endpoint} - ${error.message}`); - throw error; - } - }; - - return { - // Config endpoints - async getConfig() { - return request('/distributed/config'); - }, - - async updateWorker(workerId, data) { - return request('/distributed/config/update_worker', { - method: 'POST', - body: JSON.stringify({ worker_id: workerId, ...data }) - }); - }, - - async deleteWorker(workerId) { - return request('/distributed/config/delete_worker', { - method: 'POST', - body: JSON.stringify({ worker_id: workerId }) - }); - }, - - async updateSetting(key, value) { - return request('/distributed/config/update_setting', { - method: 'POST', - body: JSON.stringify({ key, value }) - }); - }, - - async updateMaster(data) { - return request('/distributed/config/update_master', { - method: 'POST', - body: JSON.stringify(data) - }); - }, - - // Worker management endpoints - async launchWorker(workerId) { - return request('/distributed/launch_worker', { - method: 'POST', - body: JSON.stringify({ worker_id: workerId }) - }); - }, - - async stopWorker(workerId) { - return request('/distributed/stop_worker', { - method: 'POST', - body: JSON.stringify({ worker_id: workerId }) - }); - }, - - async getManagedWorkers() { - return request('/distributed/managed_workers'); - }, - - async getWorkerLog(workerId, lines = 1000) { - return request(`/distributed/worker_log/${workerId}?lines=${lines}`); - }, - - async clearLaunchingFlag(workerId) { - return request('/distributed/worker/clear_launching', { - method: 'POST', - body: JSON.stringify({ worker_id: workerId }) - }); - }, - - // Job preparation - async prepareJob(multiJobId) { - return request('/distributed/prepare_job', { - method: 'POST', - body: JSON.stringify({ multi_job_id: multiJobId }) - }); - }, - - // Image loading - async loadImage(imagePath) { - return request('/distributed/load_image', { - method: 'POST', - body: JSON.stringify({ image_path: imagePath }) - }); - }, - - // Network info - async getNetworkInfo() { - return request('/distributed/network_info'); - }, - - // Status checking (with timeout) - async checkStatus(url, timeout = 900) { - const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), timeout); - - try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - signal: controller.signal - }); - clearTimeout(timeoutId); - - if (!response.ok) throw new Error(`HTTP ${response.status}`); - return await response.json(); - } catch (error) { - clearTimeout(timeoutId); - throw error; - } - }, - - // Batch status checking - async checkMultipleStatuses(urls) { - return Promise.allSettled( - urls.map(url => this.checkStatus(url)) - ); - } - }; - } - - injectStyles() { - const styleId = 'distributed-styles'; - if (!document.getElementById(styleId)) { - const style = document.createElement('style'); - style.id = styleId; - style.textContent = DistributedExtension.PULSE_ANIMATION_CSS; - 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; - } - - async loadConfig() { - try { - this.config = await this.api.getConfig(); - this.log("Loaded config: " + JSON.stringify(this.config)); - } catch (error) { - this.log("Failed to load config: " + error.message); - this.config = { workers: [], settings: {} }; - } - } - - async updateWorkerEnabled(workerId, enabled) { - const worker = this.config.workers.find(w => w.id === workerId); - if (worker) { - worker.enabled = enabled; - } - - try { - await this.api.updateWorker(workerId, { enabled }); - } catch (error) { - this.log("Error updating worker: " + error.message); - } - } - - 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); - - app.extensionManager.toast.add({ - severity: "success", - summary: "Setting Updated", - detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`, - life: 2000 - }); - } catch (error) { - this.log(`Error updating setting '${key}': ${error.message}`); - 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.renderSidebarContent(el) - }); - } - - _createButton(text, onClick, style) { - return this.ui.createButton(text, onClick, style); - } - - // Generic handler for parallel worker operations - async _handleWorkerOperation(button, operation, successText, errorText) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; - - button.textContent = operation.loadingText; - button.disabled = true; - - try { - const urlsToProcess = this.enabledWorkers.map(w => ({ - name: w.name, - url: this.getWorkerUrl(w) - })); - - if (urlsToProcess.length === 0) { - button.textContent = "No Workers"; - button.style.backgroundColor = "#c04c4c"; - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - button.disabled = false; - }, 3000); - return; - } - - const promises = urlsToProcess.map(target => - fetch(`${target.url}${operation.endpoint}`, { - method: 'POST', - mode: 'cors' - }) - .then(response => ({ ok: response.ok, name: target.name })) - .catch(() => ({ ok: false, name: target.name })) - ); - - const results = await Promise.all(promises); - const failures = results.filter(r => !r.ok); - - if (failures.length === 0) { - button.textContent = successText; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); - if (operation.onSuccess) operation.onSuccess(); - } else { - button.textContent = errorText; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); - this.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`); - } - - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - }, 3000); - } finally { - button.disabled = false; - } - } - - async _handleInterruptWorkers(button) { - return this._handleWorkerOperation(button, { - name: "Interrupt", - endpoint: "/interrupt", - loadingText: "Interrupting...", - onSuccess: () => setTimeout(() => this.checkAllWorkerStatuses(), 500) - }, "Interrupted!", "Error! See Console"); - } - - async _handleClearMemory(button) { - return this._handleWorkerOperation(button, { - name: "Clear memory", - endpoint: "/distributed/clear_memory", - loadingText: "Clearing..." - }, "Success!", "Error! See Console"); - } - - - async renderSidebarContent(el) { - // Panel is being opened/rendered - this.debugLog("Panel opened"); - - if (!el) { - this.debugLog("No element provided to renderSidebarContent"); - return; - } - - // Prevent infinite recursion - if (this._isRendering) { - this.debugLog("Already rendering, skipping"); - return; - } - this._isRendering = true; - - try { - // Store reference to the panel element - this.panelElement = el; - - // Reload config when panel opens to pick up any external changes - await this.loadConfig(); - - // Reload managed workers when panel opens - this.loadManagedWorkers().then(() => { - // Update all worker controls after loading - if (this.config?.workers) { - this.config.workers.forEach(w => this.updateWorkerControls(w.id)); - } - }); - - // Start checking worker statuses immediately in parallel - // Results will update the UI as they come in - setTimeout(() => this.checkAllWorkerStatuses(), 0); - - el.innerHTML = ''; - - // Create toolbar header to match ComfyUI style - const toolbar = document.createElement("div"); - toolbar.className = "p-toolbar p-component border-x-0 border-t-0 rounded-none px-2 py-1 min-h-8"; - toolbar.style.cssText = "border-bottom: 1px solid #444; background: transparent;"; - - const toolbarStart = document.createElement("div"); - toolbarStart.className = "p-toolbar-start"; - - const titleSpan = document.createElement("span"); - titleSpan.className = "text-xs 2xl:text-sm truncate"; - titleSpan.textContent = "COMFYUI DISTRIBUTED"; - titleSpan.title = "ComfyUI Distributed"; - - toolbarStart.appendChild(titleSpan); - toolbar.appendChild(toolbarStart); - - const toolbarCenter = document.createElement("div"); - toolbarCenter.className = "p-toolbar-center"; - toolbar.appendChild(toolbarCenter); - - const toolbarEnd = document.createElement("div"); - toolbarEnd.className = "p-toolbar-end"; - toolbar.appendChild(toolbarEnd); - - el.appendChild(toolbar); - - // Main container with adjusted padding - const container = document.createElement("div"); - container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: calc(100% - 32px);"; - - // Detect master info on panel open (in case CUDA info wasn't available at startup) - this.log(`Panel opened. CUDA device count: ${this.cudaDeviceCount}, Workers: ${this.config?.workers?.length || 0}`); - if (!this.cudaDeviceCount) { - await this.detectMasterIP(); - } - - // Check if we should auto-populate workers after detection - if (this.cudaDeviceCount > 0 && (!this.config.workers || this.config.workers.length === 0)) { - this.log(`Auto-populating workers based on ${this.cudaDeviceCount} CUDA devices (excluding master on CUDA ${this.masterCudaDevice})`); - - const newWorkers = []; - let workerNum = 1; - let portOffset = 0; - - for (let i = 0; i < this.cudaDeviceCount; i++) { - // Skip the CUDA device used by master - if (i === this.masterCudaDevice) { - this.log(`Skipping CUDA ${i} (used by master)`); - continue; - } - - const worker = { - id: Date.now() + i, - name: `Worker ${workerNum}`, - host: "localhost", - port: 8189 + portOffset, - cuda_device: i, - enabled: true - }; - newWorkers.push(worker); - workerNum++; - portOffset++; - } - - // Only proceed if we have workers to add - if (newWorkers.length > 0) { - this.log(`Auto-populating ${newWorkers.length} workers`); - - // Add workers to config - this.config.workers = newWorkers; - - // Save each worker using the update endpoint - for (const worker of newWorkers) { - try { - await this.api.updateWorker(worker.id, worker); - } catch (error) { - this.log(`Error saving worker ${worker.name}: ${error.message}`); - } - } - - this.log(`Auto-populated ${newWorkers.length} workers and saved config`); - - // Reload the config to include the new workers - await this.loadConfig(); - - // Continue rendering with the updated config - } else { - this.log("No additional CUDA devices available for workers (all used by master)"); - } - } - - // Master Node Section - styled exactly like a worker card - const masterDiv = this.ui.createCard('master'); - - // Left column - checkbox area (always checked and disabled) - const checkboxColumn = this.ui.createCardColumn('checkbox', { title: "Master node is always enabled" }); - - // Greyed out checkbox that's always checked - const checkbox = document.createElement("input"); - checkbox.type = "checkbox"; - checkbox.checked = true; - checkbox.disabled = true; - checkbox.style.cssText = "cursor: default; width: 16px; height: 16px; opacity: 0.6;"; - - checkboxColumn.appendChild(checkbox); - - // Right column - content area - const contentColumn = this.ui.createCardColumn('content'); - - // First row: master info - const infoRow = this.ui.createInfoRow({ onClick: null }); - infoRow.title = "Click to expand settings"; - - // Worker content - const workerContent = this.ui.createWorkerContent(); - - const statusDot = this.ui.createStatusDot("master-status", "#4CAF50", "Online"); - - const gpuInfo = document.createElement("span"); - const cudaInfo = this.masterCudaDevice !== undefined ? `CUDA ${this.masterCudaDevice} • ` : ''; - const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - gpuInfo.innerHTML = `${this.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; - - workerContent.appendChild(statusDot); - workerContent.appendChild(gpuInfo); - - // Settings arrow - const settingsArrow = document.createElement("span"); - settingsArrow.className = "settings-arrow"; - settingsArrow.innerHTML = "▶"; - settingsArrow.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;"; - - infoRow.appendChild(workerContent); - infoRow.appendChild(settingsArrow); - - // Click handler for info row - let masterSettingsExpanded = false; - infoRow.onclick = () => { - masterSettingsExpanded = !masterSettingsExpanded; - if (masterSettingsExpanded) { - masterSettingsDiv.classList.add("expanded"); - masterSettingsDiv.style.padding = "12px"; - masterSettingsDiv.style.marginTop = "8px"; - masterSettingsDiv.style.marginBottom = "8px"; - settingsArrow.style.transform = "rotate(90deg)"; - } else { - masterSettingsDiv.classList.remove("expanded"); - masterSettingsDiv.style.padding = "0 12px"; - masterSettingsDiv.style.marginTop = "0"; - masterSettingsDiv.style.marginBottom = "0"; - settingsArrow.style.transform = "rotate(0deg)"; - } - }; - - // Hover effect for entire content column - contentColumn.onmouseover = () => { - contentColumn.style.backgroundColor = "#333"; - settingsArrow.style.color = "#fff"; - }; - contentColumn.onmouseout = () => { - contentColumn.style.backgroundColor = "transparent"; - settingsArrow.style.color = "#888"; - }; - - contentColumn.appendChild(infoRow); - - // Second row: master controls (styled like remote worker) - const controlsDiv = document.createElement("div"); - controlsDiv.style.cssText = "padding: 0 8px 6px 8px; display: flex; gap: 4px;"; - - // Master info box styled like remote worker - const masterInfo = document.createElement("div"); - masterInfo.style.cssText = "background-color: #333; color: #888; padding: 4px 12px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1;"; - masterInfo.textContent = "Master"; - controlsDiv.appendChild(masterInfo); - - // Third row: expandable settings panel - const masterSettingsDiv = document.createElement("div"); - masterSettingsDiv.id = "master-settings"; - masterSettingsDiv.className = "worker-settings"; - masterSettingsDiv.style.cssText = "margin: 0 8px; padding: 0 12px; background: #1e1e1e; border-radius: 4px;"; - - // Create settings form using reusable components - const settingsForm = document.createElement("div"); - settingsForm.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; - - // Name field - const nameResult = this.ui.createFormGroup("Name:", this.config?.master?.name || "Master", "master-name"); - settingsForm.appendChild(nameResult.group); - - // Host field - const hostResult = this.ui.createFormGroup("Host:", this.config?.master?.host || "", "master-host", "text", "Auto-detect if empty"); - settingsForm.appendChild(hostResult.group); - - // Buttons row - using same style as workers - const saveBtn = this._createButton("Save", - async () => { - const nameInput = document.getElementById('master-name'); - const hostInput = document.getElementById('master-host'); - - // Update config - if (!this.config.master) this.config.master = {}; - this.config.master.name = nameInput.value.trim() || "Master"; - - // Save both host and name - await this.saveMasterConfig({ - host: hostInput.value.trim(), - name: this.config.master.name - }); - - // Update display - document.getElementById('master-name-display').textContent = this.config.master.name; - this.updateMasterDisplay(); - - saveBtn.textContent = "Saved!"; - setTimeout(() => { saveBtn.textContent = "Save"; }, 2000); - }, - "background-color: #4a7c4a;"); - saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - - const cancelBtn = this._createButton("Cancel", - () => { - // Reset inputs to original values - document.getElementById('master-name').value = this.config?.master?.name || "Master"; - document.getElementById('master-host').value = this.config?.master?.host || ""; - }, - "background-color: #555;"); - cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - - const buttonGroup = this.ui.createButtonGroup([saveBtn, cancelBtn], " margin-top: 8px;"); - settingsForm.appendChild(buttonGroup); - masterSettingsDiv.appendChild(settingsForm); - - contentColumn.appendChild(controlsDiv); - contentColumn.appendChild(masterSettingsDiv); - - masterDiv.appendChild(checkboxColumn); - masterDiv.appendChild(contentColumn); - container.appendChild(masterDiv); - - // Workers Section (no heading) - const gpuSection = document.createElement("div"); - gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; - - const gpuList = document.createElement("div"); - const workers = this.config?.workers || []; - - // If no workers exist, show a full blueprint placeholder first - if (workers.length === 0) { - const blueprintDiv = this.ui.createCard('blueprint', { - title: "Click to add your first worker", - onClick: () => this.addNewWorker() - }); - - // Left column - plus icon - const blueprintCheckbox = document.createElement("div"); - blueprintCheckbox.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 2px solid #666; color: #888; font-size: 24px; font-weight: bold;"; - blueprintCheckbox.innerHTML = "+"; - - // Right column - const blueprintContent = this.ui.createCardColumn('content'); - - // Info row with placeholder content - const blueprintInfo = this.ui.createInfoRow(); - - const blueprintWorkerContent = this.ui.createWorkerContent(); - - const blueprintDot = this.ui.createStatusDot(null, "transparent", ""); - blueprintDot.style.border = "1px solid #555"; - - const blueprintText = document.createElement("span"); - blueprintText.innerHTML = `Add New Worker
[CUDA] • [Port]`; - - blueprintWorkerContent.appendChild(blueprintDot); - blueprintWorkerContent.appendChild(blueprintText); - blueprintInfo.appendChild(blueprintWorkerContent); - - // Placeholder button row - single Launch button spanning width - const blueprintButtons = document.createElement("div"); - blueprintButtons.style.cssText = "padding: 0 8px 6px 8px;"; - - const ghostLaunch = document.createElement("div"); - ghostLaunch.style.cssText = "width: 100%; padding: 3px 8px; font-size: 11px; border: 1px solid #444; background: transparent; color: #555; border-radius: 4px; text-align: center;"; - ghostLaunch.textContent = "Launch"; - - blueprintButtons.appendChild(ghostLaunch); - - blueprintContent.appendChild(blueprintInfo); - blueprintContent.appendChild(blueprintButtons); - - blueprintDiv.appendChild(blueprintCheckbox); - blueprintDiv.appendChild(blueprintContent); - - // Hover effect - blueprintDiv.onmouseover = () => { - blueprintDiv.style.borderColor = "#888"; - blueprintDiv.style.backgroundColor = "rgba(255, 255, 255, 0.05)"; - blueprintCheckbox.style.borderColor = "#888"; - blueprintCheckbox.style.color = "#aaa"; - }; - blueprintDiv.onmouseout = () => { - blueprintDiv.style.borderColor = "#666"; - blueprintDiv.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; - blueprintCheckbox.style.borderColor = "#666"; - blueprintCheckbox.style.color = "#888"; - }; - - gpuList.appendChild(blueprintDiv); - } - - // Show existing workers - workers.forEach(worker => { - const gpuDiv = this.ui.createCard('worker'); - - // Left column - checkbox area (extends full height) - const checkboxColumn = this.ui.createCardColumn('checkbox'); - checkboxColumn.style.cursor = "pointer"; - checkboxColumn.title = "Enable/disable this worker"; - - // Standard checkbox - const checkbox = document.createElement("input"); - checkbox.type = "checkbox"; - checkbox.id = `gpu-${worker.id}`; - checkbox.checked = worker.enabled; - checkbox.style.cssText = "cursor: pointer; width: 16px; height: 16px; pointer-events: none;"; // pointer-events: none so the column click handles it - - checkboxColumn.appendChild(checkbox); - - // Click handler for checkbox column - checkboxColumn.onclick = async (e) => { - checkbox.checked = !checkbox.checked; - await this.updateWorkerEnabled(worker.id, checkbox.checked); - this.updateSummary(); - }; - - // Right column - content area - const contentColumn = this.ui.createCardColumn('content'); - - // First row: worker info - const infoRow = this.ui.createInfoRow(); - infoRow.title = "Click to expand settings"; - - // Worker content - const workerContent = this.ui.createWorkerContent(); - - const statusDot = this.ui.createStatusDot(`status-${worker.id}`, "#666", "Checking status..."); - - const gpuInfo = document.createElement("span"); - const isRemote = this.isRemoteWorker(worker); - - if (isRemote) { - // For remote workers: show host:port - gpuInfo.innerHTML = `${worker.name}
${worker.host}:${worker.port}`; - } else { - // For local workers: show CUDA and port - const cudaInfo = worker.cuda_device !== undefined ? `CUDA ${worker.cuda_device} • ` : ''; - gpuInfo.innerHTML = `${worker.name}
${cudaInfo}Port ${worker.port}`; - } - - workerContent.appendChild(statusDot); - workerContent.appendChild(gpuInfo); - - // Settings arrow - const settingsArrow = document.createElement("span"); - settingsArrow.className = "settings-arrow"; - settingsArrow.innerHTML = "▶"; - settingsArrow.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto;"; - if (this.state.isWorkerExpanded(worker.id)) { - settingsArrow.style.transform = "rotate(90deg)"; - } - - infoRow.appendChild(workerContent); - infoRow.appendChild(settingsArrow); - - - // Add content and arrow back to info row - infoRow.appendChild(workerContent); - infoRow.appendChild(settingsArrow); - - // Click handler for info row - infoRow.onclick = () => { - this.toggleWorkerExpanded(worker.id); - }; - - // Hover effect for entire content column - contentColumn.onmouseover = () => { - contentColumn.style.backgroundColor = "#333"; - settingsArrow.style.color = "#fff"; - }; - contentColumn.onmouseout = () => { - contentColumn.style.backgroundColor = "transparent"; - settingsArrow.style.color = "#888"; - }; - - contentColumn.appendChild(infoRow); - - // Second row: launch controls - const controlsDiv = document.createElement("div"); - controlsDiv.id = `controls-${worker.id}`; - controlsDiv.style.cssText = "padding: 0 8px 6px 8px; display: flex; gap: 4px;"; - - if (this.isRemoteWorker(worker)) { - // For remote workers, show styled info box - const remoteInfo = document.createElement("div"); - remoteInfo.style.cssText = "background-color: #333; color: #888; padding: 4px 12px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1;"; - remoteInfo.textContent = "Remote worker"; - controlsDiv.appendChild(remoteInfo); - } else { - // Local workers get launch/stop/log buttons - const controls = this.ui.createWorkerControls(worker.id, { - launch: () => this.launchWorker(worker.id), - stop: () => this.stopWorker(worker.id), - viewLog: () => this.viewWorkerLog(worker.id) - }); - - // Apply custom button styles - const launchBtn = controls.querySelector(`#launch-${worker.id}`); - const stopBtn = controls.querySelector(`#stop-${worker.id}`); - const logBtn = controls.querySelector(`#log-${worker.id}`); - - launchBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; background-color: #555;'; - launchBtn.title = "Launch worker (runs in background with logging)"; - - stopBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; background-color: #555;'; - stopBtn.title = "Stop worker"; - - logBtn.style.cssText += ' padding: 3px 8px; font-size: 11px; display: none; background-color: #685434;'; - - // Move controls into the controlsDiv - while (controls.firstChild) { - controlsDiv.appendChild(controls.firstChild); - } - } - - // Third row: expandable settings panel - const settingsDiv = document.createElement("div"); - settingsDiv.id = `settings-${worker.id}`; - settingsDiv.className = "worker-settings"; - if (this.state.isWorkerExpanded(worker.id)) { - settingsDiv.classList.add("expanded"); - } - settingsDiv.style.cssText = "margin: 0 8px; padding: 0 12px; background: #1e1e1e; border-radius: 4px;"; - if (this.state.isWorkerExpanded(worker.id)) { - settingsDiv.style.padding = "12px"; - settingsDiv.style.marginTop = "8px"; - settingsDiv.style.marginBottom = "8px"; - } - - // Create settings form - const settingsForm = this.createWorkerSettingsForm(worker); - settingsDiv.appendChild(settingsForm); - - contentColumn.appendChild(controlsDiv); - contentColumn.appendChild(settingsDiv); - - gpuDiv.appendChild(checkboxColumn); - gpuDiv.appendChild(contentColumn); - gpuList.appendChild(gpuDiv); - - // Update controls based on managed worker status - this.updateWorkerControls(worker.id); - }); - gpuSection.appendChild(gpuList); - - // Only show the minimal "Add Worker" box if there are existing workers - if (workers.length > 0) { - // Add Worker placeholder box - const addWorkerDiv = this.ui.createCard('add', { - title: "Click to add a new worker", - onClick: () => this.addNewWorker() - }); - - // Left column - plus icon area - const plusColumn = document.createElement("div"); - plusColumn.style.cssText = "flex: 0 0 40px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #444; color: #555; font-size: 18px;"; - plusColumn.innerHTML = "+"; - - // Right column - content area - const addContentColumn = this.ui.createCardColumn('content'); - - // Info row - const addInfoRow = this.ui.createInfoRow(); - addInfoRow.style.minHeight = "42px"; // Reduced from 60px to 42px (30% reduction) - - // Placeholder content - const addContent = this.ui.createWorkerContent(); - - // Placeholder dot - const placeholderDot = this.ui.createStatusDot(null, "transparent", ""); - placeholderDot.style.border = "1px solid #555"; - - // Placeholder text - const placeholderInfo = document.createElement("span"); - placeholderInfo.innerHTML = `Add New Worker`; - - addContent.appendChild(placeholderDot); - addContent.appendChild(placeholderInfo); - addInfoRow.appendChild(addContent); - addContentColumn.appendChild(addInfoRow); - - addWorkerDiv.appendChild(plusColumn); - addWorkerDiv.appendChild(addContentColumn); - - // Hover effects - addWorkerDiv.onmouseover = () => { - addWorkerDiv.style.borderColor = "#666"; - addWorkerDiv.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; - plusColumn.style.color = "#888"; - plusColumn.style.borderColor = "#666"; - }; - - addWorkerDiv.onmouseout = () => { - addWorkerDiv.style.borderColor = "#444"; - addWorkerDiv.style.backgroundColor = "transparent"; - plusColumn.style.color = "#555"; - plusColumn.style.borderColor = "#444"; - }; - - gpuSection.appendChild(addWorkerDiv); - } - - container.appendChild(gpuSection); - - const actionsSection = document.createElement("div"); - actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; - - // Create a row for both buttons - const buttonRow = document.createElement("div"); - buttonRow.style.cssText = "display: flex; gap: 8px;"; - - const clearMemButton = this._createButton("Clear Worker VRAM", (e) => this._handleClearMemory(e.target), DistributedExtension.BUTTON_STYLES.clearMemory); - clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; - clearMemButton.style.cssText += " flex: 1;"; - - const interruptButton = this._createButton("Interrupt Workers", (e) => this._handleInterruptWorkers(e.target), DistributedExtension.BUTTON_STYLES.clearMemory); - interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; - interruptButton.style.cssText += " flex: 1;"; - - buttonRow.appendChild(clearMemButton); - buttonRow.appendChild(interruptButton); - actionsSection.appendChild(buttonRow); - - container.appendChild(actionsSection); - - // Settings section - const settingsSection = document.createElement("div"); - settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;"; - - // Settings header with toggle - const settingsHeader = document.createElement("div"); - settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;"; - - const workerSettingsTitle = document.createElement("h4"); - workerSettingsTitle.textContent = "Settings"; - workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;"; - - const workerSettingsToggle = document.createElement("span"); - workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed - workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; - - settingsHeader.appendChild(workerSettingsTitle); - settingsHeader.appendChild(workerSettingsToggle); - - // Hover effect for header - settingsHeader.onmouseover = () => { - workerSettingsToggle.style.color = "#fff"; - }; - settingsHeader.onmouseout = () => { - workerSettingsToggle.style.color = "#888"; - }; - - // Collapsible settings content - const settingsContent = document.createElement("div"); - settingsContent.style.cssText = "max-height: 0; overflow: hidden; transition: max-height 0.3s ease;"; - - const settingsDiv = document.createElement("div"); - settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;"; - - // Toggle functionality - let settingsExpanded = false; - settingsHeader.onclick = () => { - settingsExpanded = !settingsExpanded; - if (settingsExpanded) { - settingsContent.style.maxHeight = "200px"; - workerSettingsToggle.style.transform = "rotate(90deg)"; - } else { - settingsContent.style.maxHeight = "0"; - workerSettingsToggle.style.transform = "rotate(0deg)"; - } - }; - - // Debug mode setting - const debugGroup = document.createElement("div"); - debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const debugCheckbox = document.createElement("input"); - debugCheckbox.type = "checkbox"; - debugCheckbox.id = "setting-debug"; - debugCheckbox.checked = this.config?.settings?.debug || false; - debugCheckbox.onchange = (e) => this._updateSetting('debug', e.target.checked); - - const debugLabel = document.createElement("label"); - debugLabel.htmlFor = "setting-debug"; - debugLabel.textContent = "Debug Mode"; - debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - debugGroup.appendChild(debugCheckbox); - debugGroup.appendChild(debugLabel); - - // Auto-launch workers setting - const autoLaunchGroup = document.createElement("div"); - autoLaunchGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const autoLaunchCheckbox = document.createElement("input"); - autoLaunchCheckbox.type = "checkbox"; - autoLaunchCheckbox.id = "setting-auto-launch"; - autoLaunchCheckbox.checked = this.config?.settings?.auto_launch_workers || false; - autoLaunchCheckbox.onchange = (e) => this._updateSetting('auto_launch_workers', e.target.checked); - - const autoLaunchLabel = document.createElement("label"); - autoLaunchLabel.htmlFor = "setting-auto-launch"; - autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup"; - autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - autoLaunchGroup.appendChild(autoLaunchCheckbox); - autoLaunchGroup.appendChild(autoLaunchLabel); - - // Stop workers on exit setting - const stopOnExitGroup = document.createElement("div"); - stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - - const stopOnExitCheckbox = document.createElement("input"); - stopOnExitCheckbox.type = "checkbox"; - stopOnExitCheckbox.id = "setting-stop-on-exit"; - stopOnExitCheckbox.checked = this.config?.settings?.stop_workers_on_master_exit !== false; // Default true - stopOnExitCheckbox.onchange = (e) => this._updateSetting('stop_workers_on_master_exit', e.target.checked); - - const stopOnExitLabel = document.createElement("label"); - stopOnExitLabel.htmlFor = "setting-stop-on-exit"; - stopOnExitLabel.textContent = "Stop Local Workers on Master Exit"; - stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - stopOnExitGroup.appendChild(stopOnExitCheckbox); - stopOnExitGroup.appendChild(stopOnExitLabel); - - settingsDiv.appendChild(debugGroup); - settingsDiv.appendChild(autoLaunchGroup); - settingsDiv.appendChild(stopOnExitGroup); - settingsContent.appendChild(settingsDiv); - - settingsSection.appendChild(settingsHeader); - settingsSection.appendChild(settingsContent); - container.appendChild(settingsSection); - - const summarySection = document.createElement("div"); - summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;"; - const summary = document.createElement("div"); - summary.id = "distributed-summary"; - summary.style.cssText = "font-size: 11px; color: #888;"; - summarySection.appendChild(summary); - container.appendChild(summarySection); - el.appendChild(container); - this.updateSummary(); - } finally { - // Always reset the rendering flag - this._isRendering = false; - } - } - - updateMasterDisplay() { - // Update CUDA info if element exists - const cudaInfo = document.getElementById('master-cuda-info'); - if (cudaInfo) { - const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - if (this.masterCudaDevice !== undefined && this.masterCudaDevice !== null) { - cudaInfo.textContent = `CUDA ${this.masterCudaDevice} • Port ${port}`; - } else { - cudaInfo.textContent = `Port ${port}`; - } - } - - // Update name if changed - const nameDisplay = document.getElementById('master-name-display'); - if (nameDisplay && this.config?.master?.name) { - nameDisplay.textContent = this.config.master.name; - } - } - - updateSummary() { - const summaryEl = document.getElementById('distributed-summary'); - if (summaryEl) { - const totalGPUs = this.enabledWorkers.length + 1; - if (this.isEnabled) { - summaryEl.textContent = `If Collector node is present, total generation = (${totalGPUs} GPUs × Batch Size)`; - } else { - summaryEl.textContent = "Only the master GPU will be used."; - } - } - } - - // --- Helper Methods --- - - // --- Core Logic & Execution --- - - setupInterceptor() { - api.queuePrompt = async (number, prompt) => { - if (this.isEnabled) { - const hasCollector = this.findNodesByClass(prompt.output, "DistributedCollector").length > 0; - const hasDistUpscale = this.findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0; - - if (hasCollector || hasDistUpscale) { - const result = await this.executeParallelDistributed(prompt); - // Check status after dispatching jobs - setTimeout(() => this.checkAllWorkerStatuses(), 500); - return result; - } - } - return this.originalQueuePrompt(number, prompt); - }; - } - - async executeParallelDistributed(promptWrapper) { - try { - const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution - const enabledWorkers = this.enabledWorkers; - - // Pre-flight health check on all enabled workers - const activeWorkers = await this.performPreflightCheck(enabledWorkers); - - if (activeWorkers.length === 0 && enabledWorkers.length > 0) { - this.log("No active workers found. Running on master only."); - // Fall back to master-only execution - return this.originalQueuePrompt(0, promptWrapper); - } - - this.debugLog(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`); - - // Find all distributed nodes in the workflow - const collectorNodes = this.findNodesByClass(promptWrapper.output, "DistributedCollector"); - const upscaleNodes = this.findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); - const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - - // Map original node IDs to truly unique job IDs for this specific run - const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`])); - - // Prepare a separate job queue on the backend for each unique job ID - const preparePromises = Array.from(job_id_map.values()).map(uniqueId => this._prepareDistributedJob(uniqueId)); - await Promise.all(preparePromises); - - const jobs = []; - // Use only active workers - const participants = ['master', ...activeWorkers.map(w => w.id)]; - - for (const participantId of participants) { - const options = { - enabled_worker_ids: activeWorkers.map(w => w.id), - workflow: promptWrapper.workflow, - job_id_map: job_id_map // Pass the map of unique IDs - }; - - const jobApiPrompt = this._prepareApiPromptForParticipant( - promptWrapper.output, participantId, options - ); - - if (participantId === 'master') { - jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); - } else { - const worker = activeWorkers.find(w => w.id === participantId); - if (worker) { - const job = { - type: 'worker', - worker, - prompt: jobApiPrompt, - workflow: promptWrapper.workflow - }; - - // Add image references if found for remote workers - if (options._imageReferences) { - job.imageReferences = options._imageReferences; - } - - jobs.push(job); - } - } - } - - const result = await this._executeJobs(jobs); - return result; - } catch (error) { - this.log("Parallel execution failed: " + error.message); - throw error; - } - } - - - async _executeJobs(jobs) { - let masterPromptId = null; - - // Pre-load all unique images before dispatching to workers - const allImageReferences = new Map(); - for (const job of jobs) { - if (job.type === 'worker' && job.imageReferences) { - for (const [filename, info] of job.imageReferences) { - allImageReferences.set(filename, info); - } - } - } - - if (allImageReferences.size > 0) { - this.debugLog(`Pre-loading ${allImageReferences.size} unique image(s) for all workers`); - await this.loadImagesForWorker(allImageReferences); - } - - // Now dispatch jobs in parallel - const promises = jobs.map(job => { - if (job.type === 'master') { - return this.originalQueuePrompt(0, job.promptWrapper).then(result => { - masterPromptId = result; - return result; - }); - } else { - return this._dispatchToWorker(job.worker, job.prompt, job.workflow, job.imageReferences); - } - }); - await Promise.all(promises); - return masterPromptId || { "prompt_id": "distributed-job-dispatched" }; - } - - // --- Helper Methods --- - - findNodesByClass(apiPrompt, className) { - return Object.entries(apiPrompt) - .filter(([, nodeData]) => nodeData.class_type === className) - .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); - } - - /** - * Find all image references in the workflow - * Looks for inputs named "image" that contain filename strings - */ - findImageReferences(apiPrompt) { - const images = new Map(); - // Updated regex to handle: - // - Standard files: "image.png" - // - Subfolder files: "subfolder/image.png" - // - ComfyUI special format: "clipspace/file.png [input]" - const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp)(\s*\[\w+\])?$/i; - - for (const [nodeId, node] of Object.entries(apiPrompt)) { - if (node.inputs && node.inputs.image) { - const imageValue = node.inputs.image; - // Check if it's a string filename (not an array connection) - if (typeof imageValue === 'string' && imageExtensions.test(imageValue)) { - images.set(imageValue, { - nodeId, - nodeType: node.class_type, - inputName: 'image' - }); - this.debugLog(`Found image reference: ${imageValue} in node ${nodeId} (${node.class_type})`); - } - } - } - - return images; - } - - /** - * Find all nodes connected to collector nodes (both upstream and downstream) - * @param {Object} apiPrompt - The workflow API prompt - * @param {Array} collectorIds - Array of collector node IDs - * @returns {Set} Set of node IDs connected to any collector - */ - findCollectorConnectedNodes(apiPrompt, collectorIds) { - const connected = new Set(collectorIds); // Include all collectors - const toProcess = [...collectorIds]; - - // First, build a map of which nodes output to which nodes - const outputMap = new Map(); // sourceNode -> Set of target nodes - for (const [nodeId, node] of Object.entries(apiPrompt)) { - if (node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - if (!outputMap.has(sourceNodeId)) { - outputMap.set(sourceNodeId, new Set()); - } - outputMap.get(sourceNodeId).add(nodeId); - } - } - } - } - - // Traverse both upstream and downstream - while (toProcess.length > 0) { - const nodeId = toProcess.pop(); - const node = apiPrompt[nodeId]; - - // Traverse upstream (inputs) - if (node && node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - if (!connected.has(sourceNodeId)) { - connected.add(sourceNodeId); - toProcess.push(sourceNodeId); - } - } - } - } - - // Traverse downstream (outputs) - if (outputMap.has(nodeId)) { - for (const targetNodeId of outputMap.get(nodeId)) { - if (!connected.has(targetNodeId)) { - connected.add(targetNodeId); - toProcess.push(targetNodeId); - } - } - } - } - - return connected; - } - - /** - * Find only upstream nodes (inputs) for distributed collector nodes - * This is used for workers to avoid executing downstream nodes like SaveImage - * @param {Object} apiPrompt - The API prompt containing the workflow - * @param {Array} collectorIds - Array of collector node IDs - * @returns {Set} Set of node IDs that feed into collectors - */ - findCollectorUpstreamNodes(apiPrompt, collectorIds) { - const connected = new Set(collectorIds); // Include all collectors - const toProcess = [...collectorIds]; - - // Only traverse upstream (inputs) - while (toProcess.length > 0) { - const nodeId = toProcess.pop(); - const node = apiPrompt[nodeId]; - - // Traverse upstream (inputs) only - if (node && node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - if (!connected.has(sourceNodeId)) { - connected.add(sourceNodeId); - toProcess.push(sourceNodeId); - } - } - } - } - } - - return connected; - } - - /** - * Prune workflow to only include nodes connected to distributed nodes - * @param {Object} apiPrompt - The full workflow API prompt - * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided) - * @returns {Object} Pruned API prompt with only required nodes - */ - pruneWorkflowForWorker(apiPrompt, distributedNodes = null) { - // Find all distributed nodes if not provided - if (!distributedNodes) { - const collectorNodes = this.findNodesByClass(apiPrompt, "DistributedCollector"); - const upscaleNodes = this.findNodesByClass(apiPrompt, "UltimateSDUpscaleDistributed"); - distributedNodes = [...collectorNodes, ...upscaleNodes]; - } - - if (distributedNodes.length === 0) { - // No distributed nodes, return full workflow - return apiPrompt; - } - - // Get all nodes connected to distributed nodes - const distributedIds = distributedNodes.map(node => node.id); - - // For workers, only include upstream nodes (this removes ALL downstream nodes after collectors) - const connectedNodes = this.findCollectorUpstreamNodes(apiPrompt, distributedIds); - - this.debugLog(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`); - - // Create pruned prompt with only required nodes - const prunedPrompt = {}; - for (const nodeId of connectedNodes) { - prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId])); - } - - // Check if any distributed node has downstream SaveImage nodes that were removed - // If so, add a PreviewImage node after the collector - for (const distNode of distributedNodes) { - const distNodeId = distNode.id; - - // Check if this distributed node had any downstream nodes in the original workflow - const originalOutputMap = new Map(); - for (const [nodeId, node] of Object.entries(apiPrompt)) { - if (node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2 && String(inputValue[0]) === distNodeId) { - if (!originalOutputMap.has(distNodeId)) { - originalOutputMap.set(distNodeId, []); - } - originalOutputMap.get(distNodeId).push({nodeId, inputName}); - } - } - } - } - - // If this distributed node had downstream nodes that were removed, add a PreviewImage - if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) { - // Generate a unique ID for the preview node - const previewNodeId = `preview_${distNodeId}`; - - // Add PreviewImage node connected to the distributed node - prunedPrompt[previewNodeId] = { - inputs: { - images: [distNodeId, 0] // Connect to first output of distributed node - }, - class_type: "PreviewImage", - _meta: { - title: "Preview Image (auto-added)" - } - }; - - this.debugLog(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`); - } - } - - return prunedPrompt; - } - - - _prepareApiPromptForParticipant(baseApiPrompt, participantId, options = {}) { - let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); - const isMaster = participantId === 'master'; - - // Find all distributed nodes once (before pruning) - const collectorNodes = this.findNodesByClass(jobApiPrompt, "DistributedCollector"); - const upscaleNodes = this.findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed"); - const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - - // For workers, prune the workflow to only include distributed node dependencies - if (!isMaster && allDistributedNodes.length > 0) { - jobApiPrompt = this.pruneWorkflowForWorker(jobApiPrompt, allDistributedNodes); - } - - // Handle image references for remote workers - if (!isMaster && options.enabled_worker_ids) { - // Check if this is a remote worker - const workerId = participantId; - const workerInfo = this.config.workers.find(w => w.id === workerId); - const isRemote = workerInfo && workerInfo.host; - - if (isRemote) { - // Find all image references in the pruned workflow - const imageReferences = this.findImageReferences(jobApiPrompt); - if (imageReferences.size > 0) { - this.debugLog(`Found ${imageReferences.size} image references for remote worker ${workerId}`); - // Store image references for later processing - options._imageReferences = imageReferences; - } - } - } - - // Handle Distributed seed nodes - const distributorNodes = this.findNodesByClass(jobApiPrompt, "DistributedSeed"); - if (distributorNodes.length > 0) { - this.debugLog(`Found ${distributorNodes.length} seed node(s)`); - } - - for (const seedNode of distributorNodes) { - const { inputs } = jobApiPrompt[seedNode.id]; - inputs.is_worker = !isMaster; - if (!isMaster) { - const workerIndex = options.enabled_worker_ids.indexOf(participantId); - inputs.worker_id = `worker_${workerIndex}`; - this.debugLog(`Set seed node ${seedNode.id} for worker ${workerIndex}`); - } - } - - // Handle Distributed collector nodes (already found above) - for (const collector of collectorNodes) { - const { inputs } = jobApiPrompt[collector.id]; - - // Check if this collector is downstream from a distributed upscaler - const hasUpstreamDistributedUpscaler = this._hasUpstreamNode( - jobApiPrompt, - collector.id, - 'UltimateSDUpscaleDistributed' - ); - - if (hasUpstreamDistributedUpscaler) { - // Set pass_through mode for this collector - inputs.pass_through = true; - this.debugLog(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`); - } else { - // Normal collector behavior - // Get the unique job ID from the map created for this execution - const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id; - - // Use the truly unique ID for this execution - inputs.multi_job_id = uniqueJobId; - inputs.is_worker = !isMaster; - if (isMaster) { - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - } else { - inputs.master_url = this.getMasterUrl(); - // Also make the worker_job_id unique to prevent potential caching issues - inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; - inputs.worker_id = participantId; - } - } - } - - // Handle Ultimate SD Upscale Distributed nodes - for (const upscaleNode of upscaleNodes) { - const { inputs } = jobApiPrompt[upscaleNode.id]; - - // Get the unique job ID from the map - const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id; - - inputs.multi_job_id = uniqueJobId; - inputs.is_worker = !isMaster; - - if (isMaster) { - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - } else { - inputs.master_url = this.getMasterUrl(); - inputs.worker_id = participantId; - // Workers also need the enabled_worker_ids to calculate tile distribution - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - } - } - - return jobApiPrompt; - } - - async _prepareDistributedJob(multi_job_id) { - try { - await this.api.prepareJob(multi_job_id); - } catch (error) { - this.log("Error preparing job: " + error.message); - throw error; - } - } - - /** - * Check if a node has an upstream node of a specific type - * @param {Object} apiPrompt - The workflow API prompt - * @param {string} nodeId - The node to check - * @param {string} upstreamType - The class_type to look for upstream - * @returns {boolean} True if an upstream node of the specified type exists - */ - _hasUpstreamNode(apiPrompt, nodeId, upstreamType) { - const visited = new Set(); - const toProcess = [nodeId]; - - while (toProcess.length > 0) { - const currentId = toProcess.pop(); - if (visited.has(currentId)) continue; - visited.add(currentId); - - const node = apiPrompt[currentId]; - if (!node) continue; - - // Check inputs for upstream connections - if (node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - const sourceNode = apiPrompt[sourceNodeId]; - - if (sourceNode && sourceNode.class_type === upstreamType) { - return true; - } - - if (!visited.has(sourceNodeId)) { - toProcess.push(sourceNodeId); - } - } - } - } - } - - return false; - } - - startStatusChecking() { - // Start checking every 2 seconds for more responsive updates - // Don't check immediately since panel might not be open yet - this.statusCheckInterval = setInterval(() => { - this.checkAllWorkerStatuses(); - }, 2000); - } - - async checkAllWorkerStatuses() { - // Check master status - this.checkMasterStatus(); - - if (!this.config || !this.config.workers) return; - - for (const worker of this.config.workers) { - // Check status for enabled workers OR workers that are launching - if (worker.enabled || this.state.isWorkerLaunching(worker.id)) { - this.checkWorkerStatus(worker); - } - } - } - - async checkMasterStatus() { - try { - const response = await fetch(`${window.location.origin}/prompt`, { - method: 'GET', - signal: AbortSignal.timeout(900) - }); - - if (response.ok) { - const data = await response.json(); - const queueRemaining = data.exec_info?.queue_remaining || 0; - const isProcessing = queueRemaining > 0; - - // Update master status dot - const statusDot = document.getElementById('master-status'); - if (statusDot) { - if (isProcessing) { - statusDot.style.backgroundColor = "#f0ad4e"; - statusDot.title = `Processing (${queueRemaining} in queue)`; - } else { - statusDot.style.backgroundColor = "#4CAF50"; - statusDot.title = "Online"; - } - } - } - } catch (error) { - // Master is always online (we're running on it), so keep it green - const statusDot = document.getElementById('master-status'); - if (statusDot) { - statusDot.style.backgroundColor = "#4CAF50"; - statusDot.title = "Online"; - } - } - } - - // Helper to build worker URL - getWorkerUrl(worker, endpoint = '') { - const host = worker.host || window.location.hostname; - - // Simple rule: port 443 = HTTPS, anything else = HTTP - const useHttps = worker.port === 443; - const protocol = useHttps ? 'https' : 'http'; - - // Don't add port for standard HTTPS/HTTP ports - const defaultPort = useHttps ? 443 : 80; - const port = worker.port === defaultPort ? '' : `:${worker.port}`; - - return `${protocol}://${host}${port}${endpoint}`; - } - - async checkWorkerStatus(worker) { - const url = this.getWorkerUrl(worker, '/prompt'); - const statusDot = document.getElementById(`status-${worker.id}`); - - try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - signal: AbortSignal.timeout(900) // Use your increased timeout - }); - - if (response.ok) { - const data = await response.json(); - const queueRemaining = data.exec_info?.queue_remaining || 0; - const isProcessing = queueRemaining > 0; - - // Update status - this.state.setWorkerStatus(worker.id, { - online: true, - processing: isProcessing, - queueCount: queueRemaining - }); - - // Update status dot based on processing state - if (isProcessing) { - this.updateStatusDot(worker.id, "#f0ad4e", `Online - Processing (${queueRemaining} in queue)`, false); - } else { - this.updateStatusDot(worker.id, "#3ca03c", "Online - Idle", false); - } - - // Clear launching state since worker is now online - if (this.state.isWorkerLaunching(worker.id)) { - this.state.setWorkerLaunching(worker.id, false); - this.clearLaunchingFlag(worker.id); - } - } else { - throw new Error(`HTTP ${response.status}`); - } - } catch (error) { - // Worker is offline or unreachable - this.state.setWorkerStatus(worker.id, { - online: false, - processing: false, - queueCount: 0 - }); - - // Check if worker is launching - if (this.state.isWorkerLaunching(worker.id)) { - this.updateStatusDot(worker.id, "#f0ad4e", "Launching...", true); - } else { - // Only update to red if not currently launching - this.updateStatusDot(worker.id, "#c04c4c", "Offline - Cannot connect", false); - } - - this.debugLog(`Worker ${worker.id} status check failed: ${error.message}`); - } - - // Update control buttons based on new status - this.updateWorkerControls(worker.id); - } - - async _dispatchToWorker(worker, prompt, workflow, imageReferences) { - const workerUrl = this.getWorkerUrl(worker); - - // Debug logging - always log to console for debugging - console.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`); - console.log('[Distributed] Worker URL:', workerUrl); - - // Handle image uploads for remote workers - if (imageReferences && imageReferences.size > 0) { - if (this.debugMode) { - console.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`); - } - - try { - // Load images from master - const images = await this.loadImagesForWorker(imageReferences); - - // Upload images to worker - if (images.length > 0) { - await this.uploadImagesToWorker(workerUrl, images); - if (this.debugMode) { - console.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`); - } - } - } catch (error) { - this.log(`Failed to process images for worker ${worker.name}: ${error.message}`); - // Continue with workflow execution even if image upload fails - } - } - - const promptToSend = { - prompt, - extra_data: { extra_pnginfo: { workflow } }, - client_id: api.clientId - }; - - console.log('[Distributed] Prompt data:', promptToSend); - - try { - await fetch(`${workerUrl}/prompt`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - mode: 'cors', - body: JSON.stringify(promptToSend) - }); - } catch (e) { - this.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`); - } - } - - async loadImagesForWorker(imageReferences) { - const images = []; - - // Use a cache to avoid loading the same image multiple times - if (!this._imageCache) { - this._imageCache = new Map(); - } - - for (const [filename, info] of imageReferences) { - try { - // Check cache first - if (this._imageCache.has(filename)) { - images.push(this._imageCache.get(filename)); - this.debugLog(`Using cached image: ${filename}`); - continue; - } - - // Load image from master's filesystem via API - try { - const data = await this.api.loadImage(filename); - const imageData = { - name: filename, - image: data.image_data - }; - images.push(imageData); - - // Cache the image for future use - this._imageCache.set(filename, imageData); - this.debugLog(`Loaded and cached image: ${filename}`); - } catch (loadError) { - this.log(`Failed to load image ${filename}: ${loadError.message}`); - throw loadError; - } - } catch (error) { - this.log(`Error loading image ${filename}: ${error.message}`); - } - } - - // Clear cache after a reasonable time to avoid memory issues - setTimeout(() => { - if (this._imageCache && this._imageCache.size > 0) { - this.debugLog(`Clearing image cache (${this._imageCache.size} images)`); - this._imageCache.clear(); - } - }, 30000); // Clear after 30 seconds - - return images; - } - - async uploadImagesToWorker(workerUrl, images) { - // Upload images to worker's ComfyUI instance - for (const imageData of images) { - const formData = new FormData(); - - // Convert base64 to blob - const base64Data = imageData.image.replace(/^data:image\/\w+;base64,/, ''); - const byteCharacters = atob(base64Data); - const byteNumbers = new Array(byteCharacters.length); - for (let i = 0; i < byteCharacters.length; i++) { - byteNumbers[i] = byteCharacters.charCodeAt(i); - } - const byteArray = new Uint8Array(byteNumbers); - const blob = new Blob([byteArray], { type: 'image/png' }); - - // Clean the filename - remove [input] suffix and handle subfolder - let cleanName = imageData.name; - let subfolder = ''; - - // Remove [input] or other suffixes - cleanName = cleanName.replace(/\s*\[\w+\]$/, ''); - - // Extract subfolder if present - if (cleanName.includes('/')) { - const parts = cleanName.split('/'); - subfolder = parts.slice(0, -1).join('/'); - cleanName = parts[parts.length - 1]; - } - - formData.append('image', blob, cleanName); - formData.append('type', 'input'); - formData.append('subfolder', subfolder); - formData.append('overwrite', 'true'); - - try { - const response = await fetch(`${workerUrl}/upload/image`, { - method: 'POST', - mode: 'cors', - body: formData - }); - - if (!response.ok) { - throw new Error(`Upload failed: ${response.statusText}`); - } - - this.debugLog(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`); - } catch (error) { - throw new Error(`Failed to upload ${imageData.name}: ${error.message}`); - } - } - } - - async performPreflightCheck(workers) { - if (workers.length === 0) return []; - - this.debugLog(`Performing pre-flight health check on ${workers.length} workers...`); - const startTime = Date.now(); - - const checkPromises = workers.map(async (worker) => { - const url = this.getWorkerUrl(worker, '/prompt'); - - this.debugLog(`Pre-flight checking ${worker.name} at: ${url}`); - - try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - signal: AbortSignal.timeout(900) // Use your increased timeout - }); - - if (response.ok) { - this.debugLog(`Worker ${worker.name} is active`); - return { worker, active: true }; - } else { - this.debugLog(`Worker ${worker.name} returned ${response.status}`); - return { worker, active: false }; - } - } catch (error) { - this.debugLog(`Worker ${worker.name} is offline or unreachable: ${error.message}`); - return { worker, active: false }; - } - }); - - const results = await Promise.all(checkPromises); - const activeWorkers = results.filter(r => r.active).map(r => r.worker); - - const elapsed = Date.now() - startTime; - this.debugLog(`Pre-flight check completed in ${elapsed}ms. Active workers: ${activeWorkers.length}/${workers.length}`); - - // Update UI status indicators for inactive workers - results.filter(r => !r.active).forEach(r => { - const statusDot = document.getElementById(`status-${r.worker.id}`); - if (statusDot) { - // Remove pulsing animation once status is determined - statusDot.classList.remove('status-pulsing'); - statusDot.style.backgroundColor = "#c04c4c"; // Red for offline - statusDot.title = "Offline - Cannot connect"; - } - }); - - return activeWorkers; - } - async launchWorker(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - const launchBtn = document.querySelector(`#controls-${workerId} button`); - - this.updateStatusDot(workerId, "#f0ad4e", "Launching...", true); - this.state.setWorkerLaunching(workerId, true); - - // Allow 90 seconds for worker to launch (model loading can take time) - setTimeout(() => { - this.state.setWorkerLaunching(workerId, false); - }, 90000); - - if (!launchBtn) return; - - try { - // Disable button immediately - launchBtn.disabled = true; - - const result = await this.api.launchWorker(workerId); - if (result) { - this.log(`Launched ${worker.name} (PID: ${result.pid})`); - if (result.log_file) { - this.debugLog(`Log file: ${result.log_file}`); - } - - this.state.setWorkerManaged(workerId, { - pid: result.pid, - log_file: result.log_file, - started_at: Date.now() - }); - - // Update controls immediately to hide launch button and show stop/log buttons - this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 2000); - } - } catch (error) { - // Check if worker was already running - if (error.message && error.message.includes("already running")) { - this.log(`Worker ${worker.name} is already running`); - this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 100); - } else { - this.log(`Error launching worker: ${error.message || error}`); - - // Re-enable button on error - if (launchBtn) { - launchBtn.disabled = false; - } - } - } - } - - async stopWorker(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1]; - - // Provide immediate feedback - if (stopBtn) { - stopBtn.disabled = true; - stopBtn.textContent = "Stopping..."; - stopBtn.style.backgroundColor = "#666"; - } - - try { - const result = await this.api.stopWorker(workerId); - if (result) { - this.log(`Stopped worker: ${result.message}`); - this.state.setWorkerManaged(workerId, null); - - // Flash success feedback - if (stopBtn) { - stopBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success; - stopBtn.textContent = "Stopped!"; - setTimeout(() => { - this.updateWorkerControls(workerId); - }, 1500); - } - - // Update status - setTimeout(() => this.checkWorkerStatus(worker), 500); - } else { - this.log(`Failed to stop worker: ${result.message}`); - - // Flash error feedback - if (stopBtn) { - stopBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - stopBtn.textContent = result.message.includes("already stopped") ? "Not Running" : "Failed"; - setTimeout(() => { - this.updateWorkerControls(workerId); - }, 2000); - } - } - } catch (error) { - this.log(`Error stopping worker: ${error}`); - - // Flash error feedback - if (stopBtn) { - stopBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - stopBtn.textContent = "Error"; - setTimeout(() => { - this.updateWorkerControls(workerId); - }, 2000); - } - } - } - - async clearLaunchingFlag(workerId) { - try { - await this.api.clearLaunchingFlag(workerId); - this.debugLog(`Cleared launching flag for worker ${workerId}`); - } catch (error) { - this.debugLog(`Error clearing launching flag: ${error.message || error}`); - } - } - - // Helper Functions to reduce redundancy - - /** - * Generic async button action handler - */ - async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = 3000) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; - button.disabled = true; - - try { - await action(); - button.textContent = successText; - button.style.cssText = originalStyle; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success; - return true; - } catch (error) { - button.textContent = errorText || `Error: ${error.message}`; - button.style.cssText = originalStyle; - button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - throw error; - } finally { - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - button.disabled = false; - }, resetDelay); - } - } - - /** - * Show toast notification helper - */ - showToast(severity, summary, detail, life = 3000) { - if (app.extensionManager?.toast?.add) { - app.extensionManager.toast.add({ severity, summary, detail, life }); - } - } - - /** - * Update status dot helper - */ - updateStatusDot(workerId, color, title, pulsing = false) { - const statusDot = document.getElementById(`status-${workerId}`); - if (!statusDot) return; - - statusDot.style.backgroundColor = color; - statusDot.title = title; - statusDot.classList.toggle('status-pulsing', pulsing); - } - - /** - * Cleanup method to stop intervals and listeners - */ - cleanup() { - if (this.statusCheckInterval) { - clearInterval(this.statusCheckInterval); - this.statusCheckInterval = null; - } - - if (this.logAutoRefreshInterval) { - clearInterval(this.logAutoRefreshInterval); - this.logAutoRefreshInterval = null; - } - - this.debugLog("Cleaned up intervals"); - } - - async loadManagedWorkers() { - try { - const result = await this.api.getManagedWorkers(); - // Clear all managed worker info - this.state.clearCache(); - - // Check for launching workers - for (const [workerId, info] of Object.entries(result.managed_workers)) { - this.state.setWorkerManaged(workerId, info); - - // If worker is marked as launching, add to launchingWorkers set - if (info.launching) { - this.state.setWorkerLaunching(workerId, true); - this.debugLog(`Worker ${workerId} is in launching state`); - } - } - - // Update UI for all workers - if (this.config?.workers) { - this.config.workers.forEach(w => this.updateWorkerControls(w.id)); - } - } catch (error) { - this.debugLog(`Error loading managed workers: ${error}`); - } - } - - updateWorkerControls(workerId) { - const controlsDiv = document.getElementById(`controls-${workerId}`); - - if (!controlsDiv) return; - - const worker = this.config.workers.find(w => w.id === workerId); - if (!worker) return; - - // Skip button updates for remote workers - if (this.isRemoteWorker(worker)) { - return; - } - - // Ensure we check for string ID - const managedInfo = this.state.getWorker(workerId).managed; - const status = this.state.getWorkerStatus(workerId); - - // Update button states - const buttons = controlsDiv.querySelectorAll('button'); - const launchBtn = buttons[0]; - const stopBtn = buttons[1]; - const logBtn = document.getElementById(`log-${workerId}`); - - // Show log button immediately if we have log file info (even if worker is still starting) - if (managedInfo?.log_file && logBtn) { - logBtn.style.display = 'inline-block'; - } else if (logBtn && !managedInfo) { - logBtn.style.display = 'none'; - } - - if (status?.online || managedInfo) { - // Worker is running or we just launched it - launchBtn.style.display = 'none'; // Hide launch button when running - - if (managedInfo) { - // Only show stop button if we manage this worker - stopBtn.style.display = 'inline-block'; - stopBtn.disabled = false; - stopBtn.textContent = "Stop"; - stopBtn.style.backgroundColor = "#7c4a4a"; // Red when enabled - } else { - // Hide stop button for workers launched outside UI - stopBtn.style.display = 'none'; - } - } else { - // Worker is not running - launchBtn.style.display = 'inline-block'; // Show launch button - launchBtn.disabled = false; - launchBtn.textContent = "Launch"; - launchBtn.style.backgroundColor = "#4a7c4a"; // Green when enabled - - stopBtn.style.display = 'none'; // Hide stop button when not running - } - } - - async viewWorkerLog(workerId) { - const managedInfo = this.state.getWorker(workerId).managed; - if (!managedInfo?.log_file) return; - - const logBtn = document.getElementById(`log-${workerId}`); - - // Provide immediate feedback - if (logBtn) { - logBtn.disabled = true; - logBtn.textContent = "Loading..."; - logBtn.style.backgroundColor = "#666"; - } - - try { - // Fetch log content - const data = await this.api.getWorkerLog(workerId, 1000); - - // Create modal dialog - this.showLogModal(workerId, data); - - // Restore button - if (logBtn) { - logBtn.disabled = false; - logBtn.textContent = "View Log"; - logBtn.style.backgroundColor = "#685434"; // Keep the yellow color - } - - } catch (error) { - this.log('Error viewing log: ' + error.message); - app.extensionManager.toast.add({ - severity: "error", - summary: "Error", - detail: `Failed to load log: ${error.message}`, - life: 5000 - }); - - // Flash error and restore button - if (logBtn) { - logBtn.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error; - logBtn.textContent = "Error"; - setTimeout(() => { - logBtn.disabled = false; - logBtn.textContent = "View Log"; - logBtn.style.backgroundColor = "#685434"; // Keep the yellow color - }, 2000); - } - } - } - - showLogModal(workerId, logData) { - // Remove any existing modal - const existingModal = document.getElementById('distributed-log-modal'); - if (existingModal) { - existingModal.remove(); - } - - const worker = this.config.workers.find(w => w.id === workerId); - const workerName = worker?.name || `Worker ${workerId}`; - - // Create modal container - const modal = document.createElement('div'); - modal.id = 'distributed-log-modal'; - modal.style.cssText = ` - position: fixed; - top: 0; - left: 0; - width: 100%; - height: 100%; - background: rgba(0, 0, 0, 0.8); - display: flex; - align-items: center; - justify-content: center; - z-index: 10000; - `; - - // Create modal content - const content = document.createElement('div'); - content.style.cssText = ` - background: #1e1e1e; - border-radius: 8px; - width: 90%; - max-width: 1200px; - height: 80%; - display: flex; - flex-direction: column; - box-shadow: 0 4px 20px rgba(0, 0, 0, 0.5); - `; - - // Header - const header = document.createElement('div'); - header.style.cssText = ` - padding: 15px 20px; - border-bottom: 1px solid #444; - display: flex; - justify-content: space-between; - align-items: center; - `; - - const title = document.createElement('h3'); - title.textContent = `${workerName} - Log Viewer`; - title.style.cssText = 'margin: 0; color: #fff;'; - - const headerButtons = document.createElement('div'); - headerButtons.style.cssText = 'display: flex; gap: 20px; align-items: center;'; - - // Auto-refresh container - const refreshContainer = document.createElement('div'); - refreshContainer.style.cssText = 'display: flex; align-items: center; gap: 4px;'; - - // Auto-refresh checkbox - const refreshCheckbox = document.createElement('input'); - refreshCheckbox.type = 'checkbox'; - refreshCheckbox.id = 'log-auto-refresh'; - refreshCheckbox.checked = true; // Enabled by default - refreshCheckbox.style.cssText = 'cursor: pointer;'; - refreshCheckbox.onchange = (e) => { - if (e.target.checked) { - this.startLogAutoRefresh(workerId); - } else { - this.stopLogAutoRefresh(); - } - }; - - const refreshLabel = document.createElement('label'); - refreshLabel.htmlFor = 'log-auto-refresh'; - refreshLabel.style.cssText = 'font-size: 12px; color: #ccc; cursor: pointer; white-space: nowrap;'; - refreshLabel.textContent = 'Auto-refresh'; - - // Add checkbox and label to container - refreshContainer.appendChild(refreshCheckbox); - refreshContainer.appendChild(refreshLabel); - - // Close button - const closeBtn = this._createButton('✕', - () => { - this.stopLogAutoRefresh(); - modal.remove(); - }, - 'background-color: #c04c4c;'); - closeBtn.style.cssText += ' padding: 5px 10px; font-size: 14px; font-weight: bold;'; - - headerButtons.appendChild(refreshContainer); - headerButtons.appendChild(closeBtn); - - header.appendChild(title); - header.appendChild(headerButtons); - - // Log content area - const logContainer = document.createElement('div'); - logContainer.style.cssText = ` - flex: 1; - overflow: auto; - padding: 15px; - font-family: 'Consolas', 'Monaco', 'Courier New', monospace; - font-size: 12px; - line-height: 1.4; - color: #ddd; - background: #0d0d0d; - white-space: pre-wrap; - word-wrap: break-word; - `; - logContainer.id = 'distributed-log-content'; - logContainer.textContent = logData.content; - - // Auto-scroll to bottom - setTimeout(() => { - logContainer.scrollTop = logContainer.scrollHeight; - }, 0); - - // Status bar - const statusBar = document.createElement('div'); - statusBar.style.cssText = ` - padding: 10px 20px; - border-top: 1px solid #444; - font-size: 11px; - color: #888; - `; - statusBar.textContent = `Log file: ${logData.log_file}`; - if (logData.truncated) { - statusBar.textContent += ` (showing last ${logData.lines_shown} lines of ${this.formatFileSize(logData.file_size)})`; - } - - // Assemble modal - content.appendChild(header); - content.appendChild(logContainer); - content.appendChild(statusBar); - modal.appendChild(content); - - // Close on background click - modal.onclick = (e) => { - if (e.target === modal) { - this.stopLogAutoRefresh(); - modal.remove(); - } - }; - - // Close on Escape key - const handleEscape = (e) => { - if (e.key === 'Escape') { - this.stopLogAutoRefresh(); - modal.remove(); - document.removeEventListener('keydown', handleEscape); - } - }; - document.addEventListener('keydown', handleEscape); - - document.body.appendChild(modal); - - // Start auto-refresh - this.startLogAutoRefresh(workerId); - } - - async refreshLog(workerId, silent = false) { - const logContent = document.getElementById('distributed-log-content'); - if (!logContent) return; - - try { - const data = await this.api.getWorkerLog(workerId, 1000); - - // Update content - const shouldAutoScroll = logContent.scrollTop + logContent.clientHeight >= logContent.scrollHeight - 50; - logContent.textContent = data.content; - - // Auto-scroll if was at bottom - if (shouldAutoScroll) { - logContent.scrollTop = logContent.scrollHeight; - } - - // Only show toast if not in silent mode (manual refresh) - if (!silent) { - app.extensionManager.toast.add({ - severity: "success", - summary: "Log Refreshed", - detail: "Log content updated", - life: 2000 - }); - } - - } catch (error) { - // Only show error toast if not in silent mode - if (!silent) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Refresh Failed", - detail: error.message, - life: 3000 - }); - } - } - } - - formatFileSize(bytes) { - if (bytes < 1024) return bytes + ' B'; - if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; - return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; - } - - isRemoteWorker(worker) { - const host = worker.host || window.location.hostname; - return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; - } - - getMasterUrl() { - // Always use the detected/configured master IP for consistency - if (this.config?.master?.host) { - const configuredHost = this.config.master.host; - - // If the configured host already includes protocol, use as-is - if (configuredHost.startsWith('http://') || configuredHost.startsWith('https://')) { - return configuredHost; - } - - // For domain names (not IPs), default to HTTPS - const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost); - const isLocalhost = configuredHost === 'localhost' || configuredHost === '127.0.0.1'; - - if (!isIP && !isLocalhost && configuredHost.includes('.')) { - // It's a domain name, use HTTPS - return `https://${configuredHost}`; - } else { - // For IPs and localhost, use current access method - const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - if ((window.location.protocol === 'https:' && port === '443') || - (window.location.protocol === 'http:' && port === '80')) { - return `${window.location.protocol}//${configuredHost}`; - } - return `${window.location.protocol}//${configuredHost}:${port}`; - } - } - - // If no master IP is set but we're on a network address, use it - const hostname = window.location.hostname; - if (hostname !== 'localhost' && hostname !== '127.0.0.1') { - return window.location.origin; - } - - // Fallback warning - this won't work for remote workers - this.debugLog("No master host configured - remote workers won't be able to connect. " + - "Master host should be auto-detected on startup."); - return window.location.origin; - } - - async detectMasterIP() { - try { - const data = await this.api.getNetworkInfo(); - this.log("Network info: " + JSON.stringify(data)); - - // Store CUDA device info - if (data.cuda_device !== null && data.cuda_device !== undefined) { - this.masterCudaDevice = data.cuda_device; - // Update the master display with CUDA info - this.updateMasterDisplay(); - } - - // Store CUDA device count for auto-population - if (data.cuda_device_count > 0) { - this.cudaDeviceCount = data.cuda_device_count; - this.log(`Detected ${this.cudaDeviceCount} CUDA devices`); - - // Auto-populate workers if none exist - this.log(`Current workers: ${this.config.workers ? this.config.workers.length : 'null'}`); - if (!this.config.workers || this.config.workers.length === 0) { - this.log(`Auto-populating ${this.cudaDeviceCount} workers based on CUDA devices`); - - const newWorkers = []; - for (let i = 0; i < this.cudaDeviceCount; i++) { - const worker = { - id: Date.now() + i, - name: `Worker ${i + 1}`, - host: "localhost", - port: 8189 + i, - cuda_device: i, - enabled: true - }; - newWorkers.push(worker); - } - - // Update config and save - this.config.workers = newWorkers; - - // Refresh the panel if it's already open - const panel = document.getElementById("distributed-panel"); - if (panel) { - this.renderSidebarContent(panel); - } - } - } - - // Check if we already have a master host configured - if (this.config?.master?.host) { - this.debugLog(`Master host already configured: ${this.config.master.host}`); - return; - } - - // Use the recommended IP from the backend - if (data.recommended_ip && data.recommended_ip !== '127.0.0.1') { - this.log(`Auto-detected master IP: ${data.recommended_ip}`); - - // Save the detected IP (pass true to suppress notification) - await this.saveMasterConfig({ host: data.recommended_ip }, true); - - // Show a single combined notification - if (app.extensionManager?.toast) { - app.extensionManager.toast.add({ - severity: "info", - summary: "Master IP Auto-Detected", - detail: `Master IP set to ${data.recommended_ip} for remote worker communication`, - life: 5000 - }); - } - } - } catch (error) { - this.debugLog("Error detecting master IP: " + error.message); - } - } - - async saveMasterConfig(updates, suppressNotification = false) { - try { - const response = await fetch(`${window.location.origin}/distributed/config/update_master`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(updates) - }); - - if (response.ok) { - // Update local config - if (!this.config.master) { - this.config.master = {}; - } - // Update all provided fields - Object.assign(this.config.master, updates); - - // Only show notification if not suppressed - if (!suppressNotification) { - app.extensionManager.toast.add({ - severity: "success", - summary: "Master Host Saved", - detail: masterIp.trim() ? `Master host set to ${masterIp.trim()}` : "Master host cleared (auto-detect)", - life: 3000 - }); - } - } else { - throw new Error("Failed to save master host"); - } - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Save Failed", - detail: error.message, - life: 5000 - }); - } - } - - startLogAutoRefresh(workerId) { - // Stop any existing auto-refresh - this.stopLogAutoRefresh(); - - // Refresh every 2 seconds - this.logAutoRefreshInterval = setInterval(() => { - this.refreshLog(workerId, true); // silent mode - }, 2000); - } - - stopLogAutoRefresh() { - if (this.logAutoRefreshInterval) { - clearInterval(this.logAutoRefreshInterval); - this.logAutoRefreshInterval = null; - } - } - - toggleWorkerExpanded(workerId) { - const settingsDiv = document.getElementById(`settings-${workerId}`); - const gpuDiv = settingsDiv.closest('[style*="margin-bottom: 12px"]'); - const settingsArrow = gpuDiv.querySelector('.settings-arrow'); - - if (!settingsDiv) return; - - if (this.state.isWorkerExpanded(workerId)) { - this.state.setWorkerExpanded(workerId, false); - settingsDiv.classList.remove("expanded"); - if (settingsArrow) { - settingsArrow.style.transform = "rotate(0deg)"; - } - // Animate padding to 0 - settingsDiv.style.padding = "0 12px"; - settingsDiv.style.marginTop = "0"; - settingsDiv.style.marginBottom = "0"; - } else { - this.state.setWorkerExpanded(workerId, true); - settingsDiv.classList.add("expanded"); - if (settingsArrow) { - settingsArrow.style.transform = "rotate(90deg)"; - } - // Animate padding to full - settingsDiv.style.padding = "12px"; - settingsDiv.style.marginTop = "8px"; - settingsDiv.style.marginBottom = "8px"; - } - } - - createWorkerSettingsForm(worker) { - const form = document.createElement("div"); - form.style.cssText = "display: flex; flex-direction: column; gap: 10px;"; - - // Name field - const nameGroup = this.createFormGroup("Name:", worker.name, `name-${worker.id}`); - form.appendChild(nameGroup); - - // Remote checkbox - const remoteGroup = document.createElement("div"); - remoteGroup.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; - - const remoteCheckbox = document.createElement("input"); - remoteCheckbox.type = "checkbox"; - remoteCheckbox.id = `remote-${worker.id}`; - remoteCheckbox.checked = this.isRemoteWorker(worker); - remoteCheckbox.onchange = (e) => { - const isRemote = e.target.checked; - // Show/hide relevant fields - const hostGroup = document.getElementById(`host-group-${worker.id}`); - const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); - const argsGroup = document.getElementById(`args-group-${worker.id}`); - - if (isRemote) { - hostGroup.style.display = "flex"; - cudaGroup.style.display = "none"; - argsGroup.style.display = "none"; - // If switching to remote and host is localhost, clear it - const hostInput = document.getElementById(`host-${worker.id}`); - if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { - hostInput.value = ""; - } - } else { - hostGroup.style.display = "none"; - cudaGroup.style.display = "flex"; - argsGroup.style.display = "flex"; - } - }; - - const remoteLabel = document.createElement("label"); - remoteLabel.htmlFor = `remote-${worker.id}`; - remoteLabel.textContent = "Remote Worker"; - remoteLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; - - remoteGroup.appendChild(remoteCheckbox); - remoteGroup.appendChild(remoteLabel); - form.appendChild(remoteGroup); - - // Host field (only for remote workers) - const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); - hostGroup.id = `host-group-${worker.id}`; - hostGroup.style.display = this.isRemoteWorker(worker) ? "flex" : "none"; - form.appendChild(hostGroup); - - // Port field - const portGroup = this.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number"); - form.appendChild(portGroup); - - // CUDA Device field (only for local workers) - const cudaGroup = this.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number"); - cudaGroup.id = `cuda-group-${worker.id}`; - cudaGroup.style.display = this.isRemoteWorker(worker) ? "none" : "flex"; - form.appendChild(cudaGroup); - - // Extra Args field (only for local workers) - const argsGroup = this.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`); - argsGroup.id = `args-group-${worker.id}`; - argsGroup.style.display = this.isRemoteWorker(worker) ? "none" : "flex"; - form.appendChild(argsGroup); - - // Buttons - const buttonGroup = document.createElement("div"); - buttonGroup.style.cssText = "display: flex; gap: 8px; margin-top: 8px;"; - - const saveBtn = this._createButton("Save", - () => this.saveWorkerSettings(worker.id), - "background-color: #4a7c4a;"); - saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - - const cancelBtn = this._createButton("Cancel", - () => this.cancelWorkerSettings(worker.id), - "background-color: #555;"); - cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; - - const deleteBtn = this._createButton("Delete", - () => this.deleteWorker(worker.id), - "background-color: #7c4a4a;"); - deleteBtn.style.cssText += " padding: 6px 12px; font-size: 12px; margin-left: auto;"; - - buttonGroup.appendChild(saveBtn); - buttonGroup.appendChild(cancelBtn); - buttonGroup.appendChild(deleteBtn); - - form.appendChild(buttonGroup); - - return form; - } - - createSettingsToggle() { - const settingsRow = document.createElement("div"); - settingsRow.style.cssText = "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;"; - - const settingsTitle = document.createElement("h4"); - settingsTitle.textContent = "Settings"; - settingsTitle.style.cssText = "margin: 0; font-size: 14px;"; - - const settingsToggle = document.createElement("span"); - settingsToggle.textContent = "▶"; // Right arrow when collapsed - settingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; - - settingsRow.appendChild(settingsToggle); - settingsRow.appendChild(settingsTitle); - - return { settingsRow, settingsToggle }; - } - - createFormGroup(label, value, id, type = "text", placeholder = "") { - const result = this.ui.createFormGroup(label, value, id, type, placeholder); - return result.group; - } - - async saveWorkerSettings(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - if (!worker) return; - - // Get form values - const name = document.getElementById(`name-${workerId}`).value; - const isRemote = document.getElementById(`remote-${workerId}`).checked; - const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; - const port = parseInt(document.getElementById(`port-${workerId}`).value); - const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); - const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value; - - // Validate - if (!name.trim()) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Validation Error", - detail: "Worker name is required", - life: 3000 - }); - return; - } - - if (isRemote && !host.trim()) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Validation Error", - detail: "Host is required for remote workers", - life: 3000 - }); - return; - } - - if (isNaN(port) || port < 1 || port > 65535) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Validation Error", - detail: "Port must be between 1 and 65535", - life: 3000 - }); - return; - } - - // Check for port conflicts - // Remote workers can reuse ports, but local workers cannot share ports with each other or master - if (!isRemote) { - // Check if port conflicts with master - const masterPort = parseInt(window.location.port) || (window.location.protocol === 'https:' ? 443 : 80); - if (port === masterPort) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Port Conflict", - detail: `Port ${port} is already in use by the master server`, - life: 3000 - }); - return; - } - - // Check if port conflicts with other local workers - const localPortConflict = this.config.workers.some(w => - w.id !== workerId && - w.port === port && - !w.host // local workers have no host or host is null - ); - - if (localPortConflict) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Port Conflict", - detail: `Port ${port} is already in use by another local worker`, - life: 3000 - }); - return; - } - } else { - // For remote workers, only check conflicts with other workers on the same host - const sameHostConflict = this.config.workers.some(w => - w.id !== workerId && - w.port === port && - w.host === host.trim() - ); - - if (sameHostConflict) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Port Conflict", - detail: `Port ${port} is already in use by another worker on ${host}`, - life: 3000 - }); - return; - } - } - - try { - const response = await fetch(`${window.location.origin}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - worker_id: workerId, - name: name.trim(), - host: isRemote ? host.trim() : null, - port: port, - cuda_device: isRemote ? null : cudaDevice, - extra_args: isRemote ? null : (extraArgs ? extraArgs.trim() : "") - }) - }); - - if (response.ok) { - // Update local config - worker.name = name.trim(); - if (isRemote) { - worker.host = host.trim(); - delete worker.cuda_device; - delete worker.extra_args; - } else { - delete worker.host; - worker.cuda_device = cudaDevice; - worker.extra_args = extraArgs ? extraArgs.trim() : ""; - } - worker.port = port; - - app.extensionManager.toast.add({ - severity: "success", - summary: "Settings Saved", - detail: `Worker ${name} settings updated`, - life: 3000 - }); - - // Refresh the UI - if (this.panelElement) { - this.renderSidebarContent(this.panelElement); - } - } else { - const error = await response.json(); - throw new Error(error.message || "Failed to save settings"); - } - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Save Failed", - detail: error.message, - life: 5000 - }); - } - } - - cancelWorkerSettings(workerId) { - // Collapse the settings panel - this.toggleWorkerExpanded(workerId); - - // Reset form values to original - const worker = this.config.workers.find(w => w.id === workerId); - if (worker) { - document.getElementById(`name-${workerId}`).value = worker.name; - document.getElementById(`host-${workerId}`).value = worker.host || ""; - document.getElementById(`port-${workerId}`).value = worker.port; - document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0; - document.getElementById(`args-${workerId}`).value = worker.extra_args || ""; - - // Reset remote checkbox - const remoteCheckbox = document.getElementById(`remote-${workerId}`); - if (remoteCheckbox) { - remoteCheckbox.checked = this.isRemoteWorker(worker); - } - } - } - - async deleteWorker(workerId) { - const worker = this.config.workers.find(w => w.id === workerId); - if (!worker) return; - - // Confirm deletion - if (!confirm(`Are you sure you want to delete worker "${worker.name}"?`)) { - return; - } - - try { - const response = await fetch(`${window.location.origin}/distributed/config/delete_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ worker_id: workerId }) - }); - - if (response.ok) { - // Remove from local config - const index = this.config.workers.findIndex(w => w.id === workerId); - if (index !== -1) { - this.config.workers.splice(index, 1); - } - - app.extensionManager.toast.add({ - severity: "success", - summary: "Worker Deleted", - detail: `Worker ${worker.name} has been removed`, - life: 3000 - }); - - // Refresh the UI - if (this.panelElement) { - this.renderSidebarContent(this.panelElement); - } - } else { - const error = await response.json(); - throw new Error(error.message || "Failed to delete worker"); - } - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Delete Failed", - detail: error.message, - life: 5000 - }); - } - } - - // Backward compatibility wrapper - async saveMasterIp(masterIp, suppressNotification = false) { - return this.saveMasterConfig({ host: masterIp.trim() }, suppressNotification); - } - - async addNewWorker() { - // Generate new worker ID - const newId = `worker_${Date.now()}`; - - // Find next available port - const usedPorts = this.config.workers.map(w => w.port); - let nextPort = 8189; - while (usedPorts.includes(nextPort)) { - nextPort++; - } - - // Create new worker object - const newWorker = { - id: newId, - name: `Worker ${this.config.workers.length + 1}`, - port: nextPort, - cuda_device: this.config.workers.length, - enabled: true, // Default to enabled for convenience - extra_args: "" - }; - - // Add to config - this.config.workers.push(newWorker); - - // Save immediately - try { - const response = await fetch(`${window.location.origin}/distributed/config/update_worker`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - worker_id: newId, - name: newWorker.name, - port: newWorker.port, - cuda_device: newWorker.cuda_device, - extra_args: newWorker.extra_args, - enabled: newWorker.enabled - }) - }); - - if (!response.ok) { - // Remove from local config if save failed - this.config.workers.pop(); - throw new Error("Failed to save new worker"); - } - - app.extensionManager.toast.add({ - severity: "success", - summary: "Worker Added", - detail: `New worker created on port ${nextPort}`, - life: 3000 - }); - - // Refresh UI and expand the new worker - this.state.setWorkerExpanded(newId, true); - if (this.panelElement) { - this.renderSidebarContent(this.panelElement); - } - - } catch (error) { - app.extensionManager.toast.add({ - severity: "error", - summary: "Failed to Add Worker", - detail: error.message, - life: 5000 - }); - } - } -} - -app.registerExtension({ - name: "Distributed.Panel", - async setup() { - new DistributedExtension(); - } -}); \ No newline at end of file diff --git a/web/main.js b/web/main.js new file mode 100644 index 0000000..c0310c5 --- /dev/null +++ b/web/main.js @@ -0,0 +1,1178 @@ +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 } from './sidebarRenderer.js'; +import { handleWorkerOperation, handleInterruptWorkers, handleClearMemory } from './workerUtils.js'; +import { setupInterceptor, executeParallelDistributed } from './executionUtils.js'; +import { BUTTON_STYLES, PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js'; + +class DistributedExtension { + constructor() { + this.config = null; + this.originalQueuePrompt = api.queuePrompt.bind(api); + this.statusCheckInterval = null; + this.logAutoRefreshInterval = null; + this.masterSettingsExpanded = false; + + // Initialize centralized state + this.state = createStateManager(); + + // Initialize UI component factory + this.ui = new DistributedUI(); + + // Initialize API client + this.api = createApiClient(window.location.origin); + + // Inject CSS for pulsing animation + this.injectStyles(); + + this.loadConfig().then(async () => { + this.registerSidebarTab(); + this.setupInterceptor(); + this.startStatusChecking(); + this.loadManagedWorkers(); + // Detect master IP after everything is set up + this.detectMasterIP(); + }); + } + + // 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); + } + } + + // --- State & Config Management (Single Source of Truth) --- + + get enabledWorkers() { + return this.config?.workers?.filter(w => w.enabled) || []; + } + + get isEnabled() { + return this.enabledWorkers.length > 0; + } + + 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 }); + }); + } + } catch (error) { + this.log("Failed to load config: " + error.message, "error"); + this.config = { workers: [], settings: { has_auto_populated_workers: false } }; + } + } + + async updateWorkerEnabled(workerId, enabled) { + const worker = this.config.workers.find(w => w.id === workerId); + if (worker) { + worker.enabled = enabled; + this.state.updateWorker(workerId, { enabled }); + + // 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...", true); + setTimeout(() => this.checkWorkerStatus(worker), 100); + } 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"); + } + } + + 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); + + app.extensionManager.toast.add({ + severity: "success", + summary: "Setting Updated", + detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`, + 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) => renderSidebarContent(this, el) + }); + } + + updateSummary() { + const summaryEl = document.getElementById('distributed-summary'); + if (summaryEl) { + // Count active workers from state + const activeWorkers = this.config.workers.filter(w => { + const status = this.state.getWorkerStatus(w.id); + return status.online; + }).length; + const totalGPUs = activeWorkers + 1; + if (this.isEnabled) { + summaryEl.textContent = `If Collector node is present, total generation = (${totalGPUs} GPUs × Batch Size)`; + } else { + summaryEl.textContent = "Only the master GPU will be used."; + } + } + } + + // --- Core Logic & Execution --- + + setupInterceptor() { + setupInterceptor(this); + } + + async executeParallelDistributed(promptWrapper) { + return executeParallelDistributed(this, promptWrapper); + } + + startStatusChecking() { + // Start checking every 2 seconds for more responsive updates + // Don't check immediately since panel might not be open yet + this.statusCheckInterval = setInterval(() => { + this.checkAllWorkerStatuses(); + }, 2000); + } + + async checkAllWorkerStatuses() { + // Check master status + this.checkMasterStatus(); + + if (!this.config || !this.config.workers) return; + + for (const worker of this.config.workers) { + // Check status for enabled workers OR workers that are launching + if (worker.enabled || this.state.isWorkerLaunching(worker.id)) { + this.checkWorkerStatus(worker); + } + } + } + + async checkMasterStatus() { + try { + const response = await fetch(`${window.location.origin}/prompt`, { + method: 'GET', + signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + const isProcessing = queueRemaining > 0; + + // Update master status in state + this.state.setMasterStatus(isProcessing ? 'processing' : 'online'); + + // Update master status dot + const statusDot = document.getElementById('master-status'); + if (statusDot) { + if (isProcessing) { + statusDot.style.backgroundColor = "#f0ad4e"; + statusDot.title = `Processing (${queueRemaining} in queue)`; + } else { + statusDot.style.backgroundColor = "#4CAF50"; + statusDot.title = "Online"; + } + } + } + } catch (error) { + // Master is always online (we're running on it), so keep it green + const statusDot = document.getElementById('master-status'); + if (statusDot) { + statusDot.style.backgroundColor = "#4CAF50"; + statusDot.title = "Online"; + } + } + } + + // Helper to build worker URL + getWorkerUrl(worker, endpoint = '') { + const host = worker.host || window.location.hostname; + + // Simple rule: port 443 = HTTPS, anything else = HTTP + const useHttps = worker.port === 443; + const protocol = useHttps ? 'https' : 'http'; + + // Don't add port for standard HTTPS/HTTP ports + const defaultPort = useHttps ? 443 : 80; + const port = worker.port === defaultPort ? '' : `:${worker.port}`; + + return `${protocol}://${host}${port}${endpoint}`; + } + + async checkWorkerStatus(worker) { + // Assume caller ensured enabled; proceed with check + const url = this.getWorkerUrl(worker, '/prompt'); + const statusDot = document.getElementById(`status-${worker.id}`); + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + const isProcessing = queueRemaining > 0; + + // Update status + this.state.setWorkerStatus(worker.id, { + online: true, + processing: isProcessing, + queueCount: queueRemaining + }); + + // Update status dot based on processing state + if (isProcessing) { + this.ui.updateStatusDot( + worker.id, + "#f0ad4e", + `Online - Processing (${queueRemaining} in queue)`, + false + ); + } else { + this.ui.updateStatusDot(worker.id, "#3ca03c", "Online - Idle", false); + } + + // Clear launching state since worker is now online + if (this.state.isWorkerLaunching(worker.id)) { + this.state.setWorkerLaunching(worker.id, false); + this.clearLaunchingFlag(worker.id); + } + } else { + throw new Error(`HTTP ${response.status}`); + } + } catch (error) { + // Worker is offline or unreachable + this.state.setWorkerStatus(worker.id, { + online: false, + processing: false, + queueCount: 0 + }); + + // Check if worker is launching + if (this.state.isWorkerLaunching(worker.id)) { + this.ui.updateStatusDot(worker.id, "#f0ad4e", "Launching...", true); + } else if (worker.enabled) { + // Only update to red if not currently launching AND still enabled + this.ui.updateStatusDot(worker.id, "#c04c4c", "Offline - Cannot connect", false); + } + // If disabled, don't update the dot (leave it gray) + + this.log(`Worker ${worker.id} status check failed: ${error.message}`, "debug"); + } + + // Update control buttons based on new status + this.updateWorkerControls(worker.id); + } + + async launchWorker(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + const launchBtn = document.querySelector(`#controls-${workerId} button`); + + // If worker is disabled, enable it first + if (!worker.enabled) { + await this.updateWorkerEnabled(workerId, true); + + // Update the checkbox UI + const checkbox = document.getElementById(`gpu-${workerId}`); + if (checkbox) { + checkbox.checked = true; + } + + this.updateSummary(); + } + + this.ui.updateStatusDot(workerId, "#f0ad4e", "Launching...", true); + this.state.setWorkerLaunching(workerId, true); + + // Allow 90 seconds for worker to launch (model loading can take time) + setTimeout(() => { + this.state.setWorkerLaunching(workerId, false); + }, 90000); + + if (!launchBtn) return; + + try { + // Disable button immediately + launchBtn.disabled = true; + + const result = await this.api.launchWorker(workerId); + if (result) { + this.log(`Launched ${worker.name} (PID: ${result.pid})`, "info"); + if (result.log_file) { + this.log(`Log file: ${result.log_file}`, "debug"); + } + + this.state.setWorkerManaged(workerId, { + pid: result.pid, + log_file: result.log_file, + started_at: Date.now() + }); + + // Update controls immediately to hide launch button and show stop/log buttons + this.updateWorkerControls(workerId); + setTimeout(() => this.checkWorkerStatus(worker), 2000); + } + } catch (error) { + // Check if worker was already running + if (error.message && error.message.includes("already running")) { + this.log(`Worker ${worker.name} is already running`, "info"); + this.updateWorkerControls(workerId); + setTimeout(() => this.checkWorkerStatus(worker), 100); + } else { + this.log(`Error launching worker: ${error.message || error}`, "error"); + + // Re-enable button on error + if (launchBtn) { + launchBtn.disabled = false; + } + } + } + } + + async stopWorker(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1]; + + // Provide immediate feedback + if (stopBtn) { + stopBtn.disabled = true; + stopBtn.textContent = "Stopping..."; + stopBtn.style.backgroundColor = "#666"; + } + + try { + const result = await this.api.stopWorker(workerId); + if (result) { + this.log(`Stopped worker: ${result.message}`, "info"); + this.state.setWorkerManaged(workerId, null); + + // Immediately update status to offline + this.ui.updateStatusDot(workerId, "#c04c4c", "Offline"); + this.state.setWorkerStatus(workerId, { online: false }); + + // Flash success feedback + if (stopBtn) { + stopBtn.style.backgroundColor = BUTTON_STYLES.success; + stopBtn.textContent = "Stopped!"; + setTimeout(() => { + this.updateWorkerControls(workerId); + }, 1000); + } + + // Verify status after a short delay + setTimeout(() => this.checkWorkerStatus(worker), 2000); + } else { + this.log(`Failed to stop worker: ${result.message}`, "error"); + + // Flash error feedback + if (stopBtn) { + stopBtn.style.backgroundColor = BUTTON_STYLES.error; + stopBtn.textContent = result.message.includes("already stopped") ? "Not Running" : "Failed"; + + // If already stopped, update status immediately + if (result.message.includes("already stopped")) { + this.ui.updateStatusDot(workerId, "#c04c4c", "Offline"); + this.state.setWorkerStatus(workerId, { online: false }); + } + + setTimeout(() => { + this.updateWorkerControls(workerId); + }, 1500); + } + } + } catch (error) { + this.log(`Error stopping worker: ${error}`, "error"); + + // Reset button on error + if (stopBtn) { + stopBtn.style.backgroundColor = BUTTON_STYLES.error; + stopBtn.textContent = "Error"; + setTimeout(() => { + this.updateWorkerControls(workerId); + }, 1500); + } + } + } + + async clearLaunchingFlag(workerId) { + try { + await this.api.clearLaunchingFlag(workerId); + this.log(`Cleared launching flag for worker ${workerId}`, "debug"); + } catch (error) { + this.log(`Error clearing launching flag: ${error.message || error}`, "error"); + } + } + + // Generic async button action handler + async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = 3000) { + const originalText = button.textContent; + const originalStyle = button.style.cssText; + button.disabled = true; + + try { + await action(); + button.textContent = successText; + button.style.cssText = originalStyle; + button.style.backgroundColor = BUTTON_STYLES.success; + return true; + } catch (error) { + button.textContent = errorText || `Error: ${error.message}`; + button.style.cssText = originalStyle; + button.style.backgroundColor = BUTTON_STYLES.error; + throw error; + } finally { + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + button.disabled = false; + }, resetDelay); + } + } + + /** + * Cleanup method to stop intervals and listeners + */ + cleanup() { + if (this.statusCheckInterval) { + clearInterval(this.statusCheckInterval); + this.statusCheckInterval = null; + } + + if (this.logAutoRefreshInterval) { + clearInterval(this.logAutoRefreshInterval); + this.logAutoRefreshInterval = null; + } + + this.log("Cleaned up intervals", "debug"); + } + + async loadManagedWorkers() { + try { + const result = await this.api.getManagedWorkers(); + + // Check for launching workers + for (const [workerId, info] of Object.entries(result.managed_workers)) { + this.state.setWorkerManaged(workerId, info); + + // If worker is marked as launching, add to launchingWorkers set + if (info.launching) { + this.state.setWorkerLaunching(workerId, true); + this.log(`Worker ${workerId} is in launching state`, "debug"); + } + } + + // Update UI for all workers + if (this.config?.workers) { + this.config.workers.forEach(w => this.updateWorkerControls(w.id)); + } + } catch (error) { + this.log(`Error loading managed workers: ${error}`, "error"); + } + } + + updateWorkerControls(workerId) { + const controlsDiv = document.getElementById(`controls-${workerId}`); + + if (!controlsDiv) return; + + const worker = this.config.workers.find(w => w.id === workerId); + if (!worker) return; + + // Skip button updates for remote workers + if (this.isRemoteWorker(worker)) { + return; + } + + // Ensure we check for string ID + const managedInfo = this.state.getWorker(workerId).managed; + const status = this.state.getWorkerStatus(workerId); + + // Update button states - buttons are now inside a wrapper div + const buttons = controlsDiv.querySelectorAll('button'); + const launchBtn = document.getElementById(`launch-${workerId}`); + const stopBtn = document.getElementById(`stop-${workerId}`); + const logBtn = document.getElementById(`log-${workerId}`); + + // Show log button immediately if we have log file info (even if worker is still starting) + if (managedInfo?.log_file && logBtn) { + logBtn.style.display = ''; + } else if (logBtn && !managedInfo) { + logBtn.style.display = 'none'; + } + + if (status?.online || managedInfo) { + // Worker is running or we just launched it + launchBtn.style.display = 'none'; // Hide launch button when running + + if (managedInfo) { + // Only show stop button if we manage this worker + stopBtn.style.display = ''; + stopBtn.disabled = false; + stopBtn.textContent = "Stop"; + stopBtn.style.backgroundColor = "#7c4a4a"; // Red when enabled + } else { + // Hide stop button for workers launched outside UI + stopBtn.style.display = 'none'; + } + } else { + // Worker is not running + launchBtn.style.display = ''; // Show launch button + launchBtn.disabled = false; + launchBtn.textContent = "Launch"; + launchBtn.style.backgroundColor = "#4a7c4a"; // Always green + + stopBtn.style.display = 'none'; // Hide stop button when not running + } + } + + async viewWorkerLog(workerId) { + const managedInfo = this.state.getWorker(workerId).managed; + if (!managedInfo?.log_file) return; + + const logBtn = document.getElementById(`log-${workerId}`); + + // Provide immediate feedback + if (logBtn) { + logBtn.disabled = true; + logBtn.textContent = "Loading..."; + logBtn.style.backgroundColor = "#666"; + } + + try { + // Fetch log content + const data = await this.api.getWorkerLog(workerId, 1000); + + // Create modal dialog + this.ui.showLogModal(this, workerId, data); + + // Restore button + if (logBtn) { + logBtn.disabled = false; + logBtn.textContent = "View Log"; + logBtn.style.backgroundColor = "#685434"; // Keep the yellow color + } + + } catch (error) { + this.log('Error viewing log: ' + error.message, "error"); + app.extensionManager.toast.add({ + severity: "error", + summary: "Error", + detail: `Failed to load log: ${error.message}`, + life: 5000 + }); + + // Flash error and restore button + if (logBtn) { + logBtn.style.backgroundColor = BUTTON_STYLES.error; + logBtn.textContent = "Error"; + setTimeout(() => { + logBtn.disabled = false; + logBtn.textContent = "View Log"; + logBtn.style.backgroundColor = "#685434"; // Keep the yellow color + }, 2000); + } + } + } + + async refreshLog(workerId, silent = false) { + const logContent = document.getElementById('distributed-log-content'); + if (!logContent) return; + + try { + const data = await this.api.getWorkerLog(workerId, 1000); + + // Update content + const shouldAutoScroll = logContent.scrollTop + logContent.clientHeight >= logContent.scrollHeight - 50; + logContent.textContent = data.content; + + // Auto-scroll if was at bottom + if (shouldAutoScroll) { + logContent.scrollTop = logContent.scrollHeight; + } + + // Only show toast if not in silent mode (manual refresh) + if (!silent) { + app.extensionManager.toast.add({ + severity: "success", + summary: "Log Refreshed", + detail: "Log content updated", + life: 2000 + }); + } + + } catch (error) { + // Only show error toast if not in silent mode + if (!silent) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Refresh Failed", + detail: error.message, + life: 3000 + }); + } + } + } + + isRemoteWorker(worker) { + const host = worker.host || window.location.hostname; + return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; + } + + getMasterUrl() { + // Always use the detected/configured master IP for consistency + if (this.config?.master?.host) { + const configuredHost = this.config.master.host; + + // If the configured host already includes protocol, use as-is + if (configuredHost.startsWith('http://') || configuredHost.startsWith('https://')) { + return configuredHost; + } + + // For domain names (not IPs), default to HTTPS + const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost); + const isLocalhost = configuredHost === 'localhost' || configuredHost === '127.0.0.1'; + + if (!isIP && !isLocalhost && configuredHost.includes('.')) { + // It's a domain name, use HTTPS + return `https://${configuredHost}`; + } else { + // For IPs and localhost, use current access method + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + if ((window.location.protocol === 'https:' && port === '443') || + (window.location.protocol === 'http:' && port === '80')) { + return `${window.location.protocol}//${configuredHost}`; + } + return `${window.location.protocol}//${configuredHost}:${port}`; + } + } + + // If no master IP is set but we're on a network address, use it + const hostname = window.location.hostname; + if (hostname !== 'localhost' && hostname !== '127.0.0.1') { + return window.location.origin; + } + + // Fallback warning - this won't work for remote workers + this.log("No master host configured - remote workers won't be able to connect. " + + "Master host should be auto-detected on startup.", "debug"); + return window.location.origin; + } + + async detectMasterIP() { + try { + const data = await this.api.getNetworkInfo(); + this.log("Network info: " + JSON.stringify(data), "debug"); + + // Store CUDA device info + if (data.cuda_device !== null && data.cuda_device !== undefined) { + this.masterCudaDevice = data.cuda_device; + + // Store persistently in config if not already set or changed + if (!this.config.master) this.config.master = {}; + if (this.config.master.cuda_device === undefined || this.config.master.cuda_device !== data.cuda_device) { + this.config.master.cuda_device = data.cuda_device; + try { + await this.api.updateMaster({ cuda_device: data.cuda_device }); + this.log(`Stored master CUDA device: ${data.cuda_device}`, "debug"); + } catch (error) { + this.log(`Error storing master CUDA device: ${error.message}`, "error"); + } + } + + // Update the master display with CUDA info + this.ui.updateMasterDisplay(this); + } + + // Store CUDA device count for auto-population + if (data.cuda_device_count > 0) { + this.cudaDeviceCount = data.cuda_device_count; + this.log(`Detected ${this.cudaDeviceCount} CUDA devices`, "info"); + + // Auto-populate workers if conditions are met + const shouldAutoPopulate = + !this.config.settings.has_auto_populated_workers && // Never populated before + (!this.config.workers || this.config.workers.length === 0); // No workers exist + + this.log(`Auto-population check: has_populated=${this.config.settings.has_auto_populated_workers}, workers=${this.config.workers ? this.config.workers.length : 'null'}, should_populate=${shouldAutoPopulate}`, "debug"); + + if (shouldAutoPopulate) { + this.log(`Auto-populating workers based on ${this.cudaDeviceCount} CUDA devices (excluding master on CUDA ${this.masterCudaDevice})`, "info"); + + const newWorkers = []; + let workerNum = 1; + let portOffset = 0; + + for (let i = 0; i < this.cudaDeviceCount; i++) { + // Skip the CUDA device used by master + if (i === this.masterCudaDevice) { + this.log(`Skipping CUDA ${i} (used by master)`, "debug"); + continue; + } + + const worker = { + id: crypto.randomUUID(), + name: `Worker ${workerNum}`, + host: "localhost", + port: 8189 + portOffset, + cuda_device: i, + enabled: true + }; + newWorkers.push(worker); + workerNum++; + portOffset++; + } + + // Only proceed if we have workers to add + if (newWorkers.length > 0) { + this.log(`Auto-populating ${newWorkers.length} workers`, "info"); + + // Add workers to config + this.config.workers = newWorkers; + + // Set the flag to prevent future auto-population + this.config.settings.has_auto_populated_workers = true; + + // Save each worker using the update endpoint + for (const worker of newWorkers) { + try { + await this.api.updateWorker(worker.id, worker); + } catch (error) { + this.log(`Error saving worker ${worker.name}: ${error.message}`, "error"); + } + } + + // Save the updated settings + try { + await this.api.updateSetting('has_auto_populated_workers', true); + } catch (error) { + this.log(`Error saving auto-population flag: ${error.message}`, "error"); + } + + this.log(`Auto-populated ${newWorkers.length} workers and saved config`, "info"); + + // Reload the config to include the new workers + await this.loadConfig(); + } else { + this.log("No additional CUDA devices available for workers (all used by master)", "debug"); + } + } + } + + // Check if we already have a master host configured + if (this.config?.master?.host) { + this.log(`Master host already configured: ${this.config.master.host}`, "debug"); + return; + } + + // Use the recommended IP from the backend + if (data.recommended_ip && data.recommended_ip !== '127.0.0.1') { + this.log(`Auto-detected master IP: ${data.recommended_ip}`, "info"); + + // Save the detected IP (pass true to suppress notification) + await this.api.updateMaster({ host: data.recommended_ip }); + + // Show a single combined notification + if (app.extensionManager?.toast) { + app.extensionManager.toast.add({ + severity: "info", + summary: "Master IP Auto-Detected", + detail: `Master IP set to ${data.recommended_ip} for remote worker communication`, + life: 5000 + }); + } + } + } catch (error) { + this.log("Error detecting master IP: " + error.message, "error"); + } + } + + async saveWorkerSettings(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + if (!worker) return; + + // Get form values + const name = document.getElementById(`name-${workerId}`).value; + const isRemote = document.getElementById(`remote-${workerId}`).checked; + const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; + const port = parseInt(document.getElementById(`port-${workerId}`).value); + const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); + const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value; + + // Validate + if (!name.trim()) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Worker name is required", + life: 3000 + }); + return; + } + + if (isRemote && !host.trim()) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Host is required for remote workers", + life: 3000 + }); + return; + } + + if (isNaN(port) || port < 1 || port > 65535) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Port must be between 1 and 65535", + life: 3000 + }); + return; + } + + // Check for port conflicts + // Remote workers can reuse ports, but local workers cannot share ports with each other or master + if (!isRemote) { + // Check if port conflicts with master + const masterPort = parseInt(window.location.port) || (window.location.protocol === 'https:' ? 443 : 80); + if (port === masterPort) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by the master server`, + life: 3000 + }); + return; + } + + // Check if port conflicts with other local workers + const localPortConflict = this.config.workers.some(w => + w.id !== workerId && + w.port === port && + !w.host // local workers have no host or host is null + ); + + if (localPortConflict) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by another local worker`, + life: 3000 + }); + return; + } + } else { + // For remote workers, only check conflicts with other workers on the same host + const sameHostConflict = this.config.workers.some(w => + w.id !== workerId && + w.port === port && + w.host === host.trim() + ); + + if (sameHostConflict) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by another worker on ${host}`, + life: 3000 + }); + return; + } + } + + try { + await this.api.updateWorker(workerId, { + name: name.trim(), + host: isRemote ? host.trim() : null, + port: port, + cuda_device: isRemote ? null : cudaDevice, + extra_args: isRemote ? null : (extraArgs ? extraArgs.trim() : "") + }); + + // Update local config + worker.name = name.trim(); + if (isRemote) { + worker.host = host.trim(); + delete worker.cuda_device; + delete worker.extra_args; + } else { + delete worker.host; + worker.cuda_device = cudaDevice; + worker.extra_args = extraArgs ? extraArgs.trim() : ""; + } + worker.port = port; + + // Sync to state + this.state.updateWorker(workerId, { enabled: worker.enabled }); + + app.extensionManager.toast.add({ + severity: "success", + summary: "Settings Saved", + detail: `Worker ${name} settings updated`, + life: 3000 + }); + + // Refresh the UI + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + } catch (error) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Save Failed", + detail: error.message, + life: 5000 + }); + } + } + + cancelWorkerSettings(workerId) { + // Collapse the settings panel + this.toggleWorkerExpanded(workerId); + + // Reset form values to original + const worker = this.config.workers.find(w => w.id === workerId); + if (worker) { + document.getElementById(`name-${workerId}`).value = worker.name; + document.getElementById(`host-${workerId}`).value = worker.host || ""; + document.getElementById(`port-${workerId}`).value = worker.port; + document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0; + document.getElementById(`args-${workerId}`).value = worker.extra_args || ""; + + // Reset remote checkbox + const remoteCheckbox = document.getElementById(`remote-${workerId}`); + if (remoteCheckbox) { + remoteCheckbox.checked = this.isRemoteWorker(worker); + } + } + } + + async deleteWorker(workerId) { + const worker = this.config.workers.find(w => w.id === workerId); + if (!worker) return; + + // Confirm deletion + if (!confirm(`Are you sure you want to delete worker "${worker.name}"?`)) { + return; + } + + try { + await this.api.deleteWorker(workerId); + + // Remove from local config + const index = this.config.workers.findIndex(w => w.id === workerId); + if (index !== -1) { + this.config.workers.splice(index, 1); + } + + app.extensionManager.toast.add({ + severity: "success", + summary: "Worker Deleted", + detail: `Worker ${worker.name} has been removed`, + life: 3000 + }); + + // Refresh the UI + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + } catch (error) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Delete Failed", + detail: error.message, + life: 5000 + }); + } + } + + async addNewWorker() { + // Generate new worker ID using UUID + const newId = crypto.randomUUID(); + + // Find next available port + const usedPorts = this.config.workers.map(w => w.port); + let nextPort = 8189; + while (usedPorts.includes(nextPort)) { + nextPort++; + } + + // Create new worker object + const newWorker = { + id: newId, + name: `Worker ${this.config.workers.length + 1}`, + port: nextPort, + cuda_device: this.config.workers.length, + enabled: true, // Default to enabled for convenience + extra_args: "" + }; + + // Add to config + this.config.workers.push(newWorker); + + // Save immediately + try { + await this.api.updateWorker(newId, { + name: newWorker.name, + port: newWorker.port, + cuda_device: newWorker.cuda_device, + extra_args: newWorker.extra_args, + enabled: newWorker.enabled + }); + + // Sync to state + this.state.updateWorker(newId, { enabled: true }); + + app.extensionManager.toast.add({ + severity: "success", + summary: "Worker Added", + detail: `New worker created on port ${nextPort}`, + life: 3000 + }); + + // Refresh UI and expand the new worker + this.state.setWorkerExpanded(newId, true); + if (this.panelElement) { + renderSidebarContent(this, this.panelElement); + } + + } catch (error) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Failed to Add Worker", + detail: error.message, + life: 5000 + }); + } + } + + startLogAutoRefresh(workerId) { + // Stop any existing auto-refresh + this.stopLogAutoRefresh(); + + // Refresh every 2 seconds + this.logAutoRefreshInterval = setInterval(() => { + this.refreshLog(workerId, true); // silent mode + }, 2000); + } + + stopLogAutoRefresh() { + if (this.logAutoRefreshInterval) { + clearInterval(this.logAutoRefreshInterval); + this.logAutoRefreshInterval = null; + } + } + + toggleWorkerExpanded(workerId) { + const settingsDiv = document.getElementById(`settings-${workerId}`); + const gpuDiv = settingsDiv.closest('[style*="margin-bottom: 12px"]'); + const settingsArrow = gpuDiv.querySelector('.settings-arrow'); + + if (!settingsDiv) return; + + if (this.state.isWorkerExpanded(workerId)) { + this.state.setWorkerExpanded(workerId, false); + settingsDiv.classList.remove("expanded"); + if (settingsArrow) { + settingsArrow.style.transform = "rotate(0deg)"; + } + // Animate padding to 0 + settingsDiv.style.padding = "0 12px"; + settingsDiv.style.marginTop = "0"; + settingsDiv.style.marginBottom = "0"; + } else { + this.state.setWorkerExpanded(workerId, true); + settingsDiv.classList.add("expanded"); + if (settingsArrow) { + settingsArrow.style.transform = "rotate(90deg)"; + } + // Animate padding to full + settingsDiv.style.padding = "12px"; + settingsDiv.style.marginTop = "8px"; + settingsDiv.style.marginBottom = "8px"; + } + } + + _handleInterruptWorkers(button) { + return handleInterruptWorkers(this, button); + } + + _handleClearMemory(button) { + return handleClearMemory(this, button); + } +} + +app.registerExtension({ + name: "Distributed.Panel", + async setup() { + new DistributedExtension(); + } +}); \ No newline at end of file diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js new file mode 100644 index 0000000..d88e4cd --- /dev/null +++ b/web/sidebarRenderer.js @@ -0,0 +1,273 @@ +import { BUTTON_STYLES } from './constants.js'; + +export async function renderSidebarContent(extension, el) { + // Panel is being opened/rendered + extension.log("Panel opened", "debug"); + + if (!el) { + extension.log("No element provided to renderSidebarContent", "debug"); + return; + } + + // Prevent infinite recursion + if (extension._isRendering) { + extension.log("Already rendering, skipping", "debug"); + return; + } + extension._isRendering = true; + + try { + // Store reference to the panel element + extension.panelElement = el; + + // Show loading indicator + el.innerHTML = ''; + const loadingDiv = document.createElement("div"); + loadingDiv.style.cssText = "display: flex; align-items: center; justify-content: center; height: calc(100vh - 100px); color: #888;"; + loadingDiv.innerHTML = `
◐
`; + el.appendChild(loadingDiv); + + // Add rotation animation + const style = document.createElement('style'); + style.textContent = ` + @keyframes rotate { + from { transform: rotate(0deg); } + to { transform: rotate(360deg); } + } + `; + document.head.appendChild(style); + loadingDiv.querySelector('div').style.animation = 'rotate 1s linear infinite'; + + // Preload data outside render + await extension.loadConfig(); + await extension.loadManagedWorkers(); + + el.innerHTML = ''; + + // Create toolbar header to match ComfyUI style + const toolbar = document.createElement("div"); + toolbar.className = "p-toolbar p-component border-x-0 border-t-0 rounded-none px-2 py-1 min-h-8"; + toolbar.style.cssText = "border-bottom: 1px solid #444; background: transparent; display: flex; align-items: center;"; + + const toolbarStart = document.createElement("div"); + toolbarStart.className = "p-toolbar-start"; + toolbarStart.style.cssText = "display: flex; align-items: center;"; + + const titleSpan = document.createElement("span"); + titleSpan.className = "text-xs 2xl:text-sm truncate"; + titleSpan.textContent = "COMFYUI DISTRIBUTED"; + titleSpan.title = "ComfyUI Distributed"; + + toolbarStart.appendChild(titleSpan); + toolbar.appendChild(toolbarStart); + + const toolbarCenter = document.createElement("div"); + toolbarCenter.className = "p-toolbar-center"; + toolbar.appendChild(toolbarCenter); + + const toolbarEnd = document.createElement("div"); + toolbarEnd.className = "p-toolbar-end"; + toolbar.appendChild(toolbarEnd); + + el.appendChild(toolbar); + + // Main container with adjusted padding + const container = document.createElement("div"); + container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: calc(100% - 32px);"; + + // Detect master info on panel open (in case CUDA info wasn't available at startup) + extension.log(`Panel opened. CUDA device count: ${extension.cudaDeviceCount}, Workers: ${extension.config?.workers?.length || 0}`, "debug"); + if (!extension.cudaDeviceCount) { + await extension.detectMasterIP(); + } + + + // Now render with guaranteed up-to-date config + // Master Node Section + const masterDiv = extension.ui.renderEntityCard('master', extension.config?.master, extension); + container.appendChild(masterDiv); + + // Workers Section (no heading) + const gpuSection = document.createElement("div"); + gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; + + const gpuList = document.createElement("div"); + const workers = extension.config?.workers || []; + + // If no workers exist, show a full blueprint placeholder first + if (workers.length === 0) { + const blueprintDiv = extension.ui.renderEntityCard('blueprint', { onClick: () => extension.addNewWorker() }, extension); + gpuList.appendChild(blueprintDiv); + } + + // Show existing workers + workers.forEach(worker => { + const gpuDiv = extension.ui.renderEntityCard('worker', worker, extension); + gpuList.appendChild(gpuDiv); + }); + gpuSection.appendChild(gpuList); + + // Only show the minimal "Add Worker" box if there are existing workers + if (workers.length > 0) { + const addWorkerDiv = extension.ui.renderEntityCard('add', { onClick: () => extension.addNewWorker() }, extension); + gpuSection.appendChild(addWorkerDiv); + } + + container.appendChild(gpuSection); + + const actionsSection = document.createElement("div"); + actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; + + // Create a row for both buttons + const buttonRow = document.createElement("div"); + buttonRow.style.cssText = "display: flex; gap: 8px;"; + + const clearMemButton = extension.ui.createButtonHelper( + "Clear Worker VRAM", + (e) => extension._handleClearMemory(e.target), + BUTTON_STYLES.clearMemory + ); + clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; + clearMemButton.style.cssText += " flex: 1;"; + + const interruptButton = extension.ui.createButtonHelper( + "Interrupt Workers", + (e) => extension._handleInterruptWorkers(e.target), + BUTTON_STYLES.clearMemory + ); + interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; + interruptButton.style.cssText += " flex: 1;"; + + buttonRow.appendChild(clearMemButton); + buttonRow.appendChild(interruptButton); + actionsSection.appendChild(buttonRow); + + container.appendChild(actionsSection); + + // Settings section + const settingsSection = document.createElement("div"); + settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;"; + + // Settings header with toggle + const settingsHeader = document.createElement("div"); + settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;"; + + const workerSettingsTitle = document.createElement("h4"); + workerSettingsTitle.textContent = "Settings"; + workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + + const workerSettingsToggle = document.createElement("span"); + workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed + workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + + settingsHeader.appendChild(workerSettingsTitle); + settingsHeader.appendChild(workerSettingsToggle); + + // Hover effect for header + settingsHeader.onmouseover = () => { + workerSettingsToggle.style.color = "#fff"; + }; + settingsHeader.onmouseout = () => { + workerSettingsToggle.style.color = "#888"; + }; + + // Collapsible settings content + const settingsContent = document.createElement("div"); + settingsContent.style.cssText = "max-height: 0; overflow: hidden; transition: max-height 0.3s ease;"; + + const settingsDiv = document.createElement("div"); + settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;"; + + // Toggle functionality + let settingsExpanded = false; + settingsHeader.onclick = () => { + settingsExpanded = !settingsExpanded; + if (settingsExpanded) { + settingsContent.style.maxHeight = "200px"; + workerSettingsToggle.style.transform = "rotate(90deg)"; + } else { + settingsContent.style.maxHeight = "0"; + workerSettingsToggle.style.transform = "rotate(0deg)"; + } + }; + + // Debug mode setting + const debugGroup = document.createElement("div"); + debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + + const debugCheckbox = document.createElement("input"); + debugCheckbox.type = "checkbox"; + debugCheckbox.id = "setting-debug"; + debugCheckbox.checked = extension.config?.settings?.debug || false; + debugCheckbox.onchange = (e) => extension._updateSetting('debug', e.target.checked); + + const debugLabel = document.createElement("label"); + debugLabel.htmlFor = "setting-debug"; + debugLabel.textContent = "Debug Mode"; + debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + debugGroup.appendChild(debugCheckbox); + debugGroup.appendChild(debugLabel); + + // Auto-launch workers setting + const autoLaunchGroup = document.createElement("div"); + autoLaunchGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + + const autoLaunchCheckbox = document.createElement("input"); + autoLaunchCheckbox.type = "checkbox"; + autoLaunchCheckbox.id = "setting-auto-launch"; + autoLaunchCheckbox.checked = extension.config?.settings?.auto_launch_workers || false; + autoLaunchCheckbox.onchange = (e) => extension._updateSetting('auto_launch_workers', e.target.checked); + + const autoLaunchLabel = document.createElement("label"); + autoLaunchLabel.htmlFor = "setting-auto-launch"; + autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup"; + autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + autoLaunchGroup.appendChild(autoLaunchCheckbox); + autoLaunchGroup.appendChild(autoLaunchLabel); + + // Stop workers on exit setting + const stopOnExitGroup = document.createElement("div"); + stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + + const stopOnExitCheckbox = document.createElement("input"); + stopOnExitCheckbox.type = "checkbox"; + stopOnExitCheckbox.id = "setting-stop-on-exit"; + stopOnExitCheckbox.checked = extension.config?.settings?.stop_workers_on_master_exit !== false; // Default true + stopOnExitCheckbox.onchange = (e) => extension._updateSetting('stop_workers_on_master_exit', e.target.checked); + + const stopOnExitLabel = document.createElement("label"); + stopOnExitLabel.htmlFor = "setting-stop-on-exit"; + stopOnExitLabel.textContent = "Stop Local Workers on Master Exit"; + stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + stopOnExitGroup.appendChild(stopOnExitCheckbox); + stopOnExitGroup.appendChild(stopOnExitLabel); + + settingsDiv.appendChild(debugGroup); + settingsDiv.appendChild(autoLaunchGroup); + settingsDiv.appendChild(stopOnExitGroup); + settingsContent.appendChild(settingsDiv); + + settingsSection.appendChild(settingsHeader); + settingsSection.appendChild(settingsContent); + container.appendChild(settingsSection); + + const summarySection = document.createElement("div"); + summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;"; + const summary = document.createElement("div"); + summary.id = "distributed-summary"; + summary.style.cssText = "font-size: 11px; color: #888;"; + summarySection.appendChild(summary); + container.appendChild(summarySection); + el.appendChild(container); + extension.updateSummary(); + + // Start checking worker statuses immediately in parallel + setTimeout(() => extension.checkAllWorkerStatuses(), 0); + } finally { + // Always reset the rendering flag + extension._isRendering = false; + } +} \ No newline at end of file diff --git a/web/stateManager.js b/web/stateManager.js new file mode 100644 index 0000000..f400f4a --- /dev/null +++ b/web/stateManager.js @@ -0,0 +1,61 @@ +export function createStateManager() { + const state = { + workers: new Map(), // Unified worker state: { status, managed, launching, expanded, ... } + masterStatus: 'online', + }; + + return { + // Worker state management + getWorker(workerId) { + return state.workers.get(String(workerId)) || {}; + }, + + updateWorker(workerId, updates) { + const id = String(workerId); + const current = state.workers.get(id) || {}; + state.workers.set(id, { ...current, ...updates }); + return state.workers.get(id); + }, + + setWorkerStatus(workerId, status) { + return this.updateWorker(workerId, { status }); + }, + + setWorkerManaged(workerId, info) { + return this.updateWorker(workerId, { managed: info }); + }, + + setWorkerLaunching(workerId, launching) { + return this.updateWorker(workerId, { launching }); + }, + + setWorkerExpanded(workerId, expanded) { + return this.updateWorker(workerId, { expanded }); + }, + + isWorkerLaunching(workerId) { + return this.getWorker(workerId).launching || false; + }, + + isWorkerExpanded(workerId) { + return this.getWorker(workerId).expanded || false; + }, + + isWorkerManaged(workerId) { + return !!this.getWorker(workerId).managed; + }, + + getWorkerStatus(workerId) { + return this.getWorker(workerId).status || {}; + }, + + // Master state + setMasterStatus(status) { + state.masterStatus = status; + }, + + getMasterStatus() { + return state.masterStatus; + } + }; +} \ No newline at end of file diff --git a/web/ui.js b/web/ui.js new file mode 100644 index 0000000..87bbbd9 --- /dev/null +++ b/web/ui.js @@ -0,0 +1,997 @@ +import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS } from './constants.js'; + +const cardConfigs = { + master: { + checkbox: { + enabled: false, + checked: true, + disabled: true, + opacity: 0.6, + title: "Master node is always enabled" + }, + statusDot: { + color: STATUS_COLORS.ONLINE_GREEN, + title: 'Online', + id: 'master-status', + dynamic: true + }, + infoText: (data, extension) => { + const cudaDevice = extension.config?.master?.cuda_device ?? extension.masterCudaDevice; + const cudaInfo = cudaDevice !== undefined ? `CUDA ${cudaDevice} • ` : ''; + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + return `${data?.name || extension.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; + }, + controls: { + type: 'info', + text: 'Master', + style: "background-color: #333; color: #999;" + }, + settings: { + formType: 'master', + id: 'master-settings', + expandedTracker: 'masterSettingsExpanded' + }, + hover: true, + expand: true, + border: 'solid' + }, + worker: { + checkbox: { + enabled: true, + title: "Enable/disable this worker" + }, + statusDot: { + dynamic: true, + initialColor: (data) => data.enabled ? STATUS_COLORS.OFFLINE_RED : STATUS_COLORS.DISABLED_GRAY, + initialTitle: (data) => data.enabled ? "Checking status..." : "Disabled", + pulsing: (data) => data.enabled, + id: (data) => `status-${data.id}` + }, + infoText: (data, extension) => { + const isRemote = extension.isRemoteWorker(data); + if (isRemote) { + return `${data.name}
${data.host}:${data.port}`; + } else { + const cudaInfo = data.cuda_device !== undefined ? `CUDA ${data.cuda_device} • ` : ''; + return `${data.name}
${cudaInfo}Port ${data.port}`; + } + }, + controls: { + dynamic: true + }, + settings: { + formType: 'worker', + id: (data) => `settings-${data.id}`, + expandedId: (data) => data?.id + }, + hover: true, + expand: true, + border: 'solid' + }, + blueprint: { + checkbox: { + type: 'icon', + content: '+', + width: 42, + style: 'border-right: 2px dashed #555; color: #777; font-size: 24px; font-weight: 500;' + }, + statusDot: { + color: 'transparent', + border: '1px solid #555' + }, + infoText: () => `Add New Worker
[CUDA] • [Port]`, + controls: { + type: 'ghost', + text: 'Configure', + style: 'border: 1px solid #444; background: transparent; color: #555;' + }, + hover: 'placeholder', + expand: false, + border: 'dashed' + }, + add: { + checkbox: { + type: 'icon', + content: '+', + width: 43, + style: 'border-right: 1px dashed #444; color: #555; font-size: 18px;' + }, + statusDot: { + color: 'transparent', + border: '1px solid #555' + }, + infoText: () => `Add New Worker`, + controls: null, + hover: 'placeholder', + expand: false, + border: 'dashed', + minHeight: '48px' + } +}; + +export class DistributedUI { + constructor() { + // UI element styles + this.styles = UI_STYLES; + } + + createStatusDot(id, color = "#666", title = "Status") { + const dot = document.createElement("span"); + if (id) dot.id = id; + dot.style.cssText = this.styles.statusDot + ` background-color: ${color};`; + dot.title = title; + return dot; + } + + createButton(text, onClick, customStyle = "") { + const button = document.createElement("button"); + button.textContent = text; + button.className = "distributed-button"; + button.style.cssText = BUTTON_STYLES.base + customStyle; + if (onClick) button.onclick = onClick; + return button; + } + + createButtonGroup(buttons, style = "") { + const group = document.createElement("div"); + group.style.cssText = "display: flex; gap: 4px; margin-top: 10px;" + style; + buttons.forEach(button => group.appendChild(button)); + return group; + } + + createWorkerControls(workerId, handlers = {}) { + const controlsDiv = document.createElement("div"); + controlsDiv.id = `controls-${workerId}`; + controlsDiv.style.cssText = this.styles.controlsDiv; + + const buttons = []; + + if (handlers.launch) { + const launchBtn = this.createButton('Launch', handlers.launch); + launchBtn.id = `launch-${workerId}`; + launchBtn.title = "Launch this worker instance"; + buttons.push(launchBtn); + } + + if (handlers.stop) { + const stopBtn = this.createButton('Stop', handlers.stop); + stopBtn.id = `stop-${workerId}`; + stopBtn.title = "Stop this worker instance"; + buttons.push(stopBtn); + } + + if (handlers.viewLog) { + const logBtn = this.createButton('View Log', handlers.viewLog); + logBtn.id = `log-${workerId}`; + logBtn.title = "View worker log file"; + buttons.push(logBtn); + } + + buttons.forEach(btn => controlsDiv.appendChild(btn)); + return controlsDiv; + } + + createFormGroup(label, value, id, type = "text", placeholder = "") { + const group = document.createElement("div"); + group.style.cssText = this.styles.formGroup; + + const labelEl = document.createElement("label"); + labelEl.textContent = label; + labelEl.htmlFor = id; + labelEl.style.cssText = this.styles.formLabel; + + const input = document.createElement("input"); + input.type = type; + input.id = id; + input.value = value; + input.placeholder = placeholder; + input.style.cssText = this.styles.formInput; + + group.appendChild(labelEl); + group.appendChild(input); + return { group, input }; + } + + + createInfoBox(text) { + const box = document.createElement("div"); + box.style.cssText = this.styles.infoBox; + box.textContent = text; + return box; + } + + addHoverEffect(element, onHover, onLeave) { + element.onmouseover = onHover; + element.onmouseout = onLeave; + } + + createCard(type = 'worker', options = {}) { + const card = document.createElement("div"); + const baseStyle = "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;"; + + switch(type) { + case 'master': + case 'worker': + card.style.cssText = baseStyle + "background: #2a2a2a;"; + break; + case 'blueprint': + card.style.cssText = baseStyle + "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);"; + if (options.onClick) card.onclick = options.onClick; + if (options.title) card.title = options.title; + break; + case 'add': + card.style.cssText = baseStyle + "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;"; + if (options.onClick) card.onclick = options.onClick; + if (options.title) card.title = options.title; + break; + } + + if (options.onMouseEnter) { + card.addEventListener('mouseenter', options.onMouseEnter); + } + if (options.onMouseLeave) { + card.addEventListener('mouseleave', options.onMouseLeave); + } + + return card; + } + + createCardColumn(type = 'checkbox', options = {}) { + const column = document.createElement("div"); + const baseStyle = "display: flex; align-items: center; justify-content: center;"; + + switch(type) { + case 'checkbox': + column.style.cssText = baseStyle + "flex: 0 0 44px; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);"; + if (options.title) column.title = options.title; + break; + case 'icon': + column.style.cssText = baseStyle + "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;"; + break; + case 'content': + column.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; + break; + } + + return column; + } + + createInfoRow(options = {}) { + const row = document.createElement("div"); + row.style.cssText = "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;"; + if (options.onClick) row.onclick = options.onClick; + return row; + } + + createWorkerContent() { + const content = document.createElement("div"); + content.style.cssText = "display: flex; align-items: center; gap: 10px; flex: 1;"; + return content; + } + + createSettingsForm(fields = [], options = {}) { + const form = document.createElement("div"); + form.style.cssText = "display: flex; flex-direction: column; gap: 10px;"; + + fields.forEach(field => { + if (field.type === 'checkbox') { + const group = document.createElement("div"); + group.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; + + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = field.id; + checkbox.checked = field.checked || false; + if (field.onChange) checkbox.onchange = field.onChange; + + const label = document.createElement("label"); + label.htmlFor = field.id; + label.textContent = field.label; + label.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + group.appendChild(checkbox); + group.appendChild(label); + form.appendChild(group); + } else { + const result = this.createFormGroup(field.label, field.value, field.id, field.type, field.placeholder); + if (field.groupId) result.group.id = field.groupId; + if (field.display) result.group.style.display = field.display; + form.appendChild(result.group); + } + }); + + if (options.buttons) { + const buttonGroup = this.createButtonGroup(options.buttons, options.buttonStyle || " margin-top: 8px;"); + form.appendChild(buttonGroup); + } + + return form; + } + + + createButtonHelper(text, onClick, style) { + return this.createButton(text, onClick, style); + } + + updateMasterDisplay(extension) { + // Use persistent config value as fallback + const cudaDevice = extension?.config?.master?.cuda_device ?? extension?.masterCudaDevice; + + // Update CUDA info if element exists + const cudaInfo = document.getElementById('master-cuda-info'); + if (cudaInfo) { + const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); + if (cudaDevice !== undefined && cudaDevice !== null) { + cudaInfo.textContent = `CUDA ${cudaDevice} • Port ${port}`; + } else { + cudaInfo.textContent = `Port ${port}`; + } + } + + // Update name if changed + const nameDisplay = document.getElementById('master-name-display'); + if (nameDisplay && extension?.config?.master?.name) { + nameDisplay.textContent = extension.config.master.name; + } + } + + showToast(app, severity, summary, detail, life = 3000) { + if (app.extensionManager?.toast?.add) { + app.extensionManager.toast.add({ severity, summary, detail, life }); + } + } + + updateStatusDot(workerId, color, title, pulsing = false) { + const statusDot = document.getElementById(`status-${workerId}`); + if (!statusDot) return; + + statusDot.style.backgroundColor = color; + statusDot.title = title; + statusDot.classList.toggle('status-pulsing', pulsing); + } + + showLogModal(extension, workerId, logData) { + // Remove any existing modal + const existingModal = document.getElementById('distributed-log-modal'); + if (existingModal) { + existingModal.remove(); + } + + const worker = extension.config.workers.find(w => w.id === workerId); + const workerName = worker?.name || `Worker ${workerId}`; + + // Create modal container + const modal = document.createElement('div'); + modal.id = 'distributed-log-modal'; + modal.style.cssText = ` + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background: rgba(0, 0, 0, 0.8); + display: flex; + align-items: center; + justify-content: center; + z-index: 10000; + `; + + // Create modal content + const content = document.createElement('div'); + content.style.cssText = ` + background: #1e1e1e; + border-radius: 8px; + width: 90%; + max-width: 1200px; + height: 80%; + display: flex; + flex-direction: column; + border: 1px solid #444; + `; + + // Header + const header = document.createElement('div'); + header.style.cssText = ` + padding: 15px 20px; + border-bottom: 1px solid #444; + display: flex; + justify-content: space-between; + align-items: center; + `; + + const title = document.createElement('h3'); + title.textContent = `${workerName} - Log Viewer`; + title.style.cssText = 'margin: 0; color: #fff;'; + + const headerButtons = document.createElement('div'); + headerButtons.style.cssText = 'display: flex; gap: 20px; align-items: center;'; + + // Auto-refresh container + const refreshContainer = document.createElement('div'); + refreshContainer.style.cssText = 'display: flex; align-items: center; gap: 4px;'; + + // Auto-refresh checkbox + const refreshCheckbox = document.createElement('input'); + refreshCheckbox.type = 'checkbox'; + refreshCheckbox.id = 'log-auto-refresh'; + refreshCheckbox.checked = true; // Enabled by default + refreshCheckbox.style.cssText = 'cursor: pointer;'; + refreshCheckbox.onchange = (e) => { + if (e.target.checked) { + extension.startLogAutoRefresh(workerId); + } else { + extension.stopLogAutoRefresh(); + } + }; + + const refreshLabel = document.createElement('label'); + refreshLabel.htmlFor = 'log-auto-refresh'; + refreshLabel.style.cssText = 'font-size: 12px; color: #ccc; cursor: pointer; white-space: nowrap;'; + refreshLabel.textContent = 'Auto-refresh'; + + // Add checkbox and label to container + refreshContainer.appendChild(refreshCheckbox); + refreshContainer.appendChild(refreshLabel); + + // Close button + const closeBtn = this.createButton('✕', + () => { + extension.stopLogAutoRefresh(); + modal.remove(); + }, + 'background-color: #c04c4c;'); + closeBtn.style.cssText += ' padding: 5px 10px; font-size: 14px; font-weight: bold;'; + + headerButtons.appendChild(refreshContainer); + headerButtons.appendChild(closeBtn); + + header.appendChild(title); + header.appendChild(headerButtons); + + // Log content area + const logContainer = document.createElement('div'); + logContainer.style.cssText = ` + flex: 1; + overflow: auto; + padding: 15px; + font-family: 'Consolas', 'Monaco', 'Courier New', monospace; + font-size: 12px; + line-height: 1.4; + color: #ddd; + background: #0d0d0d; + white-space: pre-wrap; + word-wrap: break-word; + `; + logContainer.id = 'distributed-log-content'; + logContainer.textContent = logData.content; + + // Auto-scroll to bottom + setTimeout(() => { + logContainer.scrollTop = logContainer.scrollHeight; + }, 0); + + // Status bar + const statusBar = document.createElement('div'); + statusBar.style.cssText = ` + padding: 10px 20px; + border-top: 1px solid #444; + font-size: 11px; + color: #888; + `; + statusBar.textContent = `Log file: ${logData.log_file}`; + if (logData.truncated) { + statusBar.textContent += ` (showing last ${logData.lines_shown} lines of ${this.formatFileSize(logData.file_size)})`; + } + + // Assemble modal + content.appendChild(header); + content.appendChild(logContainer); + content.appendChild(statusBar); + modal.appendChild(content); + + // Close on background click + modal.onclick = (e) => { + if (e.target === modal) { + extension.stopLogAutoRefresh(); + modal.remove(); + } + }; + + // Close on Escape key + const handleEscape = (e) => { + if (e.key === 'Escape') { + extension.stopLogAutoRefresh(); + modal.remove(); + document.removeEventListener('keydown', handleEscape); + } + }; + document.addEventListener('keydown', handleEscape); + + document.body.appendChild(modal); + + // Start auto-refresh + extension.startLogAutoRefresh(workerId); + } + + formatFileSize(bytes) { + if (bytes < 1024) return bytes + ' B'; + if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; + return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; + } + + createWorkerSettingsForm(extension, worker) { + const form = document.createElement("div"); + form.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; + + // Name field + const nameGroup = this.createFormGroup("Name:", worker.name, `name-${worker.id}`); + form.appendChild(nameGroup.group); + + // Remote checkbox + const remoteGroup = document.createElement("div"); + remoteGroup.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; + + const remoteCheckbox = document.createElement("input"); + remoteCheckbox.type = "checkbox"; + remoteCheckbox.id = `remote-${worker.id}`; + remoteCheckbox.checked = extension.isRemoteWorker(worker); + remoteCheckbox.onchange = (e) => { + const isRemote = e.target.checked; + // Show/hide relevant fields + const hostGroup = document.getElementById(`host-group-${worker.id}`); + const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); + const argsGroup = document.getElementById(`args-group-${worker.id}`); + + if (isRemote) { + hostGroup.style.display = "flex"; + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + // If switching to remote and host is localhost, clear it + const hostInput = document.getElementById(`host-${worker.id}`); + if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { + hostInput.value = ""; + } + } else { + hostGroup.style.display = "none"; + cudaGroup.style.display = "flex"; + argsGroup.style.display = "flex"; + } + }; + + const remoteLabel = document.createElement("label"); + remoteLabel.htmlFor = `remote-${worker.id}`; + remoteLabel.textContent = "Remote Worker"; + remoteLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + + remoteGroup.appendChild(remoteCheckbox); + remoteGroup.appendChild(remoteLabel); + form.appendChild(remoteGroup); + + // Host field (only for remote workers) + const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); + hostGroup.group.id = `host-group-${worker.id}`; + hostGroup.group.style.display = extension.isRemoteWorker(worker) ? "flex" : "none"; + form.appendChild(hostGroup.group); + + // Port field + const portGroup = this.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number"); + form.appendChild(portGroup.group); + + // CUDA Device field (only for local workers) + const cudaGroup = this.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number"); + cudaGroup.group.id = `cuda-group-${worker.id}`; + cudaGroup.group.style.display = extension.isRemoteWorker(worker) ? "none" : "flex"; + form.appendChild(cudaGroup.group); + + // Extra Args field (only for local workers) + const argsGroup = this.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`); + argsGroup.group.id = `args-group-${worker.id}`; + argsGroup.group.style.display = extension.isRemoteWorker(worker) ? "none" : "flex"; + form.appendChild(argsGroup.group); + + // Buttons + const saveBtn = this.createButton("Save", + () => extension.saveWorkerSettings(worker.id), + "background-color: #4a7c4a;"); + saveBtn.style.cssText += " padding: 6px 14px; font-size: 12px; font-weight: 500;"; + + const cancelBtn = this.createButton("Cancel", + () => extension.cancelWorkerSettings(worker.id), + "background-color: #555;"); + cancelBtn.style.cssText += " padding: 6px 14px; font-size: 12px; font-weight: 500;"; + + const deleteBtn = this.createButton("Delete", + () => extension.deleteWorker(worker.id), + "background-color: #7c4a4a;"); + deleteBtn.style.cssText += " padding: 6px 14px; font-size: 12px; font-weight: 500; margin-left: auto;"; + + const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;"); + form.appendChild(buttonGroup); + + return form; + } + + createSettingsToggle() { + const settingsRow = document.createElement("div"); + settingsRow.style.cssText = "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;"; + + const settingsTitle = document.createElement("h4"); + settingsTitle.textContent = "Settings"; + settingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + + const settingsToggle = document.createElement("span"); + settingsToggle.textContent = "▶"; // Right arrow when collapsed + settingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + + settingsRow.appendChild(settingsToggle); + settingsRow.appendChild(settingsTitle); + + return { settingsRow, settingsToggle }; + } + + + createCheckboxOrIconColumn(config, data, extension) { + const column = this.createCardColumn('checkbox'); + + if (config?.type === 'icon') { + column.style.flex = `0 0 ${config.width || 44}px`; + column.innerHTML = config.content || '+'; + if (config.style) { + const styles = config.style.split(';').filter(s => s.trim()); + styles.forEach(style => { + const [prop, value] = style.split(':').map(s => s.trim()); + if (prop && value) { + column.style[prop.replace(/-([a-z])/g, (g) => g[1].toUpperCase())] = value; + } + }); + } + } else { + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = `gpu-${data?.id || 'master'}`; + checkbox.checked = config?.checked !== undefined ? config.checked : data?.enabled; + checkbox.disabled = config?.disabled || false; + checkbox.style.cssText = `cursor: ${config?.disabled ? 'default' : 'pointer'}; width: 16px; height: 16px;`; + + if (config?.opacity) checkbox.style.opacity = config.opacity; + if (config?.title) column.title = config.title; + + if (config?.enabled && !config?.disabled && data?.id) { + checkbox.style.pointerEvents = "none"; + column.style.cursor = "pointer"; + column.onclick = async () => { + checkbox.checked = !checkbox.checked; + await extension.updateWorkerEnabled(data.id, checkbox.checked); + extension.updateSummary(); + }; + } + + column.appendChild(checkbox); + } + + return column; + } + + createStatusDotHelper(config, data, extension) { + let color = config.color || "#666"; + let title = config.title || "Status"; + let id = config.id; + + if (typeof config.initialColor === 'function') { + color = config.initialColor(data); + } + if (typeof config.initialTitle === 'function') { + title = config.initialTitle(data); + } + if (typeof config.id === 'function') { + id = config.id(data); + } + + const dot = this.createStatusDot(id, color, title); + + if (config.border) { + dot.style.border = config.border; + } + + if (config.pulsing && (typeof config.pulsing !== 'function' || config.pulsing(data))) { + dot.classList.add('status-pulsing'); + } + + return dot; + } + + createSettingsToggleHelper(expandedId, extension) { + const arrow = document.createElement("span"); + arrow.className = "settings-arrow"; + arrow.innerHTML = "▶"; + arrow.style.cssText = this.styles.settingsArrow; + + const isExpanded = typeof expandedId === 'function' ? + extension.state.isWorkerExpanded(expandedId(extension)) : + (expandedId === 'master' ? false : extension.state.isWorkerExpanded(expandedId)); + + if (isExpanded) { + arrow.style.transform = "rotate(90deg)"; + } + + return arrow; + } + + createControlsSection(config, data, extension, isRemote) { + if (!config) return null; + + const controlsDiv = document.createElement("div"); + controlsDiv.id = `controls-${data?.id || 'master'}`; + controlsDiv.style.cssText = this.styles.controlsDiv; + + // Always create a wrapper div for consistent layout + const controlsWrapper = document.createElement("div"); + controlsWrapper.style.cssText = "display: flex; gap: 6px; align-items: stretch; width: 100%;"; + + if (config.dynamic && data) { + if (isRemote) { + const remoteInfo = this.createButton("Remote worker", null, "background-color: #333;"); + remoteInfo.style.cssText += " flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; color: #999; cursor: default;"; + remoteInfo.disabled = true; + controlsWrapper.appendChild(remoteInfo); + } else { + const controls = this.createWorkerControls(data.id, { + launch: () => extension.launchWorker(data.id), + stop: () => extension.stopWorker(data.id), + viewLog: () => extension.viewWorkerLog(data.id) + }); + + const launchBtn = controls.querySelector(`#launch-${data.id}`); + const stopBtn = controls.querySelector(`#stop-${data.id}`); + const logBtn = controls.querySelector(`#log-${data.id}`); + + launchBtn.style.cssText += ` flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; ${BUTTON_STYLES.launch}`; + launchBtn.title = "Launch worker (runs in background with logging)"; + + stopBtn.style.cssText += ` flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; ${BUTTON_STYLES.stop} display: none;`; + stopBtn.title = "Stop worker"; + + logBtn.style.cssText += ' flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; display: none; background-color: #685434;'; + + while (controls.firstChild) { + controlsWrapper.appendChild(controls.firstChild); + } + } + } else if (config.type === 'info') { + const infoBtn = this.createButton(config.text, null, config.style || ""); + infoBtn.style.cssText += " flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; cursor: default;"; + infoBtn.disabled = true; + controlsWrapper.appendChild(infoBtn); + } else if (config.type === 'ghost') { + const ghostBtn = document.createElement("button"); + ghostBtn.style.cssText = `flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; border-radius: 4px; cursor: default; ${config.style || ""}`; + ghostBtn.textContent = config.text; + ghostBtn.disabled = true; + controlsWrapper.appendChild(ghostBtn); + } + + controlsDiv.appendChild(controlsWrapper); + return controlsDiv; + } + + createSettingsSection(config, data, extension) { + const settingsDiv = document.createElement("div"); + const settingsId = typeof config.id === 'function' ? config.id(data) : config.id; + settingsDiv.id = settingsId; + settingsDiv.className = "worker-settings"; + + const expandedId = typeof config.expandedId === 'function' ? config.expandedId(data) : config.expandedId; + const isExpanded = expandedId === 'master-settings' ? false : extension.state.isWorkerExpanded(expandedId); + + settingsDiv.style.cssText = this.styles.workerSettings; + + if (isExpanded) { + settingsDiv.classList.add("expanded"); + settingsDiv.style.padding = "12px"; + settingsDiv.style.marginTop = "8px"; + settingsDiv.style.marginBottom = "8px"; + } + + let settingsForm; + if (config.formType === 'master') { + settingsForm = this.createMasterSettingsForm(extension, data); + } else if (config.formType === 'worker') { + settingsForm = this.createWorkerSettingsForm(extension, data); + } + + if (settingsForm) { + settingsDiv.appendChild(settingsForm); + } + + return settingsDiv; + } + + createMasterSettingsForm(extension, data) { + const settingsForm = document.createElement("div"); + settingsForm.style.cssText = "display: flex; flex-direction: column; gap: 8px;"; + + const nameResult = this.createFormGroup("Name:", extension.config?.master?.name || "Master", "master-name"); + settingsForm.appendChild(nameResult.group); + + const hostResult = this.createFormGroup("Host:", extension.config?.master?.host || "", "master-host", "text", "Auto-detect if empty"); + settingsForm.appendChild(hostResult.group); + + const saveBtn = this.createButton("Save", async () => { + const nameInput = document.getElementById('master-name'); + const hostInput = document.getElementById('master-host'); + + if (!extension.config.master) extension.config.master = {}; + extension.config.master.name = nameInput.value.trim() || "Master"; + + await extension.api.updateMaster({ + host: hostInput.value.trim(), + name: extension.config.master.name + }); + + // Reload config to refresh any updated values + await extension.loadConfig(); + + document.getElementById('master-name-display').textContent = extension.config.master.name; + this.updateMasterDisplay(extension); + + saveBtn.textContent = "Saved!"; + setTimeout(() => { saveBtn.textContent = "Save"; }, 2000); + }, "background-color: #4a7c4a;"); + saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; + + const cancelBtn = this.createButton("Cancel", () => { + document.getElementById('master-name').value = extension.config?.master?.name || "Master"; + document.getElementById('master-host').value = extension.config?.master?.host || ""; + }, "background-color: #555;"); + cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; + + const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn], " margin-top: 8px;"); + settingsForm.appendChild(buttonGroup); + + return settingsForm; + } + + addPlaceholderHover(card, leftColumn, entityType) { + card.onmouseover = () => { + if (entityType === 'blueprint') { + card.style.borderColor = "#777"; + card.style.backgroundColor = "rgba(255, 255, 255, 0.05)"; + leftColumn.style.color = "#999"; + } else { + card.style.borderColor = "#666"; + card.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; + leftColumn.style.color = "#888"; + leftColumn.style.borderColor = "#666"; + } + }; + + card.onmouseout = () => { + if (entityType === 'blueprint') { + card.style.borderColor = "#555"; + card.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; + leftColumn.style.color = "#777"; + } else { + card.style.borderColor = "#444"; + card.style.backgroundColor = "transparent"; + leftColumn.style.color = "#555"; + leftColumn.style.borderColor = "#444"; + } + }; + } + + renderEntityCard(entityType, data, extension) { + const config = cardConfigs[entityType] || {}; + const isPlaceholder = entityType === 'blueprint' || entityType === 'add'; + const isWorker = entityType === 'worker'; + const isMaster = entityType === 'master'; + const isRemote = isWorker && extension.isRemoteWorker(data); + + const cardOptions = { + onClick: isPlaceholder ? data?.onClick : null + }; + if (isPlaceholder) { + cardOptions.title = entityType === 'blueprint' ? "Click to add your first worker" : "Click to add a new worker"; + } + const card = this.createCard(entityType, cardOptions); + + const leftColumn = this.createCheckboxOrIconColumn(config.checkbox, data, extension); + card.appendChild(leftColumn); + + const rightColumn = this.createCardColumn('content'); + + const infoRow = this.createInfoRow(); + if (config.infoRowPadding) { + infoRow.style.padding = config.infoRowPadding; + } + if (config.minHeight === 'auto') { + infoRow.style.minHeight = 'auto'; + } else if (config.minHeight) { + infoRow.style.minHeight = config.minHeight; + } + if (config.expand) { + infoRow.title = "Click to expand settings"; + infoRow.onclick = () => { + if (isMaster) { + const masterSettingsExpanded = !extension.masterSettingsExpanded; + extension.masterSettingsExpanded = masterSettingsExpanded; + const masterSettingsDiv = document.getElementById("master-settings"); + const arrow = infoRow.querySelector('.settings-arrow'); + if (masterSettingsExpanded) { + masterSettingsDiv.classList.add("expanded"); + masterSettingsDiv.style.padding = "12px"; + masterSettingsDiv.style.marginTop = "8px"; + masterSettingsDiv.style.marginBottom = "8px"; + arrow.style.transform = "rotate(90deg)"; + } else { + masterSettingsDiv.classList.remove("expanded"); + masterSettingsDiv.style.padding = "0 12px"; + masterSettingsDiv.style.marginTop = "0"; + masterSettingsDiv.style.marginBottom = "0"; + arrow.style.transform = "rotate(0deg)"; + } + } else { + extension.toggleWorkerExpanded(data.id); + } + }; + } + + const workerContent = this.createWorkerContent(); + if (entityType === 'add') { + workerContent.style.alignItems = "center"; + } + + const statusDot = this.createStatusDotHelper(config.statusDot, data, extension); + workerContent.appendChild(statusDot); + + const infoSpan = document.createElement("span"); + infoSpan.innerHTML = config.infoText(data, extension); + workerContent.appendChild(infoSpan); + + infoRow.appendChild(workerContent); + + let settingsArrow; + if (config.expand) { + const expandedId = config.settings?.expandedId || (isMaster ? 'master' : data?.id); + settingsArrow = this.createSettingsToggleHelper(expandedId, extension); + if (isMaster && !extension.masterSettingsExpanded) { + settingsArrow.style.transform = "rotate(0deg)"; + } + infoRow.appendChild(settingsArrow); + } + + rightColumn.appendChild(infoRow); + + if (config.hover === true) { + rightColumn.onmouseover = () => { + rightColumn.style.backgroundColor = "#333"; + if (settingsArrow) settingsArrow.style.color = "#fff"; + }; + rightColumn.onmouseout = () => { + rightColumn.style.backgroundColor = "transparent"; + if (settingsArrow) settingsArrow.style.color = "#888"; + }; + } + + const controlsDiv = this.createControlsSection(config.controls, data, extension, isRemote); + if (controlsDiv) { + rightColumn.appendChild(controlsDiv); + } + + if (config.settings) { + const settingsDiv = this.createSettingsSection(config.settings, data, extension); + rightColumn.appendChild(settingsDiv); + } + + card.appendChild(rightColumn); + + if (config.hover === 'placeholder') { + this.addPlaceholderHover(card, leftColumn, entityType); + } + + if (isWorker && !isRemote) { + extension.updateWorkerControls(data.id); + } + + return card; + } +} \ No newline at end of file diff --git a/web/workerUtils.js b/web/workerUtils.js new file mode 100644 index 0000000..122d59e --- /dev/null +++ b/web/workerUtils.js @@ -0,0 +1,266 @@ +import { BUTTON_STYLES } from './constants.js'; + +export async function handleWorkerOperation(extension, button, operation, successText, errorText) { + const originalText = button.textContent; + const originalStyle = button.style.cssText; + + button.textContent = operation.loadingText; + button.disabled = true; + + try { + const urlsToProcess = extension.enabledWorkers.map(w => ({ + name: w.name, + url: extension.getWorkerUrl(w) + })); + + if (urlsToProcess.length === 0) { + button.textContent = "No Workers"; + button.style.backgroundColor = "#c04c4c"; + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + button.disabled = false; + }, 3000); + return; + } + + const promises = urlsToProcess.map(target => + fetch(`${target.url}${operation.endpoint}`, { + method: 'POST', + mode: 'cors' + }) + .then(response => ({ ok: response.ok, name: target.name })) + .catch(() => ({ ok: false, name: target.name })) + ); + + const results = await Promise.all(promises); + const failures = results.filter(r => !r.ok); + + if (failures.length === 0) { + button.textContent = successText; + button.style.backgroundColor = BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + if (operation.onSuccess) operation.onSuccess(); + } else { + button.textContent = errorText; + button.style.backgroundColor = BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + extension.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`, "error"); + } + + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + }, 3000); + } finally { + button.disabled = false; + } +} + +export async function handleInterruptWorkers(extension, button) { + return handleWorkerOperation(extension, button, { + name: "Interrupt", + endpoint: "/interrupt", + loadingText: "Interrupting...", + onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), 500) + }, "Interrupted!", "Error! See Console"); +} + +export async function handleClearMemory(extension, button) { + return handleWorkerOperation(extension, button, { + name: "Clear memory", + endpoint: "/distributed/clear_memory", + loadingText: "Clearing..." + }, "Success!", "Error! See Console"); +} + +export function findNodesByClass(apiPrompt, className) { + return Object.entries(apiPrompt) + .filter(([, nodeData]) => nodeData.class_type === className) + .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); +} + +/** + * Find all image references in the workflow + * Looks for inputs named "image" that contain filename strings + */ +export function findImageReferences(extension, apiPrompt) { + const images = new Map(); + // Updated regex to handle: + // - Standard files: "image.png" + // - Subfolder files: "subfolder/image.png" + // - ComfyUI special format: "clipspace/file.png [input]" + const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp)(\s*\[\w+\])?$/i; + + for (const [nodeId, node] of Object.entries(apiPrompt)) { + if (node.inputs && node.inputs.image) { + const imageValue = node.inputs.image; + if (typeof imageValue === 'string') { + // Clean special suffixes like [input] or [output] + const cleanValue = imageValue.replace(/\s*\[\w+\]$/, '').trim(); + if (imageExtensions.test(cleanValue)) { + images.set(cleanValue, { + nodeId, + nodeType: node.class_type, + inputName: 'image' + }); + extension.log(`Found image reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug"); + } + } + } + } + + return images; +} + +/** + * Find only upstream nodes (inputs) for distributed collector nodes + * This is used for workers to avoid executing downstream nodes like SaveImage + * @param {Object} apiPrompt - The API prompt containing the workflow + * @param {Array} collectorIds - Array of collector node IDs + * @returns {Set} Set of node IDs that feed into collectors + */ +export function findCollectorUpstreamNodes(apiPrompt, collectorIds) { + const connected = new Set(collectorIds); // Include all collectors + const toProcess = [...collectorIds]; + + // Only traverse upstream (inputs) + while (toProcess.length > 0) { + const nodeId = toProcess.pop(); + const node = apiPrompt[nodeId]; + + // Traverse upstream (inputs) only + if (node && node.inputs) { + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceNodeId = String(inputValue[0]); + if (!connected.has(sourceNodeId)) { + connected.add(sourceNodeId); + toProcess.push(sourceNodeId); + } + } + } + } + } + + return connected; +} + +/** + * Prune workflow to only include nodes connected to distributed nodes + * @param {Object} apiPrompt - The full workflow API prompt + * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided) + * @returns {Object} Pruned API prompt with only required nodes + */ +export function pruneWorkflowForWorker(extension, apiPrompt, distributedNodes = null) { + // Find all distributed nodes if not provided + if (!distributedNodes) { + const collectorNodes = findNodesByClass(apiPrompt, "DistributedCollector"); + const upscaleNodes = findNodesByClass(apiPrompt, "UltimateSDUpscaleDistributed"); + distributedNodes = [...collectorNodes, ...upscaleNodes]; + } + + if (distributedNodes.length === 0) { + // No distributed nodes, return full workflow + return apiPrompt; + } + + // Get all nodes connected to distributed nodes + const distributedIds = distributedNodes.map(node => node.id); + + // For workers, only include upstream nodes (this removes ALL downstream nodes after collectors) + const connectedNodes = findCollectorUpstreamNodes(apiPrompt, distributedIds); + + extension.log(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`, "debug"); + + // Create pruned prompt with only required nodes + const prunedPrompt = {}; + for (const nodeId of connectedNodes) { + prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId])); + } + + // Check if any distributed node has downstream SaveImage nodes that were removed + // If so, add a PreviewImage node after the collector + for (const distNode of distributedNodes) { + const distNodeId = distNode.id; + + // Check if this distributed node had any downstream nodes in the original workflow + const originalOutputMap = new Map(); + for (const [nodeId, node] of Object.entries(apiPrompt)) { + if (node.inputs) { + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2 && String(inputValue[0]) === distNodeId) { + if (!originalOutputMap.has(distNodeId)) { + originalOutputMap.set(distNodeId, []); + } + originalOutputMap.get(distNodeId).push({nodeId, inputName}); + } + } + } + } + + // If this distributed node had downstream nodes that were removed, add a PreviewImage + if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) { + // Generate unique numeric ID: max existing numeric key +1 + const existingIds = Object.keys(prunedPrompt) + .filter(k => !isNaN(parseInt(k))) + .map(k => parseInt(k)); + const maxId = existingIds.length > 0 ? Math.max(...existingIds) : 0; + const previewNodeId = String(maxId + 1); + + // Add PreviewImage node connected to the distributed node + prunedPrompt[previewNodeId] = { + inputs: { + images: [distNodeId, 0] // Connect to first output of distributed node + }, + class_type: "PreviewImage", + _meta: { + title: "Preview Image (auto-added)" + } + }; + + extension.log(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`, "debug"); + } + } + + return prunedPrompt; +} + +/** + * Check if a node has an upstream node of a specific type + * @param {Object} apiPrompt - The workflow API prompt + * @param {string} nodeId - The node to check + * @param {string} upstreamType - The class_type to look for upstream + * @returns {boolean} True if an upstream node of the specified type exists + */ +export function hasUpstreamNode(apiPrompt, nodeId, upstreamType) { + const visited = new Set(); + const toProcess = [nodeId]; + + while (toProcess.length > 0) { + const currentId = toProcess.pop(); + if (visited.has(currentId)) continue; + visited.add(currentId); + + const node = apiPrompt[currentId]; + if (!node) continue; + + // Check inputs for upstream connections + if (node.inputs) { + for (const [inputName, inputValue] of Object.entries(node.inputs)) { + if (Array.isArray(inputValue) && inputValue.length === 2) { + const sourceNodeId = String(inputValue[0]); + const sourceNode = apiPrompt[sourceNodeId]; + + if (sourceNode && sourceNode.class_type === upstreamType) { + return true; + } + + if (!visited.has(sourceNodeId)) { + toProcess.push(sourceNodeId); + } + } + } + } + } + + return false; +} \ No newline at end of file From 9834a0f6e470979a2161b2ff455ac8af5b8d6bd3 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sun, 13 Jul 2025 10:00:36 +1000 Subject: [PATCH 05/93] constants --- web/apiClient.js | 8 ++++--- web/constants.js | 54 +++++++++++++++++++++++++++++++++++++++---- web/executionUtils.js | 11 ++++++--- web/main.js | 27 +++++++++++----------- web/ui.js | 54 +++++++++++++++++++++---------------------- web/workerUtils.js | 8 +++---- 6 files changed, 106 insertions(+), 56 deletions(-) diff --git a/web/apiClient.js b/web/apiClient.js index 061428d..b4fceaa 100644 --- a/web/apiClient.js +++ b/web/apiClient.js @@ -1,7 +1,9 @@ +import { TIMEOUTS } from './constants.js'; + export function createApiClient(baseUrl) { - const request = async (endpoint, options = {}, retries = 3) => { + const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => { let lastError; - let delay = 1000; // Initial delay for exponential backoff + let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff for (let attempt = 0; attempt < retries; attempt++) { try { @@ -114,7 +116,7 @@ export function createApiClient(baseUrl) { }, // Status checking (with timeout) - async checkStatus(url, timeout = 5000) { + async checkStatus(url, timeout = TIMEOUTS.DEFAULT_FETCH) { const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), timeout); diff --git a/web/constants.js b/web/constants.js index 223807d..2b3f0ad 100644 --- a/web/constants.js +++ b/web/constants.js @@ -5,7 +5,6 @@ export const BUTTON_STYLES = { error: "background-color: #c04c4c;", launch: "background-color: #4a7c4a;", stop: "background-color: #7c4a4a;", - grey: "background-color: #555;", }; export const STATUS_COLORS = { @@ -15,6 +14,18 @@ export const STATUS_COLORS = { PROCESSING_YELLOW: "#f0ad4e" }; +export const UI_COLORS = { + MUTED_TEXT: "#888", + SECONDARY_TEXT: "#ccc", + BORDER_LIGHT: "#555", + BORDER_DARK: "#444", + BORDER_DARKER: "#3a3a3a", + BACKGROUND_DARK: "#2a2a2a", + BACKGROUND_DARKER: "#1e1e1e", + ICON_COLOR: "#666", + ACCENT_COLOR: "#777" +}; + export const PULSE_ANIMATION_CSS = ` @keyframes pulse { 0% { @@ -73,11 +84,32 @@ export const UI_STYLES = { formGroup: "display: flex; flex-direction: column; gap: 5px;", formLabel: "font-size: 12px; color: #ccc; font-weight: 500;", formInput: "padding: 6px 10px; background: #2a2a2a; border: 1px solid #444; color: white; font-size: 12px; border-radius: 4px; transition: border-color 0.2s;", + + // Card styles + cardBase: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;", workerCard: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex; background: #2a2a2a;", + cardBlueprint: "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);", + cardAdd: "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;", + + // Column styles + columnBase: "display: flex; align-items: center; justify-content: center;", checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);", contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;", - infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer;", + iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;", + + // Row and content styles + infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;", workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;", + + // Form and controls styles + buttonGroup: "display: flex; gap: 4px; margin-top: 10px;", + settingsForm: "display: flex; flex-direction: column; gap: 10px;", + checkboxGroup: "display: flex; align-items: center; gap: 8px; margin: 5px 0;", + formLabelClickable: "font-size: 12px; color: #ccc; cursor: pointer;", + settingsToggle: "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;", + controlsWrapper: "display: flex; gap: 6px; align-items: stretch; width: 100%;", + + // Existing styles settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto; padding: 4px;", infoBox: "background-color: #333; color: #999; padding: 5px 14px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;", workerSettings: "margin: 0 12px; padding: 0 12px; background: #1e1e1e; border-radius: 4px; border: 1px solid #2a2a2a;" @@ -85,8 +117,22 @@ export const UI_STYLES = { export const TIMEOUTS = { DEFAULT_FETCH: 5000, // ms for general API calls - STATUS_CHECK: 2000, // ms for status checks + STATUS_CHECK: 500, // ms for status checks LAUNCH: 90000, // ms for worker launch (longer for model loading) RETRY_DELAY: 1000, // initial delay for exponential backoff - MAX_RETRIES: 3 // max retry attempts + MAX_RETRIES: 3, // max retry attempts + + // UI feedback delays + BUTTON_RESET: 3000, // button text/state reset after actions + FLASH_SHORT: 1000, // brief success feedback + FLASH_MEDIUM: 1500, // medium error feedback + FLASH_LONG: 2000, // longer error feedback + + // Operational delays + POST_ACTION_DELAY: 500, // delay after operations before status checks + STATUS_CHECK_DELAY: 100, // brief delay before status checks + + // Background tasks + LOG_REFRESH: 2000, // log auto-refresh interval + IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache }; \ No newline at end of file diff --git a/web/executionUtils.js b/web/executionUtils.js index d26fcba..9960de4 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -11,7 +11,7 @@ export function setupInterceptor(extension) { if (hasCollector || hasDistUpscale) { const result = await executeParallelDistributed(extension, prompt); // Check status after dispatching jobs - setTimeout(() => extension.checkAllWorkerStatuses(), 500); + setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); return result; } } @@ -27,8 +27,13 @@ export async function executeParallelDistributed(extension, promptWrapper) { // Pre-flight health check on all enabled workers const activeWorkers = await performPreflightCheck(extension, enabledWorkers); + // Case: Enabled workers but all offline if (activeWorkers.length === 0 && enabledWorkers.length > 0) { - extension.log("No active workers found. Running on master only."); + extension.log("No active workers found. All enabled workers are offline."); + if (extension.ui?.showToast) { + extension.ui.showToast(extension.app, "error", "All Workers Offline", + `${enabledWorkers.length} worker(s) enabled but all are offline or unreachable. Check worker connections and try again.`, 5000); + } // Fall back to master-only execution return extension.originalQueuePrompt(0, promptWrapper); } @@ -336,7 +341,7 @@ export async function loadImagesForWorker(extension, imageReferences) { extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug"); extension._imageCache.clear(); } - }, 30000); // Clear after 30 seconds + }, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds return images; } diff --git a/web/main.js b/web/main.js index c0310c5..e643d58 100644 --- a/web/main.js +++ b/web/main.js @@ -16,6 +16,7 @@ class DistributedExtension { this.statusCheckInterval = null; this.logAutoRefreshInterval = null; this.masterSettingsExpanded = false; + this.app = app; // Store app reference for toast notifications // Initialize centralized state this.state = createStateManager(); @@ -109,7 +110,7 @@ class DistributedExtension { if (enabled) { // Enabled: Start with checking state and trigger check this.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Checking status...", true); - setTimeout(() => this.checkWorkerStatus(worker), 100); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK_DELAY); } else { // Disabled: Set to gray this.ui.updateStatusDot(workerId, STATUS_COLORS.DISABLED_GRAY, "Disabled", false); @@ -192,11 +193,9 @@ class DistributedExtension { } startStatusChecking() { - // Start checking every 2 seconds for more responsive updates - // Don't check immediately since panel might not be open yet this.statusCheckInterval = setInterval(() => { this.checkAllWorkerStatuses(); - }, 2000); + }, TIMEOUTS.STATUS_CHECK); } async checkAllWorkerStatuses() { @@ -356,7 +355,7 @@ class DistributedExtension { // Allow 90 seconds for worker to launch (model loading can take time) setTimeout(() => { this.state.setWorkerLaunching(workerId, false); - }, 90000); + }, TIMEOUTS.LAUNCH); if (!launchBtn) return; @@ -379,14 +378,14 @@ class DistributedExtension { // Update controls immediately to hide launch button and show stop/log buttons this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 2000); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK); } } catch (error) { // Check if worker was already running if (error.message && error.message.includes("already running")) { this.log(`Worker ${worker.name} is already running`, "info"); this.updateWorkerControls(workerId); - setTimeout(() => this.checkWorkerStatus(worker), 100); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK_DELAY); } else { this.log(`Error launching worker: ${error.message || error}`, "error"); @@ -425,11 +424,11 @@ class DistributedExtension { stopBtn.textContent = "Stopped!"; setTimeout(() => { this.updateWorkerControls(workerId); - }, 1000); + }, TIMEOUTS.FLASH_SHORT); } // Verify status after a short delay - setTimeout(() => this.checkWorkerStatus(worker), 2000); + setTimeout(() => this.checkWorkerStatus(worker), TIMEOUTS.STATUS_CHECK); } else { this.log(`Failed to stop worker: ${result.message}`, "error"); @@ -446,7 +445,7 @@ class DistributedExtension { setTimeout(() => { this.updateWorkerControls(workerId); - }, 1500); + }, TIMEOUTS.FLASH_MEDIUM); } } } catch (error) { @@ -458,7 +457,7 @@ class DistributedExtension { stopBtn.textContent = "Error"; setTimeout(() => { this.updateWorkerControls(workerId); - }, 1500); + }, TIMEOUTS.FLASH_MEDIUM); } } } @@ -473,7 +472,7 @@ class DistributedExtension { } // Generic async button action handler - async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = 3000) { + async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = TIMEOUTS.BUTTON_RESET) { const originalText = button.textContent; const originalStyle = button.style.cssText; button.disabled = true; @@ -638,7 +637,7 @@ class DistributedExtension { logBtn.disabled = false; logBtn.textContent = "View Log"; logBtn.style.backgroundColor = "#685434"; // Keep the yellow color - }, 2000); + }, TIMEOUTS.FLASH_LONG); } } } @@ -1121,7 +1120,7 @@ class DistributedExtension { // Refresh every 2 seconds this.logAutoRefreshInterval = setInterval(() => { this.refreshLog(workerId, true); // silent mode - }, 2000); + }, TIMEOUTS.LOG_REFRESH); } stopLogAutoRefresh() { diff --git a/web/ui.js b/web/ui.js index 87bbbd9..d8a2fd6 100644 --- a/web/ui.js +++ b/web/ui.js @@ -1,4 +1,4 @@ -import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS } from './constants.js'; +import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS, UI_COLORS, TIMEOUTS } from './constants.js'; const cardConfigs = { master: { @@ -19,7 +19,7 @@ const cardConfigs = { const cudaDevice = extension.config?.master?.cuda_device ?? extension.masterCudaDevice; const cudaInfo = cudaDevice !== undefined ? `CUDA ${cudaDevice} • ` : ''; const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - return `${data?.name || extension.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; + return `${data?.name || extension.config?.master?.name || "Master"}
${cudaInfo}Port ${port}`; }, controls: { type: 'info', @@ -50,10 +50,10 @@ const cardConfigs = { infoText: (data, extension) => { const isRemote = extension.isRemoteWorker(data); if (isRemote) { - return `${data.name}
${data.host}:${data.port}`; + return `${data.name}
${data.host}:${data.port}`; } else { const cudaInfo = data.cuda_device !== undefined ? `CUDA ${data.cuda_device} • ` : ''; - return `${data.name}
${cudaInfo}Port ${data.port}`; + return `${data.name}
${cudaInfo}Port ${data.port}`; } }, controls: { @@ -73,17 +73,17 @@ const cardConfigs = { type: 'icon', content: '+', width: 42, - style: 'border-right: 2px dashed #555; color: #777; font-size: 24px; font-weight: 500;' + style: `border-right: 2px dashed ${UI_COLORS.BORDER_LIGHT}; color: ${UI_COLORS.ACCENT_COLOR}; font-size: 24px; font-weight: 500;` }, statusDot: { color: 'transparent', - border: '1px solid #555' + border: `1px solid ${UI_COLORS.BORDER_LIGHT}` }, - infoText: () => `Add New Worker
[CUDA] • [Port]`, + infoText: () => `Add New Worker
[CUDA] • [Port]`, controls: { type: 'ghost', text: 'Configure', - style: 'border: 1px solid #444; background: transparent; color: #555;' + style: `border: 1px solid ${UI_COLORS.BORDER_DARK}; background: transparent; color: ${UI_COLORS.BORDER_LIGHT};` }, hover: 'placeholder', expand: false, @@ -94,13 +94,13 @@ const cardConfigs = { type: 'icon', content: '+', width: 43, - style: 'border-right: 1px dashed #444; color: #555; font-size: 18px;' + style: `border-right: 1px dashed ${UI_COLORS.BORDER_DARK}; color: ${UI_COLORS.BORDER_LIGHT}; font-size: 18px;` }, statusDot: { color: 'transparent', - border: '1px solid #555' + border: `1px solid ${UI_COLORS.BORDER_LIGHT}` }, - infoText: () => `Add New Worker`, + infoText: () => `Add New Worker`, controls: null, hover: 'placeholder', expand: false, @@ -134,7 +134,7 @@ export class DistributedUI { createButtonGroup(buttons, style = "") { const group = document.createElement("div"); - group.style.cssText = "display: flex; gap: 4px; margin-top: 10px;" + style; + group.style.cssText = this.styles.buttonGroup + style; buttons.forEach(button => group.appendChild(button)); return group; } @@ -207,20 +207,19 @@ export class DistributedUI { createCard(type = 'worker', options = {}) { const card = document.createElement("div"); - const baseStyle = "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;"; switch(type) { case 'master': case 'worker': - card.style.cssText = baseStyle + "background: #2a2a2a;"; + card.style.cssText = this.styles.workerCard; break; case 'blueprint': - card.style.cssText = baseStyle + "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);"; + card.style.cssText = this.styles.cardBase + this.styles.cardBlueprint; if (options.onClick) card.onclick = options.onClick; if (options.title) card.title = options.title; break; case 'add': - card.style.cssText = baseStyle + "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;"; + card.style.cssText = this.styles.cardBase + this.styles.cardAdd; if (options.onClick) card.onclick = options.onClick; if (options.title) card.title = options.title; break; @@ -238,18 +237,17 @@ export class DistributedUI { createCardColumn(type = 'checkbox', options = {}) { const column = document.createElement("div"); - const baseStyle = "display: flex; align-items: center; justify-content: center;"; switch(type) { case 'checkbox': - column.style.cssText = baseStyle + "flex: 0 0 44px; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);"; + column.style.cssText = this.styles.checkboxColumn; if (options.title) column.title = options.title; break; case 'icon': - column.style.cssText = baseStyle + "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;"; + column.style.cssText = this.styles.columnBase + this.styles.iconColumn; break; case 'content': - column.style.cssText = "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;"; + column.style.cssText = this.styles.contentColumn; break; } @@ -258,25 +256,25 @@ export class DistributedUI { createInfoRow(options = {}) { const row = document.createElement("div"); - row.style.cssText = "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;"; + row.style.cssText = this.styles.infoRow; if (options.onClick) row.onclick = options.onClick; return row; } createWorkerContent() { const content = document.createElement("div"); - content.style.cssText = "display: flex; align-items: center; gap: 10px; flex: 1;"; + content.style.cssText = this.styles.workerContent; return content; } createSettingsForm(fields = [], options = {}) { const form = document.createElement("div"); - form.style.cssText = "display: flex; flex-direction: column; gap: 10px;"; + form.style.cssText = this.styles.settingsForm; fields.forEach(field => { if (field.type === 'checkbox') { const group = document.createElement("div"); - group.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; + group.style.cssText = this.styles.checkboxGroup; const checkbox = document.createElement("input"); checkbox.type = "checkbox"; @@ -287,7 +285,7 @@ export class DistributedUI { const label = document.createElement("label"); label.htmlFor = field.id; label.textContent = field.label; - label.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + label.style.cssText = this.styles.formLabelClickable; group.appendChild(checkbox); group.appendChild(label); @@ -613,7 +611,7 @@ export class DistributedUI { createSettingsToggle() { const settingsRow = document.createElement("div"); - settingsRow.style.cssText = "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;"; + settingsRow.style.cssText = this.styles.settingsToggle; const settingsTitle = document.createElement("h4"); settingsTitle.textContent = "Settings"; @@ -726,7 +724,7 @@ export class DistributedUI { // Always create a wrapper div for consistent layout const controlsWrapper = document.createElement("div"); - controlsWrapper.style.cssText = "display: flex; gap: 6px; align-items: stretch; width: 100%;"; + controlsWrapper.style.cssText = this.styles.controlsWrapper; if (config.dynamic && data) { if (isRemote) { @@ -835,7 +833,7 @@ export class DistributedUI { this.updateMasterDisplay(extension); saveBtn.textContent = "Saved!"; - setTimeout(() => { saveBtn.textContent = "Save"; }, 2000); + setTimeout(() => { saveBtn.textContent = "Save"; }, TIMEOUTS.FLASH_LONG); }, "background-color: #4a7c4a;"); saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; diff --git a/web/workerUtils.js b/web/workerUtils.js index 122d59e..745758a 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -1,4 +1,4 @@ -import { BUTTON_STYLES } from './constants.js'; +import { BUTTON_STYLES, TIMEOUTS } from './constants.js'; export async function handleWorkerOperation(extension, button, operation, successText, errorText) { const originalText = button.textContent; @@ -20,7 +20,7 @@ export async function handleWorkerOperation(extension, button, operation, succes button.textContent = originalText; button.style.cssText = originalStyle; button.disabled = false; - }, 3000); + }, TIMEOUTS.BUTTON_RESET); return; } @@ -49,7 +49,7 @@ export async function handleWorkerOperation(extension, button, operation, succes setTimeout(() => { button.textContent = originalText; button.style.cssText = originalStyle; - }, 3000); + }, TIMEOUTS.BUTTON_RESET); } finally { button.disabled = false; } @@ -60,7 +60,7 @@ export async function handleInterruptWorkers(extension, button) { name: "Interrupt", endpoint: "/interrupt", loadingText: "Interrupting...", - onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), 500) + onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY) }, "Interrupted!", "Error! See Console"); } From f1a94ee4b0a9c7668bc2e28e683adb30a3ad094f Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sun, 13 Jul 2025 13:29:45 +1000 Subject: [PATCH 06/93] ui tweaks --- web/constants.js | 2 +- web/sidebarRenderer.js | 4 +++- web/ui.js | 27 ++++++++++++++++++++++++++- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/web/constants.js b/web/constants.js index 2b3f0ad..441a766 100644 --- a/web/constants.js +++ b/web/constants.js @@ -117,7 +117,7 @@ export const UI_STYLES = { export const TIMEOUTS = { DEFAULT_FETCH: 5000, // ms for general API calls - STATUS_CHECK: 500, // ms for status checks + STATUS_CHECK: 1200, // ms for status checks LAUNCH: 90000, // ms for worker launch (longer for model loading) RETRY_DELAY: 1000, // initial delay for exponential backoff MAX_RETRIES: 3, // max retry attempts diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js index d88e4cd..3187ee4 100644 --- a/web/sidebarRenderer.js +++ b/web/sidebarRenderer.js @@ -173,7 +173,7 @@ export async function renderSidebarContent(extension, el) { // Collapsible settings content const settingsContent = document.createElement("div"); - settingsContent.style.cssText = "max-height: 0; overflow: hidden; transition: max-height 0.3s ease;"; + settingsContent.style.cssText = "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;"; const settingsDiv = document.createElement("div"); settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;"; @@ -184,9 +184,11 @@ export async function renderSidebarContent(extension, el) { settingsExpanded = !settingsExpanded; if (settingsExpanded) { settingsContent.style.maxHeight = "200px"; + settingsContent.style.opacity = "1"; workerSettingsToggle.style.transform = "rotate(90deg)"; } else { settingsContent.style.maxHeight = "0"; + settingsContent.style.opacity = "0"; workerSettingsToggle.style.transform = "rotate(0deg)"; } }; diff --git a/web/ui.js b/web/ui.js index d8a2fd6..24f0808 100644 --- a/web/ui.js +++ b/web/ui.js @@ -821,17 +821,42 @@ export class DistributedUI { if (!extension.config.master) extension.config.master = {}; extension.config.master.name = nameInput.value.trim() || "Master"; + const hostValue = hostInput.value.trim(); + await extension.api.updateMaster({ - host: hostInput.value.trim(), + host: hostValue, name: extension.config.master.name }); // Reload config to refresh any updated values await extension.loadConfig(); + // If host was emptied, trigger auto-detection + if (!hostValue) { + extension.log("Host field cleared, triggering IP auto-detection", "debug"); + await extension.detectMasterIP(); + // Reload config again to get the auto-detected IP + await extension.loadConfig(); + // Update the input field with the detected IP + document.getElementById('master-host').value = extension.config?.master?.host || ""; + } + document.getElementById('master-name-display').textContent = extension.config.master.name; this.updateMasterDisplay(extension); + // Show toast notification + if (extension.app?.extensionManager?.toast) { + const message = !hostValue ? + "Master settings saved and IP auto-detected" : + "Master settings saved successfully"; + extension.app.extensionManager.toast.add({ + severity: "success", + summary: "Master Updated", + detail: message, + life: 3000 + }); + } + saveBtn.textContent = "Saved!"; setTimeout(() => { saveBtn.textContent = "Save"; }, TIMEOUTS.FLASH_LONG); }, "background-color: #4a7c4a;"); From 606701b65bf5729cb59ed7f1745212ac0b7d990b Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 13 Jul 2025 13:31:23 +1000 Subject: [PATCH 07/93] Update README.md --- README.md | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/README.md b/README.md index 0bbf6ed..adb7d15 100644 --- a/README.md +++ b/README.md @@ -86,18 +86,6 @@ A powerful extension for ComfyUI that enables parallel and distributed processin - **Port**: The port number used when launching ComfyUI on the remote worker (e.g., 8189). - **Save** the remote worker configuration. -
-💡 Tip: Using GPUs Outside Your Local Network - -Services like [Tailscale](https://tailscale.com) create secure VPN connections that make remote computers appear as if they're on your local network. This lets you use GPUs from: -- **Other locations** (work computer, friend's gaming PC, etc.) -- **Cloud instances** (AWS, Google Cloud, etc.) -- **Any internet-connected machine** with CUDA GPUs - -Just install Tailscale on both machines and use the Tailscale IP address as the "Host" when configuring your remote worker. - -
- ### Configuration Tips | Setting | Description | Example | @@ -246,8 +234,8 @@ This project is under active development. Contributions are welcome! ### Planned Features +- [x] Support for Cloud workers (release coming soon) - [ ] Remote worker control via SSH - [ ] View remote worker logs in UI - [ ] Improve worker timeout logic -- [ ] Support for Runpod workers - [ ] ComfyUI Desktop app compatibility From eb855349fed460563ba4310fabcd4a4417a2d26b Mon Sep 17 00:00:00 2001 From: robertvoy Date: Mon, 14 Jul 2025 12:52:21 +1000 Subject: [PATCH 08/93] fixed UUID creation bug when accessing master over network --- web/main.js | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/web/main.js b/web/main.js index e643d58..ecd7efe 100644 --- a/web/main.js +++ b/web/main.js @@ -50,6 +50,19 @@ class DistributedExtension { } } + // Generate UUID with fallback for non-secure contexts + generateUUID() { + if (crypto.randomUUID) { + return crypto.randomUUID(); + } + // Fallback for non-secure contexts + return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, function(c) { + const r = Math.random() * 16 | 0; + const v = c == 'x' ? r : (r & 0x3 | 0x8); + return v.toString(16); + }); + } + injectStyles() { const styleId = 'distributed-styles'; if (!document.getElementById(styleId)) { @@ -1054,8 +1067,8 @@ class DistributedExtension { } async addNewWorker() { - // Generate new worker ID using UUID - const newId = crypto.randomUUID(); + // Generate new worker ID using UUID (fallback for non-secure contexts) + const newId = this.generateUUID(); // Find next available port const usedPorts = this.config.workers.map(w => w.port); From d561bba5aab539af87e399fd5c7f19bdcdb778f6 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Mon, 14 Jul 2025 14:14:07 +1000 Subject: [PATCH 09/93] ui tweaks --- web/constants.js | 29 ++++++++++++++++++++++------- web/sidebarRenderer.js | 12 +++++++----- web/ui.js | 22 +++++++++++----------- 3 files changed, 40 insertions(+), 23 deletions(-) diff --git a/web/constants.js b/web/constants.js index 441a766..df0d251 100644 --- a/web/constants.js +++ b/web/constants.js @@ -1,10 +1,24 @@ export const BUTTON_STYLES = { - base: "width: 100%; padding: 8px 12px; font-size: 12px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: all 0.2s; font-weight: 500;", - clearMemory: "background-color: #555;", - success: "background-color: #3ca03c;", - error: "background-color: #c04c4c;", + // Base styles with unified padding + base: "width: 100%; padding: 4px 14px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: all 0.2s; font-size: 12px; font-weight: 500;", + + // Context-specific combined styles + workerControl: "flex: 1; font-size: 11px;", + + // Layout modifiers + hidden: "display: none;", + marginLeftAuto: "margin-left: auto;", + + // Color variants + cancel: "background-color: #555;", + info: "background-color: #333;", + success: "background-color: #4a7c4a;", + error: "background-color: #7c4a4a;", launch: "background-color: #4a7c4a;", stop: "background-color: #7c4a4a;", + log: "background-color: #685434;", + clearMemory: "background-color: #555; padding: 6px 14px;", + interrupt: "background-color: #555; padding: 6px 14px;", }; export const STATUS_COLORS = { @@ -30,22 +44,23 @@ export const PULSE_ANIMATION_CSS = ` @keyframes pulse { 0% { opacity: 1; - transform: scale(1); + transform: scale(0.8); box-shadow: 0 0 0 0 rgba(240, 173, 78, 0.7); } 50% { opacity: 0.3; - transform: scale(1.3); + transform: scale(1.1); box-shadow: 0 0 0 6px rgba(240, 173, 78, 0); } 100% { opacity: 1; - transform: scale(1); + transform: scale(0.8); box-shadow: 0 0 0 0 rgba(240, 173, 78, 0); } } .status-pulsing { animation: pulse 1.2s ease-in-out infinite; + transform-origin: center; } /* Button hover effects */ diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js index 3187ee4..9fc95a0 100644 --- a/web/sidebarRenderer.js +++ b/web/sidebarRenderer.js @@ -24,7 +24,9 @@ export async function renderSidebarContent(extension, el) { el.innerHTML = ''; const loadingDiv = document.createElement("div"); loadingDiv.style.cssText = "display: flex; align-items: center; justify-content: center; height: calc(100vh - 100px); color: #888;"; - loadingDiv.innerHTML = `
◐
`; + loadingDiv.innerHTML = ` + + `; el.appendChild(loadingDiv); // Add rotation animation @@ -36,7 +38,7 @@ export async function renderSidebarContent(extension, el) { } `; document.head.appendChild(style); - loadingDiv.querySelector('div').style.animation = 'rotate 1s linear infinite'; + loadingDiv.querySelector('svg').style.animation = 'rotate 1s linear infinite'; // Preload data outside render await extension.loadConfig(); @@ -128,15 +130,15 @@ export async function renderSidebarContent(extension, el) { BUTTON_STYLES.clearMemory ); clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; - clearMemButton.style.cssText += " flex: 1;"; + clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory; const interruptButton = extension.ui.createButtonHelper( "Interrupt Workers", (e) => extension._handleInterruptWorkers(e.target), - BUTTON_STYLES.clearMemory + BUTTON_STYLES.interrupt ); interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; - interruptButton.style.cssText += " flex: 1;"; + interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt; buttonRow.appendChild(clearMemButton); buttonRow.appendChild(interruptButton); diff --git a/web/ui.js b/web/ui.js index 24f0808..de41240 100644 --- a/web/ui.js +++ b/web/ui.js @@ -591,17 +591,17 @@ export class DistributedUI { const saveBtn = this.createButton("Save", () => extension.saveWorkerSettings(worker.id), "background-color: #4a7c4a;"); - saveBtn.style.cssText += " padding: 6px 14px; font-size: 12px; font-weight: 500;"; + saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success; const cancelBtn = this.createButton("Cancel", () => extension.cancelWorkerSettings(worker.id), "background-color: #555;"); - cancelBtn.style.cssText += " padding: 6px 14px; font-size: 12px; font-weight: 500;"; + cancelBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.cancel; const deleteBtn = this.createButton("Delete", () => extension.deleteWorker(worker.id), "background-color: #7c4a4a;"); - deleteBtn.style.cssText += " padding: 6px 14px; font-size: 12px; font-weight: 500; margin-left: auto;"; + deleteBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.error + BUTTON_STYLES.marginLeftAuto; const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;"); form.appendChild(buttonGroup); @@ -728,8 +728,8 @@ export class DistributedUI { if (config.dynamic && data) { if (isRemote) { - const remoteInfo = this.createButton("Remote worker", null, "background-color: #333;"); - remoteInfo.style.cssText += " flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; color: #999; cursor: default;"; + const remoteInfo = this.createButton("Remote worker", null, BUTTON_STYLES.info); + remoteInfo.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.info + " color: #999; cursor: default;"; remoteInfo.disabled = true; controlsWrapper.appendChild(remoteInfo); } else { @@ -743,13 +743,13 @@ export class DistributedUI { const stopBtn = controls.querySelector(`#stop-${data.id}`); const logBtn = controls.querySelector(`#log-${data.id}`); - launchBtn.style.cssText += ` flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; ${BUTTON_STYLES.launch}`; + launchBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.launch; launchBtn.title = "Launch worker (runs in background with logging)"; - stopBtn.style.cssText += ` flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; ${BUTTON_STYLES.stop} display: none;`; + stopBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.stop + BUTTON_STYLES.hidden; stopBtn.title = "Stop worker"; - logBtn.style.cssText += ' flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; display: none; background-color: #685434;'; + logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.log + BUTTON_STYLES.hidden; while (controls.firstChild) { controlsWrapper.appendChild(controls.firstChild); @@ -757,7 +757,7 @@ export class DistributedUI { } } else if (config.type === 'info') { const infoBtn = this.createButton(config.text, null, config.style || ""); - infoBtn.style.cssText += " flex: 1; padding: 5px 14px; font-size: 11px; font-weight: 500; cursor: default;"; + infoBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + (config.style || BUTTON_STYLES.info) + " cursor: default;"; infoBtn.disabled = true; controlsWrapper.appendChild(infoBtn); } else if (config.type === 'ghost') { @@ -860,13 +860,13 @@ export class DistributedUI { saveBtn.textContent = "Saved!"; setTimeout(() => { saveBtn.textContent = "Save"; }, TIMEOUTS.FLASH_LONG); }, "background-color: #4a7c4a;"); - saveBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; + saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success; const cancelBtn = this.createButton("Cancel", () => { document.getElementById('master-name').value = extension.config?.master?.name || "Master"; document.getElementById('master-host').value = extension.config?.master?.host || ""; }, "background-color: #555;"); - cancelBtn.style.cssText += " padding: 6px 12px; font-size: 12px;"; + cancelBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.cancel; const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn], " margin-top: 8px;"); settingsForm.appendChild(buttonGroup); From d4ee572eccbdfbc0597e67ce283d8878d3dc3bc6 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 08:19:13 +1000 Subject: [PATCH 10/93] Runpod support --- distributed.py | 8 +++- web/main.js | 29 +++++++++--- web/ui.js | 125 ++++++++++++++++++++++++++++++++++++++----------- 3 files changed, 128 insertions(+), 34 deletions(-) diff --git a/distributed.py b/distributed.py index e9a7c0f..1a1a14e 100644 --- a/distributed.py +++ b/distributed.py @@ -297,6 +297,11 @@ async def update_worker_endpoint(request): worker.pop("extra_args", None) else: worker["extra_args"] = data["extra_args"] + + # Handle type field + if "type" in data: + worker["type"] = data["type"] + worker_found = True break @@ -310,7 +315,8 @@ async def update_worker_endpoint(request): "port": data["port"], "cuda_device": data["cuda_device"], "enabled": data.get("enabled", False), - "extra_args": data.get("extra_args", "") + "extra_args": data.get("extra_args", ""), + "type": data.get("type", "local") } if "workers" not in config: config["workers"] = [] diff --git a/web/main.js b/web/main.js index ecd7efe..efc1a63 100644 --- a/web/main.js +++ b/web/main.js @@ -266,13 +266,17 @@ class DistributedExtension { getWorkerUrl(worker, endpoint = '') { const host = worker.host || window.location.hostname; - // Simple rule: port 443 = HTTPS, anything else = HTTP - const useHttps = worker.port === 443; + // Cloud workers always use HTTPS + const isCloud = worker.type === 'cloud'; + // Simple rule: port 443 = HTTPS, cloud workers = HTTPS, anything else = HTTP + const useHttps = isCloud || worker.port === 443; const protocol = useHttps ? 'https' : 'http'; + // Cloud workers use default HTTPS port (443) + const effectivePort = isCloud ? 443 : worker.port; // Don't add port for standard HTTPS/HTTP ports const defaultPort = useHttps ? 443 : 80; - const port = worker.port === defaultPort ? '' : `:${worker.port}`; + const port = effectivePort === defaultPort ? '' : `:${effectivePort}`; return `${protocol}://${host}${port}${endpoint}`; } @@ -695,10 +699,19 @@ class DistributedExtension { } isRemoteWorker(worker) { + // Check if explicitly marked as cloud worker + if (worker.type === "cloud") { + return true; + } + // Otherwise check by host (backward compatibility) const host = worker.host || window.location.hostname; return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; } + isCloudWorker(worker) { + return worker.type === "cloud"; + } + getMasterUrl() { // Always use the detected/configured master IP for consistency if (this.config?.master?.host) { @@ -873,7 +886,9 @@ class DistributedExtension { // Get form values const name = document.getElementById(`name-${workerId}`).value; - const isRemote = document.getElementById(`remote-${workerId}`).checked; + const workerType = document.getElementById(`worker-type-${workerId}`).value; + const isRemote = workerType === 'remote' || workerType === 'cloud'; + const isCloud = workerType === 'cloud'; const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; const port = parseInt(document.getElementById(`port-${workerId}`).value); const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); @@ -890,7 +905,7 @@ class DistributedExtension { return; } - if (isRemote && !host.trim()) { + if ((workerType === 'remote' || workerType === 'cloud') && !host.trim()) { app.extensionManager.toast.add({ severity: "error", summary: "Validation Error", @@ -900,7 +915,7 @@ class DistributedExtension { return; } - if (isNaN(port) || port < 1 || port > 65535) { + if (!isCloud && (isNaN(port) || port < 1 || port > 65535)) { app.extensionManager.toast.add({ severity: "error", summary: "Validation Error", @@ -963,6 +978,7 @@ class DistributedExtension { try { await this.api.updateWorker(workerId, { name: name.trim(), + type: workerType, host: isRemote ? host.trim() : null, port: port, cuda_device: isRemote ? null : cudaDevice, @@ -971,6 +987,7 @@ class DistributedExtension { // Update local config worker.name = name.trim(); + worker.type = workerType; if (isRemote) { worker.host = host.trim(); delete worker.cuda_device; diff --git a/web/ui.js b/web/ui.js index de41240..2421cdf 100644 --- a/web/ui.js +++ b/web/ui.js @@ -49,7 +49,12 @@ const cardConfigs = { }, infoText: (data, extension) => { const isRemote = extension.isRemoteWorker(data); - if (isRemote) { + const isCloud = data.type === 'cloud'; + + if (isCloud) { + // For cloud workers, don't show port (it's always 443) + return `${data.name}
${data.host}`; + } else if (isRemote) { return `${data.name}
${data.host}:${data.port}`; } else { const cudaInfo = data.cuda_device !== undefined ? `CUDA ${data.cuda_device} • ` : ''; @@ -525,66 +530,130 @@ export class DistributedUI { const nameGroup = this.createFormGroup("Name:", worker.name, `name-${worker.id}`); form.appendChild(nameGroup.group); - // Remote checkbox - const remoteGroup = document.createElement("div"); - remoteGroup.style.cssText = "display: flex; align-items: center; gap: 8px; margin: 5px 0;"; + // Worker type dropdown + const typeGroup = document.createElement("div"); + typeGroup.style.cssText = "display: flex; flex-direction: column; gap: 4px; margin: 5px 0;"; - const remoteCheckbox = document.createElement("input"); - remoteCheckbox.type = "checkbox"; - remoteCheckbox.id = `remote-${worker.id}`; - remoteCheckbox.checked = extension.isRemoteWorker(worker); - remoteCheckbox.onchange = (e) => { - const isRemote = e.target.checked; + const typeLabel = document.createElement("label"); + typeLabel.htmlFor = `worker-type-${worker.id}`; + typeLabel.textContent = "Worker Type:"; + typeLabel.style.cssText = "font-size: 12px; color: #ccc;"; + + const typeSelect = document.createElement("select"); + typeSelect.id = `worker-type-${worker.id}`; + typeSelect.style.cssText = "padding: 4px 8px; background: #333; color: #fff; border: 1px solid #555; border-radius: 4px; font-size: 12px;"; + + // Create options + const localOption = document.createElement("option"); + localOption.value = "local"; + localOption.textContent = "Local"; + + const remoteOption = document.createElement("option"); + remoteOption.value = "remote"; + remoteOption.textContent = "Remote"; + + const cloudOption = document.createElement("option"); + cloudOption.value = "cloud"; + cloudOption.textContent = "Cloud"; + + typeSelect.appendChild(localOption); + typeSelect.appendChild(remoteOption); + typeSelect.appendChild(cloudOption); + + // Create powered by Runpod text (initially hidden) + const runpodText = document.createElement("a"); + runpodText.id = `runpod-text-${worker.id}`; + runpodText.href = "https://www.runpod.io"; + runpodText.target = "_blank"; + runpodText.textContent = "Deploy cloud worker with Runpod"; + runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;"; + + typeSelect.onchange = (e) => { + const workerType = e.target.value; // Show/hide relevant fields const hostGroup = document.getElementById(`host-group-${worker.id}`); + const hostInput = document.getElementById(`host-${worker.id}`); + const portGroup = document.getElementById(`port-group-${worker.id}`); + const portInput = document.getElementById(`port-${worker.id}`); const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); const argsGroup = document.getElementById(`args-group-${worker.id}`); + const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`); - if (isRemote) { + if (workerType === "local") { + hostGroup.style.display = "none"; + portGroup.style.display = "flex"; + cudaGroup.style.display = "flex"; + argsGroup.style.display = "flex"; + runpodTextElem.style.display = "none"; + } else if (workerType === "remote") { hostGroup.style.display = "flex"; + portGroup.style.display = "flex"; cudaGroup.style.display = "none"; argsGroup.style.display = "none"; + runpodTextElem.style.display = "none"; + // Update placeholder for remote workers + hostInput.placeholder = "e.g., 192.168.1.100"; // If switching to remote and host is localhost, clear it - const hostInput = document.getElementById(`host-${worker.id}`); if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { hostInput.value = ""; } - } else { - hostGroup.style.display = "none"; - cudaGroup.style.display = "flex"; - argsGroup.style.display = "flex"; + } else if (workerType === "cloud") { + hostGroup.style.display = "flex"; + portGroup.style.display = "flex"; // Keep port visible for cloud workers + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + runpodTextElem.style.display = "block"; + // Update placeholder for cloud workers + hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com"; + // Set port to 443 for cloud workers + portInput.value = "443"; + // If switching to cloud and host is localhost, clear it + if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { + hostInput.value = ""; + } } }; - const remoteLabel = document.createElement("label"); - remoteLabel.htmlFor = `remote-${worker.id}`; - remoteLabel.textContent = "Remote Worker"; - remoteLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + typeGroup.appendChild(typeLabel); + typeGroup.appendChild(typeSelect); + typeGroup.appendChild(runpodText); + form.appendChild(typeGroup); - remoteGroup.appendChild(remoteCheckbox); - remoteGroup.appendChild(remoteLabel); - form.appendChild(remoteGroup); + // Set initial value based on worker type + if (worker.type === "cloud") { + typeSelect.value = "cloud"; + // Show Runpod text immediately for cloud workers + runpodText.style.display = "block"; + } else if (extension.isRemoteWorker(worker)) { + typeSelect.value = "remote"; + } else { + typeSelect.value = "local"; + } + + // Trigger initial state + typeSelect.dispatchEvent(new Event('change')); // Host field (only for remote workers) const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); hostGroup.group.id = `host-group-${worker.id}`; - hostGroup.group.style.display = extension.isRemoteWorker(worker) ? "flex" : "none"; + hostGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "flex" : "none"; form.appendChild(hostGroup.group); // Port field const portGroup = this.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number"); + portGroup.group.id = `port-group-${worker.id}`; form.appendChild(portGroup.group); // CUDA Device field (only for local workers) const cudaGroup = this.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number"); cudaGroup.group.id = `cuda-group-${worker.id}`; - cudaGroup.group.style.display = extension.isRemoteWorker(worker) ? "none" : "flex"; + cudaGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex"; form.appendChild(cudaGroup.group); // Extra Args field (only for local workers) const argsGroup = this.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`); argsGroup.group.id = `args-group-${worker.id}`; - argsGroup.group.style.display = extension.isRemoteWorker(worker) ? "none" : "flex"; + argsGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex"; form.appendChild(argsGroup.group); // Buttons @@ -728,7 +797,9 @@ export class DistributedUI { if (config.dynamic && data) { if (isRemote) { - const remoteInfo = this.createButton("Remote worker", null, BUTTON_STYLES.info); + const isCloud = data.type === 'cloud'; + const workerTypeText = isCloud ? "Cloud worker" : "Remote worker"; + const remoteInfo = this.createButton(workerTypeText, null, BUTTON_STYLES.info); remoteInfo.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.info + " color: #999; cursor: default;"; remoteInfo.disabled = true; controlsWrapper.appendChild(remoteInfo); From 9650280c7f50898720bf9773e8874d75e7c20846 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 08:20:05 +1000 Subject: [PATCH 11/93] Runpod support --- web/main.js | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/web/main.js b/web/main.js index efc1a63..4cce62b 100644 --- a/web/main.js +++ b/web/main.js @@ -268,17 +268,29 @@ class DistributedExtension { // Cloud workers always use HTTPS const isCloud = worker.type === 'cloud'; - // Simple rule: port 443 = HTTPS, cloud workers = HTTPS, anything else = HTTP - const useHttps = isCloud || worker.port === 443; + + // Detect if we're running on Runpod (for local workers on Runpod infrastructure) + const isRunpodProxy = host.endsWith('.proxy.runpod.net'); + + // For local workers on Runpod, construct the port-specific proxy URL + let finalHost = host; + if (!worker.host && isRunpodProxy) { + // Extract pod ID from current host (format: podId-xxxx.proxy.runpod.net) + const podId = host.split('-')[0]; + const domain = host.substring(host.indexOf('-') + 1); + finalHost = `${podId}-${worker.port}.${domain}`; + } + + // Determine protocol: HTTPS for cloud, Runpod proxies, or port 443 + const useHttps = isCloud || isRunpodProxy || worker.port === 443; const protocol = useHttps ? 'https' : 'http'; - // Cloud workers use default HTTPS port (443) - const effectivePort = isCloud ? 443 : worker.port; - // Don't add port for standard HTTPS/HTTP ports + // Only add port if non-standard const defaultPort = useHttps ? 443 : 80; - const port = effectivePort === defaultPort ? '' : `:${effectivePort}`; + const needsPort = !isRunpodProxy && worker.port !== defaultPort; + const portStr = needsPort ? `:${worker.port}` : ''; - return `${protocol}://${host}${port}${endpoint}`; + return `${protocol}://${finalHost}${portStr}${endpoint}`; } async checkWorkerStatus(worker) { From fa09a2da10dfc63cecfbec1f8d4a0516e923b911 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 11:16:34 +1000 Subject: [PATCH 12/93] Image Batch Divider --- __init__.py | 30 +++++++++++ distributed.py | 80 ++++++++++++++++++++++++++++- utils/constants.py | 4 +- web/ui.js | 123 ++++++++++++++++++++++++--------------------- 4 files changed, 177 insertions(+), 60 deletions(-) diff --git a/__init__.py b/__init__.py index 9bc2dde..18d9765 100644 --- a/__init__.py +++ b/__init__.py @@ -4,6 +4,36 @@ import sys # Add the directory to Python path sys.path.append(os.path.dirname(__file__)) +# Patch ComfyUI execution validation for DistributedVideoCollector +try: + import execution + from .distributed import ImageBatchDivider + + # Store original validate function if it exists + if hasattr(execution, 'validate_outputs'): + original_validate_outputs = execution.validate_outputs + + def patched_validate_outputs(executor, node_id, result, node_class): + if node_class == ImageBatchDivider: + return # Skip validation for our dynamic output node + return original_validate_outputs(executor, node_id, result, node_class) + + execution.validate_outputs = patched_validate_outputs + else: + # Fallback: patch the main execute method if validate_outputs doesn't exist + if hasattr(execution.PromptExecutor, 'execute'): + original_execute = execution.PromptExecutor.execute + + def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]): + # This is a more complex patch - for now just call original + # The ByPassTypeTuple should handle most validation issues + return original_execute(self, prompt, prompt_id, extra_data, execute_outputs) + + execution.PromptExecutor.execute = patched_execute + +except ImportError: + pass # ComfyUI execution module not available during import + # Import everything needed from the main module from .distributed import ( NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, diff --git a/distributed.py b/distributed.py index 1a1a14e..bf2b2c1 100644 --- a/distributed.py +++ b/distributed.py @@ -1638,11 +1638,87 @@ class DistributedSeed: # Fallback: return original seed return (seed,) +# Define ByPassTypeTuple for flexible return types +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + +any_type = AnyType("*") + +class ByPassTypeTuple(tuple): + def __getitem__(self, index): + if index > 0: + index = 0 + item = super().__getitem__(index) + if isinstance(item, str): + return any_type + return item + +class ImageBatchDivider: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "divide_by": ("INT", { + "default": 2, + "min": 1, + "max": 10, + "step": 1, + "display": "number", + "tooltip": "Number of parts to divide the batch into" + }), + } + } + + RETURN_TYPES = ByPassTypeTuple(("IMAGE", )) # Flexible for variable outputs + RETURN_NAMES = ByPassTypeTuple(tuple([f"batch_{i+1}" for i in range(10)])) + FUNCTION = "divide_batch" + OUTPUT_NODE = True + CATEGORY = "image" + + def divide_batch(self, images, divide_by): + import torch + + # Use divide_by directly + total_splits = divide_by + + if total_splits > 10: + total_splits = 10 # Cap to max + + # Get total number of frames + total_frames = images.shape[0] + frames_per_split = total_frames // total_splits + remainder = total_frames % total_splits + + outputs = [] + start_idx = 0 + + for i in range(total_splits): + current_frames = frames_per_split + (1 if i < remainder else 0) + end_idx = start_idx + current_frames + split_frames = images[start_idx:end_idx] + outputs.append(split_frames) + start_idx = end_idx + + # Pad with empty tensors up to max (10) to match potential RETURN_TYPES len + empty_shape = (1, images.shape[1], images.shape[2], images.shape[3]) if total_frames > 0 else (1, 512, 512, 3) + empty_tensor = torch.zeros(empty_shape, dtype=images.dtype if total_frames > 0 else torch.float32, + device=images.device if total_frames > 0 else 'cpu') + + while len(outputs) < 10: + outputs.append(empty_tensor) + + return tuple(outputs) + + NODE_CLASS_MAPPINGS = { "DistributedCollector": DistributedCollectorNode, - "DistributedSeed": DistributedSeed + "DistributedSeed": DistributedSeed, + "ImageBatchDivider": ImageBatchDivider } NODE_DISPLAY_NAME_MAPPINGS = { "DistributedCollector": "Distributed Collector", - "DistributedSeed": "Distributed Seed" + "DistributedSeed": "Distributed Seed", + "ImageBatchDivider": "Image Batch Divider" } diff --git a/utils/constants.py b/utils/constants.py index a199f24..ff1869d 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -3,8 +3,8 @@ Shared constants for ComfyUI-Distributed. """ # Timeouts (in seconds) -WORKER_JOB_TIMEOUT = 10.0 -TILE_COLLECTION_TIMEOUT = 60.0 +WORKER_JOB_TIMEOUT = 30.0 +TILE_COLLECTION_TIMEOUT = 30.0 TILE_WAIT_TIMEOUT = 30.0 PROCESS_TERMINATION_TIMEOUT = 5.0 diff --git a/web/ui.js b/web/ui.js index 2421cdf..bdbc566 100644 --- a/web/ui.js +++ b/web/ui.js @@ -568,50 +568,58 @@ export class DistributedUI { runpodText.textContent = "Deploy cloud worker with Runpod"; runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;"; - typeSelect.onchange = (e) => { - const workerType = e.target.value; - // Show/hide relevant fields - const hostGroup = document.getElementById(`host-group-${worker.id}`); - const hostInput = document.getElementById(`host-${worker.id}`); - const portGroup = document.getElementById(`port-group-${worker.id}`); - const portInput = document.getElementById(`port-${worker.id}`); - const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); - const argsGroup = document.getElementById(`args-group-${worker.id}`); - const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`); - - if (workerType === "local") { - hostGroup.style.display = "none"; - portGroup.style.display = "flex"; - cudaGroup.style.display = "flex"; - argsGroup.style.display = "flex"; - runpodTextElem.style.display = "none"; - } else if (workerType === "remote") { - hostGroup.style.display = "flex"; - portGroup.style.display = "flex"; - cudaGroup.style.display = "none"; - argsGroup.style.display = "none"; - runpodTextElem.style.display = "none"; - // Update placeholder for remote workers - hostInput.placeholder = "e.g., 192.168.1.100"; - // If switching to remote and host is localhost, clear it - if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { - hostInput.value = ""; + // Store the onchange function to be assigned later + const createOnChangeHandler = () => { + return (e) => { + const workerType = e.target.value; + // Show/hide relevant fields + const hostGroup = document.getElementById(`host-group-${worker.id}`); + const hostInput = document.getElementById(`host-${worker.id}`); + const portGroup = document.getElementById(`port-group-${worker.id}`); + const portInput = document.getElementById(`port-${worker.id}`); + const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); + const argsGroup = document.getElementById(`args-group-${worker.id}`); + const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`); + + // Check if elements exist before accessing them + if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput) { + return; // Elements not ready yet } - } else if (workerType === "cloud") { - hostGroup.style.display = "flex"; - portGroup.style.display = "flex"; // Keep port visible for cloud workers - cudaGroup.style.display = "none"; - argsGroup.style.display = "none"; - runpodTextElem.style.display = "block"; - // Update placeholder for cloud workers - hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com"; - // Set port to 443 for cloud workers - portInput.value = "443"; - // If switching to cloud and host is localhost, clear it - if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { - hostInput.value = ""; + + if (workerType === "local") { + hostGroup.style.display = "none"; + portGroup.style.display = "flex"; + cudaGroup.style.display = "flex"; + argsGroup.style.display = "flex"; + runpodTextElem.style.display = "none"; + } else if (workerType === "remote") { + hostGroup.style.display = "flex"; + portGroup.style.display = "flex"; + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + runpodTextElem.style.display = "none"; + // Update placeholder for remote workers + hostInput.placeholder = "e.g., 192.168.1.100"; + // If switching to remote and host is localhost, clear it + if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { + hostInput.value = ""; + } + } else if (workerType === "cloud") { + hostGroup.style.display = "flex"; + portGroup.style.display = "flex"; // Keep port visible for cloud workers + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + runpodTextElem.style.display = "block"; + // Update placeholder for cloud workers + hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com"; + // Set port to 443 for cloud workers + portInput.value = "443"; + // If switching to cloud and host is localhost, clear it + if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { + hostInput.value = ""; + } } - } + }; }; typeGroup.appendChild(typeLabel); @@ -619,20 +627,6 @@ export class DistributedUI { typeGroup.appendChild(runpodText); form.appendChild(typeGroup); - // Set initial value based on worker type - if (worker.type === "cloud") { - typeSelect.value = "cloud"; - // Show Runpod text immediately for cloud workers - runpodText.style.display = "block"; - } else if (extension.isRemoteWorker(worker)) { - typeSelect.value = "remote"; - } else { - typeSelect.value = "local"; - } - - // Trigger initial state - typeSelect.dispatchEvent(new Event('change')); - // Host field (only for remote workers) const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); hostGroup.group.id = `host-group-${worker.id}`; @@ -675,6 +669,23 @@ export class DistributedUI { const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;"); form.appendChild(buttonGroup); + // Assign the onchange handler now that all elements are created + typeSelect.onchange = createOnChangeHandler(); + + // Set initial value and trigger state after all DOM elements are created + if (worker.type === "cloud") { + typeSelect.value = "cloud"; + // Show Runpod text immediately for cloud workers + runpodText.style.display = "block"; + } else if (extension.isRemoteWorker(worker)) { + typeSelect.value = "remote"; + } else { + typeSelect.value = "local"; + } + + // Trigger initial state now that all elements exist + typeSelect.dispatchEvent(new Event('change')); + return form; } From b771e963a51231a84ef7969a36609f7f0491c328 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 12:22:15 +1000 Subject: [PATCH 13/93] optimised large batch collection --- distributed.py | 215 +++++++++++++++++++++++++++++++++++++------------ 1 file changed, 162 insertions(+), 53 deletions(-) diff --git a/distributed.py b/distributed.py index bf2b2c1..aceab71 100644 --- a/distributed.py +++ b/distributed.py @@ -14,6 +14,7 @@ import platform import time import atexit import signal +from comfy.utils import ProgressBar # Import shared utilities from .utils.logging import debug_log, log @@ -1299,38 +1300,76 @@ async def job_complete_endpoint(request): try: data = await request.post() multi_job_id = data.get('multi_job_id') - image_file = data.get('image') worker_id = data.get('worker_id') - image_index = data.get('image_index') is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - debug_log(f"job_complete received - job_id: {multi_job_id}, worker: {worker_id}, index: {image_index}, is_last: {is_last}") - - if not all([multi_job_id, image_file]): - return await handle_api_error(request, "Missing job_id or image data", 400) - - # Process image with error handling - try: - img_data = image_file.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - # Convert to tensor using utility function - tensor = pil_to_tensor(img) - # Ensure tensor is contiguous - tensor = ensure_contiguous(tensor) - except Exception as e: - log(f"Error processing image from worker {worker_id}: {e}") - return await handle_api_error(request, f"Image processing error: {e}", 400) + # Check for batch mode + batch_size = int(data.get('batch_size', 0)) + tensors = [] + + if batch_size > 0: + # Batch mode: Extract multiple images + debug_log(f"job_complete received batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") + + for i in range(batch_size): + image_field = data.get(f'image_{i}') + if image_field is None: + return await handle_api_error(request, f"Missing image_{i}", 400) + + try: + img_data = image_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + tensor = pil_to_tensor(img) + tensor = ensure_contiguous(tensor) + tensors.append(tensor) + except Exception as e: + log(f"Error processing image {i} from worker {worker_id}: {e}") + return await handle_api_error(request, f"Image processing error: {e}", 400) + else: + # Fallback for single-image mode (backward compatibility) + image_file = data.get('image') + image_index = data.get('image_index') + + debug_log(f"job_complete received single - job_id: {multi_job_id}, worker: {worker_id}, index: {image_index}, is_last: {is_last}") + + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + try: + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + tensor = pil_to_tensor(img) + tensor = ensure_contiguous(tensor) + tensors = [tensor] + except Exception as e: + log(f"Error processing image from worker {worker_id}: {e}") + return await handle_api_error(request, f"Image processing error: {e}", 400) + # Put batch into queue async with prompt_server.distributed_jobs_lock: debug_log(f"Current pending jobs: {list(prompt_server.distributed_pending_jobs.keys())}") if multi_job_id in prompt_server.distributed_pending_jobs: - await prompt_server.distributed_pending_jobs[multi_job_id].put({ - 'tensor': tensor, - 'worker_id': worker_id, - 'image_index': int(image_index) if image_index else 0, - 'is_last': is_last - }) - debug_log(f"Received result for job {multi_job_id} from worker {worker_id} (last: {is_last})") + if batch_size > 0: + # Put batch as single item + await prompt_server.distributed_pending_jobs[multi_job_id].put({ + 'worker_id': worker_id, + 'tensors': tensors, + 'is_last': is_last + }) + debug_log(f"Received batch result for job {multi_job_id} from worker {worker_id}, size={len(tensors)}") + else: + # Put single image (backward compat) + await prompt_server.distributed_pending_jobs[multi_job_id].put({ + 'tensor': tensors[0], + 'worker_id': worker_id, + 'image_index': int(image_index) if image_index else 0, + 'is_last': is_last + }) + debug_log(f"Received single result for job {multi_job_id} from worker {worker_id}") + debug_log(f"Queue size after put: {prompt_server.distributed_pending_jobs[multi_job_id].qsize()}") return web.json_response({"status": "success"}) else: @@ -1373,6 +1412,33 @@ class DistributedCollectorNode: ) return result + async def send_batch_to_master(self, image_batch, multi_job_id, master_url, worker_id): + """Send entire image batch in one POST using multi-part FormData""" + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', 'True') # Full batch, so last for this worker + data.add_field('batch_size', str(image_batch.shape[0])) # Number of images in batch + + # Loop to add each image as a separate field + for i in range(image_batch.shape[0]): + # Convert tensor slice to PIL + img = tensor_to_pil(image_batch[i:i+1], 0) # tensor_to_pil handles single-image batch + byte_io = io.BytesIO() + img.save(byte_io, format='PNG', compress_level=0) # Lossless, no compression for speed + byte_io.seek(0) + data.add_field(f'image_{i}', byte_io, filename=f'image_{i}.png', content_type='image/png') + + try: + session = await get_client_session() + url = f"{master_url}/distributed/job_complete" + debug_log(f"Worker - Sending batch of {image_batch.shape[0]} images to URL: {url}") + async with session.post(url, data=data) as response: + response.raise_for_status() + except Exception as e: + log(f"Worker - Failed to send batch to master: {e}") + debug_log(f"Worker - Full error details: URL={url}") + async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index, worker_id, is_last=False): """Helper method to send a single image to the master""" # Ensure we handle the tensor shape correctly (H, W, C) @@ -1405,14 +1471,9 @@ class DistributedCollectorNode: async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): if is_worker: - # Worker mode: send images to master + # Worker mode: send images to master in a single batch debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") - - # Send all images, marking the last one - for i in range(images.shape[0]): - is_last = (i == images.shape[0] - 1) - await self.send_image_to_master(images[i], multi_job_id, master_url, i, worker_id, is_last) - + await self.send_batch_to_master(images, multi_job_id, master_url, worker_id) return (images,) else: # Master mode: collect images from workers @@ -1457,7 +1518,12 @@ class DistributedCollectorNode: q = prompt_server.distributed_pending_jobs[multi_job_id] initial_size = q.qsize() debug_log(f"Master - Queue size before collection: {initial_size}") - + + # NEW: Initialize progress bar with estimated total images + expected_total = master_batch_size + (num_workers * worker_batch_size) # Estimate; adjust if batch sizes vary + p = ProgressBar(expected_total) + p.update(master_batch_size) # Start with master's contribution (already "collected") + while len(workers_done) < num_workers: try: # Get the queue again each time to ensure we have the right reference @@ -1469,21 +1535,48 @@ class DistributedCollectorNode: result = await asyncio.wait_for(q.get(), timeout=timeout) worker_id = result['worker_id'] - image_index = result['image_index'] - tensor = result['tensor'] is_last = result.get('is_last', False) - debug_log(f"Master - Got result from worker {worker_id}, image {image_index}, is_last={is_last}") - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - # Debug: Check worker tensor properties - if collected_count == 0: # Only print once - debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") - - collected_count += 1 + # Check if batch mode + tensors = result.get('tensors', []) + if tensors: + # Batch mode + debug_log(f"Master - Got batch from worker {worker_id}, size={len(tensors)}, is_last={is_last}") + + if worker_id not in worker_images: + worker_images[worker_id] = {} + + # Add all tensors with sequential indices + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + + # Debug: Check worker tensor properties + if collected_count == 0: # Only print once + debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") + + collected_count += len(tensors) + + # Update progress for batch + p.update(len(tensors)) + else: + # Single image mode (backward compat) + image_index = result['image_index'] + tensor = result['tensor'] + + debug_log(f"Master - Got single result from worker {worker_id}, image {image_index}, is_last={is_last}") + + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor + + # Debug: Check worker tensor properties + if collected_count == 0: # Only print once + debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") + + collected_count += 1 + + # Update progress for single image + p.update(1) # Once we start receiving images, use shorter timeout timeout = WORKER_JOB_TIMEOUT @@ -1519,15 +1612,31 @@ class DistributedCollectorNode: # Process them for item in remaining_items: worker_id = item['worker_id'] - image_index = item['image_index'] - tensor = item['tensor'] is_last = item.get('is_last', False) - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - collected_count += 1 + # Check if batch mode + tensors = item.get('tensors', []) + if tensors: + # Batch mode + if worker_id not in worker_images: + worker_images[worker_id] = {} + + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + + collected_count += len(tensors) + p.update(len(tensors)) + else: + # Single image mode + image_index = item['image_index'] + tensor = item['tensor'] + + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor + + collected_count += 1 + p.update(1) if is_last: workers_done.add(worker_id) From 5a2bde287ea309496c3066a4aeea6e764fedcde5 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 12:34:32 +1000 Subject: [PATCH 14/93] distributed upscaler speed boost --- distributed_upscale.py | 226 +++++++++++++++++++++++++++++++++-------- 1 file changed, 182 insertions(+), 44 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 7bf702b..4a40dde 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -174,10 +174,10 @@ class UltimateSDUpscaleDistributed: 'extracted_height': eh }) - # Send results back using the server's event loop + # Send results back using the server's event loop in a single batch try: run_async_in_server_loop( - self._async_send_tiles_to_master(processed_tiles, multi_job_id, master_url, + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), timeout=TILE_SEND_TIMEOUT ) @@ -186,18 +186,6 @@ class UltimateSDUpscaleDistributed: return (upscaled_image,) - async def _async_send_tiles_to_master(self, processed_tiles, multi_job_id, master_url, - padding, worker_id): - """Async helper to send tiles to master.""" - for i, tile_data in enumerate(processed_tiles): - is_last = (i == len(processed_tiles) - 1) - await self.send_tile_to_master( - tile_data['tile'], multi_job_id, master_url, - tile_data['tile_idx'], tile_data['x'], tile_data['y'], - tile_data['extracted_width'], tile_data['extracted_height'], - padding, worker_id, is_last - ) - def process_master(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, @@ -433,12 +421,33 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Loop status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected, workers_done set: {workers_done}") try: result = await asyncio.wait_for(q.get(), timeout=timeout) - tile_idx = result['tile_idx'] worker_id = result['worker_id'] is_last = result.get('is_last', False) - collected_tiles[tile_idx] = result - debug_log(f"UltimateSDUpscale Master - Received tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + # Check if batch mode + tiles = result.get('tiles', []) + if tiles: + # Batch mode + debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") + + for tile_data in tiles: + tile_idx = tile_data['tile_idx'] + # Store the full tile data including metadata + collected_tiles[tile_idx] = { + 'tensor': tile_data['tensor'], + 'tile_idx': tile_idx, + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': tile_data['padding'], + 'worker_id': worker_id + } + else: + # Single tile mode (backward compat) + tile_idx = result['tile_idx'] + collected_tiles[tile_idx] = result + debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") if is_last: workers_done.add(worker_id) @@ -603,6 +612,59 @@ class UltimateSDUpscaleDistributed: return result.convert('RGB') + async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, + padding, worker_id): + """Send all processed tiles to master in a single batch request.""" + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', 'True') # Full batch, so last for this worker + data.add_field('batch_size', str(len(processed_tiles))) + data.add_field('padding', str(padding)) + + # Add each tile as a separate field with metadata + for i, tile_data in enumerate(processed_tiles): + # Convert tensor to PIL + img = tensor_to_pil(tile_data['tile'], 0) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + + # Add image field + data.add_field(f'tile_{i}', byte_io, filename=f'tile_{i}.png', content_type='image/png') + + # Add metadata for this tile + data.add_field(f'tile_{i}_idx', str(tile_data['tile_idx'])) + data.add_field(f'tile_{i}_x', str(tile_data['x'])) + data.add_field(f'tile_{i}_y', str(tile_data['y'])) + data.add_field(f'tile_{i}_width', str(tile_data['extracted_width'])) + data.add_field(f'tile_{i}_height', str(tile_data['extracted_height'])) + + # Retry logic with exponential backoff + max_retries = 5 + retry_delay = 0.5 + + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/tile_complete" + + debug_log(f"UltimateSDUpscale Worker - Sending batch of {len(processed_tiles)} tiles to {url}, attempt {attempt + 1}") + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"UltimateSDUpscale Worker - Successfully sent batch of {len(processed_tiles)} tiles") + return + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"UltimateSDUpscale Worker - Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 # Exponential backoff + else: + log(f"UltimateSDUpscale Worker - Failed to send batch after {max_retries} attempts: {e}") + raise + async def send_tile_to_master(self, tile_tensor, multi_job_id, master_url, tile_idx, x, y, extracted_width, extracted_height, padding, worker_id, is_last=False): @@ -714,43 +776,119 @@ async def tile_complete_endpoint(request): try: data = await request.post() multi_job_id = data.get('multi_job_id') - image_file = data.get('image') worker_id = data.get('worker_id') - tile_idx = int(data.get('tile_idx', 0)) - x = int(data.get('x', 0)) - y = int(data.get('y', 0)) - extracted_width = int(data.get('extracted_width', 512)) - extracted_height = int(data.get('extracted_height', 512)) - padding = int(data.get('padding', 32)) is_last = data.get('is_last', 'False').lower() == 'true' - - if not all([multi_job_id, image_file]): - return await handle_api_error(request, "Missing job_id or image data", 400) - - # Process image - img_data = image_file.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale - tile_complete: Checking distributed_pending_tile_jobs for job {multi_job_id}") - debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + + # Check for batch mode + batch_size = int(data.get('batch_size', 0)) + tiles = [] + + if batch_size > 0: + # Batch mode: Extract multiple tiles + padding = int(data.get('padding', 32)) + debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + return await handle_api_error(request, f"Missing tile_{i}", 400) + + # Get metadata for this tile + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + try: + # Process image + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + except Exception as e: + log(f"Error processing tile {i} from worker {worker_id}: {e}") + return await handle_api_error(request, f"Tile processing error: {e}", 400) + else: + # Single tile mode (backward compatibility) + image_file = data.get('image') + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + tile_idx = int(data.get('tile_idx', 0)) + x = int(data.get('x', 0)) + y = int(data.get('y', 0)) + extracted_width = int(data.get('extracted_width', 512)) + extracted_height = int(data.get('extracted_height', 512)) + padding = int(data.get('padding', 32)) + + debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}") + + try: + # Process image + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles = [{ 'tensor': tensor, - 'worker_id': worker_id, 'tile_idx': tile_idx, 'x': x, 'y': y, 'extracted_width': extracted_width, 'extracted_height': extracted_height, - 'padding': padding, - 'is_last': is_last - }) - debug_log(f"UltimateSDUpscale - Received tile {tile_idx} for job {multi_job_id} from worker {worker_id} (last: {is_last})") - debug_log(f"UltimateSDUpscale - tile_complete: Successfully queued tile {tile_idx}") + 'padding': padding + }] + except Exception as e: + log(f"Error processing tile from worker {worker_id}: {e}") + return await handle_api_error(request, f"Tile processing error: {e}", 400) + + # Put tiles into queue + async with prompt_server.distributed_tile_jobs_lock: + debug_log(f"UltimateSDUpscale - tile_complete: Checking distributed_pending_tile_jobs for job {multi_job_id}") + debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") + + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + if batch_size > 0: + # Put batch as single item + await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last + }) + debug_log(f"UltimateSDUpscale - Received batch of {len(tiles)} tiles for job {multi_job_id} from worker {worker_id}") + else: + # Put single tile (backward compat) + tile_data = tiles[0] + await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + 'tensor': tile_data['tensor'], + 'worker_id': worker_id, + 'tile_idx': tile_data['tile_idx'], + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': tile_data['padding'], + 'is_last': is_last + }) + debug_log(f"UltimateSDUpscale - Received single tile {tile_data['tile_idx']} for job {multi_job_id} from worker {worker_id}") + return web.json_response({"status": "success"}) else: debug_log(f"UltimateSDUpscale - API Error: Job {multi_job_id} not found in distributed_pending_tile_jobs") From b4ed92a363fab14d944b87c610aa9a0b4c87c085 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 13:28:08 +1000 Subject: [PATCH 15/93] handle large batches better --- distributed.py | 185 +++++++++++++++--------- distributed_upscale.py | 315 +++++++++++++++++++++++------------------ utils/constants.py | 5 +- 3 files changed, 298 insertions(+), 207 deletions(-) diff --git a/distributed.py b/distributed.py index aceab71..5dc0b1e 100644 --- a/distributed.py +++ b/distributed.py @@ -26,7 +26,7 @@ from .utils.async_helpers import run_async_in_server_loop from .utils.constants import ( WORKER_JOB_TIMEOUT, PROCESS_TERMINATION_TIMEOUT, WORKER_CHECK_INTERVAL, STATUS_CHECK_INTERVAL, CHUNK_SIZE, LOG_TAIL_BYTES, WORKER_LOG_PATTERN, - WORKER_STARTUP_DELAY, PROCESS_WAIT_TIMEOUT, MEMORY_CLEAR_DELAY + WORKER_STARTUP_DELAY, PROCESS_WAIT_TIMEOUT, MEMORY_CLEAR_DELAY, MAX_BATCH ) # Try to import psutil for better process management @@ -1314,20 +1314,74 @@ async def job_complete_endpoint(request): # Batch mode: Extract multiple images debug_log(f"job_complete received batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") + # Check for JSON metadata + metadata_field = data.get('images_metadata') + metadata = None + if metadata_field: + # New JSON metadata format + try: + # Handle different types of metadata field + if hasattr(metadata_field, 'file'): + # File-like object + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + # Direct bytes/bytearray + metadata_str = metadata_field.decode('utf-8') + else: + # String + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + debug_log(f"Using JSON metadata for batch from worker {worker_id}") + except Exception as e: + log(f"Error parsing JSON metadata from worker {worker_id}: {e}") + return await handle_api_error(request, f"Metadata parsing error: {e}", 400) + else: + # Legacy format - log deprecation warning + debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") + + # Process images with per-item error handling + image_data_list = [] # List of tuples: (index, tensor) + for i in range(batch_size): image_field = data.get(f'image_{i}') if image_field is None: - return await handle_api_error(request, f"Missing image_{i}", 400) + log(f"Missing image_{i} from worker {worker_id}, skipping") + continue try: img_data = image_field.file.read() img = Image.open(io.BytesIO(img_data)).convert("RGB") tensor = pil_to_tensor(img) tensor = ensure_contiguous(tensor) - tensors.append(tensor) + + # Get expected index from metadata or use position + if metadata and i < len(metadata): + expected_idx = metadata[i].get('index', i) + if i != expected_idx: + debug_log(f"Warning: Image order mismatch at position {i}, expected index {expected_idx}") + else: + expected_idx = i + + image_data_list.append((expected_idx, tensor)) except Exception as e: - log(f"Error processing image {i} from worker {worker_id}: {e}") - return await handle_api_error(request, f"Image processing error: {e}", 400) + log(f"Error processing image {i} from worker {worker_id}: {e}, skipping this image") + # Continue processing other images instead of failing entire batch + continue + + # Check if we got any valid images + if not image_data_list: + return await handle_api_error(request, "No valid images in batch", 400) + + # Sort by index to ensure correct order + image_data_list.sort(key=lambda x: x[0]) + tensors = [tensor for _, tensor in image_data_list] + + # Validate final order + if metadata: + debug_log(f"Reordered {len(tensors)} images based on metadata indices") else: # Fallback for single-image mode (backward compatibility) image_file = data.get('image') @@ -1413,61 +1467,60 @@ class DistributedCollectorNode: return result async def send_batch_to_master(self, image_batch, multi_job_id, master_url, worker_id): - """Send entire image batch in one POST using multi-part FormData""" - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('is_last', 'True') # Full batch, so last for this worker - data.add_field('batch_size', str(image_batch.shape[0])) # Number of images in batch - - # Loop to add each image as a separate field - for i in range(image_batch.shape[0]): - # Convert tensor slice to PIL - img = tensor_to_pil(image_batch[i:i+1], 0) # tensor_to_pil handles single-image batch - byte_io = io.BytesIO() - img.save(byte_io, format='PNG', compress_level=0) # Lossless, no compression for speed - byte_io.seek(0) - data.add_field(f'image_{i}', byte_io, filename=f'image_{i}.png', content_type='image/png') - - try: - session = await get_client_session() - url = f"{master_url}/distributed/job_complete" - debug_log(f"Worker - Sending batch of {image_batch.shape[0]} images to URL: {url}") - async with session.post(url, data=data) as response: - response.raise_for_status() - except Exception as e: - log(f"Worker - Failed to send batch to master: {e}") - debug_log(f"Worker - Full error details: URL={url}") - - async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index, worker_id, is_last=False): - """Helper method to send a single image to the master""" - # Ensure we handle the tensor shape correctly (H, W, C) - if image_tensor.dim() == 4: # Has batch dimension - img = tensor_to_pil(image_tensor, 0) - else: # Single image without batch dimension - # Add batch dimension temporarily for tensor_to_pil - img = tensor_to_pil(image_tensor.unsqueeze(0), 0) - byte_io = io.BytesIO() - # Use PNG with no compression for lossless transfer - img.save(byte_io, format='PNG', compress_level=0) - byte_io.seek(0) + """Send image batch to master, chunked if large.""" + batch_size = image_batch.shape[0] + if batch_size == 0: + return - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('image_index', str(image_index)) - data.add_field('is_last', str(is_last)) - data.add_field('image', byte_io, filename=f'image_{image_index}.png', content_type='image/png') + debug_log(f"Worker - Sending {batch_size} images in chunks of max {MAX_BATCH}") + + for start in range(0, batch_size, MAX_BATCH): + chunk = image_batch[start:start + MAX_BATCH] + chunk_size = chunk.shape[0] + is_chunk_last = (start + chunk_size == batch_size) # True only for final chunk + + debug_log(f"Worker - Processing chunk {start//MAX_BATCH + 1}, images {start} to {start + chunk_size - 1}") + + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', str(is_chunk_last)) + data.add_field('batch_size', str(chunk_size)) + + # Chunk metadata: Absolute index from full batch + metadata = [{'index': start + j} for j in range(chunk_size)] + data.add_field('images_metadata', json.dumps(metadata), content_type='application/json') + + # Add chunk images + for j in range(chunk_size): + # Convert tensor slice to PIL + img = tensor_to_pil(chunk[j:j+1], 0) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + data.add_field(f'image_{j}', byte_io, filename=f'image_{j}.png', content_type='image/png') + + try: + # Estimate payload size for logging + payload_size = 0 + for field in data._fields: + if hasattr(field[2], 'read'): + field[2].seek(0) + payload_size += len(field[2].read()) + field[2].seek(0) + debug_log(f"Worker - Chunk payload size: {payload_size:,} bytes") + + session = await get_client_session() + url = f"{master_url}/distributed/job_complete" + debug_log(f"Worker - Sending chunk of {chunk_size} images to URL: {url}") + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Worker - Successfully sent chunk of {chunk_size} images") + except Exception as e: + log(f"Worker - Failed to send chunk to master: {e}") + debug_log(f"Worker - Full error details: URL={url}") + raise # Re-raise to handle at caller level - try: - session = await get_client_session() - url = f"{master_url}/distributed/job_complete" - debug_log(f"Worker - Sending image to URL: {url}") - async with session.post(url, data=data) as response: - response.raise_for_status() - except Exception as e: - log(f"Worker - Failed to send image {image_index+1} to master: {e}") - debug_log(f"Worker - Full error details: URL={master_url}/distributed/job_complete") async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): if is_worker: @@ -1519,10 +1572,8 @@ class DistributedCollectorNode: initial_size = q.qsize() debug_log(f"Master - Queue size before collection: {initial_size}") - # NEW: Initialize progress bar with estimated total images - expected_total = master_batch_size + (num_workers * worker_batch_size) # Estimate; adjust if batch sizes vary - p = ProgressBar(expected_total) - p.update(master_batch_size) # Start with master's contribution (already "collected") + # NEW: Initialize progress bar for workers (total = num_workers) + p = ProgressBar(num_workers) while len(workers_done) < num_workers: try: @@ -1555,9 +1606,6 @@ class DistributedCollectorNode: debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") collected_count += len(tensors) - - # Update progress for batch - p.update(len(tensors)) else: # Single image mode (backward compat) image_index = result['image_index'] @@ -1574,15 +1622,13 @@ class DistributedCollectorNode: debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") collected_count += 1 - - # Update progress for single image - p.update(1) # Once we start receiving images, use shorter timeout timeout = WORKER_JOB_TIMEOUT if is_last: workers_done.add(worker_id) + p.update(1) # +1 per completed worker debug_log(f"Master - Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") else: debug_log(f"Master - Collected image {image_index + 1} from worker {worker_id}") @@ -1625,7 +1671,6 @@ class DistributedCollectorNode: worker_images[worker_id][idx] = tensor collected_count += len(tensors) - p.update(len(tensors)) else: # Single image mode image_index = item['image_index'] @@ -1636,10 +1681,10 @@ class DistributedCollectorNode: worker_images[worker_id][image_index] = tensor collected_count += 1 - p.update(1) if is_last: workers_done.add(worker_id) + p.update(1) # +1 here too debug_log(f"Master - Worker {worker_id} done (found in timeout drain)") else: log(f"Master - Queue {multi_job_id} no longer exists!") diff --git a/distributed_upscale.py b/distributed_upscale.py index 4a40dde..6be5de0 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -20,7 +20,7 @@ from .utils.network import get_server_port, get_server_loop, get_client_session, from .utils.async_helpers import run_async_in_server_loop from .utils.constants import ( TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_TRANSFER_TIMEOUT, - QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT + QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH ) # Helper function to ensure persistent state is initialized @@ -614,119 +614,86 @@ class UltimateSDUpscaleDistributed: async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, padding, worker_id): - """Send all processed tiles to master in a single batch request.""" - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('is_last', 'True') # Full batch, so last for this worker - data.add_field('batch_size', str(len(processed_tiles))) - data.add_field('padding', str(padding)) + """Send all processed tiles to master, chunked if large.""" + if not processed_tiles: + return # Early exit if empty - # Add each tile as a separate field with metadata - for i, tile_data in enumerate(processed_tiles): - # Convert tensor to PIL - img = tensor_to_pil(tile_data['tile'], 0) - byte_io = io.BytesIO() - img.save(byte_io, format='PNG', compress_level=0) - byte_io.seek(0) + total_tiles = len(processed_tiles) + debug_log(f"UltimateSDUpscale Worker - Sending {total_tiles} tiles in chunks of max {MAX_BATCH}") + + for start in range(0, total_tiles, MAX_BATCH): + chunk = processed_tiles[start:start + MAX_BATCH] + chunk_size = len(chunk) + is_chunk_last = (start + chunk_size == total_tiles) # True only for final chunk - # Add image field - data.add_field(f'tile_{i}', byte_io, filename=f'tile_{i}.png', content_type='image/png') + debug_log(f"UltimateSDUpscale Worker - Processing chunk {start//MAX_BATCH + 1}, tiles {start} to {start + chunk_size - 1}") - # Add metadata for this tile - data.add_field(f'tile_{i}_idx', str(tile_data['tile_idx'])) - data.add_field(f'tile_{i}_x', str(tile_data['x'])) - data.add_field(f'tile_{i}_y', str(tile_data['y'])) - data.add_field(f'tile_{i}_width', str(tile_data['extracted_width'])) - data.add_field(f'tile_{i}_height', str(tile_data['extracted_height'])) - - # Retry logic with exponential backoff - max_retries = 5 - retry_delay = 0.5 - - for attempt in range(max_retries): - try: - session = await get_client_session() - url = f"{master_url}/distributed/tile_complete" + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', str(is_chunk_last)) + data.add_field('batch_size', str(chunk_size)) + data.add_field('padding', str(padding)) + + # Chunk metadata: Keep absolute tile_idx + metadata = [] + for tile_data in chunk: + metadata.append({ + 'tile_idx': tile_data['tile_idx'], # Original/absolute idx + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'] + }) + + # Add JSON metadata + data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') + + # Add chunk images as 'tile_{j}' where j=0 to chunk_size-1 + for j, tile_data in enumerate(chunk): + # Convert tensor to PIL + img = tensor_to_pil(tile_data['tile'], 0) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) - debug_log(f"UltimateSDUpscale Worker - Sending batch of {len(processed_tiles)} tiles to {url}, attempt {attempt + 1}") - - async with session.post(url, data=data) as response: - response.raise_for_status() - debug_log(f"UltimateSDUpscale Worker - Successfully sent batch of {len(processed_tiles)} tiles") - return + # Add image field + data.add_field(f'tile_{j}', byte_io, filename=f'tile_{j}.png', content_type='image/png') + + # Estimate payload size for logging + payload_size = 0 + for field in data._fields: + if hasattr(field[2], 'read'): + field[2].seek(0) + payload_size += len(field[2].read()) + field[2].seek(0) + debug_log(f"UltimateSDUpscale Worker - Chunk payload size: {payload_size:,} bytes") + + # Retry logic with exponential backoff + max_retries = 5 + retry_delay = 0.5 + + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/tile_complete" - except Exception as e: - if attempt < max_retries - 1: - debug_log(f"UltimateSDUpscale Worker - Retry {attempt + 1}/{max_retries} after error: {e}") - await asyncio.sleep(retry_delay) - retry_delay *= 2 # Exponential backoff - else: - log(f"UltimateSDUpscale Worker - Failed to send batch after {max_retries} attempts: {e}") - raise - - async def send_tile_to_master(self, tile_tensor, multi_job_id, master_url, - tile_idx, x, y, extracted_width, extracted_height, - padding, worker_id, is_last=False): - """Send processed tile to master with retry logic.""" - # Convert tensor to PIL image - img = tensor_to_pil(tile_tensor, 0) - - # Store image bytes for retry - img_bytes = io.BytesIO() - img.save(img_bytes, format='PNG', compress_level=0) - img_data = img_bytes.getvalue() - - # Retry logic with exponential backoff - max_retries = 5 - retry_delay = 0.5 # Start with 500ms - - for attempt in range(max_retries): - try: - # Create fresh FormData for each attempt - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('tile_idx', str(tile_idx)) - data.add_field('x', str(x)) - data.add_field('y', str(y)) - data.add_field('extracted_width', str(extracted_width)) - data.add_field('extracted_height', str(extracted_height)) - data.add_field('padding', str(padding)) - data.add_field('is_last', str(is_last)) - data.add_field('image', io.BytesIO(img_data), filename=f'tile_{tile_idx}.png', content_type='image/png') - - session = await get_client_session() - timeout = aiohttp.ClientTimeout(total=TILE_TRANSFER_TIMEOUT) - async with session.post(f"{master_url}/distributed/tile_complete", data=data, timeout=timeout) as response: - if response.status == 404 and attempt < max_retries - 1: - # Queue not ready yet, wait and retry - debug_log(f"UltimateSDUpscale Worker - Queue not ready for job {multi_job_id}, attempt {attempt + 1}/{max_retries}, retrying in {retry_delay}s...") + debug_log(f"UltimateSDUpscale Worker - Sending chunk of {chunk_size} tiles to {url}, attempt {attempt + 1}") + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"UltimateSDUpscale Worker - Successfully sent chunk of {chunk_size} tiles") + break # Success, move to next chunk + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"UltimateSDUpscale Worker - Retry {attempt + 1}/{max_retries} after error: {e}") await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) # Exponential backoff with cap at 5s - continue - response.raise_for_status() - debug_log(f"UltimateSDUpscale Worker - Successfully sent tile {tile_idx} on attempt {attempt + 1}") - return # Success - except aiohttp.ClientResponseError as e: - if e.status == 404 and attempt < max_retries - 1: - debug_log(f"UltimateSDUpscale Worker - Got 404 for job {multi_job_id}, attempt {attempt + 1}/{max_retries}, retrying in {retry_delay}s...") - await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) - continue - elif attempt == max_retries - 1: - log(f"UltimateSDUpscale Worker - Failed to send tile {tile_idx} after {max_retries} attempts: {e}") - raise - except Exception as e: - if attempt == max_retries - 1: - log(f"UltimateSDUpscale Worker - Failed to send tile {tile_idx} after {max_retries} attempts: {e}") - raise - else: - debug_log(f"UltimateSDUpscale Worker - Error on attempt {attempt + 1}: {e}, retrying...") - await asyncio.sleep(retry_delay) - retry_delay = min(retry_delay * 2, 5.0) - - + retry_delay *= 2 # Exponential backoff + else: + log(f"UltimateSDUpscale Worker - Failed to send chunk after {max_retries} attempts: {e}") + raise + def process_single_gpu(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles): @@ -793,37 +760,113 @@ async def tile_complete_endpoint(request): padding = int(data.get('padding', 32)) debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - return await handle_api_error(request, f"Missing tile_{i}", 400) - - # Get metadata for this tile - tile_idx = int(data.get(f'tile_{i}_idx', i)) - x = int(data.get(f'tile_{i}_x', 0)) - y = int(data.get(f'tile_{i}_y', 0)) - extracted_width = int(data.get(f'tile_{i}_width', 512)) - extracted_height = int(data.get(f'tile_{i}_height', 512)) - + # Check for JSON metadata + metadata_field = data.get('tiles_metadata') + if metadata_field: + # New JSON metadata format try: - # Process image - img_data = tile_field.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] + # Handle different types of metadata field + if hasattr(metadata_field, 'file'): + # File-like object + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + # Direct bytes/bytearray + metadata_str = metadata_field.decode('utf-8') + else: + # String + metadata_str = str(metadata_field) - tiles.append({ - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': x, - 'y': y, - 'extracted_width': extracted_width, - 'extracted_height': extracted_height, - 'padding': padding - }) + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + + tile_data_list = [] # Temporary list to collect and sort tiles + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + log(f"Missing tile_{i} from worker {worker_id}, skipping") + continue + + try: + # Process image + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + # Get metadata for this tile + if i < len(metadata): + tile_data = metadata[i] + tile_idx = tile_data.get('tile_idx', i) + + # Validate order + if i != tile_idx % batch_size: + debug_log(f"Warning: Tile order mismatch at position {i}, expected tile_idx {tile_idx}") + + tile_info = { + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': padding + } + tile_data_list.append(tile_info) + else: + log(f"Missing metadata for tile {i} from worker {worker_id}, skipping") + continue + + except Exception as e: + log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") + continue + + # Sort tiles by tile_idx to ensure correct order + tile_data_list.sort(key=lambda x: x['tile_idx']) + tiles.extend(tile_data_list) + + if tile_data_list: + debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx") except Exception as e: - log(f"Error processing tile {i} from worker {worker_id}: {e}") - return await handle_api_error(request, f"Tile processing error: {e}", 400) + log(f"Error processing JSON metadata from worker {worker_id}: {e}") + return await handle_api_error(request, f"Metadata processing error: {e}", 400) + else: + # Legacy format: individual fields per tile (backward compatibility) + debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + log(f"Missing tile_{i} from worker {worker_id}, skipping") + continue + + # Get metadata for this tile + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + try: + # Process image + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + except Exception as e: + log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") + continue else: # Single tile mode (backward compatibility) image_file = data.get('image') diff --git a/utils/constants.py b/utils/constants.py index ff1869d..2ce53c4 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -31,4 +31,7 @@ QUEUE_INIT_TIMEOUT = 5.0 TILE_SEND_TIMEOUT = 60.0 # Memory operations -MEMORY_CLEAR_DELAY = 0.5 \ No newline at end of file +MEMORY_CLEAR_DELAY = 0.5 + +# Batch processing +MAX_BATCH = 20 # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs) \ No newline at end of file From 189165f2ec948a7d52b26e271a76afb00d089abc Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 15 Jul 2025 19:19:00 +1000 Subject: [PATCH 16/93] cross-platform paths --- distributed.py | 18 +++++++++++ web/executionUtils.js | 69 +++++++++++++++++++++++++++++++++++++++---- web/workerUtils.js | 50 +++++++++++++++++++++++++++++++ 3 files changed, 131 insertions(+), 6 deletions(-) diff --git a/distributed.py b/distributed.py index 5dc0b1e..2b4eaca 100644 --- a/distributed.py +++ b/distributed.py @@ -1295,6 +1295,24 @@ async def load_image_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) +@server.PromptServer.instance.routes.get("/distributed/system_info") +async def system_info_endpoint(request): + """Return system information including platform details""" + import platform + + return web.json_response({ + "platform": { + "os_name": os.name, # 'nt' for Windows, 'posix' for Unix/Linux + "path_separator": os.sep, # '\\' on Windows, '/' on Linux + "system": platform.system(), # 'Windows', 'Linux', 'Darwin' + "machine": platform.machine(), # 'x86_64', 'AMD64', etc. + }, + "comfyui": { + # Can add ComfyUI-specific info here + "version": "distributed-extension", + } + }) + @server.PromptServer.instance.routes.post("/distributed/job_complete") async def job_complete_endpoint(request): try: diff --git a/web/executionUtils.js b/web/executionUtils.js index 9960de4..f7a40f6 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -1,7 +1,42 @@ import { api } from "../../scripts/api.js"; -import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker } from './workerUtils.js'; +import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo } from './workerUtils.js'; import { TIMEOUTS } from './constants.js'; +/** + * Convert paths in the API prompt to match the target platform's separator + * @param {Object} apiPrompt - The workflow API prompt + * @param {string} targetSeparator - The target path separator ('\\' or '/') + * @returns {Object} The converted API prompt + */ +function convertPathsForPlatform(apiPrompt, targetSeparator) { + // Regex to identify likely file paths with extensions + const isLikelyFilename = (value) => { + return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i); + }; + + function convert(obj) { + if (typeof obj === 'string') { + // Only convert strings that look like file paths + if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) { + // Replace any path separator with the target one + return obj.replace(/[\\\/]/g, targetSeparator); + } + return obj; + } else if (Array.isArray(obj)) { + return obj.map(convert); + } else if (typeof obj === 'object' && obj !== null) { + const newObj = {}; + for (const [key, value] of Object.entries(obj)) { + newObj[key] = convert(value); + } + return newObj; + } + return obj; + } + + return convert(apiPrompt); +} + export function setupInterceptor(extension) { api.queuePrompt = async (number, prompt) => { if (extension.isEnabled) { @@ -63,7 +98,7 @@ export async function executeParallelDistributed(extension, promptWrapper) { job_id_map: job_id_map // Pass the map of unique IDs }; - const jobApiPrompt = prepareApiPromptForParticipant( + const jobApiPrompt = await prepareApiPromptForParticipant( extension, promptWrapper.output, participantId, options ); @@ -97,7 +132,7 @@ export async function executeParallelDistributed(extension, promptWrapper) { } } -export function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) { +export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) { let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); const isMaster = participantId === 'master'; @@ -106,9 +141,31 @@ export function prepareApiPromptForParticipant(extension, baseApiPrompt, partici const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed"); const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - // For workers, prune the workflow to only include distributed node dependencies - if (!isMaster && allDistributedNodes.length > 0) { - jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes); + // For workers, handle platform-specific path conversion + if (!isMaster) { + const workerInfo = extension.config.workers.find(w => w.id === participantId); + + if (workerInfo && workerInfo.host) { + // Remote or cloud worker - needs path translation + try { + const workerUrl = extension.getWorkerUrl(workerInfo); + const systemInfo = await getCachedWorkerSystemInfo(workerUrl); + const targetSeparator = systemInfo.platform.path_separator; + + // Convert paths to match worker's platform + jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator); + + extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug"); + } catch (e) { + extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn"); + // Continue without path conversion + } + } + + // Prune the workflow to only include distributed node dependencies + if (allDistributedNodes.length > 0) { + jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes); + } } // Handle image references for remote workers diff --git a/web/workerUtils.js b/web/workerUtils.js index 745758a..63280f7 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -263,4 +263,54 @@ export function hasUpstreamNode(apiPrompt, nodeId, upstreamType) { } return false; +} + +/** + * Get system information from a worker + * @param {string} workerUrl - The worker URL + * @returns {Promise} System information including platform details + */ +export async function getWorkerSystemInfo(workerUrl) { + try { + const response = await fetch(`${workerUrl}/distributed/system_info`); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + return await response.json(); + } catch (error) { + console.warn(`Failed to get system info from ${workerUrl}:`, error); + // Return sensible defaults + return { + platform: { + os_name: 'posix', // Assume Linux + path_separator: '/', + system: 'Linux' + } + }; + } +} + +// Cache system info to avoid repeated calls +const systemInfoCache = new Map(); + +/** + * Get cached system information from a worker + * @param {string} workerUrl - The worker URL + * @returns {Promise} Cached or fresh system information + */ +export async function getCachedWorkerSystemInfo(workerUrl) { + if (systemInfoCache.has(workerUrl)) { + return systemInfoCache.get(workerUrl); + } + + const info = await getWorkerSystemInfo(workerUrl); + systemInfoCache.set(workerUrl, info); + return info; +} + +/** + * Clear the system info cache + */ +export function clearSystemInfoCache() { + systemInfoCache.clear(); } \ No newline at end of file From 026c1de6435cd9a96a550aef32ec59726152992b Mon Sep 17 00:00:00 2001 From: robertvoy Date: Wed, 16 Jul 2025 12:30:51 +1000 Subject: [PATCH 17/93] image batch divider js --- web/image_batch_divider.js | 86 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 web/image_batch_divider.js diff --git a/web/image_batch_divider.js b/web/image_batch_divider.js new file mode 100644 index 0000000..522530e --- /dev/null +++ b/web/image_batch_divider.js @@ -0,0 +1,86 @@ +import { app } from "/scripts/app.js"; + +app.registerExtension({ + name: "Distributed.ImageBatchDivider", + async nodeCreated(node) { + if (node.comfyClass === "ImageBatchDivider") { + try { + const updateOutputs = () => { + if (!node.widgets) return; + + const divideByWidget = node.widgets.find(w => w.name === "divide_by"); + if (!divideByWidget) return; + + const divideBy = parseInt(divideByWidget.value, 10) || 1; + const totalOutputs = divideBy; // Direct divide by value + + // Ensure outputs array exists + if (!node.outputs) node.outputs = []; + + // Remove excess outputs + while (node.outputs.length > totalOutputs) { + node.removeOutput(node.outputs.length - 1); + } + + // Add missing outputs + while (node.outputs.length < totalOutputs) { + const outputIndex = node.outputs.length + 1; + node.addOutput(`batch_${outputIndex}`, "IMAGE"); + } + + if (node.setDirty) node.setDirty(true); // Refresh canvas + }; + + // Initial update with delay to allow workflow loading + setTimeout(updateOutputs, 200); + + // Find the widget and set up responsive handlers + const divideByWidget = node.widgets.find(w => w.name === "divide_by"); + if (divideByWidget) { + // Override callback for immediate trigger on value set + const originalCallback = divideByWidget.callback; + divideByWidget.callback = (value) => { + updateOutputs(); + if (originalCallback) originalCallback.call(divideByWidget, value); // Preserve 'this' context + }; + + // Add event listener for real-time input changes (e.g., typing/dragging) + if (divideByWidget.inputEl) { + divideByWidget.inputEl.addEventListener('input', updateOutputs); + } + + // Lightweight MutationObserver as fallback (observe attributes on widget element if available) + const observer = new MutationObserver(updateOutputs); + if (divideByWidget.element) { + observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true }); + } + + // Store cleanup function + node._batchDividerCleanup = () => { + observer.disconnect(); + if (divideByWidget.inputEl) { + divideByWidget.inputEl.removeEventListener('input', updateOutputs); + } + divideByWidget.callback = originalCallback; // Restore original + }; + } + + // Add post-configure hook for reliable workflow loading + const originalConfigure = node.configure; + node.configure = function(data) { + const result = originalConfigure ? originalConfigure.call(this, data) : undefined; + updateOutputs(); // Re-run after config load + return result; + }; + } catch (error) { + console.error("Error in ImageBatchDivider extension:", error); + } + } + }, + + nodeBeforeRemove(node) { + if (node.comfyClass === "ImageBatchDivider" && node._batchDividerCleanup) { + node._batchDividerCleanup(); + } + } +}); \ No newline at end of file From 24c6b41a73ea21c84e8b127dd4681f20ee517517 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Wed, 16 Jul 2025 12:33:24 +1000 Subject: [PATCH 18/93] runpod address bug fix --- web/main.js | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/web/main.js b/web/main.js index 4cce62b..7d00c38 100644 --- a/web/main.js +++ b/web/main.js @@ -275,10 +275,15 @@ class DistributedExtension { // For local workers on Runpod, construct the port-specific proxy URL let finalHost = host; if (!worker.host && isRunpodProxy) { - // Extract pod ID from current host (format: podId-xxxx.proxy.runpod.net) - const podId = host.split('-')[0]; - const domain = host.substring(host.indexOf('-') + 1); - finalHost = `${podId}-${worker.port}.${domain}`; + const match = host.match(/^(.*)\.proxy\.runpod\.net$/); + if (match) { + const podId = match[1]; + const domain = 'proxy.runpod.net'; + finalHost = `${podId}-${worker.port}.${domain}`; + } else { + // Fallback or log error if no match (shouldn't happen) + console.error(`[Distributed] Failed to parse Runpod proxy host: ${host}`); + } } // Determine protocol: HTTPS for cloud, Runpod proxies, or port 443 From 338456b79de2e039e6faec649c490464da64f00a Mon Sep 17 00:00:00 2001 From: robertvoy Date: Wed, 16 Jul 2025 13:01:05 +1000 Subject: [PATCH 19/93] auto detect runpod host --- web/main.js | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/web/main.js b/web/main.js index 7d00c38..7a5be5d 100644 --- a/web/main.js +++ b/web/main.js @@ -771,6 +771,12 @@ class DistributedExtension { async detectMasterIP() { try { + // Detect if we're running on Runpod + const isRunpod = window.location.hostname.endsWith('.proxy.runpod.net'); + if (isRunpod) { + this.log("Detected Runpod environment", "info"); + } + const data = await this.api.getNetworkInfo(); this.log("Network info: " + JSON.stringify(data), "debug"); @@ -826,7 +832,8 @@ class DistributedExtension { host: "localhost", port: 8189 + portOffset, cuda_device: i, - enabled: true + enabled: true, + extra_args: isRunpod ? "--listen" : "" }; newWorkers.push(worker); workerNum++; @@ -875,6 +882,30 @@ class DistributedExtension { return; } + // For Runpod, use the proxy hostname as master host + if (isRunpod) { + const runpodHost = window.location.hostname; + this.log(`Setting Runpod master host: ${runpodHost}`, "info"); + + // Save the Runpod host + await this.api.updateMaster({ host: runpodHost }); + + // Update local config + if (!this.config.master) this.config.master = {}; + this.config.master.host = runpodHost; + + // Show notification + if (app.extensionManager?.toast) { + app.extensionManager.toast.add({ + severity: "info", + summary: "Runpod Auto-Configuration", + detail: `Master host set to ${runpodHost} with --listen flag for workers`, + life: 5000 + }); + } + return; // Skip regular IP detection for Runpod + } + // Use the recommended IP from the backend if (data.recommended_ip && data.recommended_ip !== '127.0.0.1') { this.log(`Auto-detected master IP: ${data.recommended_ip}`, "info"); From 3a6c1f9d04b6f98de433df9253fcde7580691030 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Wed, 16 Jul 2025 13:26:11 +1000 Subject: [PATCH 20/93] better status checks --- web/executionUtils.js | 8 +++++++- web/main.js | 29 ++++++++++++++++++++++++++--- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/web/executionUtils.js b/web/executionUtils.js index f7a40f6..bbf700b 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -45,7 +45,9 @@ export function setupInterceptor(extension) { if (hasCollector || hasDistUpscale) { const result = await executeParallelDistributed(extension, prompt); - // Check status after dispatching jobs + // Immediate status check for instant feedback + extension.checkAllWorkerStatuses(); + // Another check after a short delay to catch state changes setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); return result; } @@ -298,6 +300,10 @@ export async function executeJobs(extension, jobs) { } }); await Promise.all(promises); + + // Trigger immediate status check for instant feedback + extension.checkAllWorkerStatuses(); + return masterPromptId || { "prompt_id": "distributed-job-dispatched" }; } diff --git a/web/main.js b/web/main.js index 7a5be5d..f59ce4e 100644 --- a/web/main.js +++ b/web/main.js @@ -26,6 +26,9 @@ class DistributedExtension { // Initialize API client this.api = createApiClient(window.location.origin); + + // Initialize status check timeout reference + this.statusCheckTimeout = null; // Inject CSS for pulsing animation this.injectStyles(); @@ -206,9 +209,7 @@ class DistributedExtension { } startStatusChecking() { - this.statusCheckInterval = setInterval(() => { - this.checkAllWorkerStatuses(); - }, TIMEOUTS.STATUS_CHECK); + this.checkAllWorkerStatuses(); } async checkAllWorkerStatuses() { @@ -223,6 +224,23 @@ class DistributedExtension { this.checkWorkerStatus(worker); } } + + // Determine next interval based on current state + let isActive = this.state.getMasterStatus() === 'processing'; // Master is busy + + // Check workers for activity + this.config.workers.forEach(worker => { + const ws = this.state.getWorker(worker.id); // Get worker state + if (ws.launching || ws.status?.processing) { // Launching or processing + isActive = true; + } + }); + + // Set next delay: 1s if active, 5s if idle + const nextInterval = isActive ? 1000 : 5000; + + // Schedule the next check + this.statusCheckTimeout = setTimeout(() => this.checkAllWorkerStatuses(), nextInterval); } async checkMasterStatus() { @@ -545,6 +563,11 @@ class DistributedExtension { this.logAutoRefreshInterval = null; } + if (this.statusCheckTimeout) { + clearTimeout(this.statusCheckTimeout); + this.statusCheckTimeout = null; + } + this.log("Cleaned up intervals", "debug"); } From dd23503883fdf319e8beb6e7a190445ecf89973c Mon Sep 17 00:00:00 2001 From: robertvoy Date: Wed, 16 Jul 2025 14:07:52 +1000 Subject: [PATCH 21/93] better shutdown and security --- distributed.py | 126 +++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 102 insertions(+), 24 deletions(-) diff --git a/distributed.py b/distributed.py index 2b4eaca..9664862 100644 --- a/distributed.py +++ b/distributed.py @@ -14,6 +14,8 @@ import platform import time import atexit import signal +import sys +import shlex from comfy.utils import ProgressBar # Import shared utilities @@ -781,9 +783,20 @@ class WorkerProcessManager: raise RuntimeError(error_msg) - # Add any extra arguments + # Add any extra arguments safely if worker_config.get('extra_args'): - cmd.extend(worker_config['extra_args'].split()) + raw_args = worker_config['extra_args'].strip() + if raw_args: + # Safely split using shlex to handle quotes and spaces + extra_args_list = shlex.split(raw_args) + + # Validate: Block dangerous shell meta-characters (e.g., ; | & > < ` $) + forbidden_chars = set(';|>&<`$()[]{}*!?') + for arg in extra_args_list: + if any(c in forbidden_chars for c in arg): + raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {''.join(forbidden_chars)}") + + cmd.extend(extra_args_list) return cmd @@ -1210,15 +1223,9 @@ def delayed_auto_launch(): timer.daemon = True timer.start() -# Call delayed auto-launch only if we're the master (not a worker) -if not os.environ.get('COMFYUI_MASTER_PID'): - delayed_auto_launch() -else: - debug_log("Running as worker, skipping auto-launch") - -# Register cleanup on exit - only clean up if setting is enabled -def cleanup_on_exit(signum=None, frame=None): - """Handle cleanup on exit or signal""" +# Async cleanup function for proper shutdown +async def async_cleanup_and_exit(signum=None): + """Async-friendly cleanup and exit.""" try: config = load_config() if config.get('settings', {}).get('stop_workers_on_master_exit', True): @@ -1226,24 +1233,95 @@ def cleanup_on_exit(signum=None, frame=None): worker_manager.cleanup_all() else: print("\n[Distributed] Master shutting down, workers will continue running") - # Still save the current state + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + + # On Windows, we need to exit differently + if platform.system() == "Windows": + # Force exit on Windows + sys.exit(0) + else: + # On Unix, stop the event loop gracefully + loop = asyncio.get_running_loop() + loop.stop() + +def register_async_signals(): + """Register async signal handlers for graceful shutdown.""" + if platform.system() == "Windows": + # Windows doesn't support add_signal_handler, use traditional signal handling + def signal_handler(signum, frame): + # Schedule the async cleanup in the event loop + loop = server.PromptServer.instance.loop + if loop and loop.is_running(): + asyncio.run_coroutine_threadsafe(async_cleanup_and_exit(signum), loop) + else: + # Fallback to sync cleanup if loop isn't running + try: + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + print("\n[Distributed] Master shutting down, stopping all managed workers...") + worker_manager.cleanup_all() + else: + print("\n[Distributed] Master shutting down, workers will continue running") + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + else: + # Unix-like systems support add_signal_handler + loop = server.PromptServer.instance.loop + + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, lambda s=sig: asyncio.create_task(async_cleanup_and_exit(s))) + + # SIGHUP is Unix-only + loop.add_signal_handler(signal.SIGHUP, lambda: asyncio.create_task(async_cleanup_and_exit(signal.SIGHUP))) + +def sync_cleanup(): + """Synchronous wrapper for atexit.""" + try: + # For atexit, we don't want to stop the loop or exit + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + print("\n[Distributed] Master shutting down, stopping all managed workers...") + worker_manager.cleanup_all() + else: + print("\n[Distributed] Master shutting down, workers will continue running") worker_manager.save_processes() except Exception as e: print(f"[Distributed] Error during cleanup: {e}") -# Register cleanup handlers -atexit.register(cleanup_on_exit) +# Register atexit handler for normal exits +atexit.register(sync_cleanup) -# Handle terminal window closing and Ctrl+C -try: - signal.signal(signal.SIGINT, cleanup_on_exit) - signal.signal(signal.SIGTERM, cleanup_on_exit) - - if platform.system() != "Windows": - # SIGHUP is sent when terminal closes on Unix - signal.signal(signal.SIGHUP, cleanup_on_exit) -except Exception as e: - print(f"[Distributed] Warning: Could not set signal handlers: {e}") +# Call delayed auto-launch only if we're the master (not a worker) +if not os.environ.get('COMFYUI_MASTER_PID'): + delayed_auto_launch() + try: + register_async_signals() # Register async signal handlers + except Exception as e: + print(f"[Distributed] Warning: Could not register async signal handlers: {e}") + # Fall back to basic signal handling + def basic_signal_handler(signum, frame): + print("\n[Distributed] Received signal, shutting down...") + try: + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + worker_manager.cleanup_all() + else: + worker_manager.save_processes() + except Exception as cleanup_error: + print(f"[Distributed] Error during cleanup: {cleanup_error}") + sys.exit(0) + + signal.signal(signal.SIGINT, basic_signal_handler) + signal.signal(signal.SIGTERM, basic_signal_handler) +else: + debug_log("Running as worker, skipping auto-launch") # --- Persistent State Storage --- # Store job queue on the persistent server instance to survive script reloads From 5703b6206227c97b9bfe11d540951435a659efaf Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 09:53:05 +1000 Subject: [PATCH 22/93] distributed upscale video support --- distributed_upscale.py | 1079 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 991 insertions(+), 88 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 6be5de0..371050c 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1,3 +1,15 @@ +""" +ComfyUI Distributed Upscale Module + +This module implements a distributed version of Ultimate SD Upscale with multi-mode +batch handling for efficient processing of single images and video frames. + +Processing Modes: +- Single GPU: Process all tiles locally when no workers are available +- Static Mode: For small batches (≤ dynamic_threshold), flattens tiles across batch +- Dynamic Mode: For large batches, assigns whole images to workers dynamically +""" + import torch import numpy as np from PIL import Image, ImageFilter, ImageDraw @@ -7,6 +19,9 @@ import aiohttp from aiohttp import web import io import server +import math +import os +import time from typing import List, Tuple # Import ComfyUI modules @@ -20,9 +35,15 @@ from .utils.network import get_server_port, get_server_loop, get_client_session, from .utils.async_helpers import run_async_in_server_loop from .utils.constants import ( TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_TRANSFER_TIMEOUT, - QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH + QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT ) +# Make MAX_BATCH configurable +MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) + +# Configure maximum payload size (50MB default, configurable via environment variable) +MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) + # Helper function to ensure persistent state is initialized def ensure_tile_jobs_initialized(): """Ensure tile job storage is initialized on the server instance.""" @@ -31,6 +52,15 @@ def ensure_tile_jobs_initialized(): debug_log("Initializing persistent tile job queue on server instance.") prompt_server.distributed_pending_tile_jobs = {} prompt_server.distributed_tile_jobs_lock = asyncio.Lock() + else: + # Clean up any legacy queue structures that don't have the 'mode' field + to_remove = [] + for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items(): + if not isinstance(job_data, dict) or 'mode' not in job_data: + debug_log(f"Removing legacy queue structure for job {job_id}") + to_remove.append(job_id) + for job_id in to_remove: + del prompt_server.distributed_pending_tile_jobs[job_id] return prompt_server # Note: tensor_to_pil and pil_to_tensor are imported from utils.image @@ -38,7 +68,23 @@ def ensure_tile_jobs_initialized(): class UltimateSDUpscaleDistributed: """ Distributed version of Ultimate SD Upscale (No Upscale). - Distributes tile processing across multiple workers. + + Supports three processing modes: + 1. Single GPU: No workers available, process everything locally + 2. Static Mode: Small batches, distributes tiles across workers (flattened) + 3. Dynamic Mode: Large batches, assigns whole images to workers dynamically + + Features: + - Multi-mode batch handling for efficient video/image upscaling + - Tiled VAE support for memory efficiency + - Dynamic load balancing for large batches + - Backward compatible with single-image workflows + + Environment Variables: + - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20) + - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB) + + Threshold: dynamic_threshold input controls mode switch (default 8) """ def __init__(self): @@ -67,6 +113,7 @@ class UltimateSDUpscaleDistributed: "padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}), "mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}), "force_uniform_tiles": ("BOOLEAN", {"default": True}), + "tiled_decode": ("BOOLEAN", {"default": False}), }, "hidden": { "multi_job_id": ("STRING", {"default": ""}), @@ -75,6 +122,7 @@ class UltimateSDUpscaleDistributed: "enabled_worker_ids": ("STRING", {"default": "[]"}), "worker_id": ("STRING", {"default": ""}), "tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility + "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}), }, } @@ -106,67 +154,106 @@ class UltimateSDUpscaleDistributed: def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, - mask_blur, force_uniform_tiles, multi_job_id="", is_worker=False, - master_url="", enabled_worker_ids="[]", worker_id="", tile_indices=""): + mask_blur, force_uniform_tiles, tiled_decode, multi_job_id="", is_worker=False, + master_url="", enabled_worker_ids="[]", worker_id="", tile_indices="", dynamic_threshold=8): """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler.""" if not multi_job_id: # No distributed processing, run single GPU version return self.process_single_gpu(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles) + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode) if is_worker: # Worker mode: process tiles synchronously return self.process_worker_tiles(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, master_url, - worker_id, enabled_worker_ids) + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold) else: # Master mode: distribute and collect synchronously return self.process_master(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, enabled_worker_ids) + force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold) def process_worker_tiles(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, master_url, - worker_id, enabled_worker_ids): - """Worker processing - SYNCHRONOUS tile processing with async communication.""" + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold): + """Worker processing - determines mode and processes accordingly.""" + # Get batch size to determine mode + batch_size = upscaled_image.shape[0] + + # Ensure mode consistency across master/workers via shared threshold + # Determine mode (must match master's logic) + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) + + if batch_size == 1 or batch_size <= dynamic_threshold: + mode = "static" + else: + mode = "dynamic" + + debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}") + + if mode == "dynamic": + return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold) + + # Static mode - continue with existing tile-based processing # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) # Calculate all tile positions _, height, width, _ = upscaled_image.shape - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - # Get assigned tiles for this worker - assigned_tiles = self._get_worker_tiles(all_tiles, enabled_worker_ids, worker_id) - if not assigned_tiles: + # Calculate total tiles across all images + num_tiles_per_image = len(all_tiles) + total_tiles = batch_size * num_tiles_per_image + + # Get assigned global tile indices for this worker + assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_worker_ids, worker_id) + if not assigned_global_indices: return (upscaled_image,) - debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(assigned_tiles)} tiles") + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(assigned_global_indices)} tiles from batch of {batch_size}") # Process tiles SYNCHRONOUSLY processed_tiles = [] - for tile_idx in assigned_tiles: + for global_idx in assigned_global_indices: + # Calculate which image and tile this corresponds to + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + # Skip if batch_idx is out of range + if batch_idx >= upscaled_image.shape[0]: + debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") + continue + x, y = all_tiles[tile_idx] - # Extract tile + # Extract tile from the specific image in the batch tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( - upscaled_image, x, y, tile_width, tile_height, padding + upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding ) - # Process tile through SD (SYNCHRONOUS - no async!) + # Process tile through SD with unique seed + image_seed = seed + batch_idx * 1000 processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, - seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise) + image_seed + tile_idx, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode) processed_tiles.append({ 'tile': processed_tile, + 'global_idx': global_idx, + 'batch_idx': batch_idx, 'tile_idx': tile_idx, 'x': x1, # Use extraction position 'y': y1, @@ -189,17 +276,17 @@ class UltimateSDUpscaleDistributed: def process_master(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, multi_job_id, enabled_worker_ids): + force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold): """Master processing - SYNCHRONOUS with async collection.""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) - # Get image dimensions - _, height, width, _ = upscaled_image.shape + # Get image dimensions and batch size + batch_size, height, width, _ = upscaled_image.shape # Calculate all tiles - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) total_tiles = len(all_tiles) debug_log(f"UltimateSDUpscale Master - Total tiles: {total_tiles}") @@ -208,11 +295,44 @@ class UltimateSDUpscaleDistributed: enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) + # Synchronous queue initialization before mode determination + import threading + def sync_init(): + prompt_server = ensure_tile_jobs_initialized() + if multi_job_id not in prompt_server.distributed_pending_tile_jobs: + prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() + debug_log(f"Synchronous queue init for {multi_job_id}") + thread = threading.Thread(target=sync_init) + thread.start() + thread.join(timeout=5.0) # Block with timeout + if thread.is_alive(): + log(f"Queue init timed out for {multi_job_id}") + raise RuntimeError("Queue init failed") + + # Determine processing mode if num_workers == 0: + mode = "single_gpu" + elif batch_size == 1 or batch_size <= dynamic_threshold: + mode = "static" + else: + mode = "dynamic" + + debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}") + + if mode == "single_gpu": # No workers, process all tiles locally return self.process_single_gpu(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles) + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode) + + elif mode == "dynamic": + # Dynamic mode for large batches + return self.process_master_dynamic(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers) + + # Static mode - continue with existing implementation # Initialize job queue FIRST before any processing # Initialize the queue directly in the persistent server storage @@ -235,33 +355,51 @@ class UltimateSDUpscaleDistributed: # loop = asyncio.new_event_loop() # asyncio.set_event_loop(loop) - # Prepare result image - image_pil = tensor_to_pil(upscaled_image, 0) - result_image = image_pil.copy() + # Convert batch to PIL list for processing + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + result_images.append(image_pil.copy()) # Mask will be created per tile for proper blending - # Calculate master's tiles - master_tiles_indices = self._get_master_tiles(all_tiles, num_workers) - master_tiles = len(master_tiles_indices) + # Calculate total tiles across all images + num_tiles_per_image = len(all_tiles) + total_tiles_all_images = batch_size * num_tiles_per_image - debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally") + # Calculate master's global tile indices + master_global_indices = self._get_master_global_indices(total_tiles_all_images, num_workers) + master_tiles = len(master_global_indices) + + debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally from batch of {batch_size}") # Process master tiles SYNCHRONOUSLY - for tile_idx in master_tiles_indices: - result_image = self._process_and_blend_tile( - tile_idx, all_tiles[tile_idx], upscaled_image, result_image, - model, positive, negative, vae, seed, steps, cfg, + for global_idx in master_global_indices: + # Calculate which image and tile this corresponds to + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + # Skip if batch_idx is out of range + if batch_idx >= len(result_images): + debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds result_images length {len(result_images)}") + continue + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + result_images[batch_idx] = self._process_and_blend_tile( + tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive, negative, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height + padding, mask_blur, width, height, tiled_decode ) # Collect worker tiles using async - worker_tiles_expected = total_tiles - master_tiles + worker_tiles_expected = total_tiles_all_images - master_tiles if worker_tiles_expected > 0: # Calculate tile distribution for active worker calculation - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) + tiles_per_participant = total_tiles_all_images // (num_workers + 1) + remainder = total_tiles_all_images % (num_workers + 1) # Calculate which workers will actually have tiles active_workers = [] @@ -274,7 +412,7 @@ class UltimateSDUpscaleDistributed: worker_start_idx += remainder - 1 if remainder > 0 else 0 worker_tile_count = tiles_per_participant - if worker_start_idx < total_tiles and worker_tile_count > 0: + if worker_start_idx < total_tiles_all_images and worker_tile_count > 0: active_workers.append(enabled_workers[i]) debug_log(f"UltimateSDUpscale Master - Waiting for {worker_tiles_expected} tiles from {len(active_workers)} active workers") @@ -296,6 +434,14 @@ class UltimateSDUpscaleDistributed: extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) + # Get batch index if available (for batch processing) + batch_idx = tile_data.get('batch_idx', 0) + + # Skip if batch_idx is out of range + if batch_idx >= len(result_images): + debug_log(f"Warning: Received tile for batch_idx {batch_idx} but only have {len(result_images)} images") + continue + # Convert and blend tile_pil = tensor_to_pil(tile_tensor, 0) # Get the original tile position from the tile index @@ -303,17 +449,53 @@ class UltimateSDUpscaleDistributed: # Create mask for this specific tile tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) # Use extraction position and size for blending - result_image = self.blend_tile(result_image, tile_pil, + result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, x, y, (extracted_width, extracted_height), tile_mask, padding) # Convert back to tensor - result_tensor = pil_to_tensor(result_image) + # For batch processing, combine all result images + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + # Convert all images and concatenate + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + if upscaled_image.is_cuda: result_tensor = result_tensor.cuda() debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") return (result_tensor,) + def _get_worker_global_indices(self, total_tiles, enabled_worker_ids, worker_id): + """Calculate which global tile indices are assigned to a specific worker in flattened mode.""" + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) + + debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}") + + try: + worker_index = enabled_workers.index(worker_id) + except ValueError: + log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}") + return [] + + # Calculate tile distribution + tiles_per_participant = total_tiles // (num_workers + 1) + remainder = total_tiles % (num_workers + 1) + master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) + + start_idx = master_tiles + (worker_index * tiles_per_participant) + if worker_index < remainder - 1: + start_idx += worker_index + end_idx = start_idx + tiles_per_participant + 1 + else: + start_idx += remainder - 1 if remainder > 0 else 0 + end_idx = start_idx + tiles_per_participant + + end_idx = min(end_idx, total_tiles) + return list(range(start_idx, end_idx)) + def _get_worker_tiles(self, all_tiles, enabled_worker_ids, worker_id): """Calculate which tiles are assigned to a specific worker.""" enabled_workers = json.loads(enabled_worker_ids) @@ -344,6 +526,13 @@ class UltimateSDUpscaleDistributed: end_idx = min(end_idx, total_tiles) return list(range(start_idx, end_idx)) + def _get_master_global_indices(self, total_tiles, num_workers): + """Calculate which global tile indices are assigned to the master in flattened mode.""" + tiles_per_participant = total_tiles // (num_workers + 1) + remainder = total_tiles % (num_workers + 1) + master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) + return list(range(master_tiles)) + def _get_master_tiles(self, all_tiles, num_workers): """Calculate which tiles are assigned to the master.""" total_tiles = len(all_tiles) @@ -355,7 +544,7 @@ class UltimateSDUpscaleDistributed: def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, image_width, image_height): + padding, mask_blur, image_width, image_height, tiled_decode): """Process a single tile and blend it into the result image.""" x, y = tile_pos @@ -366,7 +555,7 @@ class UltimateSDUpscaleDistributed: processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise) + scheduler, denoise, tiled_decode) # Convert and blend processed_pil = tensor_to_pil(processed_tile, 0) @@ -389,6 +578,29 @@ class UltimateSDUpscaleDistributed: text = await response.text() raise RuntimeError(f"Failed to prepare job: {response.status} - {text}") + async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, worker_indices=None): + """Initialize the job queue for dynamic mode with pending images.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_tile_jobs: + prompt_server.distributed_pending_tile_jobs[multi_job_id] = { + 'queue': asyncio.Queue(), + 'mode': 'dynamic', + 'pending_images': asyncio.Queue(), + 'completed_images': {}, + 'completed_tiles': {}, # For uniformity + 'batch_size': batch_size, + 'assigned_to_workers': assigned_to_workers or {}, + 'worker_status': worker_status or {} + } + # Initialize pending images queue with worker indices only + pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] + for i in worker_indices: + await pending_images.put(i) + debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {len(worker_indices)} pending images for workers") + else: + debug_log(f"UltimateSDUpscale Master Dynamic - Queue already exists for {multi_job_id}") + async def _init_job_queue(self, multi_job_id): """Initialize the job queue for collecting tiles.""" prompt_server = ensure_tile_jobs_initialized() @@ -396,7 +608,11 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - _init_job_queue: Checking if {multi_job_id} exists in distributed_pending_tile_jobs") debug_log(f"UltimateSDUpscale Master - _init_job_queue: Current jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() + prompt_server.distributed_pending_tile_jobs[multi_job_id] = { + 'queue': asyncio.Queue(), + 'mode': 'static', + 'completed_tiles': {} # For static mode data + } debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") debug_log(f"UltimateSDUpscale Master - _init_job_queue: Created new queue for {multi_job_id}") else: @@ -409,7 +625,12 @@ class UltimateSDUpscaleDistributed: async with prompt_server.distributed_tile_jobs_lock: if multi_job_id not in prompt_server.distributed_pending_tile_jobs: raise RuntimeError(f"Job queue not initialized for {multi_job_id}") - q = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + raise RuntimeError("Invalid job data structure") + if job_data['mode'] != 'static': + raise RuntimeError(f"Mode mismatch: expected static, got {job_data['mode']}") + q = job_data['queue'] collected_tiles = {} workers_done = set() @@ -431,9 +652,18 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") for tile_data in tiles: + # Validate required fields + if 'batch_idx' not in tile_data: + log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") + continue + batch_idx = tile_data['batch_idx'] + tile_idx = tile_data['tile_idx'] + # Use global_idx as key if available (for batch processing) + key = tile_data.get('global_idx', tile_idx) + # Store the full tile data including metadata - collected_tiles[tile_idx] = { + collected_tiles[key] = { 'tensor': tile_data['tensor'], 'tile_idx': tile_idx, 'x': tile_data['x'], @@ -441,7 +671,9 @@ class UltimateSDUpscaleDistributed: 'extracted_width': tile_data['extracted_width'], 'extracted_height': tile_data['extracted_height'], 'padding': tile_data['padding'], - 'worker_id': worker_id + 'worker_id': worker_id, + 'batch_idx': tile_data.get('batch_idx', 0), + 'global_idx': tile_data.get('global_idx', tile_idx) } else: # Single tile mode (backward compat) @@ -470,21 +702,160 @@ class UltimateSDUpscaleDistributed: return collected_tiles + async def _async_process_dynamic(self, multi_job_id, expected_from_workers, num_workers): + """Async helper for dynamic mode - collect processed whole images from workers. + + This method implements dynamic load balancing where workers request images + as they become available, ensuring optimal utilization of GPU resources.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_tile_jobs: + raise RuntimeError(f"Job queue not initialized for {multi_job_id}") + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + raise RuntimeError("Invalid job data structure") + if job_data['mode'] != 'dynamic': + raise RuntimeError(f"Mode mismatch: expected dynamic, got {job_data['mode']}") + q = job_data['queue'] + completed_images = job_data['completed_images'] + + workers_done = set() + timeout = TILE_WAIT_TIMEOUT * 2 # Longer timeout for full images + + debug_log(f"UltimateSDUpscale Master Dynamic - Waiting for {expected_from_workers} images from {num_workers} workers") + + last_heartbeat_check = time.time() + + while len(completed_images) < expected_from_workers and len(workers_done) < num_workers: + try: + result = await asyncio.wait_for(q.get(), timeout=10.0) # Shorter timeout for regular checks + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + completed_images[image_idx] = image_pil + debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}") + + if is_last: + workers_done.add(worker_id) + debug_log(f"UltimateSDUpscale Master Dynamic - Worker {worker_id} completed") + + except asyncio.TimeoutError: + # Check for worker timeouts every 10 seconds + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data_check = prompt_server.distributed_pending_tile_jobs[multi_job_id] + # Check for timed out workers (60 second timeout) + for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): + if current_time - last_heartbeat > 60: + log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") + # Requeue assigned images from this worker + for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): + if idx not in completed_images: + await job_data_check['pending_images'].put(idx) + debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx} from timed out worker {worker}") + # Mark worker as done and clean up + workers_done.add(worker) + if 'worker_status' in job_data_check: + del job_data_check['worker_status'][worker] + if 'assigned_to_workers' in job_data_check: + job_data_check['assigned_to_workers'][worker] = [] + last_heartbeat_check = current_time + + # Check if we've been waiting too long overall + if current_time - last_heartbeat_check > timeout: + log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") + # Requeue all unfinished images + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + pending_queue = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] + for idx in range(job_data_check.get('batch_size', 0)): + if idx not in completed_images: + await pending_queue.put(idx) + debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx}") + break + + debug_log(f"UltimateSDUpscale Master Dynamic - Collection complete. Got {len(completed_images)} images") + + # Clean up job queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + del prompt_server.distributed_pending_tile_jobs[multi_job_id] + + return completed_images + def round_to_multiple(self, value: int, multiple: int = 8) -> int: """Round value to nearest multiple.""" return round(value / multiple) * multiple def calculate_tiles(self, image_width: int, image_height: int, - tile_width: int, tile_height: int) -> List[Tuple[int, int]]: + tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]: """Calculate tile positions for the image. Tiles are placed at grid positions without overlap in their placement. - The overlap happens during extraction with padding.""" - tiles = [] - for y in range(0, image_height, tile_height): - for x in range(0, image_width, tile_width): - tiles.append((x, y)) - return tiles + The overlap happens during extraction with padding. + + If force_uniform_tiles is True, tiles are distributed evenly across the image.""" + if force_uniform_tiles: + # Calculate number of tiles needed + num_tiles_x = math.ceil(image_width / tile_width) + num_tiles_y = math.ceil(image_height / tile_height) + + # Calculate uniform spacing + tiles = [] + + # Handle X dimension + if num_tiles_x == 1: + tiles_x = [0] + else: + # Calculate step to distribute tiles evenly + step_x = (image_width - tile_width) / (num_tiles_x - 1) + # Check for excessive overlap + if step_x < tile_width / 2: + debug_log(f"Warning: Uniform tiles would have excessive overlap (step_x={step_x}, tile_width={tile_width}). Falling back to non-uniform mode.") + force_uniform_tiles = False + else: + tiles_x = [round(i * step_x) for i in range(num_tiles_x)] + # Ensure tiles are aligned to 8-pixel boundaries + tiles_x = [self.round_to_multiple(x, 8) for x in tiles_x] + + # Handle Y dimension (if still using uniform tiles) + if force_uniform_tiles: + if num_tiles_y == 1: + tiles_y = [0] + else: + # Calculate step to distribute tiles evenly + step_y = (image_height - tile_height) / (num_tiles_y - 1) + # Check for excessive overlap + if step_y < tile_height / 2: + debug_log(f"Warning: Uniform tiles would have excessive overlap (step_y={step_y}, tile_height={tile_height}). Falling back to non-uniform mode.") + force_uniform_tiles = False + else: + tiles_y = [round(i * step_y) for i in range(num_tiles_y)] + # Ensure tiles are aligned to 8-pixel boundaries + tiles_y = [self.round_to_multiple(y, 8) for y in tiles_y] + + # If still using uniform tiles, generate positions + if force_uniform_tiles: + # Generate all tile positions + for y in tiles_y: + for x in tiles_x: + tiles.append((x, y)) + + return tiles + + # Fall through to non-uniform mode if needed + if not force_uniform_tiles: + # Original non-uniform tile calculation + tiles = [] + for y in range(0, image_height, tile_height): + for x in range(0, image_width, tile_width): + tiles.append((x, y)) + return tiles def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int, tile_width: int, tile_height: int, padding: int) -> Tuple[torch.Tensor, int, int, int, int]: @@ -524,11 +895,20 @@ class UltimateSDUpscaleDistributed: def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae, seed: int, steps: int, cfg: float, sampler_name: str, - scheduler: str, denoise: float) -> torch.Tensor: + scheduler: str, denoise: float, tiled_decode: bool = False) -> torch.Tensor: """Process a single tile through SD sampling.""" # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode + # Try to import tiled VAE nodes if available + try: + from nodes import VAEEncodeTiled, VAEDecodeTiled + tiled_vae_available = True + except ImportError: + tiled_vae_available = False + if tiled_decode: + debug_log("Tiled VAE nodes not available, falling back to standard VAE") + # Convert to PIL and back to ensure clean tensor without gradient tracking tile_pil = tensor_to_pil(tile_tensor, 0) clean_tensor = pil_to_tensor(tile_pil) @@ -538,14 +918,20 @@ class UltimateSDUpscaleDistributed: clean_tensor = clean_tensor.cuda() # Encode to latent - latent = VAEEncode().encode(vae, clean_tensor)[0] + if tiled_decode and tiled_vae_available: + latent = VAEEncodeTiled().encode(vae, clean_tensor)[0] + else: + latent = VAEEncode().encode(vae, clean_tensor)[0] # Sample samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=denoise)[0] # Decode back to image - image = VAEDecode().decode(vae, samples)[0] + if tiled_decode and tiled_vae_available: + image = VAEDecodeTiled().decode(vae, samples)[0] + else: + image = VAEDecode().decode(vae, samples)[0] return image @@ -572,7 +958,7 @@ class UltimateSDUpscaleDistributed: def blend_tile(self, base_image: Image.Image, tile_image: Image.Image, x: int, y: int, extracted_size: Tuple[int, int], - mask: Image.Image, padding: int) -> Image.Image: + mask: int, padding: int) -> Image.Image: """Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach. This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py @@ -635,16 +1021,22 @@ class UltimateSDUpscaleDistributed: data.add_field('batch_size', str(chunk_size)) data.add_field('padding', str(padding)) - # Chunk metadata: Keep absolute tile_idx + # Chunk metadata: Keep absolute tile_idx and batch_idx if present metadata = [] for tile_data in chunk: - metadata.append({ + meta = { 'tile_idx': tile_data['tile_idx'], # Original/absolute idx 'x': tile_data['x'], 'y': tile_data['y'], 'extracted_width': tile_data['extracted_width'], 'extracted_height': tile_data['extracted_height'] - }) + } + # Add batch_idx and global_idx if present (for batch processing) + if 'batch_idx' in tile_data: + meta['batch_idx'] = tile_data['batch_idx'] + if 'global_idx' in tile_data: + meta['global_idx'] = tile_data['global_idx'] + metadata.append(meta) # Add JSON metadata data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') @@ -696,51 +1088,493 @@ class UltimateSDUpscaleDistributed: def process_single_gpu(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, force_uniform_tiles): + tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode): """Process all tiles on a single GPU (no distribution).""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) - # Get image dimensions - _, height, width, _ = upscaled_image.shape + # Get image dimensions and batch size + batch_size, height, width, _ = upscaled_image.shape # Calculate all tiles - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally") + debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally for batch of {batch_size}") - # Convert to PIL - image_pil = tensor_to_pil(upscaled_image, 0) - result_image = image_pil.copy() + # Convert batch to PIL list + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + # Ensure RGB mode for consistency + if image_pil.mode != 'RGB': + image_pil = image_pil.convert('RGB') + result_images.append(image_pil.copy()) - # Mask will be created per tile for proper blending - - # Process each tile - for idx, tile_pos in enumerate(all_tiles): - result_image = self._process_and_blend_tile( - idx, tile_pos, upscaled_image, result_image, - model, positive, negative, vae, seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height - ) + # Process each image in the batch + for batch_idx in range(batch_size): + # Process each tile for this image + for tile_idx, tile_pos in enumerate(all_tiles): + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + result_images[batch_idx] = self._process_and_blend_tile( + tile_idx, tile_pos, upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive, negative, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode + ) # Convert back to tensor - result_tensor = pil_to_tensor(result_image) + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + if upscaled_image.is_cuda: result_tensor = result_tensor.cuda() return (result_tensor,) + + def process_master_dynamic(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers): + """Dynamic mode for large batches - assigns whole images to workers dynamically.""" + # Get batch size and dimensions + batch_size, height, width, _ = upscaled_image.shape + num_workers = len(enabled_workers) + + debug_log(f"UltimateSDUpscale Master Dynamic - Processing batch of {batch_size} images with {num_workers} workers") + + # Calculate master's share similar to static mode + num_participants = num_workers + 1 + images_per_participant = batch_size // num_participants + remainder = batch_size % num_participants + master_images_count = images_per_participant + (1 if remainder > 0 else 0) + master_indices = list(range(master_images_count)) + worker_indices = list(range(master_images_count, batch_size)) + expected_from_workers = len(worker_indices) + + # Calculate tiles for processing + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + # Initialize job queue for communication - synchronous approach + try: + # Add worker tracking data + assigned_to_workers = {w: [] for w in enabled_workers} + worker_status = {w: time.time() for w in enabled_workers} + + # Use short timeout for async init, with threading fallback + try: + run_async_in_server_loop( + self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, worker_indices), + timeout=2.0 # Short timeout + ) + except Exception as init_error: + debug_log(f"UltimateSDUpscale Master Dynamic - Async init failed, using threading: {init_error}") + # Threading fallback for dynamic mode + import threading + def sync_dynamic_init(): + prompt_server = ensure_tile_jobs_initialized() + if multi_job_id not in prompt_server.distributed_pending_tile_jobs: + prompt_server.distributed_pending_tile_jobs[multi_job_id] = { + 'queue': asyncio.Queue(), + 'mode': 'dynamic', + 'pending_images': asyncio.Queue(), + 'completed_images': {}, + 'completed_tiles': {}, + 'batch_size': batch_size, + 'assigned_to_workers': assigned_to_workers, + 'worker_status': worker_status + } + # Initialize pending images queue with worker indices only + pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] + for i in worker_indices: + pending_images.put_nowait(i) + debug_log(f"Synchronous dynamic init for {multi_job_id} with {len(worker_indices)} pending images for workers") + thread = threading.Thread(target=sync_dynamic_init) + thread.start() + thread.join(timeout=3.0) + if thread.is_alive(): + raise RuntimeError("Dynamic queue init failed") + except Exception as e: + debug_log(f"UltimateSDUpscale Master Dynamic - Queue initialization error: {e}") + raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}") + + # Convert batch to PIL list + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + # Ensure RGB mode for consistency + if image_pil.mode != 'RGB': + image_pil = image_pil.convert('RGB') + result_images.append(image_pil.copy()) + + # Process master's share locally + debug_log(f"UltimateSDUpscale Master Dynamic - Processing {master_images_count} images locally") + for idx in master_indices: + single_tensor = upscaled_image[idx:idx+1] + local_image = result_images[idx] + image_seed = seed + idx * 1000 + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive, negative, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode + ) + result_images[idx] = local_image + + # Process with workers using dynamic assignment + collected_images = run_async_in_server_loop( + self._async_process_dynamic(multi_job_id, expected_from_workers, num_workers), + timeout=TILE_COLLECTION_TIMEOUT * 2 # Longer timeout for full images + ) + + # Update result images with processed ones + for idx, processed_img in collected_images.items(): + if idx < batch_size: + result_images[idx] = processed_img + + # Convert back to tensor + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + debug_log(f"UltimateSDUpscale Master Dynamic - Job {multi_job_id} complete") + return (result_tensor,) + + def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold): + """Worker processing in dynamic mode - processes whole images.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get dimensions + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} ready to process images") + + # Keep track of processed images for is_last detection + processed_count = 0 + + # Poll for job readiness to avoid races during master init + max_poll_attempts = 20 # ~20s at 1s sleep + for attempt in range(max_poll_attempts): + ready = run_async_in_server_loop( + self._check_job_status(multi_job_id, master_url), + timeout=5.0 + ) + if ready: + debug_log(f"Job {multi_job_id} ready after {attempt} attempts") + break + time.sleep(1.0) # Poll every 1s + else: + log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") + return (upscaled_image,) + + # Loop to request and process images + while True: + # Request an image to process + image_idx, estimated_remaining = run_async_in_server_loop( + self._request_image_from_master(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT + ) + + if image_idx is None: + debug_log(f"UltimateSDUpscale Worker Dynamic - No more images to process") + break + + debug_log(f"UltimateSDUpscale Worker Dynamic - Mode: dynamic, assigned image {image_idx} to worker {worker_id}") + processed_count += 1 + + # Determine if this should be marked as last for this worker + is_last_for_worker = (estimated_remaining == 0) + + # Extract single image tensor + single_tensor = upscaled_image[image_idx:image_idx+1] + + # Convert to PIL for processing + local_image = tensor_to_pil(single_tensor, 0).copy() + + # Process all tiles for this image + image_seed = seed + image_idx * 1000 + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive, negative, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode + ) + + # Send processed image back to master + try: + # Use the estimated remaining to determine if this is the last image + is_last = is_last_for_worker + run_async_in_server_loop( + self._send_full_image_to_master(local_image, image_idx, multi_job_id, + master_url, worker_id, is_last), + timeout=TILE_SEND_TIMEOUT + ) + # Send heartbeat after processing + run_async_in_server_loop( + self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + if is_last: + break + except Exception as e: + log(f"UltimateSDUpscale Worker Dynamic - Error sending image {image_idx}: {e}") + # Continue processing other images + + # Send final is_last signal + debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} processed {processed_count} images, sending completion signal") + try: + run_async_in_server_loop( + self._send_worker_complete_signal(multi_job_id, master_url, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + except Exception as e: + log(f"UltimateSDUpscale Worker Dynamic - Error sending completion signal: {e}") + + return (upscaled_image,) + + async def _request_image_from_master(self, multi_job_id, master_url, worker_id): + """Request an image index to process from master in dynamic mode.""" + # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races + max_retries = 10 + retry_delay = 0.5 + start_time = time.time() + + for attempt in range(max_retries): + # Check total timeout + if time.time() - start_time > 30: + log(f"Total request timeout after 30s for worker {worker_id}") + return None, 0 + + try: + session = await get_client_session() + url = f"{master_url}/distributed/request_image" + + async with session.post(url, json={ + 'worker_id': str(worker_id), + 'multi_job_id': multi_job_id + }) as response: + if response.status == 200: + data = await response.json() + image_idx = data.get('image_idx') + estimated_remaining = data.get('estimated_remaining', 0) + return image_idx, estimated_remaining + elif response.status == 404: + # Special handling for 404 - job not found yet + text = await response.text() + debug_log(f"Job not found (404), will retry: {text}") + await asyncio.sleep(1.0) + else: + text = await response.text() + debug_log(f"Request image failed: {response.status} - {text}") + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + else: + log(f"Failed to request image after {max_retries} attempts: {e}") + raise + + return None, 0 + + async def _check_job_status(self, multi_job_id, master_url): + """Check if job is ready on the master.""" + try: + session = await get_client_session() + url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}" + async with session.get(url) as response: + if response.status == 200: + data = await response.json() + return data.get('ready', False) + return False + except Exception as e: + debug_log(f"Job status check failed: {e}") + return False + + async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, + master_url, worker_id, is_last): + """Send a processed full image back to master in dynamic mode.""" + # Serialize image to PNG + byte_io = io.BytesIO() + image_pil.save(byte_io, format='PNG', compress_level=0) + byte_io.seek(0) + + # Prepare form data + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('image_idx', str(image_idx)) + data.add_field('is_last', str(is_last)) + data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', + content_type='image/png') + + # Retry logic + max_retries = 5 + retry_delay = 0.5 + + for attempt in range(max_retries): + try: + session = await get_client_session() + url = f"{master_url}/distributed/tile_complete" + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Successfully sent image {image_idx} to master") + return + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + else: + log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}") + raise + + async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id): + """Send heartbeat signal to master.""" + try: + data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)} + session = await get_client_session() + url = f"{master_url}/distributed/heartbeat" + + async with session.post(url, json=data) as response: + response.raise_for_status() + debug_log(f"Worker {worker_id} sent heartbeat") + except Exception as e: + debug_log(f"Worker {worker_id} heartbeat failed: {e}") + + async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): + """Send completion signal to master in dynamic mode.""" + # Send a dummy request with is_last=True + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', 'true') + # Send empty tile data to signal completion + data.add_field('batch_size', '0') + + session = await get_client_session() + url = f"{master_url}/distributed/tile_complete" + + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Worker {worker_id} sent completion signal") # Ensure initialization before registering routes ensure_tile_jobs_initialized() # API Endpoint for tile completion +@server.PromptServer.instance.routes.post("/distributed/request_image") +async def request_image_endpoint(request): + """Endpoint for workers to request images in dynamic mode.""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'dynamic': + return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) + if 'pending_images' in job_data: + pending_queue = job_data['pending_images'] + try: + image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + # Track assigned image + if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: + job_data['assigned_to_workers'][worker_id].append(image_idx) + # Update worker heartbeat + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + # Get estimated remaining count + remaining = pending_queue.qsize() # Approximate + debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") + return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) + except asyncio.TimeoutError: + return web.json_response({"image_idx": None}) + else: + return await handle_api_error(request, "Invalid dynamic mode configuration", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.get("/distributed/job_status") +async def job_status_endpoint(request): + """Endpoint to check if a job is ready.""" + multi_job_id = request.query.get('multi_job_id') + if not multi_job_id: + return web.json_response({"ready": False}) + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) + return web.json_response({"ready": ready}) + +@server.PromptServer.instance.routes.post("/distributed/heartbeat") +async def heartbeat_endpoint(request): + """Endpoint for workers to send heartbeat signals.""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Worker status tracking not available", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + @server.PromptServer.instance.routes.post("/distributed/tile_complete") async def tile_complete_endpoint(request): """Endpoint for receiving completed tiles from workers.""" try: + # Check payload size before processing + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413) + data = await request.post() multi_job_id = data.get('multi_job_id') worker_id = data.get('worker_id') @@ -751,10 +1585,68 @@ async def tile_complete_endpoint(request): prompt_server = ensure_tile_jobs_initialized() + # Check if this is a full image submission (dynamic mode) + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + is_last = data.get('is_last', 'False').lower() == 'true' + + try: + # Process full image + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}") + + # Put into dynamic mode queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'dynamic': + return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) + # Check for duplicate image submission + if 'completed_images' in job_data and image_idx in job_data['completed_images']: + log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring") + return web.json_response({"status": "duplicate"}) + if 'queue' in job_data: + await job_data['queue'].put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Invalid dynamic mode configuration", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + log(f"Error processing full image from worker {worker_id}: {e}") + return await handle_api_error(request, f"Image processing error: {e}", 400) + # Check for batch mode batch_size = int(data.get('batch_size', 0)) tiles = [] + # Handle completion signals (batch_size=0 with is_last=True) + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if isinstance(job_data, dict) and 'queue' in job_data: + await job_data['queue'].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Invalid job structure", 500) + else: + return await handle_api_error(request, "Job not found", 404) + if batch_size > 0: # Batch mode: Extract multiple tiles padding = int(data.get('padding', 32)) @@ -811,7 +1703,9 @@ async def tile_complete_endpoint(request): 'y': tile_data['y'], 'extracted_width': tile_data['extracted_width'], 'extracted_height': tile_data['extracted_height'], - 'padding': padding + 'padding': padding, + 'batch_idx': tile_data.get('batch_idx', 0), + 'global_idx': tile_data.get('global_idx', tile_idx) } tile_data_list.append(tile_info) else: @@ -896,7 +1790,9 @@ async def tile_complete_endpoint(request): 'y': y, 'extracted_width': extracted_width, 'extracted_height': extracted_height, - 'padding': padding + 'padding': padding, + 'batch_idx': 0, # Default for legacy single-tile + 'global_idx': tile_idx # Default for legacy single-tile }] except Exception as e: log(f"Error processing tile from worker {worker_id}: {e}") @@ -908,9 +1804,16 @@ async def tile_complete_endpoint(request): debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'static': + return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400) + + q = job_data['queue'] if batch_size > 0: # Put batch as single item - await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + await q.put({ 'worker_id': worker_id, 'tiles': tiles, 'is_last': is_last @@ -919,7 +1822,7 @@ async def tile_complete_endpoint(request): else: # Put single tile (backward compat) tile_data = tiles[0] - await prompt_server.distributed_pending_tile_jobs[multi_job_id].put({ + await q.put({ 'tensor': tile_data['tensor'], 'worker_id': worker_id, 'tile_idx': tile_data['tile_idx'], From abcbab757422f5ff0fee99da382ebc1b9d0ce878 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 10:23:27 +1000 Subject: [PATCH 23/93] support sending videos to remote workers --- distributed.py | 193 ++++++++++++++++++++++++++++++++++++++---- web/executionUtils.js | 31 ++++--- web/workerUtils.js | 21 +++-- 3 files changed, 210 insertions(+), 35 deletions(-) diff --git a/distributed.py b/distributed.py index 9664862..6e7bd2e 100644 --- a/distributed.py +++ b/distributed.py @@ -16,6 +16,8 @@ import atexit import signal import sys import shlex +import uuid +from multiprocessing import Queue from comfy.utils import ProgressBar # Import shared utilities @@ -258,6 +260,33 @@ async def get_network_info_endpoint(request): "recommended_ip": None }) +@server.PromptServer.instance.routes.get("/distributed/system_info") +async def get_system_info_endpoint(request): + """Get system information including machine ID for local worker detection.""" + try: + import socket + + return web.json_response({ + "status": "success", + "hostname": socket.gethostname(), + "machine_id": get_machine_id(), + "platform": { + "system": platform.system(), + "machine": platform.machine(), + "node": platform.node(), + "path_separator": os.sep, # Add path separator + "os_name": os.name # Add OS name (posix, nt, etc.) + }, + "is_docker": is_docker_environment(), + "is_runpod": is_runpod_environment(), + "runpod_pod_id": os.environ.get('RUNPOD_POD_ID') + }) + except Exception as e: + return web.json_response({ + "status": "error", + "message": str(e) + }, status=500) + @server.PromptServer.instance.routes.post("/distributed/config/update_worker") async def update_worker_endpoint(request): try: @@ -1156,6 +1185,106 @@ class WorkerProcessManager: # Create global instance worker_manager = WorkerProcessManager() +# Add IPC queue storage +worker_manager.queues = {} + +# Helper to get aiohttp session +async def get_client_session(): + """Get or create aiohttp client session.""" + if not hasattr(server.PromptServer.instance, '_distributed_session'): + server.PromptServer.instance._distributed_session = aiohttp.ClientSession() + return server.PromptServer.instance._distributed_session + +# Local worker detection functions +async def is_local_worker(worker_config): + """Check if a worker is running on the same machine as the master.""" + host = worker_config.get('host', 'localhost') + if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local': + return True + + # For cloud workers, check if on same physical host + if worker_config.get('type') == 'cloud': + return await is_same_physical_host(worker_config) + + return False + +async def is_same_physical_host(worker_config): + """Compare machine IDs to determine if worker is on same physical host.""" + try: + # Get master machine ID + master_machine_id = get_machine_id() + + # Fetch worker's machine ID via API + host = worker_config.get('host', 'localhost') + port = worker_config.get('port', 8188) + + session = await get_client_session() + async with session.get( + f"http://{host}:{port}/distributed/system_info", + timeout=aiohttp.ClientTimeout(total=5) + ) as resp: + if resp.status == 200: + data = await resp.json() + worker_machine_id = data.get('machine_id') + return worker_machine_id == master_machine_id + else: + debug_log(f"Failed to get system info from worker: HTTP {resp.status}") + return False + except Exception as e: + debug_log(f"Error checking same physical host: {e}") + return False + +def get_machine_id(): + """Get a unique identifier for this machine.""" + # Try multiple methods to get a stable machine ID + try: + # Method 1: MAC address-based UUID + return str(uuid.getnode()) + except: + try: + # Method 2: Platform + hostname + import socket + return f"{platform.machine()}_{socket.gethostname()}" + except: + # Fallback + return platform.machine() + +def is_docker_environment(): + """Check if running inside Docker container.""" + return (os.path.exists('/.dockerenv') or + os.environ.get('DOCKER_CONTAINER', False) or + 'docker' in platform.node().lower()) + +def is_runpod_environment(): + """Check if running in Runpod environment.""" + return (os.environ.get('RUNPOD_POD_ID') is not None or + os.environ.get('RUNPOD_API_KEY') is not None) + +async def get_comms_channel(worker_id, worker_config): + """Get communication channel for a worker (HTTP URL or IPC Queue).""" + if await is_local_worker(worker_config): + comm_mode = worker_config.get('communicationMode', 'http') + + if comm_mode == 'ipc': + # Use multiprocessing Queue for IPC + if worker_id not in worker_manager.queues: + worker_manager.queues[worker_id] = Queue() + return worker_manager.queues[worker_id] + elif is_docker_environment(): + # Docker: use host.docker.internal + return f"http://host.docker.internal:{worker_config['port']}" + else: + # Local: use loopback + return f"http://127.0.0.1:{worker_config['port']}" + elif worker_config.get('type') == 'cloud' and is_runpod_environment(): + # Runpod same-host optimization (if detected) + host = worker_config.get('host', 'localhost') + return f"http://{host}:{worker_config['port']}" + else: + # Remote worker: use configured endpoint + host = worker_config.get('host', 'localhost') + return f"http://{host}:{worker_config['port']}" + # Auto-launch workers if enabled def auto_launch_workers(): """Launch enabled workers if auto_launch_workers is set to true.""" @@ -1335,7 +1464,7 @@ if not hasattr(prompt_server, 'distributed_pending_jobs'): @server.PromptServer.instance.routes.post("/distributed/load_image") async def load_image_endpoint(request): - """Load an image file and return it as base64 data.""" + """Load an image or video file and return it as base64 data.""" try: data = await request.json() image_path = data.get("image_path") @@ -1345,30 +1474,58 @@ async def load_image_endpoint(request): import folder_paths import base64 - from PIL import Image import io - # Use ComfyUI's folder paths to find the image + # Use ComfyUI's folder paths to find the file full_path = folder_paths.get_annotated_filepath(image_path) if not os.path.exists(full_path): - return await handle_api_error(request, f"Image not found: {image_path}", 404) + return await handle_api_error(request, f"File not found: {image_path}", 404) - # Load and convert to base64 - with Image.open(full_path) as img: - # Convert to RGB if needed - if img.mode not in ('RGB', 'RGBA'): - img = img.convert('RGB') + # Check if it's a video file + video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'} + file_ext = os.path.splitext(full_path)[1].lower() + + if file_ext in video_extensions: + # For video files, read the raw bytes + with open(full_path, 'rb') as f: + file_data = f.read() - # Save to bytes - buffer = io.BytesIO() - img.save(buffer, format='PNG', compress_level=1) # Fast compression - img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') - - return web.json_response({ - "status": "success", - "image_data": f"data:image/png;base64,{img_base64}" - }) + # Determine MIME type + mime_types = { + '.mp4': 'video/mp4', + '.avi': 'video/x-msvideo', + '.mov': 'video/quicktime', + '.mkv': 'video/x-matroska', + '.webm': 'video/webm' + } + mime_type = mime_types.get(file_ext, 'video/mp4') + + # Return base64 encoded video with data URL + video_base64 = base64.b64encode(file_data).decode('utf-8') + return web.json_response({ + "status": "success", + "image_data": f"data:{mime_type};base64,{video_base64}" + }) + else: + # For images, use PIL + from PIL import Image + + # Load and convert to base64 + with Image.open(full_path) as img: + # Convert to RGB if needed + if img.mode not in ('RGB', 'RGBA'): + img = img.convert('RGB') + + # Save to bytes + buffer = io.BytesIO() + img.save(buffer, format='PNG', compress_level=1) # Fast compression + img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8') + + return web.json_response({ + "status": "success", + "image_data": f"data:image/png;base64,{img_base64}" + }) except Exception as e: return await handle_api_error(request, e, 500) diff --git a/web/executionUtils.js b/web/executionUtils.js index bbf700b..56ad8c2 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -9,6 +9,12 @@ import { TIMEOUTS } from './constants.js'; * @returns {Object} The converted API prompt */ function convertPathsForPlatform(apiPrompt, targetSeparator) { + // Validate target separator + if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) { + console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion'); + return apiPrompt; + } + // Regex to identify likely file paths with extensions const isLikelyFilename = (value) => { return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i); @@ -152,12 +158,15 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p try { const workerUrl = extension.getWorkerUrl(workerInfo); const systemInfo = await getCachedWorkerSystemInfo(workerUrl); - const targetSeparator = systemInfo.platform.path_separator; + const targetSeparator = systemInfo?.platform?.path_separator; - // Convert paths to match worker's platform - jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator); - - extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug"); + if (targetSeparator) { + // Convert paths to match worker's platform + jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator); + extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug"); + } else { + extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug"); + } } catch (e) { extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn"); // Continue without path conversion @@ -178,10 +187,10 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p const isRemote = workerInfo && workerInfo.host; if (isRemote) { - // Find all image references in the pruned workflow + // Find all image/video references in the pruned workflow const imageReferences = findImageReferences(extension, jobApiPrompt); if (imageReferences.size > 0) { - extension.log(`Found ${imageReferences.size} image references for remote worker ${workerId}`, "debug"); + extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug"); // Store image references for later processing options._imageReferences = imageReferences; } @@ -284,7 +293,7 @@ export async function executeJobs(extension, jobs) { } if (allImageReferences.size > 0) { - extension.log(`Pre-loading ${allImageReferences.size} unique image(s) for all workers`, "debug"); + extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug"); await loadImagesForWorker(extension, allImageReferences); } @@ -414,12 +423,12 @@ export async function uploadImagesToWorker(extension, workerUrl, images) { for (const imageData of images) { const formData = new FormData(); - // Detect MIME type from base64 header - const mimeMatch = imageData.image.match(/^data:(image\/\w+);base64,/); + // Detect MIME type from base64 header (supports both image and video) + const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\w+);base64,/); const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected // Convert base64 to blob - const base64Data = imageData.image.replace(/^data:image\/\w+;base64,/, ''); + const base64Data = imageData.image.replace(/^data:(?:image|video)\/\w+;base64,/, ''); const byteCharacters = atob(base64Data); const byteNumbers = new Array(byteCharacters.length); for (let i = 0; i < byteCharacters.length; i++) { diff --git a/web/workerUtils.js b/web/workerUtils.js index 63280f7..ef87542 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -88,21 +88,30 @@ export function findImageReferences(extension, apiPrompt) { // - Standard files: "image.png" // - Subfolder files: "subfolder/image.png" // - ComfyUI special format: "clipspace/file.png [input]" - const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp)(\s*\[\w+\])?$/i; + // - Video files: "video.mp4", "animation.avi", etc. + const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i; for (const [nodeId, node] of Object.entries(apiPrompt)) { + // Check for both 'image' and 'video' inputs + const mediaInputs = []; if (node.inputs && node.inputs.image) { - const imageValue = node.inputs.image; - if (typeof imageValue === 'string') { + mediaInputs.push(node.inputs.image); + } + if (node.inputs && node.inputs.video) { + mediaInputs.push(node.inputs.video); + } + + for (const mediaValue of mediaInputs) { + if (typeof mediaValue === 'string') { // Clean special suffixes like [input] or [output] - const cleanValue = imageValue.replace(/\s*\[\w+\]$/, '').trim(); + const cleanValue = mediaValue.replace(/\s*\[\w+\]$/, '').trim(); if (imageExtensions.test(cleanValue)) { images.set(cleanValue, { nodeId, nodeType: node.class_type, - inputName: 'image' + inputName: 'image' // Keep as 'image' for compatibility }); - extension.log(`Found image reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug"); + extension.log(`Found media reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug"); } } } From 45e66bd26b81fd4d998f839f5e19bdc5da0baf40 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 10:29:18 +1000 Subject: [PATCH 24/93] fix race condition in upscaler --- distributed_upscale.py | 48 ++++++------------------------------------ 1 file changed, 6 insertions(+), 42 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 371050c..66d9e3d 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -54,12 +54,11 @@ def ensure_tile_jobs_initialized(): prompt_server.distributed_tile_jobs_lock = asyncio.Lock() else: # Clean up any legacy queue structures that don't have the 'mode' field - to_remove = [] - for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items(): - if not isinstance(job_data, dict) or 'mode' not in job_data: - debug_log(f"Removing legacy queue structure for job {job_id}") - to_remove.append(job_id) + # (Should be rare after fixes, but keep for safety) + to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() + if not isinstance(job_data, dict) or 'mode' not in job_data] for job_id in to_remove: + debug_log(f"Removing legacy queue structure for job {job_id}") del prompt_server.distributed_pending_tile_jobs[job_id] return prompt_server @@ -132,24 +131,7 @@ class UltimateSDUpscaleDistributed: @classmethod def IS_CHANGED(cls, **kwargs): - """Force re-execution and ensure queues exist for the job.""" - multi_job_id = kwargs.get('multi_job_id', '') - if multi_job_id: - # Initialize queue for this job immediately - prompt_server = ensure_tile_jobs_initialized() - loop = get_server_loop() - - async def init_queue(): - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() - debug_log(f"UltimateSDUpscaleDistributed - Pre-initialized queue for job {multi_job_id}") - - try: - asyncio.run_coroutine_threadsafe(init_queue(), loop).result(timeout=1.0) - except Exception as e: - debug_log(f"UltimateSDUpscaleDistributed - Error pre-initializing queue: {e}") - + """Force re-execution.""" return float("nan") # Always re-execute def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, @@ -295,20 +277,6 @@ class UltimateSDUpscaleDistributed: enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) - # Synchronous queue initialization before mode determination - import threading - def sync_init(): - prompt_server = ensure_tile_jobs_initialized() - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() - debug_log(f"Synchronous queue init for {multi_job_id}") - thread = threading.Thread(target=sync_init) - thread.start() - thread.join(timeout=5.0) # Block with timeout - if thread.is_alive(): - log(f"Queue init timed out for {multi_job_id}") - raise RuntimeError("Queue init failed") - # Determine processing mode if num_workers == 0: mode = "single_gpu" @@ -345,11 +313,7 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Queue initialization complete") except Exception as e: debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}") - # Try synchronous initialization as fallback - prompt_server = ensure_tile_jobs_initialized() - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = asyncio.Queue() - debug_log(f"UltimateSDUpscale Master - Queue initialized via fallback") + raise RuntimeError(f"Failed to initialize static mode queue: {e}") # Don't try to prepare via API since we're using local queues # loop = asyncio.new_event_loop() From 963ff1d9de1d4e0555c4e4281dfce1698d14e140 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 11:04:06 +1000 Subject: [PATCH 25/93] fixed incorrect tile index warning --- distributed_upscale.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 66d9e3d..5fd51b8 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1656,10 +1656,6 @@ async def tile_complete_endpoint(request): tile_data = metadata[i] tile_idx = tile_data.get('tile_idx', i) - # Validate order - if i != tile_idx % batch_size: - debug_log(f"Warning: Tile order mismatch at position {i}, expected tile_idx {tile_idx}") - tile_info = { 'tensor': tensor, 'tile_idx': tile_idx, From e9e03998126831793d70e0cb3dc30c2cb1071d0f Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 14:18:17 +1000 Subject: [PATCH 26/93] only upload images/videos to workers if they dont already have them --- distributed.py | 58 ++++++++++++++++++++++++++++++++++++++++--- web/executionUtils.js | 58 ++++++++++++++++++++++++++++++++++--------- 2 files changed, 101 insertions(+), 15 deletions(-) diff --git a/distributed.py b/distributed.py index 6e7bd2e..7b53c41 100644 --- a/distributed.py +++ b/distributed.py @@ -1464,7 +1464,7 @@ if not hasattr(prompt_server, 'distributed_pending_jobs'): @server.PromptServer.instance.routes.post("/distributed/load_image") async def load_image_endpoint(request): - """Load an image or video file and return it as base64 data.""" + """Load an image or video file and return it as base64 data with hash.""" try: data = await request.json() image_path = data.get("image_path") @@ -1475,6 +1475,7 @@ async def load_image_endpoint(request): import folder_paths import base64 import io + import hashlib # Use ComfyUI's folder paths to find the file full_path = folder_paths.get_annotated_filepath(image_path) @@ -1482,6 +1483,13 @@ async def load_image_endpoint(request): if not os.path.exists(full_path): return await handle_api_error(request, f"File not found: {image_path}", 404) + # Calculate file hash + hash_md5 = hashlib.md5() + with open(full_path, 'rb') as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_md5.update(chunk) + file_hash = hash_md5.hexdigest() + # Check if it's a video file video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'} file_ext = os.path.splitext(full_path)[1].lower() @@ -1505,7 +1513,8 @@ async def load_image_endpoint(request): video_base64 = base64.b64encode(file_data).decode('utf-8') return web.json_response({ "status": "success", - "image_data": f"data:{mime_type};base64,{video_base64}" + "image_data": f"data:{mime_type};base64,{video_base64}", + "hash": file_hash }) else: # For images, use PIL @@ -1524,12 +1533,55 @@ async def load_image_endpoint(request): return web.json_response({ "status": "success", - "image_data": f"data:image/png;base64,{img_base64}" + "image_data": f"data:image/png;base64,{img_base64}", + "hash": file_hash }) except Exception as e: return await handle_api_error(request, e, 500) +@server.PromptServer.instance.routes.post("/distributed/check_file") +async def check_file_endpoint(request): + """Check if a file exists and matches the given hash.""" + try: + data = await request.json() + filename = data.get("filename") + expected_hash = data.get("hash") + + if not filename or not expected_hash: + return await handle_api_error(request, "Missing filename or hash", 400) + + import folder_paths + import hashlib + + # Use ComfyUI's folder paths to find the file + full_path = folder_paths.get_annotated_filepath(filename) + + if not os.path.exists(full_path): + return web.json_response({ + "status": "success", + "exists": False + }) + + # Calculate file hash + hash_md5 = hashlib.md5() + with open(full_path, 'rb') as f: + for chunk in iter(lambda: f.read(4096), b""): + hash_md5.update(chunk) + file_hash = hash_md5.hexdigest() + + # Check if hash matches + hash_matches = file_hash == expected_hash + + return web.json_response({ + "status": "success", + "exists": True, + "hash_matches": hash_matches + }) + + except Exception as e: + return await handle_api_error(request, e, 500) + @server.PromptServer.instance.routes.get("/distributed/system_info") async def system_info_endpoint(request): """Return system information including platform details""" diff --git a/web/executionUtils.js b/web/executionUtils.js index 56ad8c2..4c116cd 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -325,20 +325,27 @@ async function dispatchToWorker(extension, worker, prompt, workflow, imageRefere // Handle image uploads for remote workers if (imageReferences && imageReferences.size > 0) { - extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug"); + // Check if this is a local worker (same host as master) + const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost'); - try { - // Load images from master - const images = await loadImagesForWorker(extension, imageReferences); + if (isLocalWorker) { + extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug"); + } else { + extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug"); - // Upload images to worker - if (images.length > 0) { - await uploadImagesToWorker(extension, workerUrl, images); - extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug"); + try { + // Load images from master + const images = await loadImagesForWorker(extension, imageReferences); + + // Upload images to worker + if (images.length > 0) { + await uploadImagesToWorker(extension, workerUrl, images); + extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug"); + } + } catch (error) { + extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error"); + // Continue with workflow execution even if image upload fails } - } catch (error) { - extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error"); - // Continue with workflow execution even if image upload fails } } @@ -391,7 +398,8 @@ export async function loadImagesForWorker(extension, imageReferences) { const data = await extension.api.loadImage(filename); const imageData = { name: filename, - image: data.image_data + image: data.image_data, + hash: data.hash // Include hash from the response }; images.push(imageData); @@ -421,6 +429,32 @@ export async function loadImagesForWorker(extension, imageReferences) { export async function uploadImagesToWorker(extension, workerUrl, images) { // Upload images to worker's ComfyUI instance for (const imageData of images) { + // Check if file already exists with matching hash + if (imageData.hash) { + try { + const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify({ + filename: imageData.name, + hash: imageData.hash + }) + }); + + if (checkResponse.ok) { + const result = await checkResponse.json(); + if (result.exists && result.hash_matches) { + extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug"); + continue; + } + } + } catch (error) { + // If check fails, proceed with upload + extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug"); + } + } + const formData = new FormData(); // Detect MIME type from base64 header (supports both image and video) From 35a502eb681ea7e9d5d29306e218691157b05211 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 15:50:11 +1000 Subject: [PATCH 27/93] master shares images dynamicially --- distributed_upscale.py | 195 +++++++++++++++++++---------------------- 1 file changed, 92 insertions(+), 103 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 5fd51b8..5e5c5a2 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -542,7 +542,7 @@ class UltimateSDUpscaleDistributed: text = await response.text() raise RuntimeError(f"Failed to prepare job: {response.status} - {text}") - async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, worker_indices=None): + async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None): """Initialize the job queue for dynamic mode with pending images.""" prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: @@ -557,11 +557,11 @@ class UltimateSDUpscaleDistributed: 'assigned_to_workers': assigned_to_workers or {}, 'worker_status': worker_status or {} } - # Initialize pending images queue with worker indices only + # Initialize pending images queue with all indices pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] - for i in worker_indices: + for i in all_indices: await pending_images.put(i) - debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {len(worker_indices)} pending images for workers") + debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {len(all_indices)} pending images for all participants") else: debug_log(f"UltimateSDUpscale Master Dynamic - Queue already exists for {multi_job_id}") @@ -666,11 +666,16 @@ class UltimateSDUpscaleDistributed: return collected_tiles - async def _async_process_dynamic(self, multi_job_id, expected_from_workers, num_workers): - """Async helper for dynamic mode - collect processed whole images from workers. - - This method implements dynamic load balancing where workers request images - as they become available, ensuring optimal utilization of GPU resources.""" + async def _mark_image_completed(self, multi_job_id, image_idx, image_pil): + """Mark an image as completed in the job data.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + job_data['completed_images'][image_idx] = image_pil + + async def _async_collect_dynamic_images(self, multi_job_id, expected_from_workers, num_workers): + """Collect processed images from workers only.""" prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id not in prompt_server.distributed_pending_tile_jobs: @@ -684,13 +689,13 @@ class UltimateSDUpscaleDistributed: completed_images = job_data['completed_images'] workers_done = set() + collected_count = 0 timeout = TILE_WAIT_TIMEOUT * 2 # Longer timeout for full images + last_heartbeat_check = time.time() debug_log(f"UltimateSDUpscale Master Dynamic - Waiting for {expected_from_workers} images from {num_workers} workers") - last_heartbeat_check = time.time() - - while len(completed_images) < expected_from_workers and len(workers_done) < num_workers: + while collected_count < expected_from_workers and len(workers_done) < num_workers: try: result = await asyncio.wait_for(q.get(), timeout=10.0) # Shorter timeout for regular checks worker_id = result['worker_id'] @@ -700,6 +705,7 @@ class UltimateSDUpscaleDistributed: image_idx = result['image_idx'] image_pil = result['image'] completed_images[image_idx] = image_pil + collected_count += 1 debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}") if is_last: @@ -707,7 +713,7 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master Dynamic - Worker {worker_id} completed") except asyncio.TimeoutError: - # Check for worker timeouts every 10 seconds + # Check for worker timeouts current_time = time.time() if current_time - last_heartbeat_check >= 10.0: async with prompt_server.distributed_tile_jobs_lock: @@ -717,33 +723,17 @@ class UltimateSDUpscaleDistributed: for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): if current_time - last_heartbeat > 60: log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") - # Requeue assigned images from this worker - for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): - if idx not in completed_images: - await job_data_check['pending_images'].put(idx) - debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx} from timed out worker {worker}") - # Mark worker as done and clean up workers_done.add(worker) if 'worker_status' in job_data_check: del job_data_check['worker_status'][worker] - if 'assigned_to_workers' in job_data_check: - job_data_check['assigned_to_workers'][worker] = [] last_heartbeat_check = current_time # Check if we've been waiting too long overall if current_time - last_heartbeat_check > timeout: log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") - # Requeue all unfinished images - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - pending_queue = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] - for idx in range(job_data_check.get('batch_size', 0)): - if idx not in completed_images: - await pending_queue.put(idx) - debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx}") break - debug_log(f"UltimateSDUpscale Master Dynamic - Collection complete. Got {len(completed_images)} images") + debug_log(f"UltimateSDUpscale Master Dynamic - Collection complete. Got {collected_count} images from workers") # Clean up job queue async with prompt_server.distributed_tile_jobs_lock: @@ -1104,83 +1094,60 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers): - """Dynamic mode for large batches - assigns whole images to workers dynamically.""" + """Dynamic mode for large batches - assigns whole images to workers dynamically, including master.""" # Get batch size and dimensions batch_size, height, width, _ = upscaled_image.shape num_workers = len(enabled_workers) debug_log(f"UltimateSDUpscale Master Dynamic - Processing batch of {batch_size} images with {num_workers} workers") - # Calculate master's share similar to static mode - num_participants = num_workers + 1 - images_per_participant = batch_size // num_participants - remainder = batch_size % num_participants - master_images_count = images_per_participant + (1 if remainder > 0 else 0) - master_indices = list(range(master_images_count)) - worker_indices = list(range(master_images_count, batch_size)) - expected_from_workers = len(worker_indices) + # No fixed share - all images are dynamic + all_indices = list(range(batch_size)) + + log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.") # Calculate tiles for processing all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - # Initialize job queue for communication - synchronous approach + # Initialize job queue for communication try: - # Add worker tracking data assigned_to_workers = {w: [] for w in enabled_workers} worker_status = {w: time.time() for w in enabled_workers} - # Use short timeout for async init, with threading fallback - try: - run_async_in_server_loop( - self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, worker_indices), - timeout=2.0 # Short timeout - ) - except Exception as init_error: - debug_log(f"UltimateSDUpscale Master Dynamic - Async init failed, using threading: {init_error}") - # Threading fallback for dynamic mode - import threading - def sync_dynamic_init(): - prompt_server = ensure_tile_jobs_initialized() - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = { - 'queue': asyncio.Queue(), - 'mode': 'dynamic', - 'pending_images': asyncio.Queue(), - 'completed_images': {}, - 'completed_tiles': {}, - 'batch_size': batch_size, - 'assigned_to_workers': assigned_to_workers, - 'worker_status': worker_status - } - # Initialize pending images queue with worker indices only - pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] - for i in worker_indices: - pending_images.put_nowait(i) - debug_log(f"Synchronous dynamic init for {multi_job_id} with {len(worker_indices)} pending images for workers") - thread = threading.Thread(target=sync_dynamic_init) - thread.start() - thread.join(timeout=3.0) - if thread.is_alive(): - raise RuntimeError("Dynamic queue init failed") + run_async_in_server_loop( + self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, all_indices), + timeout=2.0 + ) except Exception as e: debug_log(f"UltimateSDUpscale Master Dynamic - Queue initialization error: {e}") raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}") # Convert batch to PIL list - result_images = [] - for b in range(batch_size): - image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) - # Ensure RGB mode for consistency - if image_pil.mode != 'RGB': - image_pil = image_pil.convert('RGB') - result_images.append(image_pil.copy()) + result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)] - # Process master's share locally - debug_log(f"UltimateSDUpscale Master Dynamic - Processing {master_images_count} images locally") - for idx in master_indices: - single_tensor = upscaled_image[idx:idx+1] - local_image = result_images[idx] - image_seed = seed + idx * 1000 + # Process images dynamically with master participating + prompt_server = ensure_tile_jobs_initialized() + processed_count = 0 + + # Process loop - master pulls from queue and processes synchronously + while True: + # Try to get an image to process using async helper + image_idx = run_async_in_server_loop( + self._get_next_image_index(multi_job_id), + timeout=5.0 + ) + + if image_idx is None: + break + + log(f"Master processing image {image_idx} dynamically") + processed_count += 1 + + # Process locally + single_tensor = upscaled_image[image_idx:image_idx+1] + local_image = result_images[image_idx] + image_seed = seed + image_idx * 1000 + for tile_idx, pos in enumerate(all_tiles): local_image = self._process_and_blend_tile( tile_idx, pos, single_tensor, local_image, @@ -1188,32 +1155,54 @@ class UltimateSDUpscaleDistributed: sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, tiled_decode ) - result_images[idx] = local_image + + result_images[image_idx] = local_image + + # Mark as completed using async helper + run_async_in_server_loop( + self._mark_image_completed(multi_job_id, image_idx, local_image), + timeout=5.0 + ) - # Process with workers using dynamic assignment - collected_images = run_async_in_server_loop( - self._async_process_dynamic(multi_job_id, expected_from_workers, num_workers), - timeout=TILE_COLLECTION_TIMEOUT * 2 # Longer timeout for full images - ) + debug_log(f"Master processed {processed_count} images locally") - # Update result images with processed ones - for idx, processed_img in collected_images.items(): - if idx < batch_size: - result_images[idx] = processed_img + # Collect remaining images from workers + expected_from_workers = batch_size - processed_count + if expected_from_workers > 0: + collected_images = run_async_in_server_loop( + self._async_collect_dynamic_images(multi_job_id, expected_from_workers, num_workers), + timeout=TILE_COLLECTION_TIMEOUT * 2 + ) + + # Update result images with worker results + for idx, processed_img in collected_images.items(): + if idx < batch_size: + result_images[idx] = processed_img # Convert back to tensor - if batch_size == 1: - result_tensor = pil_to_tensor(result_images[0]) - else: - result_tensors = [pil_to_tensor(img) for img in result_images] - result_tensor = torch.cat(result_tensors, dim=0) - + result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0]) if upscaled_image.is_cuda: result_tensor = result_tensor.cuda() debug_log(f"UltimateSDUpscale Master Dynamic - Job {multi_job_id} complete") + log(f"Completed processing all {batch_size} images") return (result_tensor,) + + async def _get_next_image_index(self, multi_job_id): + """Get next image index from pending queue for master.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or 'pending_images' not in job_data: + return None + try: + image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=0.1) + return image_idx + except asyncio.TimeoutError: + return None + + def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, From a77af79edae165ecbf32c13fce4247550faddc49 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 16:25:59 +1000 Subject: [PATCH 28/93] upscaler bug fixes --- distributed_upscale.py | 216 +++++++++++++++++++++++++++++++++-------- 1 file changed, 174 insertions(+), 42 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 5e5c5a2..d22fb5a 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -674,8 +674,8 @@ class UltimateSDUpscaleDistributed: if job_data and 'completed_images' in job_data: job_data['completed_images'][image_idx] = image_pil - async def _async_collect_dynamic_images(self, multi_job_id, expected_from_workers, num_workers): - """Collect processed images from workers only.""" + async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): + """Collect remaining processed images from workers.""" prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id not in prompt_server.distributed_pending_tile_jobs: @@ -693,9 +693,9 @@ class UltimateSDUpscaleDistributed: timeout = TILE_WAIT_TIMEOUT * 2 # Longer timeout for full images last_heartbeat_check = time.time() - debug_log(f"UltimateSDUpscale Master Dynamic - Waiting for {expected_from_workers} images from {num_workers} workers") + debug_log(f"UltimateSDUpscale Master Dynamic - Waiting for {remaining_to_collect} more images from {num_workers} workers") - while collected_count < expected_from_workers and len(workers_done) < num_workers: + while collected_count < remaining_to_collect and len(workers_done) < num_workers: try: result = await asyncio.wait_for(q.get(), timeout=10.0) # Shorter timeout for regular checks worker_id = result['worker_id'] @@ -723,14 +723,30 @@ class UltimateSDUpscaleDistributed: for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): if current_time - last_heartbeat > 60: log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") + # Requeue assigned images from this worker + for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): + if idx not in completed_images: + await job_data_check['pending_images'].put(idx) + debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx} from timed out worker {worker}") + # Mark worker as done and clean up workers_done.add(worker) if 'worker_status' in job_data_check: del job_data_check['worker_status'][worker] + if 'assigned_to_workers' in job_data_check: + job_data_check['assigned_to_workers'][worker] = [] last_heartbeat_check = current_time # Check if we've been waiting too long overall if current_time - last_heartbeat_check > timeout: log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") + # Requeue all unfinished images + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + pending_queue = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] + for idx in range(batch_size): + if idx not in completed_images: + await pending_queue.put(idx) + debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx}") break debug_log(f"UltimateSDUpscale Master Dynamic - Collection complete. Got {collected_count} images from workers") @@ -1128,56 +1144,106 @@ class UltimateSDUpscaleDistributed: # Process images dynamically with master participating prompt_server = ensure_tile_jobs_initialized() processed_count = 0 + consecutive_retries = 0 + max_consecutive_retries = 10 # Process loop - master pulls from queue and processes synchronously - while True: - # Try to get an image to process using async helper + while processed_count < batch_size: + # Try to get an image to process image_idx = run_async_in_server_loop( self._get_next_image_index(multi_job_id), - timeout=5.0 + timeout=5.0 # Short timeout to allow frequent checks ) - - if image_idx is None: - break - - log(f"Master processing image {image_idx} dynamically") - processed_count += 1 - - # Process locally - single_tensor = upscaled_image[image_idx:image_idx+1] - local_image = result_images[image_idx] - image_seed = seed + image_idx * 1000 - - for tile_idx, pos in enumerate(all_tiles): - local_image = self._process_and_blend_tile( - tile_idx, pos, single_tensor, local_image, - model, positive, negative, vae, image_seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode + + if image_idx is not None: + # Reset retry counter and process locally + consecutive_retries = 0 + debug_log(f"Master processing image {image_idx} dynamically") + processed_count += 1 + + # Get total completed count for accurate progress + total_completed = run_async_in_server_loop( + self._get_total_completed_count(multi_job_id), + timeout=1.0 ) - - result_images[image_idx] = local_image - - # Mark as completed using async helper - run_async_in_server_loop( - self._mark_image_completed(multi_job_id, image_idx, local_image), - timeout=5.0 - ) + log(f"{total_completed}/{batch_size} completed images") + + # Process locally + single_tensor = upscaled_image[image_idx:image_idx+1] + local_image = result_images[image_idx] + image_seed = seed + image_idx * 1000 + + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive, negative, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode + ) + + result_images[image_idx] = local_image + + # Mark as completed + run_async_in_server_loop( + self._mark_image_completed(multi_job_id, image_idx, local_image), + timeout=5.0 + ) + else: + # Queue empty: collect any queued worker results to update progress + drained_count = run_async_in_server_loop( + self._drain_worker_results_queue(multi_job_id), + timeout=5.0 + ) + + # Now check total completed (includes newly collected) + completed_now = run_async_in_server_loop( + self._get_total_completed_count(multi_job_id), + timeout=1.0 + ) + + # Show accurate progress if we drained any new results + if drained_count > 0: + log(f"{completed_now}/{batch_size} completed images") + + debug_log(f"Queue empty, completed count: {completed_now}/{batch_size}") + + if completed_now >= batch_size: + debug_log("All images completed, exiting master loop") + break + + consecutive_retries += 1 + if consecutive_retries >= max_consecutive_retries: + log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") + break # Force exit to collection phase + + log("Queue empty temporarily, waiting 2s before retry") + time.sleep(2) debug_log(f"Master processed {processed_count} images locally") - # Collect remaining images from workers - expected_from_workers = batch_size - processed_count - if expected_from_workers > 0: + # Get all completed images to check what needs to be collected + all_completed = run_async_in_server_loop( + self._get_all_completed_images(multi_job_id), + timeout=5.0 + ) + + # Calculate how many we still need to collect + remaining_to_collect = batch_size - len(all_completed) + + if remaining_to_collect > 0: + debug_log(f"Waiting for {remaining_to_collect} more images from workers") collected_images = run_async_in_server_loop( - self._async_collect_dynamic_images(multi_job_id, expected_from_workers, num_workers), + self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count), timeout=TILE_COLLECTION_TIMEOUT * 2 ) - # Update result images with worker results - for idx, processed_img in collected_images.items(): - if idx < batch_size: - result_images[idx] = processed_img + # Merge collected with already completed + all_completed.update(collected_images) + + # Update result images with all completed images + for idx, processed_img in all_completed.items(): + if idx < batch_size: + result_images[idx] = processed_img # Convert back to tensor result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0]) @@ -1197,11 +1263,77 @@ class UltimateSDUpscaleDistributed: if not job_data or 'pending_images' not in job_data: return None try: - image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=0.1) + image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=1.0) return image_idx except asyncio.TimeoutError: return None + async def _get_completed_count(self, multi_job_id): + """Get count of completed images from workers.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + # Count only worker-completed images (not master's) + completed = job_data['completed_images'] + worker_count = len([idx for idx in completed if isinstance(completed[idx], Image.Image)]) + return worker_count + return 0 + + async def _get_total_completed_count(self, multi_job_id): + """Get total count of all completed images (master + workers).""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + return len(job_data['completed_images']) + return 0 + + async def _get_all_completed_images(self, multi_job_id): + """Get all completed images.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'completed_images' in job_data: + return job_data['completed_images'].copy() + return {} + + async def _drain_worker_results_queue(self, multi_job_id): + """Drain pending worker results from queue and update completed_images. Returns count of drained images.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or 'queue' not in job_data or 'completed_images' not in job_data: + return 0 + q = job_data['queue'] + completed_images = job_data['completed_images'] + + collected = 0 + while not q.empty(): + try: + result = await asyncio.wait_for(q.get(), timeout=0.1) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + if image_idx not in completed_images: + completed_images[image_idx] = image_pil + collected += 1 + debug_log(f"Drained image {image_idx} from worker {worker_id}") + + if is_last: + # Optional: track worker completion if needed + pass + except asyncio.TimeoutError: + break # No more immediately available + + if collected > 0: + debug_log(f"Drained {collected} worker images during retry") + + return collected + def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, From 1cb193de32e89e4abd42b560daf270ef36ce0c13 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 21:32:33 +1000 Subject: [PATCH 29/93] optimisations --- distributed_upscale.py | 369 +++++++++++++++++++++++++++-------------- 1 file changed, 244 insertions(+), 125 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index d22fb5a..2908d68 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1,15 +1,4 @@ -""" -ComfyUI Distributed Upscale Module - -This module implements a distributed version of Ultimate SD Upscale with multi-mode -batch handling for efficient processing of single images and video frames. - -Processing Modes: -- Single GPU: Process all tiles locally when no workers are available -- Static Mode: For small batches (≤ dynamic_threshold), flattens tiles across batch -- Dynamic Mode: For large batches, assigns whole images to workers dynamically -""" - +# Import statements (original + new) import torch import numpy as np from PIL import Image, ImageFilter, ImageDraw @@ -38,6 +27,12 @@ from .utils.constants import ( QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT ) +# Import for controller support +from .utils.usdu_utils import crop_cond + +# Additional import for deepcopy +import copy + # Make MAX_BATCH configurable MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) @@ -134,6 +129,7 @@ class UltimateSDUpscaleDistributed: """Force re-execution.""" return float("nan") # Always re-execute + def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id="", is_worker=False, @@ -173,10 +169,7 @@ class UltimateSDUpscaleDistributed: enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) - if batch_size == 1 or batch_size <= dynamic_threshold: - mode = "static" - else: - mode = "dynamic" + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}") @@ -201,7 +194,7 @@ class UltimateSDUpscaleDistributed: total_tiles = batch_size * num_tiles_per_image # Get assigned global tile indices for this worker - assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_worker_ids, worker_id) + assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) if not assigned_global_indices: return (upscaled_image,) @@ -209,6 +202,9 @@ class UltimateSDUpscaleDistributed: # Process tiles SYNCHRONOUSLY processed_tiles = [] + # Track which batch indices we've already pre-sliced conditioning for + sliced_conditioning_cache = {} + for global_idx in assigned_global_indices: # Calculate which image and tile this corresponds to batch_idx = global_idx // num_tiles_per_image @@ -218,6 +214,22 @@ class UltimateSDUpscaleDistributed: if batch_idx >= upscaled_image.shape[0]: debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") continue + + # Pre-slice conditioning once per image (cache it) + if batch_idx not in sliced_conditioning_cache: + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] x, y = all_tiles[tile_idx] @@ -228,9 +240,10 @@ class UltimateSDUpscaleDistributed: # Process tile through SD with unique seed image_seed = seed + batch_idx * 1000 - processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, + processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, image_seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode) + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) processed_tiles.append({ 'tile': processed_tile, @@ -278,12 +291,7 @@ class UltimateSDUpscaleDistributed: num_workers = len(enabled_workers) # Determine processing mode - if num_workers == 0: - mode = "single_gpu" - elif batch_size == 1 or batch_size <= dynamic_threshold: - mode = "static" - else: - mode = "dynamic" + mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}") @@ -315,10 +323,6 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}") raise RuntimeError(f"Failed to initialize static mode queue: {e}") - # Don't try to prepare via API since we're using local queues - # loop = asyncio.new_event_loop() - # asyncio.set_event_loop(loop) - # Convert batch to PIL list for processing result_images = [] for b in range(batch_size): @@ -338,6 +342,9 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally from batch of {batch_size}") # Process master tiles SYNCHRONOUSLY + # Track which batch indices we've already pre-sliced conditioning for + sliced_conditioning_cache = {} + for global_idx in master_global_indices: # Calculate which image and tile this corresponds to batch_idx = global_idx // num_tiles_per_image @@ -347,15 +354,31 @@ class UltimateSDUpscaleDistributed: if batch_idx >= len(result_images): debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds result_images length {len(result_images)}") continue + + # Pre-slice conditioning once per image (cache it) + if batch_idx not in sliced_conditioning_cache: + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] # Use unique seed per image image_seed = seed + batch_idx * 1000 result_images[batch_idx] = self._process_and_blend_tile( tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], - model, positive, negative, vae, image_seed, steps, cfg, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx ) # Collect worker tiles using async @@ -431,9 +454,8 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") return (result_tensor,) - def _get_worker_global_indices(self, total_tiles, enabled_worker_ids, worker_id): + def _get_worker_global_indices(self, total_tiles, enabled_workers, worker_id): """Calculate which global tile indices are assigned to a specific worker in flattened mode.""" - enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}") @@ -444,71 +466,20 @@ class UltimateSDUpscaleDistributed: log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}") return [] - # Calculate tile distribution - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - - start_idx = master_tiles + (worker_index * tiles_per_participant) - if worker_index < remainder - 1: - start_idx += worker_index - end_idx = start_idx + tiles_per_participant + 1 - else: - start_idx += remainder - 1 if remainder > 0 else 0 - end_idx = start_idx + tiles_per_participant - - end_idx = min(end_idx, total_tiles) - return list(range(start_idx, end_idx)) - - def _get_worker_tiles(self, all_tiles, enabled_worker_ids, worker_id): - """Calculate which tiles are assigned to a specific worker.""" - enabled_workers = json.loads(enabled_worker_ids) - num_workers = len(enabled_workers) - total_tiles = len(all_tiles) - - debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}") - - try: - worker_index = enabled_workers.index(worker_id) - except ValueError: - log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}") - return [] - - # Calculate tile distribution - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - - start_idx = master_tiles + (worker_index * tiles_per_participant) - if worker_index < remainder - 1: - start_idx += worker_index - end_idx = start_idx + tiles_per_participant + 1 - else: - start_idx += remainder - 1 if remainder > 0 else 0 - end_idx = start_idx + tiles_per_participant - - end_idx = min(end_idx, total_tiles) - return list(range(start_idx, end_idx)) + global_indices = list(range(total_tiles)) + all_assignments = self._distribute_items(global_indices, num_workers + 1) + return all_assignments[worker_index + 1] # +1 because 0 is master def _get_master_global_indices(self, total_tiles, num_workers): """Calculate which global tile indices are assigned to the master in flattened mode.""" - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - return list(range(master_tiles)) - - def _get_master_tiles(self, all_tiles, num_workers): - """Calculate which tiles are assigned to the master.""" - total_tiles = len(all_tiles) - tiles_per_participant = total_tiles // (num_workers + 1) - remainder = total_tiles % (num_workers + 1) - master_tiles = tiles_per_participant + (1 if remainder > 0 else 0) - return list(range(master_tiles)) + global_indices = list(range(total_tiles)) + all_assignments = self._distribute_items(global_indices, num_workers + 1) + return all_assignments[0] def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, image_width, image_height, tiled_decode): + padding, mask_blur, image_width, image_height, tiled_decode, batch_idx: int = 0): """Process a single tile and blend it into the result image.""" x, y = tile_pos @@ -519,7 +490,8 @@ class UltimateSDUpscaleDistributed: processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode) + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height)) # Convert and blend processed_pil = tensor_to_pil(processed_tile, 0) @@ -531,17 +503,6 @@ class UltimateSDUpscaleDistributed: return result_image - async def _prepare_multigpu_job(self, multi_job_id): - """Prepare job via API endpoint to ensure server is ready.""" - session = await get_client_session() - # Use the actual server port (get from server instance) - port = get_server_port() - async with session.post(f"http://localhost:{port}/distributed/prepare_job", - json={"multi_job_id": multi_job_id}) as response: - if response.status != 200: - text = await response.text() - raise RuntimeError(f"Failed to prepare job: {response.status} - {text}") - async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None): """Initialize the job queue for dynamic mode with pending images.""" prompt_server = ensure_tile_jobs_initialized() @@ -864,9 +825,31 @@ class UltimateSDUpscaleDistributed: return tile_tensor, x1, y1, extracted_width, extracted_height def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae, - seed: int, steps: int, cfg: float, sampler_name: str, - scheduler: str, denoise: float, tiled_decode: bool = False) -> torch.Tensor: - """Process a single tile through SD sampling.""" + seed: int, steps: int, cfg: float, sampler_name: str, + scheduler: str, denoise: float, tiled_decode: bool = False, + batch_idx: int = 0, region: Tuple[int, int, int, int] = None, + image_size: Tuple[int, int] = None) -> torch.Tensor: + """Process a single tile through SD sampling. + Note: positive and negative should already be pre-sliced for the current batch_idx.""" + debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}") + + # Log conditioning details + debug_log(f"[process_tile] Positive conditioning items: {len(positive)}") + debug_log(f"[process_tile] Negative conditioning items: {len(negative)}") + + # Check for ControlNet in conditioning + for i, (emb, cond_dict) in enumerate(positive): + if 'control' in cond_dict: + control = cond_dict['control'] + debug_log(f"[process_tile] Found ControlNet in positive[{i}]") + debug_log(f"[process_tile] Control id: {id(control)}") + if hasattr(control, 'control_model'): + debug_log(f"[process_tile] control_model id: {id(control.control_model)}") + if hasattr(control, 'cond_hint_original'): + debug_log(f"[process_tile] cond_hint_original shape: {control.cond_hint_original.shape}") + debug_log(f"[process_tile] Expected to use hint at index {batch_idx} from batch") + debug_log(f"[process_tile] Embedding shape: {emb.shape}") + # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode @@ -887,15 +870,38 @@ class UltimateSDUpscaleDistributed: if tile_tensor.is_cuda: clean_tensor = clean_tensor.cuda() + # Crop conditioning to tile region if provided (assumes hints at image resolution) + if region is not None and image_size is not None: + positive_cropped = copy.deepcopy(positive) + negative_cropped = copy.deepcopy(negative) + init_size = image_size # (width, height) of full image + canvas_size = image_size + tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height) + w_pad = 0 # No extra pad needed; region already includes padding + h_pad = 0 + positive_cropped = crop_cond(positive_cropped, region, init_size, canvas_size, tile_size, w_pad, h_pad) + negative_cropped = crop_cond(negative_cropped, region, init_size, canvas_size, tile_size, w_pad, h_pad) + else: + # No region cropping needed, use conditioning as-is + positive_cropped = positive + negative_cropped = negative + # Encode to latent if tiled_decode and tiled_vae_available: latent = VAEEncodeTiled().encode(vae, clean_tensor)[0] else: latent = VAEEncode().encode(vae, clean_tensor)[0] + # Log latent shape + if 'samples' in latent: + debug_log(f"[process_tile] Latent shape: {latent['samples'].shape}") + # Sample + debug_log(f"[process_tile] About to call common_ksampler with batch_idx={batch_idx}") + debug_log(f"[process_tile] Model id: {id(model)}") samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive, negative, latent, denoise=denoise)[0] + positive_cropped, negative_cropped, latent, denoise=denoise)[0] + debug_log(f"[process_tile] common_ksampler completed") # Decode back to image if tiled_decode and tiled_vae_available: @@ -928,7 +934,7 @@ class UltimateSDUpscaleDistributed: def blend_tile(self, base_image: Image.Image, tile_image: Image.Image, x: int, y: int, extracted_size: Tuple[int, int], - mask: int, padding: int) -> Image.Image: + mask: Image.Image, padding: int) -> Image.Image: """Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach. This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py @@ -1083,15 +1089,28 @@ class UltimateSDUpscaleDistributed: # Process each image in the batch for batch_idx in range(batch_size): - # Process each tile for this image + debug_log(f"[process_single_gpu] Processing batch_idx {batch_idx}") + # Pre-slice conditioning once per image (not per tile) + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + + # Process each tile for this image with pre-sliced conditioning for tile_idx, tile_pos in enumerate(all_tiles): # Use unique seed per image image_seed = seed + batch_idx * 1000 result_images[batch_idx] = self._process_and_blend_tile( tile_idx, tile_pos, upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], - model, positive, negative, vae, image_seed, steps, cfg, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx ) # Convert back to tensor @@ -1160,25 +1179,30 @@ class UltimateSDUpscaleDistributed: consecutive_retries = 0 debug_log(f"Master processing image {image_idx} dynamically") processed_count += 1 - - # Get total completed count for accurate progress - total_completed = run_async_in_server_loop( - self._get_total_completed_count(multi_job_id), - timeout=1.0 - ) - log(f"{total_completed}/{batch_size} completed images") # Process locally single_tensor = upscaled_image[image_idx:image_idx+1] local_image = result_images[image_idx] image_seed = seed + image_idx * 1000 + # Pre-slice conditioning once per image (not per tile) + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[image_idx:image_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + for tile_idx, pos in enumerate(all_tiles): local_image = self._process_and_blend_tile( tile_idx, pos, single_tensor, local_image, - model, positive, negative, vae, image_seed, steps, cfg, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode + padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) result_images[image_idx] = local_image @@ -1194,6 +1218,16 @@ class UltimateSDUpscaleDistributed: self._drain_worker_results_queue(multi_job_id), timeout=5.0 ) + + # Check for timed out workers and requeue their images + requeued_count = run_async_in_server_loop( + self._check_and_requeue_timed_out_workers(multi_job_id, batch_size), + timeout=5.0 + ) + if requeued_count > 0: + log(f"Requeued {requeued_count} images from timed out workers") + consecutive_retries = 0 # Reset since we have work to do + continue # Now check total completed (includes newly collected) completed_now = run_async_in_server_loop( @@ -1201,9 +1235,6 @@ class UltimateSDUpscaleDistributed: timeout=1.0 ) - # Show accurate progress if we drained any new results - if drained_count > 0: - log(f"{completed_now}/{batch_size} completed images") debug_log(f"Queue empty, completed count: {completed_now}/{batch_size}") @@ -1211,6 +1242,17 @@ class UltimateSDUpscaleDistributed: debug_log("All images completed, exiting master loop") break + # Check if there are pending images in the queue (could be requeued) + pending_count = run_async_in_server_loop( + self._get_pending_count(multi_job_id), + timeout=1.0 + ) + + if pending_count > 0: + debug_log(f"Found {pending_count} pending images (possibly requeued), continuing to process") + consecutive_retries = 0 # Reset retries since there's work to do + continue + consecutive_retries += 1 if consecutive_retries >= max_consecutive_retries: log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") @@ -1298,6 +1340,15 @@ class UltimateSDUpscaleDistributed: return job_data['completed_images'].copy() return {} + async def _get_pending_count(self, multi_job_id): + """Get count of pending images in the queue.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and 'pending_images' in job_data: + return job_data['pending_images'].qsize() + return 0 + async def _drain_worker_results_queue(self, multi_job_id): """Drain pending worker results from queue and update completed_images. Returns count of drained images.""" prompt_server = ensure_tile_jobs_initialized() @@ -1334,6 +1385,36 @@ class UltimateSDUpscaleDistributed: return collected + async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): + """Check for timed out workers and requeue their assigned images. Returns count of requeued images.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data: + return 0 + + current_time = time.time() + requeued_count = 0 + completed_images = job_data.get('completed_images', {}) + + # Check for timed out workers (60 second timeout) + for worker, last_heartbeat in list(job_data.get('worker_status', {}).items()): + if current_time - last_heartbeat > 60: + log(f"Worker {worker} timed out") + # Requeue assigned images from this worker + for idx in job_data.get('assigned_to_workers', {}).get(worker, []): + if idx not in completed_images: + await job_data['pending_images'].put(idx) + requeued_count += 1 + debug_log(f"Requeued image {idx} from timed out worker {worker}") + # Clean up worker tracking + if 'worker_status' in job_data: + del job_data['worker_status'][worker] + if 'assigned_to_workers' in job_data: + job_data['assigned_to_workers'][worker] = [] + + return requeued_count + def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, @@ -1395,12 +1476,25 @@ class UltimateSDUpscaleDistributed: # Process all tiles for this image image_seed = seed + image_idx * 1000 + + # Pre-slice conditioning once per image (not per tile) + positive_sliced = copy.deepcopy(positive) + negative_sliced = copy.deepcopy(negative) + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[image_idx:image_idx+1] + # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + for tile_idx, pos in enumerate(all_tiles): local_image = self._process_and_blend_tile( tile_idx, pos, single_tensor, local_image, - model, positive, negative, vae, image_seed, steps, cfg, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode + padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) # Send processed image back to master @@ -1566,6 +1660,31 @@ class UltimateSDUpscaleDistributed: response.raise_for_status() debug_log(f"Worker {worker_id} sent completion signal") + def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str: + """Determines the processing mode based on batch size and worker count.""" + if num_workers == 0: + return "single_gpu" + if batch_size <= dynamic_threshold: + return "static" + return "dynamic" + + def _distribute_items(self, items: list, num_participants: int) -> List[List[any]]: + """Distributes a list of items among N participants.""" + if num_participants == 0: + return [items] # All items for the single participant (master) + + items_per_participant = len(items) // num_participants + remainder = len(items) % num_participants + + assignments = [] + start_idx = 0 + for i in range(num_participants): + count = items_per_participant + (1 if i < remainder else 0) + end_idx = start_idx + count + assignments.append(items[start_idx:end_idx]) + start_idx = end_idx + + return assignments # Ensure initialization before registering routes ensure_tile_jobs_initialized() From fbfc605709256fb0b8fb72d046d6495926de5d52 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Thu, 17 Jul 2025 21:41:42 +1000 Subject: [PATCH 30/93] working batch controlnet --- distributed_upscale.py | 68 ++++++++++++++++++++++-------------------- 1 file changed, 35 insertions(+), 33 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 2908d68..20cd86b 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -129,7 +129,6 @@ class UltimateSDUpscaleDistributed: """Force re-execution.""" return float("nan") # Always re-execute - def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id="", is_worker=False, @@ -224,7 +223,13 @@ class UltimateSDUpscaleDistributed: emb, cond_dict = cond_list[i] if emb.shape[0] > 1: cond_list[i][0] = emb[batch_idx:batch_idx+1] - # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[batch_idx:batch_idx+1] + control = control.previous_controlnet if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) @@ -364,7 +369,13 @@ class UltimateSDUpscaleDistributed: emb, cond_dict = cond_list[i] if emb.shape[0] > 1: cond_list[i][0] = emb[batch_idx:batch_idx+1] - # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[batch_idx:batch_idx+1] + control = control.previous_controlnet if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) @@ -831,25 +842,6 @@ class UltimateSDUpscaleDistributed: image_size: Tuple[int, int] = None) -> torch.Tensor: """Process a single tile through SD sampling. Note: positive and negative should already be pre-sliced for the current batch_idx.""" - debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}") - - # Log conditioning details - debug_log(f"[process_tile] Positive conditioning items: {len(positive)}") - debug_log(f"[process_tile] Negative conditioning items: {len(negative)}") - - # Check for ControlNet in conditioning - for i, (emb, cond_dict) in enumerate(positive): - if 'control' in cond_dict: - control = cond_dict['control'] - debug_log(f"[process_tile] Found ControlNet in positive[{i}]") - debug_log(f"[process_tile] Control id: {id(control)}") - if hasattr(control, 'control_model'): - debug_log(f"[process_tile] control_model id: {id(control.control_model)}") - if hasattr(control, 'cond_hint_original'): - debug_log(f"[process_tile] cond_hint_original shape: {control.cond_hint_original.shape}") - debug_log(f"[process_tile] Expected to use hint at index {batch_idx} from batch") - debug_log(f"[process_tile] Embedding shape: {emb.shape}") - # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode @@ -892,16 +884,9 @@ class UltimateSDUpscaleDistributed: else: latent = VAEEncode().encode(vae, clean_tensor)[0] - # Log latent shape - if 'samples' in latent: - debug_log(f"[process_tile] Latent shape: {latent['samples'].shape}") - # Sample - debug_log(f"[process_tile] About to call common_ksampler with batch_idx={batch_idx}") - debug_log(f"[process_tile] Model id: {id(model)}") samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive_cropped, negative_cropped, latent, denoise=denoise)[0] - debug_log(f"[process_tile] common_ksampler completed") # Decode back to image if tiled_decode and tiled_vae_available: @@ -1089,7 +1074,6 @@ class UltimateSDUpscaleDistributed: # Process each image in the batch for batch_idx in range(batch_size): - debug_log(f"[process_single_gpu] Processing batch_idx {batch_idx}") # Pre-slice conditioning once per image (not per tile) positive_sliced = copy.deepcopy(positive) negative_sliced = copy.deepcopy(negative) @@ -1098,7 +1082,13 @@ class UltimateSDUpscaleDistributed: emb, cond_dict = cond_list[i] if emb.shape[0] > 1: cond_list[i][0] = emb[batch_idx:batch_idx+1] - # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[batch_idx:batch_idx+1] + control = control.previous_controlnet if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] @@ -1193,7 +1183,13 @@ class UltimateSDUpscaleDistributed: emb, cond_dict = cond_list[i] if emb.shape[0] > 1: cond_list[i][0] = emb[image_idx:image_idx+1] - # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[image_idx:image_idx+1] + control = control.previous_controlnet if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] @@ -1485,7 +1481,13 @@ class UltimateSDUpscaleDistributed: emb, cond_dict = cond_list[i] if emb.shape[0] > 1: cond_list[i][0] = emb[image_idx:image_idx+1] - # Don't slice ControlNet - crop_cond will handle spatial cropping + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[image_idx:image_idx+1] + control = control.previous_controlnet if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] From 4447998de02ec2fe7274933d5ddaf98393739778 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 07:46:10 +1000 Subject: [PATCH 31/93] fixed controlnet loading for every tile --- distributed_upscale.py | 94 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 77 insertions(+), 17 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 20cd86b..dc3356c 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -39,6 +39,36 @@ MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Configure maximum payload size (50MB default, configurable via environment variable) MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) +# Helper functions for shallow copying conditioning without duplicating models +def clone_control_chain(control, clone_hint=True): + """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" + if control is None: + return None + new_control = copy.copy(control) # Shallow copy (shares model) + if clone_hint and hasattr(control, 'cond_hint_original'): + new_control.cond_hint_original = control.cond_hint_original.clone() + if hasattr(control, 'previous_controlnet'): + new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) + return new_control + +def clone_conditioning(cond_list, clone_hints=True): + """Clone conditioning without duplicating ControlNet models.""" + new_cond = [] + for emb, cond_dict in cond_list: + new_emb = emb.clone() if emb is not None else None + new_dict = cond_dict.copy() + if 'control' in new_dict: + new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) + if 'mask' in new_dict: + new_dict['mask'] = new_dict['mask'].clone() + # Handle other potential fields if needed + if 'pooled_output' in new_dict: + new_dict['pooled_output'] = new_dict['pooled_output'].clone() + if 'area' in new_dict: + new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple + new_cond.append([new_emb, new_dict]) + return new_cond + # Helper function to ensure persistent state is initialized def ensure_tile_jobs_initialized(): """Ensure tile job storage is initialized on the server instance.""" @@ -129,6 +159,7 @@ class UltimateSDUpscaleDistributed: """Force re-execution.""" return float("nan") # Always re-execute + def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id="", is_worker=False, @@ -216,8 +247,8 @@ class UltimateSDUpscaleDistributed: # Pre-slice conditioning once per image (cache it) if batch_idx not in sliced_conditioning_cache: - positive_sliced = copy.deepcopy(positive) - negative_sliced = copy.deepcopy(negative) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) for cond_list in [positive_sliced, negative_sliced]: for i in range(len(cond_list)): emb, cond_dict = cond_list[i] @@ -362,8 +393,8 @@ class UltimateSDUpscaleDistributed: # Pre-slice conditioning once per image (cache it) if batch_idx not in sliced_conditioning_cache: - positive_sliced = copy.deepcopy(positive) - negative_sliced = copy.deepcopy(negative) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) for cond_list in [positive_sliced, negative_sliced]: for i in range(len(cond_list)): emb, cond_dict = cond_list[i] @@ -842,6 +873,25 @@ class UltimateSDUpscaleDistributed: image_size: Tuple[int, int] = None) -> torch.Tensor: """Process a single tile through SD sampling. Note: positive and negative should already be pre-sliced for the current batch_idx.""" + debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}") + + # Log conditioning details + debug_log(f"[process_tile] Positive conditioning items: {len(positive)}") + debug_log(f"[process_tile] Negative conditioning items: {len(negative)}") + + # Check for ControlNet in conditioning + for i, (emb, cond_dict) in enumerate(positive): + if 'control' in cond_dict: + control = cond_dict['control'] + debug_log(f"[process_tile] Found ControlNet in positive[{i}]") + debug_log(f"[process_tile] Control id: {id(control)}") + if hasattr(control, 'control_model'): + debug_log(f"[process_tile] control_model id: {id(control.control_model)}") + if hasattr(control, 'cond_hint_original'): + debug_log(f"[process_tile] cond_hint_original shape: {control.cond_hint_original.shape}") + debug_log(f"[process_tile] Expected to use hint at index {batch_idx} from batch") + debug_log(f"[process_tile] Embedding shape: {emb.shape}") + # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode @@ -862,21 +912,23 @@ class UltimateSDUpscaleDistributed: if tile_tensor.is_cuda: clean_tensor = clean_tensor.cuda() + # Clone conditioning per tile (shares models, clones hints for cropping) + positive_tile = clone_conditioning(positive, clone_hints=True) + negative_tile = clone_conditioning(negative, clone_hints=True) + # Crop conditioning to tile region if provided (assumes hints at image resolution) if region is not None and image_size is not None: - positive_cropped = copy.deepcopy(positive) - negative_cropped = copy.deepcopy(negative) init_size = image_size # (width, height) of full image canvas_size = image_size tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height) w_pad = 0 # No extra pad needed; region already includes padding h_pad = 0 - positive_cropped = crop_cond(positive_cropped, region, init_size, canvas_size, tile_size, w_pad, h_pad) - negative_cropped = crop_cond(negative_cropped, region, init_size, canvas_size, tile_size, w_pad, h_pad) + positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) + negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad) else: - # No region cropping needed, use conditioning as-is - positive_cropped = positive - negative_cropped = negative + # No region cropping needed, use cloned conditioning as-is + positive_cropped = positive_tile + negative_cropped = negative_tile # Encode to latent if tiled_decode and tiled_vae_available: @@ -884,9 +936,16 @@ class UltimateSDUpscaleDistributed: else: latent = VAEEncode().encode(vae, clean_tensor)[0] + # Log latent shape + if 'samples' in latent: + debug_log(f"[process_tile] Latent shape: {latent['samples'].shape}") + # Sample + debug_log(f"[process_tile] About to call common_ksampler with batch_idx={batch_idx}") + debug_log(f"[process_tile] Model id: {id(model)}") samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive_cropped, negative_cropped, latent, denoise=denoise)[0] + debug_log(f"[process_tile] common_ksampler completed") # Decode back to image if tiled_decode and tiled_vae_available: @@ -1074,9 +1133,10 @@ class UltimateSDUpscaleDistributed: # Process each image in the batch for batch_idx in range(batch_size): + debug_log(f"[process_single_gpu] Processing batch_idx {batch_idx}") # Pre-slice conditioning once per image (not per tile) - positive_sliced = copy.deepcopy(positive) - negative_sliced = copy.deepcopy(negative) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) for cond_list in [positive_sliced, negative_sliced]: for i in range(len(cond_list)): emb, cond_dict = cond_list[i] @@ -1176,8 +1236,8 @@ class UltimateSDUpscaleDistributed: image_seed = seed + image_idx * 1000 # Pre-slice conditioning once per image (not per tile) - positive_sliced = copy.deepcopy(positive) - negative_sliced = copy.deepcopy(negative) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) for cond_list in [positive_sliced, negative_sliced]: for i in range(len(cond_list)): emb, cond_dict = cond_list[i] @@ -1474,8 +1534,8 @@ class UltimateSDUpscaleDistributed: image_seed = seed + image_idx * 1000 # Pre-slice conditioning once per image (not per tile) - positive_sliced = copy.deepcopy(positive) - negative_sliced = copy.deepcopy(negative) + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) for cond_list in [positive_sliced, negative_sliced]: for i in range(len(cond_list)): emb, cond_dict = cond_list[i] From 71f6bcbfa35b6710b10cc29071266e553906b3d7 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 09:13:14 +1000 Subject: [PATCH 32/93] remove exessive logging, fixed tiled decode, reduce wait fo requeue --- distributed.py | 32 --------------- distributed_upscale.py | 93 ++++++++++++------------------------------ 2 files changed, 25 insertions(+), 100 deletions(-) diff --git a/distributed.py b/distributed.py index 7b53c41..46df518 100644 --- a/distributed.py +++ b/distributed.py @@ -1777,14 +1777,12 @@ class DistributedCollectorNode: if batch_size == 0: return - debug_log(f"Worker - Sending {batch_size} images in chunks of max {MAX_BATCH}") for start in range(0, batch_size, MAX_BATCH): chunk = image_batch[start:start + MAX_BATCH] chunk_size = chunk.shape[0] is_chunk_last = (start + chunk_size == batch_size) # True only for final chunk - debug_log(f"Worker - Processing chunk {start//MAX_BATCH + 1}, images {start} to {start + chunk_size - 1}") data = aiohttp.FormData() data.add_field('multi_job_id', multi_job_id) @@ -1806,21 +1804,11 @@ class DistributedCollectorNode: data.add_field(f'image_{j}', byte_io, filename=f'image_{j}.png', content_type='image/png') try: - # Estimate payload size for logging - payload_size = 0 - for field in data._fields: - if hasattr(field[2], 'read'): - field[2].seek(0) - payload_size += len(field[2].read()) - field[2].seek(0) - debug_log(f"Worker - Chunk payload size: {payload_size:,} bytes") session = await get_client_session() url = f"{master_url}/distributed/job_complete" - debug_log(f"Worker - Sending chunk of {chunk_size} images to URL: {url}") async with session.post(url, data=data) as response: response.raise_for_status() - debug_log(f"Worker - Successfully sent chunk of {chunk_size} images") except Exception as e: log(f"Worker - Failed to send chunk to master: {e}") debug_log(f"Worker - Full error details: URL={url}") @@ -1847,8 +1835,6 @@ class DistributedCollectorNode: # Ensure master images are contiguous images_on_cpu = ensure_contiguous(images_on_cpu) - # Debug: Check master tensor properties - debug_log(f"Master tensor - shape: {images_on_cpu.shape}, dtype: {images_on_cpu.dtype}, device: {images_on_cpu.device}") # Initialize storage for collected images worker_images = {} # Dict to store images by worker_id and index @@ -1869,13 +1855,11 @@ class DistributedCollectorNode: # Use a reasonable timeout for the first image timeout = WORKER_JOB_TIMEOUT - debug_log(f"Master - Starting collection loop, expecting {num_workers} workers") # Get queue size before starting async with prompt_server.distributed_jobs_lock: q = prompt_server.distributed_pending_jobs[multi_job_id] initial_size = q.qsize() - debug_log(f"Master - Queue size before collection: {initial_size}") # NEW: Initialize progress bar for workers (total = num_workers) p = ProgressBar(num_workers) @@ -1887,7 +1871,6 @@ class DistributedCollectorNode: q = prompt_server.distributed_pending_jobs[multi_job_id] current_size = q.qsize() - debug_log(f"Master - Waiting for queue item, timeout={timeout}s, queue size={current_size}") result = await asyncio.wait_for(q.get(), timeout=timeout) worker_id = result['worker_id'] @@ -1906,9 +1889,6 @@ class DistributedCollectorNode: for idx, tensor in enumerate(tensors): worker_images[worker_id][idx] = tensor - # Debug: Check worker tensor properties - if collected_count == 0: # Only print once - debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") collected_count += len(tensors) else: @@ -1922,9 +1902,6 @@ class DistributedCollectorNode: worker_images[worker_id] = {} worker_images[worker_id][image_index] = tensor - # Debug: Check worker tensor properties - if collected_count == 0: # Only print once - debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") collected_count += 1 @@ -1934,9 +1911,7 @@ class DistributedCollectorNode: if is_last: workers_done.add(worker_id) p.update(1) # +1 per completed worker - debug_log(f"Master - Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") else: - debug_log(f"Master - Collected image {image_index + 1} from worker {worker_id}") except asyncio.TimeoutError: missing_workers = set(str(w) for w in enabled_workers) - workers_done @@ -1947,7 +1922,6 @@ class DistributedCollectorNode: if multi_job_id in prompt_server.distributed_pending_jobs: final_q = prompt_server.distributed_pending_jobs[multi_job_id] final_size = final_q.qsize() - debug_log(f"Master - Queue size at timeout: {final_size}") # Try to drain any remaining items remaining_items = [] @@ -1959,7 +1933,6 @@ class DistributedCollectorNode: break if remaining_items: - debug_log(f"Master - Found {len(remaining_items)} items in queue after timeout!") # Process them for item in remaining_items: worker_id = item['worker_id'] @@ -1990,15 +1963,11 @@ class DistributedCollectorNode: if is_last: workers_done.add(worker_id) p.update(1) # +1 here too - debug_log(f"Master - Worker {worker_id} done (found in timeout drain)") else: log(f"Master - Queue {multi_job_id} no longer exists!") break total_collected = sum(len(imgs) for imgs in worker_images.values()) - debug_log(f"Master - Collection complete. Received {total_collected} images from {len(workers_done)} workers") - debug_log(f"Master - Workers done: {workers_done}, Enabled workers: {enabled_workers}") - debug_log(f"Master - Worker images keys: {list(worker_images.keys())}") # Clean up job queue async with prompt_server.distributed_jobs_lock: @@ -2038,7 +2007,6 @@ class DistributedCollectorNode: return (combined,) except Exception as e: log(f"Master - Error combining images: {e}") - debug_log(f"Master - Tensor shapes: {[t.shape for t in cpu_tensors]}") # Return just the master images as fallback return (images,) diff --git a/distributed_upscale.py b/distributed_upscale.py index dc3356c..fefa0cd 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -572,8 +572,6 @@ class UltimateSDUpscaleDistributed: """Initialize the job queue for collecting tiles.""" prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Checking if {multi_job_id} exists in distributed_pending_tile_jobs") - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Current jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") if multi_job_id not in prompt_server.distributed_pending_tile_jobs: prompt_server.distributed_pending_tile_jobs[multi_job_id] = { 'queue': asyncio.Queue(), @@ -581,9 +579,6 @@ class UltimateSDUpscaleDistributed: 'completed_tiles': {} # For static mode data } debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Created new queue for {multi_job_id}") - else: - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Queue already exists for {multi_job_id}") async def _async_collect_worker_tiles(self, multi_job_id, num_workers): """Async helper to collect tiles from workers.""" @@ -606,7 +601,6 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {num_workers} workers to complete") while len(workers_done) < num_workers: - debug_log(f"UltimateSDUpscale Master - Loop status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected, workers_done set: {workers_done}") try: result = await asyncio.wait_for(q.get(), timeout=timeout) worker_id = result['worker_id'] @@ -651,12 +645,9 @@ class UltimateSDUpscaleDistributed: if is_last: workers_done.add(worker_id) debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed") - debug_log(f"UltimateSDUpscale Master - Worker '{worker_id}' marked as done. Total workers done: {len(workers_done)}/{num_workers}") except asyncio.TimeoutError: log(f"UltimateSDUpscale Master - Timeout waiting for tiles") - debug_log(f"UltimateSDUpscale Master - Final status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected") - debug_log(f"UltimateSDUpscale Master - Workers that completed: {workers_done}") break debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_tiles)} tiles from {len(workers_done)} workers") @@ -724,7 +715,7 @@ class UltimateSDUpscaleDistributed: job_data_check = prompt_server.distributed_pending_tile_jobs[multi_job_id] # Check for timed out workers (60 second timeout) for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): - if current_time - last_heartbeat > 60: + if current_time - last_heartbeat > 30: log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") # Requeue assigned images from this worker for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): @@ -875,22 +866,6 @@ class UltimateSDUpscaleDistributed: Note: positive and negative should already be pre-sliced for the current batch_idx.""" debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}") - # Log conditioning details - debug_log(f"[process_tile] Positive conditioning items: {len(positive)}") - debug_log(f"[process_tile] Negative conditioning items: {len(negative)}") - - # Check for ControlNet in conditioning - for i, (emb, cond_dict) in enumerate(positive): - if 'control' in cond_dict: - control = cond_dict['control'] - debug_log(f"[process_tile] Found ControlNet in positive[{i}]") - debug_log(f"[process_tile] Control id: {id(control)}") - if hasattr(control, 'control_model'): - debug_log(f"[process_tile] control_model id: {id(control.control_model)}") - if hasattr(control, 'cond_hint_original'): - debug_log(f"[process_tile] cond_hint_original shape: {control.cond_hint_original.shape}") - debug_log(f"[process_tile] Expected to use hint at index {batch_idx} from batch") - debug_log(f"[process_tile] Embedding shape: {emb.shape}") # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode @@ -930,26 +905,16 @@ class UltimateSDUpscaleDistributed: positive_cropped = positive_tile negative_cropped = negative_tile - # Encode to latent - if tiled_decode and tiled_vae_available: - latent = VAEEncodeTiled().encode(vae, clean_tensor)[0] - else: - latent = VAEEncode().encode(vae, clean_tensor)[0] - - # Log latent shape - if 'samples' in latent: - debug_log(f"[process_tile] Latent shape: {latent['samples'].shape}") + # Encode to latent (always non-tiled, matching original node) + latent = VAEEncode().encode(vae, clean_tensor)[0] # Sample - debug_log(f"[process_tile] About to call common_ksampler with batch_idx={batch_idx}") - debug_log(f"[process_tile] Model id: {id(model)}") samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive_cropped, negative_cropped, latent, denoise=denoise)[0] - debug_log(f"[process_tile] common_ksampler completed") # Decode back to image if tiled_decode and tiled_vae_available: - image = VAEDecodeTiled().decode(vae, samples)[0] + image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] else: image = VAEDecode().decode(vae, samples)[0] @@ -1032,8 +997,6 @@ class UltimateSDUpscaleDistributed: chunk_size = len(chunk) is_chunk_last = (start + chunk_size == total_tiles) # True only for final chunk - debug_log(f"UltimateSDUpscale Worker - Processing chunk {start//MAX_BATCH + 1}, tiles {start} to {start + chunk_size - 1}") - data = aiohttp.FormData() data.add_field('multi_job_id', multi_job_id) data.add_field('worker_id', str(worker_id)) @@ -1072,14 +1035,6 @@ class UltimateSDUpscaleDistributed: # Add image field data.add_field(f'tile_{j}', byte_io, filename=f'tile_{j}.png', content_type='image/png') - # Estimate payload size for logging - payload_size = 0 - for field in data._fields: - if hasattr(field[2], 'read'): - field[2].seek(0) - payload_size += len(field[2].read()) - field[2].seek(0) - debug_log(f"UltimateSDUpscale Worker - Chunk payload size: {payload_size:,} bytes") # Retry logic with exponential backoff max_retries = 5 @@ -1090,16 +1045,12 @@ class UltimateSDUpscaleDistributed: session = await get_client_session() url = f"{master_url}/distributed/tile_complete" - debug_log(f"UltimateSDUpscale Worker - Sending chunk of {chunk_size} tiles to {url}, attempt {attempt + 1}") - async with session.post(url, data=data) as response: response.raise_for_status() - debug_log(f"UltimateSDUpscale Worker - Successfully sent chunk of {chunk_size} tiles") break # Success, move to next chunk except Exception as e: if attempt < max_retries - 1: - debug_log(f"UltimateSDUpscale Worker - Retry {attempt + 1}/{max_retries} after error: {e}") await asyncio.sleep(retry_delay) retry_delay *= 2 # Exponential backoff else: @@ -1260,6 +1211,7 @@ class UltimateSDUpscaleDistributed: sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) + # Note: No per-tile drain here – that's what makes this "per-image" result_images[image_idx] = local_image @@ -1268,6 +1220,22 @@ class UltimateSDUpscaleDistributed: self._mark_image_completed(multi_job_id, image_idx, local_image), timeout=5.0 ) + + # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing) + drained_count = run_async_in_server_loop( + self._drain_worker_results_queue(multi_job_id), + timeout=5.0 + ) + + if drained_count > 0: + debug_log(f"Drained {drained_count} worker images after master's image {image_idx}") + + # NEW: Log overall progress (includes master's image + any drained workers) + completed_now = run_async_in_server_loop( + self._get_total_completed_count(multi_job_id), + timeout=1.0 + ) + log(f"Progress: {completed_now}/{batch_size} images completed") else: # Queue empty: collect any queued worker results to update progress drained_count = run_async_in_server_loop( @@ -1291,11 +1259,9 @@ class UltimateSDUpscaleDistributed: timeout=1.0 ) + log(f"Progress: {completed_now}/{batch_size} images completed") - debug_log(f"Queue empty, completed count: {completed_now}/{batch_size}") - if completed_now >= batch_size: - debug_log("All images completed, exiting master loop") break # Check if there are pending images in the queue (could be requeued) @@ -1305,7 +1271,6 @@ class UltimateSDUpscaleDistributed: ) if pending_count > 0: - debug_log(f"Found {pending_count} pending images (possibly requeued), continuing to process") consecutive_retries = 0 # Reset retries since there's work to do continue @@ -1314,7 +1279,7 @@ class UltimateSDUpscaleDistributed: log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") break # Force exit to collection phase - log("Queue empty temporarily, waiting 2s before retry") + debug_log("Waiting for workers") time.sleep(2) debug_log(f"Master processed {processed_count} images locally") @@ -1455,7 +1420,7 @@ class UltimateSDUpscaleDistributed: # Check for timed out workers (60 second timeout) for worker, last_heartbeat in list(job_data.get('worker_status', {}).items()): - if current_time - last_heartbeat > 60: + if current_time - last_heartbeat > 30: log(f"Worker {worker} timed out") # Requeue assigned images from this worker for idx in job_data.get('assigned_to_workers', {}).get(worker, []): @@ -1701,9 +1666,8 @@ class UltimateSDUpscaleDistributed: async with session.post(url, json=data) as response: response.raise_for_status() - debug_log(f"Worker {worker_id} sent heartbeat") except Exception as e: - debug_log(f"Worker {worker_id} heartbeat failed: {e}") + pass async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): """Send completion signal to master in dynamic mode.""" @@ -2062,9 +2026,6 @@ async def tile_complete_endpoint(request): # Put tiles into queue async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale - tile_complete: Checking distributed_pending_tile_jobs for job {multi_job_id}") - debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") - if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: @@ -2080,7 +2041,6 @@ async def tile_complete_endpoint(request): 'tiles': tiles, 'is_last': is_last }) - debug_log(f"UltimateSDUpscale - Received batch of {len(tiles)} tiles for job {multi_job_id} from worker {worker_id}") else: # Put single tile (backward compat) tile_data = tiles[0] @@ -2095,12 +2055,9 @@ async def tile_complete_endpoint(request): 'padding': tile_data['padding'], 'is_last': is_last }) - debug_log(f"UltimateSDUpscale - Received single tile {tile_data['tile_idx']} for job {multi_job_id} from worker {worker_id}") return web.json_response({"status": "success"}) else: - debug_log(f"UltimateSDUpscale - API Error: Job {multi_job_id} not found in distributed_pending_tile_jobs") - debug_log(f"UltimateSDUpscale - API Error: Available jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") return await handle_api_error(request, "Job not found or already complete", 404) except Exception as e: return await handle_api_error(request, e, 500) From 84a28d71224b5f05bd1df8ba9c16e83d844a3976 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 10:30:29 +1000 Subject: [PATCH 33/93] remove redundant endpoint and warn users when cloudflare tunnel is down --- distributed.py | 17 ------ web/executionUtils.js | 32 +++++++++++ web/ui.js | 125 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 157 insertions(+), 17 deletions(-) diff --git a/distributed.py b/distributed.py index 46df518..aee2359 100644 --- a/distributed.py +++ b/distributed.py @@ -1582,23 +1582,6 @@ async def check_file_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) -@server.PromptServer.instance.routes.get("/distributed/system_info") -async def system_info_endpoint(request): - """Return system information including platform details""" - import platform - - return web.json_response({ - "platform": { - "os_name": os.name, # 'nt' for Windows, 'posix' for Unix/Linux - "path_separator": os.sep, # '\\' on Windows, '/' on Linux - "system": platform.system(), # 'Windows', 'Linux', 'Darwin' - "machine": platform.machine(), # 'x86_64', 'AMD64', etc. - }, - "comfyui": { - # Can add ComfyUI-specific info here - "version": "distributed-extension", - } - }) @server.PromptServer.instance.routes.post("/distributed/job_complete") async def job_complete_endpoint(request): diff --git a/web/executionUtils.js b/web/executionUtils.js index 4c116cd..7c493c9 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -83,6 +83,38 @@ export async function executeParallelDistributed(extension, promptWrapper) { extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug"); + // Check if master host might be unreachable by workers (cloudflare tunnel down) + const masterHost = extension.config?.master?.host || ''; + const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost); + + if (isCloudflareHost && activeWorkers.length > 0) { + // Try to verify if the cloudflare tunnel is actually up + try { + const testUrl = `${window.location.protocol}//${masterHost}/distributed/status`; + const response = await fetch(testUrl, { + method: 'GET', + mode: 'cors', + cache: 'no-cache', + signal: AbortSignal.timeout(3000) // 3 second timeout + }); + + if (!response.ok) { + throw new Error('Master not reachable'); + } + } catch (error) { + // Cloudflare tunnel appears to be down + extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error"); + + if (extension.ui?.showCloudflareWarning) { + extension.ui.showCloudflareWarning(extension, masterHost); + } + + // Stop execution - workers won't be able to send results back + extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error"); + return null; // This will prevent the workflow from running + } + } + // Find all distributed nodes in the workflow const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector"); const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); diff --git a/web/ui.js b/web/ui.js index bdbc566..9bc0411 100644 --- a/web/ui.js +++ b/web/ui.js @@ -344,6 +344,131 @@ export class DistributedUI { } } + showCloudflareWarning(extension, masterHost) { + // Remove any existing banner first + const existingBanner = document.getElementById('cloudflare-warning-banner'); + if (existingBanner) { + existingBanner.remove(); + } + + // Create warning banner + const banner = document.createElement('div'); + banner.id = 'cloudflare-warning-banner'; + banner.style.cssText = ` + position: fixed; + top: 0; + left: 0; + right: 0; + background: #ff9800; + color: #333; + padding: 8px 16px; + text-align: center; + z-index: 10000; + display: flex; + align-items: center; + justify-content: center; + gap: 16px; + box-shadow: 0 2px 5px rgba(0,0,0,0.2); + font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; + `; + + const messageSpan = document.createElement('span'); + messageSpan.innerHTML = `Connection issue: Master address ${masterHost} is not reachable. The cloudflare tunnel may be offline.`; + messageSpan.style.fontSize = '13px'; + + const resetButton = document.createElement('button'); + resetButton.textContent = 'Reset Master Address'; + resetButton.style.cssText = ` + background: #333; + color: white; + border: none; + padding: 6px 14px; + border-radius: 4px; + cursor: pointer; + font-weight: 500; + font-size: 13px; + transition: background 0.2s; + `; + resetButton.onmouseover = () => resetButton.style.background = '#555'; + resetButton.onmouseout = () => resetButton.style.background = '#333'; + + const dismissButton = document.createElement('button'); + dismissButton.textContent = 'Dismiss'; + dismissButton.style.cssText = ` + background: transparent; + color: #333; + border: 1px solid #333; + padding: 6px 14px; + border-radius: 4px; + cursor: pointer; + font-size: 13px; + transition: opacity 0.2s; + `; + dismissButton.onmouseover = () => dismissButton.style.opacity = '0.7'; + dismissButton.onmouseout = () => dismissButton.style.opacity = '1'; + + // Add click handlers + resetButton.onclick = async () => { + resetButton.disabled = true; + resetButton.textContent = 'Resetting...'; + + try { + // Save with empty host - this will trigger auto-detection + await extension.api.updateMaster({ + name: extension.config?.master?.name || "Master", + host: "" + }); + + // The API call above doesn't trigger auto-detection, so we need to do it + await extension.detectMasterIP(); + + // Reload config to get the new detected IP + await extension.loadConfig(); + + // Update UI if sidebar is open + if (extension.panelElement) { + const hostInput = document.getElementById('master-host'); + if (hostInput) { + hostInput.value = extension.config?.master?.host || ""; + } + } + + // Show success message + extension.app.extensionManager.toast.add({ + severity: "success", + summary: "Master Host Reset", + detail: `IP auto-detected: ${extension.config?.master?.host || 'detecting...'}`, + life: 4000 + }); + + banner.remove(); + } catch (error) { + resetButton.disabled = false; + resetButton.textContent = 'Reset Master Host'; + extension.log(`Failed to reset master host: ${error.message}`, "error"); + } + }; + + dismissButton.onclick = () => banner.remove(); + + // Assemble banner + banner.appendChild(messageSpan); + banner.appendChild(resetButton); + banner.appendChild(dismissButton); + + // Add to page + document.body.prepend(banner); + + // Auto-dismiss after 30 seconds + setTimeout(() => { + if (document.getElementById('cloudflare-warning-banner')) { + banner.style.transition = 'opacity 0.5s'; + banner.style.opacity = '0'; + setTimeout(() => banner.remove(), 500); + } + }, 30000); + } + updateStatusDot(workerId, color, title, pulsing = false) { const statusDot = document.getElementById(`status-${workerId}`); if (!statusDot) return; From 1a80cf90b839f5ee3b378bde2d4b0f9c0b4db956 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 10:35:15 +1000 Subject: [PATCH 34/93] fix: empty else block --- distributed.py | 1 + 1 file changed, 1 insertion(+) diff --git a/distributed.py b/distributed.py index aee2359..347fcc2 100644 --- a/distributed.py +++ b/distributed.py @@ -1895,6 +1895,7 @@ class DistributedCollectorNode: workers_done.add(worker_id) p.update(1) # +1 per completed worker else: + pass # Continue waiting for more results except asyncio.TimeoutError: missing_workers = set(str(w) for w in enabled_workers) - workers_done From 882dcf3df5aa6aff82a9bb39fb744a02b3bc2cb0 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 10:48:00 +1000 Subject: [PATCH 35/93] usdu util --- utils/usdu_utils.py | 455 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 455 insertions(+) create mode 100644 utils/usdu_utils.py diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py new file mode 100644 index 0000000..a783f0a --- /dev/null +++ b/utils/usdu_utils.py @@ -0,0 +1,455 @@ +import numpy as np +from PIL import Image, ImageFilter +import torch +import torch.nn.functional as F +from torchvision.transforms import GaussianBlur +import math + +if (not hasattr(Image, 'Resampling')): # For older versions of Pillow + Image.Resampling = Image + +BLUR_KERNEL_SIZE = 15 + + +def tensor_to_pil(img_tensor, batch_index=0): + # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels] + # and returns an RGB PIL Image. Assumes channels=3 + return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8)) + + +def pil_to_tensor(image): + # Takes a PIL image and returns a tensor of shape [1, height, width, channels] + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image).unsqueeze(0) + if len(image.shape) == 3: # If the image is grayscale, add a channel dimension + image = image.unsqueeze(-1) + return image + + +def controlnet_hint_to_pil(tensor, batch_index=0): + return tensor_to_pil(tensor.movedim(1, -1), batch_index) + + +def pil_to_controlnet_hint(img): + return pil_to_tensor(img).movedim(-1, 1) + + +def crop_tensor(tensor, region): + # Takes a tensor of shape [batch_size, height, width, channels] and crops it to the given region + x1, y1, x2, y2 = region + return tensor[:, y1:y2, x1:x2, :] + + +def resize_tensor(tensor, size, mode="nearest-exact"): + # Takes a tensor of shape [B, C, H, W] and resizes + # it to a shape of [B, C, size[0], size[1]] using the given mode + return torch.nn.functional.interpolate(tensor, size=size, mode=mode) + + +def get_crop_region(mask, pad=0): + # Takes a black and white PIL image in 'L' mode and returns the coordinates of the white rectangular mask region + # Should be equivalent to the get_crop_region function from https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/masking.py + coordinates = mask.getbbox() + if coordinates is not None: + x1, y1, x2, y2 = coordinates + else: + x1, y1, x2, y2 = mask.width, mask.height, 0, 0 + # Apply padding + x1 = max(x1 - pad, 0) + y1 = max(y1 - pad, 0) + x2 = min(x2 + pad, mask.width) + y2 = min(y2 + pad, mask.height) + return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height)) + + +def fix_crop_region(region, image_size): + # Remove the extra pixel added by the get_crop_region function + image_width, image_height = image_size + x1, y1, x2, y2 = region + if x2 < image_width: + x2 -= 1 + if y2 < image_height: + y2 -= 1 + return x1, y1, x2, y2 + + +def expand_crop(region, width, height, target_width, target_height): + ''' + Expands a crop region to a specified target size. + :param region: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points + of the rectangular region. Expected to have x2 > x1 and y2 > y1. + :param width: The width of the image the crop region is from. + :param height: The height of the image the crop region is from. + :param target_width: The desired width of the crop region. + :param target_height: The desired height of the crop region. + ''' + x1, y1, x2, y2 = region + actual_width = x2 - x1 + actual_height = y2 - y1 + # target_width = math.ceil(actual_width / 8) * 8 + # target_height = math.ceil(actual_height / 8) * 8 + + # Try to expand region to the right of half the difference + width_diff = target_width - actual_width + x2 = min(x2 + width_diff // 2, width) + # Expand region to the left of the difference including the pixels that could not be expanded to the right + width_diff = target_width - (x2 - x1) + x1 = max(x1 - width_diff, 0) + # Try the right again + width_diff = target_width - (x2 - x1) + x2 = min(x2 + width_diff, width) + + # Try to expand region to the bottom of half the difference + height_diff = target_height - actual_height + y2 = min(y2 + height_diff // 2, height) + # Expand region to the top of the difference including the pixels that could not be expanded to the bottom + height_diff = target_height - (y2 - y1) + y1 = max(y1 - height_diff, 0) + # Try the bottom again + height_diff = target_height - (y2 - y1) + y2 = min(y2 + height_diff, height) + + return (x1, y1, x2, y2), (target_width, target_height) + + +def resize_region(region, init_size, resize_size): + # Resize a crop so that it fits an image that was resized to the given width and height + x1, y1, x2, y2 = region + init_width, init_height = init_size + resize_width, resize_height = resize_size + x1 = math.floor(x1 * resize_width / init_width) + x2 = math.ceil(x2 * resize_width / init_width) + y1 = math.floor(y1 * resize_height / init_height) + y2 = math.ceil(y2 * resize_height / init_height) + return (x1, y1, x2, y2) + + +def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): + ''' + Pads an image with the given number of pixels on each side and fills the padding with data from the edges. + :param image: A PIL image + :param left_pad: The number of pixels to pad on the left side + :param right_pad: The number of pixels to pad on the right side + :param top_pad: The number of pixels to pad on the top side + :param bottom_pad: The number of pixels to pad on the bottom side + :param blur: Whether to blur the padded edges + :return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad) + ''' + left_edge = image.crop((0, 1, 1, image.height - 1)) + right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1)) + top_edge = image.crop((1, 0, image.width - 1, 1)) + bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height)) + new_width = image.width + left_pad + right_pad + new_height = image.height + top_pad + bottom_pad + padded_image = Image.new(image.mode, (new_width, new_height)) + padded_image.paste(image, (left_pad, top_pad)) + if fill: + for i in range(left_pad): + edge = left_edge.resize( + (1, new_height - i * (top_pad + bottom_pad) // left_pad), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (i, i * top_pad // left_pad)) + for i in range(right_pad): + edge = right_edge.resize( + (1, new_height - i * (top_pad + bottom_pad) // right_pad), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (new_width - 1 - i, i * top_pad // right_pad)) + for i in range(top_pad): + edge = top_edge.resize( + (new_width - i * (left_pad + right_pad) // top_pad, 1), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (i * left_pad // top_pad, i)) + for i in range(bottom_pad): + edge = bottom_edge.resize( + (new_width - i * (left_pad + right_pad) // bottom_pad, 1), resample=Image.Resampling.NEAREST) + padded_image.paste(edge, (i * left_pad // bottom_pad, new_height - 1 - i)) + if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0): + padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE)) + padded_image.paste(image, (left_pad, top_pad)) + return padded_image + + +def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): + ''' + Pads an image with the given number of pixels on each side and fills the padding with data from the edges. + Faster than pad_image, but only pads with edge data in straight lines. + :param image: A PIL image + :param left_pad: The number of pixels to pad on the left side + :param right_pad: The number of pixels to pad on the right side + :param top_pad: The number of pixels to pad on the top side + :param bottom_pad: The number of pixels to pad on the bottom side + :param blur: Whether to blur the padded edges + :return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad) + ''' + left_edge = image.crop((0, 1, 1, image.height - 1)) + right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1)) + top_edge = image.crop((1, 0, image.width - 1, 1)) + bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height)) + new_width = image.width + left_pad + right_pad + new_height = image.height + top_pad + bottom_pad + padded_image = Image.new(image.mode, (new_width, new_height)) + padded_image.paste(image, (left_pad, top_pad)) + if fill: + if left_pad > 0: + padded_image.paste(left_edge.resize((left_pad, new_height), resample=Image.Resampling.NEAREST), (0, 0)) + if right_pad > 0: + padded_image.paste(right_edge.resize((right_pad, new_height), + resample=Image.Resampling.NEAREST), (new_width - right_pad, 0)) + if top_pad > 0: + padded_image.paste(top_edge.resize((new_width, top_pad), resample=Image.Resampling.NEAREST), (0, 0)) + if bottom_pad > 0: + padded_image.paste(bottom_edge.resize((new_width, bottom_pad), + resample=Image.Resampling.NEAREST), (0, new_height - bottom_pad)) + if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0): + padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE)) + padded_image.paste(image, (left_pad, top_pad)) + return padded_image + + +def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False): + ''' + Pads an image tensor with the given number of pixels on each side and fills the padding with data from the edges. + :param tensor: A tensor of shape [B, H, W, C] + :param left_pad: The number of pixels to pad on the left side + :param right_pad: The number of pixels to pad on the right side + :param top_pad: The number of pixels to pad on the top side + :param bottom_pad: The number of pixels to pad on the bottom side + :param blur: Whether to blur the padded edges + :return: A tensor of shape [B, H + top_pad + bottom_pad, W + left_pad + right_pad, C] + ''' + batch_size, channels, height, width = tensor.shape + h_pad = left_pad + right_pad + v_pad = top_pad + bottom_pad + new_width = width + h_pad + new_height = height + v_pad + + # Create empty image + padded = torch.zeros((batch_size, channels, new_height, new_width), dtype=tensor.dtype) + + # Copy the original image into the centor of the padded tensor + padded[:, :, top_pad:top_pad + height, left_pad:left_pad + width] = tensor + + # Duplicate the edges of the original image into the padding + if top_pad > 0: + padded[:, :, :top_pad, :] = padded[:, :, top_pad:top_pad + 1, :] # Top edge + if bottom_pad > 0: + padded[:, :, -bottom_pad:, :] = padded[:, :, -bottom_pad - 1:-bottom_pad, :] # Bottom edge + if left_pad > 0: + padded[:, :, :, :left_pad] = padded[:, :, :, left_pad:left_pad + 1] # Left edge + if right_pad > 0: + padded[:, :, :, -right_pad:] = padded[:, :, :, -right_pad - 1:-right_pad] # Right edge + + return padded + + +def resize_and_pad_image(image, width, height, fill=False, blur=False): + ''' + Resizes an image to the given width and height and pads it to the given width and height. + :param image: A PIL image + :param width: The width of the resized image + :param height: The height of the resized image + :param fill: Whether to fill the padding with data from the edges + :param blur: Whether to blur the padded edges + :return: A PIL image of size (width, height) + ''' + width_ratio = width / image.width + height_ratio = height / image.height + if height_ratio > width_ratio: + resize_ratio = width_ratio + else: + resize_ratio = height_ratio + resize_width = round(image.width * resize_ratio) + resize_height = round(image.height * resize_ratio) + resized = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS) + # Pad the sides of the image to get the image to the desired size that wasn't covered by the resize + horizontal_pad = (width - resize_width) // 2 + vertical_pad = (height - resize_height) // 2 + result = pad_image2(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur) + result = result.resize((width, height), resample=Image.Resampling.LANCZOS) + return result, (horizontal_pad, vertical_pad) + + +def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False): + ''' + Resizes an image tensor to the given width and height and pads it to the given width and height. + :param tensor: A tensor of shape [B, H, W, C] + :param width: The width of the resized image + :param height: The height of the resized image + :param fill: Whether to fill the padding with data from the edges + :param blur: Whether to blur the padded edges + :return: A tensor of shape [B, height, width, C] + ''' + # Resize the image to the closest size that maintains the aspect ratio + width_ratio = width / tensor.shape[3] + height_ratio = height / tensor.shape[2] + if height_ratio > width_ratio: + resize_ratio = width_ratio + else: + resize_ratio = height_ratio + resize_width = round(tensor.shape[3] * resize_ratio) + resize_height = round(tensor.shape[2] * resize_ratio) + resized = F.interpolate(tensor, size=(resize_height, resize_width), mode='nearest-exact') + # Pad the sides of the image to get the image to the desired size that wasn't covered by the resize + horizontal_pad = (width - resize_width) // 2 + vertical_pad = (height - resize_height) // 2 + result = pad_tensor(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur) + result = F.interpolate(result, size=(height, width), mode='nearest-exact') + return result + + +def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "control" not in cond_dict: + return + c = cond_dict["control"] + controlnet = c.copy() + cond_dict["control"] = controlnet + while c is not None: + # hint is shape (B, C, H, W) + hint = controlnet.cond_hint_original + resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1]) + hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1) + hint = resize_tensor(hint, tile_size[::-1]) + controlnet.cond_hint_original = hint + c = c.previous_controlnet + controlnet.set_previous_controlnet(c.copy() if c is not None else None) + controlnet = controlnet.previous_controlnet + + +def region_intersection(region1, region2): + """ + Returns the coordinates of the intersection of two rectangular regions. + :param region1: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points + of the first rectangular region. Expected to have x2 > x1 and y2 > y1. + :param region2: The second rectangular region with the same format as the first. + :return: A tuple of the form (x1, y1, x2, y2) denoting the rectangular intersection. + None if there is no intersection. + """ + x1, y1, x2, y2 = region1 + x1_, y1_, x2_, y2_ = region2 + x1 = max(x1, x1_) + y1 = max(y1, y1_) + x2 = min(x2, x2_) + y2 = min(y2, y2_) + if x1 >= x2 or y1 >= y2: + return None + return (x1, y1, x2, y2) + + +def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "gligen" not in cond_dict: + return + type, model, cond = cond_dict["gligen"] + if type != "position": + from warnings import warn + warn(f"Unknown gligen type {type}") + return + cropped = [] + for c in cond: + emb, h, w, y, x = c + # Get the coordinates of the box in the upscaled image + x1 = x * 8 + y1 = y * 8 + x2 = x1 + w * 8 + y2 = y1 + h * 8 + gligen_upscaled_box = resize_region((x1, y1, x2, y2), init_size, canvas_size) + + # Calculate the intersection of the gligen box and the region + intersection = region_intersection(gligen_upscaled_box, region) + if intersection is None: + continue + x1, y1, x2, y2 = intersection + + # Offset the gligen box so that the origin is at the top left of the tile region + x1 -= region[0] + y1 -= region[1] + x2 -= region[0] + y2 -= region[1] + + # Add the padding + x1 += w_pad + y1 += h_pad + x2 += w_pad + y2 += h_pad + + # Set the new position params + h = (y2 - y1) // 8 + w = (x2 - x1) // 8 + x = x1 // 8 + y = y1 // 8 + cropped.append((emb, h, w, y, x)) + + cond_dict["gligen"] = (type, model, cropped) + + +def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "area" not in cond_dict: + return + + # Resize the area conditioning to the canvas size and confine it to the tile region + h, w, y, x = cond_dict["area"] + w, h, x, y = 8 * w, 8 * h, 8 * x, 8 * y + x1, y1, x2, y2 = resize_region((x, y, x + w, y + h), init_size, canvas_size) + intersection = region_intersection((x1, y1, x2, y2), region) + if intersection is None: + del cond_dict["area"] + del cond_dict["strength"] + return + x1, y1, x2, y2 = intersection + + # Offset origin to the top left of the tile + x1 -= region[0] + y1 -= region[1] + x2 -= region[0] + y2 -= region[1] + + # Add the padding + x1 += w_pad + y1 += h_pad + x2 += w_pad + y2 += h_pad + + # Set the params for tile + w, h = (x2 - x1) // 8, (y2 - y1) // 8 + x, y = x1 // 8, y1 // 8 + + cond_dict["area"] = (h, w, y, x) + + +def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): + if "mask" not in cond_dict: + return + mask_tensor = cond_dict["mask"] # (B, H, W) + masks = [] + for i in range(mask_tensor.shape[0]): + # Convert to PIL image + mask = tensor_to_pil(mask_tensor, i) # W x H + + # Resize the mask to the canvas size + mask = mask.resize(canvas_size, Image.Resampling.BICUBIC) + + # Crop the mask to the region + mask = mask.crop(region) + + # Add padding + mask, _ = resize_and_pad_image(mask, tile_size[0], tile_size[1], fill=True) + + # Resize the mask to the tile size + if tile_size != mask.size: + mask = mask.resize(tile_size, Image.Resampling.BICUBIC) + + # Convert back to tensor + mask = pil_to_tensor(mask) # (1, H, W, 1) + mask = mask.squeeze(-1) # (1, H, W) + masks.append(mask) + + cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W) + + +def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0): + cropped = [] + for emb, x in cond: + cond_dict = x.copy() + n = [emb, cond_dict] + crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + cropped.append(n) + return cropped From 9505d2a29eb8699e2503b35698a29bc846c6f80d Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 13:11:11 +1000 Subject: [PATCH 36/93] fixed cloudflare detction bug --- web/executionUtils.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/executionUtils.js b/web/executionUtils.js index 7c493c9..6e2c583 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -90,7 +90,7 @@ export async function executeParallelDistributed(extension, promptWrapper) { if (isCloudflareHost && activeWorkers.length > 0) { // Try to verify if the cloudflare tunnel is actually up try { - const testUrl = `${window.location.protocol}//${masterHost}/distributed/status`; + const testUrl = `${window.location.protocol}//${masterHost}/prompt`; const response = await fetch(testUrl, { method: 'GET', mode: 'cors', From 319781bb475aeab6f08e8cb9f11b5f3d9a180aba Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 13:12:35 +1000 Subject: [PATCH 37/93] heartbeat timeout variable added --- distributed_upscale.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index fefa0cd..47bdd02 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -39,6 +39,9 @@ MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Configure maximum payload size (50MB default, configurable via environment variable) MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) +# Heartbeat timeout in seconds (configurable via env var, default 120s to cover typical image processing) +HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '120')) + # Helper functions for shallow copying conditioning without duplicating models def clone_control_chain(control, clone_hint=True): """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" @@ -715,7 +718,7 @@ class UltimateSDUpscaleDistributed: job_data_check = prompt_server.distributed_pending_tile_jobs[multi_job_id] # Check for timed out workers (60 second timeout) for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): - if current_time - last_heartbeat > 30: + if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") # Requeue assigned images from this worker for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): @@ -1420,7 +1423,7 @@ class UltimateSDUpscaleDistributed: # Check for timed out workers (60 second timeout) for worker, last_heartbeat in list(job_data.get('worker_status', {}).items()): - if current_time - last_heartbeat > 30: + if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: log(f"Worker {worker} timed out") # Requeue assigned images from this worker for idx in job_data.get('assigned_to_workers', {}).get(worker, []): @@ -1523,6 +1526,10 @@ class UltimateSDUpscaleDistributed: sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) + run_async_in_server_loop( + self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) # Send processed image back to master try: From 60eab42174671ced855779f69670196fe46049aa Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sat, 19 Jul 2025 10:13:13 +1000 Subject: [PATCH 38/93] strategic yields in dynamic mode --- distributed_upscale.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 47bdd02..5184785 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -1214,6 +1214,9 @@ class UltimateSDUpscaleDistributed: sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) + + # Yield after each tile to minimize worker downtime + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Note: No per-tile drain here – that's what makes this "per-image" result_images[image_idx] = local_image @@ -1239,18 +1242,24 @@ class UltimateSDUpscaleDistributed: timeout=1.0 ) log(f"Progress: {completed_now}/{batch_size} images completed") + + # Yield to allow workers to get new images after completing one + run_async_in_server_loop(self._async_yield(), timeout=0.1) else: # Queue empty: collect any queued worker results to update progress drained_count = run_async_in_server_loop( self._drain_worker_results_queue(multi_job_id), timeout=5.0 ) + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain # Check for timed out workers and requeue their images requeued_count = run_async_in_server_loop( self._check_and_requeue_timed_out_workers(multi_job_id, batch_size), timeout=5.0 ) + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue + if requeued_count > 0: log(f"Requeued {requeued_count} images from timed out workers") consecutive_retries = 0 # Reset since we have work to do @@ -1267,6 +1276,8 @@ class UltimateSDUpscaleDistributed: if completed_now >= batch_size: break + run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check + # Check if there are pending images in the queue (could be requeued) pending_count = run_async_in_server_loop( self._get_pending_count(multi_job_id), @@ -1283,7 +1294,8 @@ class UltimateSDUpscaleDistributed: break # Force exit to collection phase debug_log("Waiting for workers") - time.sleep(2) + # Use async sleep to allow event loop to process worker requests + run_async_in_server_loop(asyncio.sleep(2), timeout=3.0) debug_log(f"Master processed {processed_count} images locally") @@ -1321,6 +1333,10 @@ class UltimateSDUpscaleDistributed: return (result_tensor,) + async def _async_yield(self): + """Simple async yield to allow event loop processing.""" + await asyncio.sleep(0) + async def _get_next_image_index(self, multi_job_id): """Get next image index from pending queue for master.""" prompt_server = ensure_tile_jobs_initialized() From fa2dff641276390e59a15756ced19c1ad72e2f77 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sat, 19 Jul 2025 11:23:00 +1000 Subject: [PATCH 39/93] continuous working polling only when panel open --- web/main.js | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 3 deletions(-) diff --git a/web/main.js b/web/main.js index f59ce4e..f200ad4 100644 --- a/web/main.js +++ b/web/main.js @@ -29,6 +29,9 @@ class DistributedExtension { // Initialize status check timeout reference this.statusCheckTimeout = null; + + // Initialize abort controller for status checks + this.statusCheckAbortController = null; // Inject CSS for pulsing animation this.injectStyles(); @@ -36,7 +39,8 @@ class DistributedExtension { this.loadConfig().then(async () => { this.registerSidebarTab(); this.setupInterceptor(); - this.startStatusChecking(); + // Don't start polling until panel opens + // this.startStatusChecking(); this.loadManagedWorkers(); // Detect master IP after everything is set up this.detectMasterIP(); @@ -177,9 +181,41 @@ class DistributedExtension { title: "Distributed", tooltip: "Distributed Control Panel", type: "custom", - render: (el) => renderSidebarContent(this, el) + render: (el) => { + this.panelElement = el; + this.onPanelOpen(); + return renderSidebarContent(this, el); + }, + destroy: () => { + this.onPanelClose(); + } }); } + + onPanelOpen() { + this.log("Panel opened - starting status polling", "debug"); + if (!this.statusCheckTimeout) { + this.checkAllWorkerStatuses(); + } + } + + 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.panelElement = null; + } updateSummary() { const summaryEl = document.getElementById('distributed-summary'); @@ -213,6 +249,13 @@ class DistributedExtension { } async checkAllWorkerStatuses() { + // Don't continue if panel is closed + if (!this.panelElement) return; + + // Create new abort controller for this round of checks + this.statusCheckAbortController = new AbortController(); + + // Check master status this.checkMasterStatus(); @@ -322,10 +365,16 @@ class DistributedExtension { const statusDot = document.getElementById(`status-${worker.id}`); try { + // Combine timeout with abort controller signal + const timeoutSignal = AbortSignal.timeout(TIMEOUTS.STATUS_CHECK); + const signal = this.statusCheckAbortController + ? AbortSignal.any([timeoutSignal, this.statusCheckAbortController.signal]) + : timeoutSignal; + const response = await fetch(url, { method: 'GET', mode: 'cors', - signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + signal: signal }); if (response.ok) { @@ -361,6 +410,11 @@ class DistributedExtension { throw new Error(`HTTP ${response.status}`); } } catch (error) { + // Don't process aborted requests + if (error.name === 'AbortError') { + return; + } + // Worker is offline or unreachable this.state.setWorkerStatus(worker.id, { online: false, From 526c603ecb1b0faef1d6c8e425bbc8c8bf8e5bda Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sat, 19 Jul 2025 21:13:16 +1000 Subject: [PATCH 40/93] fixed cloudflare banner button not resetting host address --- web/main.js | 12 +++--------- web/ui.js | 13 +++++++++++-- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/web/main.js b/web/main.js index f200ad4..4ae885c 100644 --- a/web/main.js +++ b/web/main.js @@ -990,15 +990,9 @@ class DistributedExtension { // Save the detected IP (pass true to suppress notification) await this.api.updateMaster({ host: data.recommended_ip }); - // Show a single combined notification - if (app.extensionManager?.toast) { - app.extensionManager.toast.add({ - severity: "info", - summary: "Master IP Auto-Detected", - detail: `Master IP set to ${data.recommended_ip} for remote worker communication`, - life: 5000 - }); - } + // Update local config immediately + if (!this.config.master) this.config.master = {}; + this.config.master.host = data.recommended_ip; } } catch (error) { this.log("Error detecting master IP: " + error.message, "error"); diff --git a/web/ui.js b/web/ui.js index 9bc0411..93b6274 100644 --- a/web/ui.js +++ b/web/ui.js @@ -419,12 +419,21 @@ export class DistributedUI { host: "" }); + // Clear the local config host so detectMasterIP() doesn't skip + if (extension.config?.master) { + extension.config.master.host = ""; + } + // The API call above doesn't trigger auto-detection, so we need to do it await extension.detectMasterIP(); // Reload config to get the new detected IP await extension.loadConfig(); + // Log the new master URL for debugging + const newMasterUrl = extension.getMasterUrl(); + extension.log(`Master host reset. New URL: ${newMasterUrl}`, "info"); + // Update UI if sidebar is open if (extension.panelElement) { const hostInput = document.getElementById('master-host'); @@ -433,11 +442,11 @@ export class DistributedUI { } } - // Show success message + // Show success message with the actual URL that will be used extension.app.extensionManager.toast.add({ severity: "success", summary: "Master Host Reset", - detail: `IP auto-detected: ${extension.config?.master?.host || 'detecting...'}`, + detail: `New address: ${newMasterUrl}`, life: 4000 }); From fcdf6ff28de17991eff0ebb58d1477e6b4e2231d Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sun, 20 Jul 2025 18:06:11 +1000 Subject: [PATCH 41/93] upscaler refactoring: moved functions and endpoints to utils --- distributed_upscale.py | 403 +--------------------------------------- utils/usdu_utils.py | 411 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 412 insertions(+), 402 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 5184785..bb22970 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -28,68 +28,15 @@ from .utils.constants import ( ) # Import for controller support -from .utils.usdu_utils import crop_cond +from .utils.usdu_utils import crop_cond, clone_control_chain, clone_conditioning, ensure_tile_jobs_initialized -# Additional import for deepcopy -import copy # Make MAX_BATCH configurable MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) -# Configure maximum payload size (50MB default, configurable via environment variable) -MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) - # Heartbeat timeout in seconds (configurable via env var, default 120s to cover typical image processing) HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '120')) -# Helper functions for shallow copying conditioning without duplicating models -def clone_control_chain(control, clone_hint=True): - """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" - if control is None: - return None - new_control = copy.copy(control) # Shallow copy (shares model) - if clone_hint and hasattr(control, 'cond_hint_original'): - new_control.cond_hint_original = control.cond_hint_original.clone() - if hasattr(control, 'previous_controlnet'): - new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) - return new_control - -def clone_conditioning(cond_list, clone_hints=True): - """Clone conditioning without duplicating ControlNet models.""" - new_cond = [] - for emb, cond_dict in cond_list: - new_emb = emb.clone() if emb is not None else None - new_dict = cond_dict.copy() - if 'control' in new_dict: - new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) - if 'mask' in new_dict: - new_dict['mask'] = new_dict['mask'].clone() - # Handle other potential fields if needed - if 'pooled_output' in new_dict: - new_dict['pooled_output'] = new_dict['pooled_output'].clone() - if 'area' in new_dict: - new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple - new_cond.append([new_emb, new_dict]) - return new_cond - -# Helper function to ensure persistent state is initialized -def ensure_tile_jobs_initialized(): - """Ensure tile job storage is initialized on the server instance.""" - prompt_server = server.PromptServer.instance - if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): - debug_log("Initializing persistent tile job queue on server instance.") - prompt_server.distributed_pending_tile_jobs = {} - prompt_server.distributed_tile_jobs_lock = asyncio.Lock() - else: - # Clean up any legacy queue structures that don't have the 'mode' field - # (Should be rare after fixes, but keep for safety) - to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() - if not isinstance(job_data, dict) or 'mode' not in job_data] - for job_id in to_remove: - debug_log(f"Removing legacy queue structure for job {job_id}") - del prompt_server.distributed_pending_tile_jobs[job_id] - return prompt_server - # Note: tensor_to_pil and pil_to_tensor are imported from utils.image class UltimateSDUpscaleDistributed: @@ -1738,354 +1685,6 @@ class UltimateSDUpscaleDistributed: # Ensure initialization before registering routes ensure_tile_jobs_initialized() -# API Endpoint for tile completion -@server.PromptServer.instance.routes.post("/distributed/request_image") -async def request_image_endpoint(request): - """Endpoint for workers to request images in dynamic mode.""" - try: - data = await request.json() - worker_id = data.get('worker_id') - multi_job_id = data.get('multi_job_id') - - if not worker_id or not multi_job_id: - return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'dynamic': - return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) - if 'pending_images' in job_data: - pending_queue = job_data['pending_images'] - try: - image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) - # Track assigned image - if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: - job_data['assigned_to_workers'][worker_id].append(image_idx) - # Update worker heartbeat - if 'worker_status' in job_data: - job_data['worker_status'][worker_id] = time.time() - # Get estimated remaining count - remaining = pending_queue.qsize() # Approximate - debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") - return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) - except asyncio.TimeoutError: - return web.json_response({"image_idx": None}) - else: - return await handle_api_error(request, "Invalid dynamic mode configuration", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.get("/distributed/job_status") -async def job_status_endpoint(request): - """Endpoint to check if a job is ready.""" - multi_job_id = request.query.get('multi_job_id') - if not multi_job_id: - return web.json_response({"ready": False}) - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) - return web.json_response({"ready": ready}) - -@server.PromptServer.instance.routes.post("/distributed/heartbeat") -async def heartbeat_endpoint(request): - """Endpoint for workers to send heartbeat signals.""" - try: - data = await request.json() - worker_id = data.get('worker_id') - multi_job_id = data.get('multi_job_id') - - if not worker_id or not multi_job_id: - return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if 'worker_status' in job_data: - job_data['worker_status'][worker_id] = time.time() - debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}") - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Worker status tracking not available", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.post("/distributed/tile_complete") -async def tile_complete_endpoint(request): - """Endpoint for receiving completed tiles from workers.""" - try: - # Check payload size before processing - content_length = request.headers.get('content-length') - if content_length and int(content_length) > MAX_PAYLOAD_SIZE: - return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413) - - data = await request.post() - multi_job_id = data.get('multi_job_id') - worker_id = data.get('worker_id') - is_last = data.get('is_last', 'False').lower() == 'true' - - if multi_job_id is None or worker_id is None: - return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - - # Check if this is a full image submission (dynamic mode) - if 'full_image' in data and 'image_idx' in data: - image_idx = int(data.get('image_idx')) - is_last = data.get('is_last', 'False').lower() == 'true' - - try: - # Process full image - img_data = data['full_image'].file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - - debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}") - - # Put into dynamic mode queue - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'dynamic': - return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) - # Check for duplicate image submission - if 'completed_images' in job_data and image_idx in job_data['completed_images']: - log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring") - return web.json_response({"status": "duplicate"}) - if 'queue' in job_data: - await job_data['queue'].put({ - 'worker_id': worker_id, - 'image_idx': image_idx, - 'image': img, - 'is_last': is_last - }) - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Invalid dynamic mode configuration", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - log(f"Error processing full image from worker {worker_id}: {e}") - return await handle_api_error(request, f"Image processing error: {e}", 400) - - # Check for batch mode - batch_size = int(data.get('batch_size', 0)) - tiles = [] - - # Handle completion signals (batch_size=0 with is_last=True) - if batch_size == 0 and is_last: - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if isinstance(job_data, dict) and 'queue' in job_data: - await job_data['queue'].put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Invalid job structure", 500) - else: - return await handle_api_error(request, "Job not found", 404) - - if batch_size > 0: - # Batch mode: Extract multiple tiles - padding = int(data.get('padding', 32)) - debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") - - # Check for JSON metadata - metadata_field = data.get('tiles_metadata') - if metadata_field: - # New JSON metadata format - try: - # Handle different types of metadata field - if hasattr(metadata_field, 'file'): - # File-like object - metadata_str = metadata_field.file.read().decode('utf-8') - elif isinstance(metadata_field, (bytes, bytearray)): - # Direct bytes/bytearray - metadata_str = metadata_field.decode('utf-8') - else: - # String - metadata_str = str(metadata_field) - - metadata = json.loads(metadata_str) - if len(metadata) != batch_size: - return await handle_api_error(request, "Metadata length mismatch", 400) - - tile_data_list = [] # Temporary list to collect and sort tiles - - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - log(f"Missing tile_{i} from worker {worker_id}, skipping") - continue - - try: - # Process image - img_data = tile_field.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] - - # Get metadata for this tile - if i < len(metadata): - tile_data = metadata[i] - tile_idx = tile_data.get('tile_idx', i) - - tile_info = { - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - 'padding': padding, - 'batch_idx': tile_data.get('batch_idx', 0), - 'global_idx': tile_data.get('global_idx', tile_idx) - } - tile_data_list.append(tile_info) - else: - log(f"Missing metadata for tile {i} from worker {worker_id}, skipping") - continue - - except Exception as e: - log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") - continue - - # Sort tiles by tile_idx to ensure correct order - tile_data_list.sort(key=lambda x: x['tile_idx']) - tiles.extend(tile_data_list) - - if tile_data_list: - debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx") - except Exception as e: - log(f"Error processing JSON metadata from worker {worker_id}: {e}") - return await handle_api_error(request, f"Metadata processing error: {e}", 400) - else: - # Legacy format: individual fields per tile (backward compatibility) - debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") - - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - log(f"Missing tile_{i} from worker {worker_id}, skipping") - continue - - # Get metadata for this tile - tile_idx = int(data.get(f'tile_{i}_idx', i)) - x = int(data.get(f'tile_{i}_x', 0)) - y = int(data.get(f'tile_{i}_y', 0)) - extracted_width = int(data.get(f'tile_{i}_width', 512)) - extracted_height = int(data.get(f'tile_{i}_height', 512)) - - try: - # Process image - img_data = tile_field.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] - - tiles.append({ - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': x, - 'y': y, - 'extracted_width': extracted_width, - 'extracted_height': extracted_height, - 'padding': padding - }) - except Exception as e: - log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") - continue - else: - # Single tile mode (backward compatibility) - image_file = data.get('image') - if not image_file: - return await handle_api_error(request, "Missing image data", 400) - - tile_idx = int(data.get('tile_idx', 0)) - x = int(data.get('x', 0)) - y = int(data.get('y', 0)) - extracted_width = int(data.get('extracted_width', 512)) - extracted_height = int(data.get('extracted_height', 512)) - padding = int(data.get('padding', 32)) - - debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}") - - try: - # Process image - img_data = image_file.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] - - tiles = [{ - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': x, - 'y': y, - 'extracted_width': extracted_width, - 'extracted_height': extracted_height, - 'padding': padding, - 'batch_idx': 0, # Default for legacy single-tile - 'global_idx': tile_idx # Default for legacy single-tile - }] - except Exception as e: - log(f"Error processing tile from worker {worker_id}: {e}") - return await handle_api_error(request, f"Tile processing error: {e}", 400) - - # Put tiles into queue - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'static': - return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400) - - q = job_data['queue'] - if batch_size > 0: - # Put batch as single item - await q.put({ - 'worker_id': worker_id, - 'tiles': tiles, - 'is_last': is_last - }) - else: - # Put single tile (backward compat) - tile_data = tiles[0] - await q.put({ - 'tensor': tile_data['tensor'], - 'worker_id': worker_id, - 'tile_idx': tile_data['tile_idx'], - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - 'padding': tile_data['padding'], - 'is_last': is_last - }) - - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Job not found or already complete", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - - # Node registration NODE_CLASS_MAPPINGS = { "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed, diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py index a783f0a..23693fd 100644 --- a/utils/usdu_utils.py +++ b/utils/usdu_utils.py @@ -4,12 +4,27 @@ import torch import torch.nn.functional as F from torchvision.transforms import GaussianBlur import math +import copy +import asyncio +import json +import io +import time +import os +import server +from aiohttp import web + +# Import other utilities from the same package +from .logging import debug_log, log +from .network import handle_api_error if (not hasattr(Image, 'Resampling')): # For older versions of Pillow Image.Resampling = Image BLUR_KERNEL_SIZE = 15 +# Configure maximum payload size (50MB default, configurable via environment variable) +MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) + def tensor_to_pil(img_tensor, batch_index=0): # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels] @@ -453,3 +468,399 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0) crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) cropped.append(n) return cropped + + +# Helper functions for shallow copying conditioning without duplicating models +def clone_control_chain(control, clone_hint=True): + """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" + if control is None: + return None + new_control = copy.copy(control) # Shallow copy (shares model) + if clone_hint and hasattr(control, 'cond_hint_original'): + new_control.cond_hint_original = control.cond_hint_original.clone() + if hasattr(control, 'previous_controlnet'): + new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) + return new_control + +def clone_conditioning(cond_list, clone_hints=True): + """Clone conditioning without duplicating ControlNet models.""" + new_cond = [] + for emb, cond_dict in cond_list: + new_emb = emb.clone() if emb is not None else None + new_dict = cond_dict.copy() + if 'control' in new_dict: + new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) + if 'mask' in new_dict: + new_dict['mask'] = new_dict['mask'].clone() + # Handle other potential fields if needed + if 'pooled_output' in new_dict: + new_dict['pooled_output'] = new_dict['pooled_output'].clone() + if 'area' in new_dict: + new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple + new_cond.append([new_emb, new_dict]) + return new_cond + +# Helper function to ensure persistent state is initialized +def ensure_tile_jobs_initialized(): + """Ensure tile job storage is initialized on the server instance.""" + prompt_server = server.PromptServer.instance + if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): + debug_log("Initializing persistent tile job queue on server instance.") + prompt_server.distributed_pending_tile_jobs = {} + prompt_server.distributed_tile_jobs_lock = asyncio.Lock() + else: + # Clean up any legacy queue structures that don't have the 'mode' field + # (Should be rare after fixes, but keep for safety) + to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() + if not isinstance(job_data, dict) or 'mode' not in job_data] + for job_id in to_remove: + debug_log(f"Removing legacy queue structure for job {job_id}") + del prompt_server.distributed_pending_tile_jobs[job_id] + return prompt_server + +# API Endpoint for tile completion +@server.PromptServer.instance.routes.post("/distributed/request_image") +async def request_image_endpoint(request): + """Endpoint for workers to request images in dynamic mode.""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'dynamic': + return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) + if 'pending_images' in job_data: + pending_queue = job_data['pending_images'] + try: + image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + # Track assigned image + if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: + job_data['assigned_to_workers'][worker_id].append(image_idx) + # Update worker heartbeat + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + # Get estimated remaining count + remaining = pending_queue.qsize() # Approximate + debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") + return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) + except asyncio.TimeoutError: + return web.json_response({"image_idx": None}) + else: + return await handle_api_error(request, "Invalid dynamic mode configuration", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.get("/distributed/job_status") +async def job_status_endpoint(request): + """Endpoint to check if a job is ready.""" + multi_job_id = request.query.get('multi_job_id') + if not multi_job_id: + return web.json_response({"ready": False}) + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) + return web.json_response({"ready": ready}) + +@server.PromptServer.instance.routes.post("/distributed/heartbeat") +async def heartbeat_endpoint(request): + """Endpoint for workers to send heartbeat signals.""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Worker status tracking not available", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.post("/distributed/tile_complete") +async def tile_complete_endpoint(request): + """Endpoint for receiving completed tiles from workers.""" + try: + # Check payload size before processing + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + # Check if this is a full image submission (dynamic mode) + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + is_last = data.get('is_last', 'False').lower() == 'true' + + try: + # Process full image + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}") + + # Put into dynamic mode queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'dynamic': + return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) + # Check for duplicate image submission + if 'completed_images' in job_data and image_idx in job_data['completed_images']: + log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring") + return web.json_response({"status": "duplicate"}) + if 'queue' in job_data: + await job_data['queue'].put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Invalid dynamic mode configuration", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + log(f"Error processing full image from worker {worker_id}: {e}") + return await handle_api_error(request, f"Image processing error: {e}", 400) + + # Check for batch mode + batch_size = int(data.get('batch_size', 0)) + tiles = [] + + # Handle completion signals (batch_size=0 with is_last=True) + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if isinstance(job_data, dict) and 'queue' in job_data: + await job_data['queue'].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Invalid job structure", 500) + else: + return await handle_api_error(request, "Job not found", 404) + + if batch_size > 0: + # Batch mode: Extract multiple tiles + padding = int(data.get('padding', 32)) + debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") + + # Check for JSON metadata + metadata_field = data.get('tiles_metadata') + if metadata_field: + # New JSON metadata format + try: + # Handle different types of metadata field + if hasattr(metadata_field, 'file'): + # File-like object + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + # Direct bytes/bytearray + metadata_str = metadata_field.decode('utf-8') + else: + # String + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + + tile_data_list = [] # Temporary list to collect and sort tiles + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + log(f"Missing tile_{i} from worker {worker_id}, skipping") + continue + + try: + # Process image + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + # Get metadata for this tile + if i < len(metadata): + tile_data = metadata[i] + tile_idx = tile_data.get('tile_idx', i) + + tile_info = { + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': padding, + 'batch_idx': tile_data.get('batch_idx', 0), + 'global_idx': tile_data.get('global_idx', tile_idx) + } + tile_data_list.append(tile_info) + else: + log(f"Missing metadata for tile {i} from worker {worker_id}, skipping") + continue + + except Exception as e: + log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") + continue + + # Sort tiles by tile_idx to ensure correct order + tile_data_list.sort(key=lambda x: x['tile_idx']) + tiles.extend(tile_data_list) + + if tile_data_list: + debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx") + except Exception as e: + log(f"Error processing JSON metadata from worker {worker_id}: {e}") + return await handle_api_error(request, f"Metadata processing error: {e}", 400) + else: + # Legacy format: individual fields per tile (backward compatibility) + debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + log(f"Missing tile_{i} from worker {worker_id}, skipping") + continue + + # Get metadata for this tile + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + try: + # Process image + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + except Exception as e: + log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") + continue + else: + # Single tile mode (backward compatibility) + image_file = data.get('image') + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + tile_idx = int(data.get('tile_idx', 0)) + x = int(data.get('x', 0)) + y = int(data.get('y', 0)) + extracted_width = int(data.get('extracted_width', 512)) + extracted_height = int(data.get('extracted_height', 512)) + padding = int(data.get('padding', 32)) + + debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}") + + try: + # Process image + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles = [{ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding, + 'batch_idx': 0, # Default for legacy single-tile + 'global_idx': tile_idx # Default for legacy single-tile + }] + except Exception as e: + log(f"Error processing tile from worker {worker_id}: {e}") + return await handle_api_error(request, f"Tile processing error: {e}", 400) + + # Put tiles into queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'static': + return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400) + + q = job_data['queue'] + if batch_size > 0: + # Put batch as single item + await q.put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last + }) + else: + # Put single tile (backward compat) + tile_data = tiles[0] + await q.put({ + 'tensor': tile_data['tensor'], + 'worker_id': worker_id, + 'tile_idx': tile_data['tile_idx'], + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': tile_data['padding'], + 'is_last': is_last + }) + + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Job not found or already complete", 404) + except Exception as e: + return await handle_api_error(request, e, 500) \ No newline at end of file From f0b4078f0004d88821ec9538d74268968d9b99b7 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Mon, 21 Jul 2025 12:58:01 +1000 Subject: [PATCH 42/93] upscaler: dynamic/static refactoring --- distributed.py | 73 +++ distributed_upscale.py | 1274 ++++++++++++++++++++++++--------------- utils/constants.py | 6 +- utils/usdu_managment.py | 803 ++++++++++++++++++++++++ utils/usdu_utils.py | 412 +------------ 5 files changed, 1656 insertions(+), 912 deletions(-) create mode 100644 utils/usdu_managment.py diff --git a/distributed.py b/distributed.py index 347fcc2..c1d5c55 100644 --- a/distributed.py +++ b/distributed.py @@ -611,6 +611,79 @@ async def get_managed_workers_endpoint(request): return await handle_api_error(request, e, 500) +@server.PromptServer.instance.routes.get("/distributed/local-worker-status") +async def get_local_worker_status_endpoint(request): + """Check status of all local workers (localhost/no host specified).""" + try: + config = load_config() + worker_statuses = {} + + for worker in config.get("workers", []): + # Only check local workers + if not worker.get("host") or worker.get("host") in ["localhost", "127.0.0.1"]: + worker_id = worker["id"] + port = worker["port"] + + # Check if worker is enabled + if not worker.get("enabled", False): + worker_statuses[worker_id] = { + "online": False, + "enabled": False, + "processing": False, + "queue_count": 0 + } + continue + + # Try to connect to worker + try: + session = await get_client_session() + async with session.get( + f"http://localhost:{port}/prompt", + timeout=aiohttp.ClientTimeout(total=2) + ) as resp: + if resp.status == 200: + data = await resp.json() + queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0) + worker_statuses[worker_id] = { + "online": True, + "enabled": True, + "processing": queue_remaining > 0, + "queue_count": queue_remaining + } + else: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": f"HTTP {resp.status}" + } + except asyncio.TimeoutError: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": "Timeout" + } + except Exception as e: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": str(e) + } + + return web.json_response({ + "status": "success", + "worker_statuses": worker_statuses + }) + except Exception as e: + debug_log(f"Error checking local worker status: {e}") + return await handle_api_error(request, e, 500) + + @server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}") async def get_worker_log_endpoint(request): """Get log content for a specific worker.""" diff --git a/distributed_upscale.py b/distributed_upscale.py index bb22970..32e6546 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -12,9 +12,11 @@ import math import os import time from typing import List, Tuple +from functools import wraps # Import ComfyUI modules import comfy.samplers +import comfy.model_management # Import shared utilities from .utils.logging import debug_log, log @@ -24,18 +26,38 @@ from .utils.network import get_server_port, get_server_loop, get_client_session, from .utils.async_helpers import run_async_in_server_loop from .utils.constants import ( TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_TRANSFER_TIMEOUT, - QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT + QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH, HEARTBEAT_TIMEOUT ) # Import for controller support -from .utils.usdu_utils import crop_cond, clone_control_chain, clone_conditioning, ensure_tile_jobs_initialized +from .utils.usdu_utils import crop_cond +from .utils.usdu_managment import ( + clone_conditioning, ensure_tile_jobs_initialized, + # Job management functions + _init_job_queue, _distribute_tasks, _get_next_task, _drain_results_queue, + _check_and_requeue_timed_out_workers, _get_completed_count, _mark_task_completed, + _send_heartbeat_to_master, _cleanup_job, + # Constants + JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS +) -# Make MAX_BATCH configurable -MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) +# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants +# They can be overridden via environment variables: +# - COMFYUI_MAX_BATCH (default: 20) +# - COMFYUI_HEARTBEAT_TIMEOUT (default: 120) -# Heartbeat timeout in seconds (configurable via env var, default 120s to cover typical image processing) -HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '120')) +# Sync wrapper decorator for async methods +def sync_wrapper(async_func): + """Decorator to wrap async methods for synchronous execution.""" + @wraps(async_func) + def sync_func(self, *args, **kwargs): + # Use run_async_in_server_loop for ComfyUI compatibility + return run_async_in_server_loop( + async_func(self, *args, **kwargs), + timeout=600.0 # 10 minute timeout for long operations + ) + return sync_func # Note: tensor_to_pil and pil_to_tensor are imported from utils.image @@ -123,11 +145,11 @@ class UltimateSDUpscaleDistributed: if is_worker: # Worker mode: process tiles synchronously - return self.process_worker_tiles(upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_worker_ids, dynamic_threshold) + return self.process_worker(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold) else: # Master mode: distribute and collect synchronously return self.process_master(upscaled_image, model, positive, negative, vae, @@ -135,12 +157,12 @@ class UltimateSDUpscaleDistributed: tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold) - def process_worker_tiles(self, upscaled_image, model, positive, negative, vae, - seed, steps, cfg, sampler_name, scheduler, denoise, - tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_worker_ids, dynamic_threshold): - """Worker processing - determines mode and processes accordingly.""" + def process_worker(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_worker_ids, dynamic_threshold): + """Unified worker processing - handles both static and dynamic modes.""" # Get batch size to determine mode batch_size = upscaled_image.shape[0] @@ -160,105 +182,18 @@ class UltimateSDUpscaleDistributed: force_uniform_tiles, tiled_decode, multi_job_id, master_url, worker_id, enabled_worker_ids, dynamic_threshold) - # Static mode - continue with existing tile-based processing - # Round tile dimensions - tile_width = self.round_to_multiple(tile_width) - tile_height = self.round_to_multiple(tile_height) - - # Calculate all tile positions - _, height, width, _ = upscaled_image.shape - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - - # Calculate total tiles across all images - num_tiles_per_image = len(all_tiles) - total_tiles = batch_size * num_tiles_per_image - - # Get assigned global tile indices for this worker - assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) - if not assigned_global_indices: - return (upscaled_image,) - - debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(assigned_global_indices)} tiles from batch of {batch_size}") - - # Process tiles SYNCHRONOUSLY - processed_tiles = [] - # Track which batch indices we've already pre-sliced conditioning for - sliced_conditioning_cache = {} - - for global_idx in assigned_global_indices: - # Calculate which image and tile this corresponds to - batch_idx = global_idx // num_tiles_per_image - tile_idx = global_idx % num_tiles_per_image - - # Skip if batch_idx is out of range - if batch_idx >= upscaled_image.shape[0]: - debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") - continue - - # Pre-slice conditioning once per image (cache it) - if batch_idx not in sliced_conditioning_cache: - positive_sliced = clone_conditioning(positive) - negative_sliced = clone_conditioning(negative) - for cond_list in [positive_sliced, negative_sliced]: - for i in range(len(cond_list)): - emb, cond_dict = cond_list[i] - if emb.shape[0] > 1: - cond_list[i][0] = emb[batch_idx:batch_idx+1] - if 'control' in cond_dict: - control = cond_dict['control'] - while control is not None: - hint = control.cond_hint_original - if hint.shape[0] > 1: - control.cond_hint_original = hint[batch_idx:batch_idx+1] - control = control.previous_controlnet - if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: - cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] - sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) - else: - positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] - - x, y = all_tiles[tile_idx] - - # Extract tile from the specific image in the batch - tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( - upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding - ) - - # Process tile through SD with unique seed - image_seed = seed + batch_idx * 1000 - processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, - image_seed + tile_idx, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode, batch_idx=batch_idx, - region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) - - processed_tiles.append({ - 'tile': processed_tile, - 'global_idx': global_idx, - 'batch_idx': batch_idx, - 'tile_idx': tile_idx, - 'x': x1, # Use extraction position - 'y': y1, - 'extracted_width': ew, - 'extracted_height': eh - }) - - # Send results back using the server's event loop in a single batch - try: - run_async_in_server_loop( - self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, - padding, worker_id), - timeout=TILE_SEND_TIMEOUT - ) - except Exception as e: - log(f"UltimateSDUpscale Worker - Error sending tiles: {e}") - - return (upscaled_image,) + # Static mode - enhanced with health monitoring and retry logic + return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_workers) def process_master(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold): - """Master processing - SYNCHRONOUS with async collection.""" + """Unified master processing with enhanced monitoring and failure handling.""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) @@ -268,9 +203,9 @@ class UltimateSDUpscaleDistributed: # Calculate all tiles all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) - total_tiles = len(all_tiles) + num_tiles_per_image = len(all_tiles) - debug_log(f"UltimateSDUpscale Master - Total tiles: {total_tiles}") + debug_log(f"UltimateSDUpscale Master - Tiles per image: {num_tiles_per_image}") # Parse enabled workers enabled_workers = json.loads(enabled_worker_ids) @@ -294,157 +229,12 @@ class UltimateSDUpscaleDistributed: tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers) - # Static mode - continue with existing implementation - - # Initialize job queue FIRST before any processing - # Initialize the queue directly in the persistent server storage - # This ensures it's accessible from the web server's event loop - try: - run_async_in_server_loop( - self._init_job_queue(multi_job_id), - timeout=QUEUE_INIT_TIMEOUT - ) - debug_log(f"UltimateSDUpscale Master - Queue initialization complete") - except Exception as e: - debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}") - raise RuntimeError(f"Failed to initialize static mode queue: {e}") - - # Convert batch to PIL list for processing - result_images = [] - for b in range(batch_size): - image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) - result_images.append(image_pil.copy()) - - # Mask will be created per tile for proper blending - - # Calculate total tiles across all images - num_tiles_per_image = len(all_tiles) - total_tiles_all_images = batch_size * num_tiles_per_image - - # Calculate master's global tile indices - master_global_indices = self._get_master_global_indices(total_tiles_all_images, num_workers) - master_tiles = len(master_global_indices) - - debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally from batch of {batch_size}") - - # Process master tiles SYNCHRONOUSLY - # Track which batch indices we've already pre-sliced conditioning for - sliced_conditioning_cache = {} - - for global_idx in master_global_indices: - # Calculate which image and tile this corresponds to - batch_idx = global_idx // num_tiles_per_image - tile_idx = global_idx % num_tiles_per_image - - # Skip if batch_idx is out of range - if batch_idx >= len(result_images): - debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds result_images length {len(result_images)}") - continue - - # Pre-slice conditioning once per image (cache it) - if batch_idx not in sliced_conditioning_cache: - positive_sliced = clone_conditioning(positive) - negative_sliced = clone_conditioning(negative) - for cond_list in [positive_sliced, negative_sliced]: - for i in range(len(cond_list)): - emb, cond_dict = cond_list[i] - if emb.shape[0] > 1: - cond_list[i][0] = emb[batch_idx:batch_idx+1] - if 'control' in cond_dict: - control = cond_dict['control'] - while control is not None: - hint = control.cond_hint_original - if hint.shape[0] > 1: - control.cond_hint_original = hint[batch_idx:batch_idx+1] - control = control.previous_controlnet - if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: - cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] - sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) - else: - positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] - - # Use unique seed per image - image_seed = seed + batch_idx * 1000 - - result_images[batch_idx] = self._process_and_blend_tile( - tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], - model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx - ) - - # Collect worker tiles using async - worker_tiles_expected = total_tiles_all_images - master_tiles - if worker_tiles_expected > 0: - # Calculate tile distribution for active worker calculation - tiles_per_participant = total_tiles_all_images // (num_workers + 1) - remainder = total_tiles_all_images % (num_workers + 1) - - # Calculate which workers will actually have tiles - active_workers = [] - for i in range(num_workers): - worker_start_idx = master_tiles + (i * tiles_per_participant) - if i < remainder - 1: - worker_start_idx += i - worker_tile_count = tiles_per_participant + 1 - else: - worker_start_idx += remainder - 1 if remainder > 0 else 0 - worker_tile_count = tiles_per_participant - - if worker_start_idx < total_tiles_all_images and worker_tile_count > 0: - active_workers.append(enabled_workers[i]) - - debug_log(f"UltimateSDUpscale Master - Waiting for {worker_tiles_expected} tiles from {len(active_workers)} active workers") - debug_log(f"UltimateSDUpscale Master - Active worker IDs: {active_workers}") - debug_log(f"UltimateSDUpscale Master - All enabled worker IDs: {enabled_workers}") - - # Collect tiles from workers using the server's event loop - collected_tiles = run_async_in_server_loop( - self._async_collect_worker_tiles(multi_job_id, len(active_workers)), - timeout=TILE_COLLECTION_TIMEOUT - ) - - # Blend collected tiles SYNCHRONOUSLY - for tile_data in collected_tiles.values(): - x = tile_data['x'] - y = tile_data['y'] - tile_tensor = tile_data['tensor'] - tile_idx = tile_data['tile_idx'] - extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) - extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) - - # Get batch index if available (for batch processing) - batch_idx = tile_data.get('batch_idx', 0) - - # Skip if batch_idx is out of range - if batch_idx >= len(result_images): - debug_log(f"Warning: Received tile for batch_idx {batch_idx} but only have {len(result_images)} images") - continue - - # Convert and blend - tile_pil = tensor_to_pil(tile_tensor, 0) - # Get the original tile position from the tile index - orig_x, orig_y = all_tiles[tile_idx] - # Create mask for this specific tile - tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) - # Use extraction position and size for blending - result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, - x, y, (extracted_width, extracted_height), tile_mask, padding) - - # Convert back to tensor - # For batch processing, combine all result images - if batch_size == 1: - result_tensor = pil_to_tensor(result_images[0]) - else: - # Convert all images and concatenate - result_tensors = [pil_to_tensor(img) for img in result_images] - result_tensor = torch.cat(result_tensors, dim=0) - - if upscaled_image.is_cuda: - result_tensor = result_tensor.cuda() - - debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") - return (result_tensor,) + # Static mode - enhanced with unified job management + return self._process_master_static_sync(upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, + all_tiles, num_tiles_per_image) def _get_worker_global_indices(self, total_tiles, enabled_workers, worker_id): """Calculate which global tile indices are assigned to a specific worker in flattened mode.""" @@ -459,13 +249,13 @@ class UltimateSDUpscaleDistributed: return [] global_indices = list(range(total_tiles)) - all_assignments = self._distribute_items(global_indices, num_workers + 1) + all_assignments = _distribute_tasks(global_indices, num_workers + 1) return all_assignments[worker_index + 1] # +1 because 0 is master def _get_master_global_indices(self, total_tiles, num_workers): """Calculate which global tile indices are assigned to the master in flattened mode.""" global_indices = list(range(total_tiles)) - all_assignments = self._distribute_items(global_indices, num_workers + 1) + all_assignments = _distribute_tasks(global_indices, num_workers + 1) return all_assignments[0] def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image, @@ -495,43 +285,62 @@ class UltimateSDUpscaleDistributed: return result_image + async def _init_job_queue_unified(self, multi_job_id, mode='static', batch_size=None, + total_tiles=None, enabled_workers=None, num_tiles_per_image=None, + assigned_to_workers=None, all_indices=None, task_assignments=None): + """Unified job queue initialization for both static and dynamic modes.""" + if mode == 'dynamic': + # Dynamic mode initialization + await _init_job_queue(multi_job_id, 'dynamic', batch_size=batch_size, + all_indices=all_indices, + enabled_workers=list(assigned_to_workers.keys()) if assigned_to_workers else []) + # Add additional fields for backward compatibility + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data['completed_images'] = {} + job_data['completed_tiles'] = {} # For uniformity + # Fix: The endpoint expects 'pending_images', not 'pending_tasks' + job_data['pending_images'] = job_data[JOB_PENDING_TASKS] + debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {batch_size} pending images for all participants") + else: + # Static mode initialization + # If num_tiles_per_image not provided, calculate it + if num_tiles_per_image is None and total_tiles is not None and batch_size is not None: + num_tiles_per_image = total_tiles // batch_size if batch_size > 0 else total_tiles + + await _init_job_queue(multi_job_id, 'static', + batch_size=batch_size, + num_tiles_per_image=num_tiles_per_image, + all_indices=None, + enabled_workers=enabled_workers, + task_assignments=task_assignments) + + # Add completed_tiles and total_items for backward compatibility + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + job_data['completed_tiles'] = {} + if total_tiles is not None: + job_data['total_items'] = total_tiles + debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") + + # Keep compatibility wrappers async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None): """Initialize the job queue for dynamic mode with pending images.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = { - 'queue': asyncio.Queue(), - 'mode': 'dynamic', - 'pending_images': asyncio.Queue(), - 'completed_images': {}, - 'completed_tiles': {}, # For uniformity - 'batch_size': batch_size, - 'assigned_to_workers': assigned_to_workers or {}, - 'worker_status': worker_status or {} - } - # Initialize pending images queue with all indices - pending_images = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] - for i in all_indices: - await pending_images.put(i) - debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {len(all_indices)} pending images for all participants") - else: - debug_log(f"UltimateSDUpscale Master Dynamic - Queue already exists for {multi_job_id}") + await self._init_job_queue_unified(multi_job_id, mode='dynamic', batch_size=batch_size, + assigned_to_workers=assigned_to_workers, all_indices=all_indices) - async def _init_job_queue(self, multi_job_id): + async def _init_job_queue(self, multi_job_id, total_tiles, enabled_workers, batch_size=1, num_tiles_per_image=None, task_assignments=None): """Initialize the job queue for collecting tiles.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - prompt_server.distributed_pending_tile_jobs[multi_job_id] = { - 'queue': asyncio.Queue(), - 'mode': 'static', - 'completed_tiles': {} # For static mode data - } - debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") + await self._init_job_queue_unified(multi_job_id, mode='static', batch_size=batch_size, + total_tiles=total_tiles, enabled_workers=enabled_workers, + num_tiles_per_image=num_tiles_per_image, task_assignments=task_assignments) - async def _async_collect_worker_tiles(self, multi_job_id, num_workers): - """Async helper to collect tiles from workers.""" + async def _async_collect_results(self, multi_job_id, num_workers, mode='static', + remaining_to_collect=None, batch_size=None): + """Unified async helper to collect results from workers (tiles or images).""" # Get the already initialized queue prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: @@ -540,78 +349,133 @@ class UltimateSDUpscaleDistributed: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: raise RuntimeError("Invalid job data structure") - if job_data['mode'] != 'static': - raise RuntimeError(f"Mode mismatch: expected static, got {job_data['mode']}") + if job_data['mode'] != mode: + raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}") q = job_data['queue'] + + # For dynamic mode, get reference to completed_images + if mode == 'dynamic': + completed_images = job_data['completed_images'] + # Calculate expected count for logging + expected_count = remaining_to_collect or batch_size + else: + # Calculate expected count from job data for static mode + expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set())) - collected_tiles = {} + mode_str = "Dynamic" if mode == 'dynamic' else "Static" + item_type = "images" if mode == 'dynamic' else "tiles" + debug_log(f"UltimateSDUpscale Master {mode_str} - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers") + + collected_results = {} workers_done = set() - timeout = TILE_WAIT_TIMEOUT - - debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {num_workers} workers to complete") + timeout = 120.0 if mode == 'dynamic' else 60.0 # Longer timeout for full images + last_heartbeat_check = time.time() + collected_count = 0 while len(workers_done) < num_workers: + # Check for user interruption + if comfy.model_management.processing_interrupted(): + log("Processing interrupted by user") + raise comfy.model_management.InterruptProcessingException() + + # For dynamic mode with remaining_to_collect, check if we've collected enough + if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect: + break + try: - result = await asyncio.wait_for(q.get(), timeout=timeout) + # Shorter timeout for regular checks in dynamic mode + wait_timeout = 10.0 if mode == 'dynamic' else timeout + result = await asyncio.wait_for(q.get(), timeout=wait_timeout) worker_id = result['worker_id'] is_last = result.get('is_last', False) - # Check if batch mode - tiles = result.get('tiles', []) - if tiles: - # Batch mode - debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") - - for tile_data in tiles: - # Validate required fields - if 'batch_idx' not in tile_data: - log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") - continue - batch_idx = tile_data['batch_idx'] + if mode == 'static': + # Handle tiles + tiles = result.get('tiles', []) + if tiles: + # Batch mode + debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})") - tile_idx = tile_data['tile_idx'] - # Use global_idx as key if available (for batch processing) - key = tile_data.get('global_idx', tile_idx) - - # Store the full tile data including metadata - collected_tiles[key] = { - 'tensor': tile_data['tensor'], - 'tile_idx': tile_idx, - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - 'padding': tile_data['padding'], - 'worker_id': worker_id, - 'batch_idx': tile_data.get('batch_idx', 0), - 'global_idx': tile_data.get('global_idx', tile_idx) - } - else: - # Single tile mode (backward compat) - tile_idx = result['tile_idx'] - collected_tiles[tile_idx] = result - debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + for tile_data in tiles: + # Validate required fields + if 'batch_idx' not in tile_data: + log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping") + continue + + tile_idx = tile_data['tile_idx'] + # Use global_idx as key if available (for batch processing) + key = tile_data.get('global_idx', tile_idx) + + # Store the full tile data including metadata + collected_results[key] = { + 'tensor': tile_data['tensor'], + 'tile_idx': tile_idx, + 'x': tile_data['x'], + 'y': tile_data['y'], + 'extracted_width': tile_data['extracted_width'], + 'extracted_height': tile_data['extracted_height'], + 'padding': tile_data['padding'], + 'worker_id': worker_id, + 'batch_idx': tile_data.get('batch_idx', 0), + 'global_idx': tile_data.get('global_idx', tile_idx) + } + else: + # Single tile mode (backward compat) + tile_idx = result['tile_idx'] + collected_results[tile_idx] = result + debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})") + + elif mode == 'dynamic': + # Handle full images + if 'image_idx' in result and 'image' in result: + image_idx = result['image_idx'] + image_pil = result['image'] + completed_images[image_idx] = image_pil + collected_results[image_idx] = image_pil + collected_count += 1 + debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}") if is_last: workers_done.add(worker_id) - debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed") - + debug_log(f"UltimateSDUpscale Master {mode_str} - Worker {worker_id} completed") + except asyncio.TimeoutError: - log(f"UltimateSDUpscale Master - Timeout waiting for tiles") - break + if mode == 'dynamic': + # Check for worker timeouts periodically + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + # Use the class method to check and requeue + requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size) + if requeued > 0: + log(f"UltimateSDUpscale Master Dynamic - Requeued {requeued} images from timed out workers") + last_heartbeat_check = current_time + + # Check if we've been waiting too long overall + if current_time - last_heartbeat_check > timeout: + log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") + break + else: + log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}") + break - debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_tiles)} tiles from {len(workers_done)} workers") + debug_log(f"UltimateSDUpscale Master {mode_str} - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers") # Clean up job queue - prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: if multi_job_id in prompt_server.distributed_pending_tile_jobs: del prompt_server.distributed_pending_tile_jobs[multi_job_id] - return collected_tiles + return collected_results if mode == 'static' else completed_images + + # Keep compatibility wrappers for existing code + async def _async_collect_worker_tiles(self, multi_job_id, num_workers): + """Async helper to collect tiles from workers.""" + return await self._async_collect_results(multi_job_id, num_workers, mode='static') async def _mark_image_completed(self, multi_job_id, image_idx, image_pil): """Mark an image as completed in the job data.""" + # Mark the image as completed with the image data + await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil}) prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) @@ -620,87 +484,9 @@ class UltimateSDUpscaleDistributed: async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): """Collect remaining processed images from workers.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_tile_jobs: - raise RuntimeError(f"Job queue not initialized for {multi_job_id}") - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - raise RuntimeError("Invalid job data structure") - if job_data['mode'] != 'dynamic': - raise RuntimeError(f"Mode mismatch: expected dynamic, got {job_data['mode']}") - q = job_data['queue'] - completed_images = job_data['completed_images'] - - workers_done = set() - collected_count = 0 - timeout = TILE_WAIT_TIMEOUT * 2 # Longer timeout for full images - last_heartbeat_check = time.time() - - debug_log(f"UltimateSDUpscale Master Dynamic - Waiting for {remaining_to_collect} more images from {num_workers} workers") - - while collected_count < remaining_to_collect and len(workers_done) < num_workers: - try: - result = await asyncio.wait_for(q.get(), timeout=10.0) # Shorter timeout for regular checks - worker_id = result['worker_id'] - is_last = result.get('is_last', False) - - if 'image_idx' in result and 'image' in result: - image_idx = result['image_idx'] - image_pil = result['image'] - completed_images[image_idx] = image_pil - collected_count += 1 - debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}") - - if is_last: - workers_done.add(worker_id) - debug_log(f"UltimateSDUpscale Master Dynamic - Worker {worker_id} completed") - - except asyncio.TimeoutError: - # Check for worker timeouts - current_time = time.time() - if current_time - last_heartbeat_check >= 10.0: - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data_check = prompt_server.distributed_pending_tile_jobs[multi_job_id] - # Check for timed out workers (60 second timeout) - for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): - if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: - log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") - # Requeue assigned images from this worker - for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): - if idx not in completed_images: - await job_data_check['pending_images'].put(idx) - debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx} from timed out worker {worker}") - # Mark worker as done and clean up - workers_done.add(worker) - if 'worker_status' in job_data_check: - del job_data_check['worker_status'][worker] - if 'assigned_to_workers' in job_data_check: - job_data_check['assigned_to_workers'][worker] = [] - last_heartbeat_check = current_time - - # Check if we've been waiting too long overall - if current_time - last_heartbeat_check > timeout: - log(f"UltimateSDUpscale Master Dynamic - Overall timeout waiting for images") - # Requeue all unfinished images - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - pending_queue = prompt_server.distributed_pending_tile_jobs[multi_job_id]['pending_images'] - for idx in range(batch_size): - if idx not in completed_images: - await pending_queue.put(idx) - debug_log(f"UltimateSDUpscale Master Dynamic - Requeued image {idx}") - break - - debug_log(f"UltimateSDUpscale Master Dynamic - Collection complete. Got {collected_count} images from workers") - - # Clean up job queue - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - del prompt_server.distributed_pending_tile_jobs[multi_job_id] - - return completed_images + return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic', + remaining_to_collect=remaining_to_collect, + batch_size=batch_size) def round_to_multiple(self, value: int, multiple: int = 8) -> int: """Round value to nearest multiple.""" @@ -833,9 +619,15 @@ class UltimateSDUpscaleDistributed: tile_pil = tensor_to_pil(tile_tensor, 0) clean_tensor = pil_to_tensor(tile_pil) + # Ensure tensor is detached and doesn't require gradients + clean_tensor = clean_tensor.detach() + if hasattr(clean_tensor, 'requires_grad_'): + clean_tensor.requires_grad_(False) + # Move to correct device if tile_tensor.is_cuda: clean_tensor = clean_tensor.cuda() + clean_tensor = clean_tensor.detach() # Detach again after device transfer # Clone conditioning per tile (shares models, clones hints for cropping) positive_tile = clone_conditioning(positive, clone_hints=True) @@ -932,6 +724,551 @@ class UltimateSDUpscaleDistributed: # Convert back to RGB return result.convert('RGB') + def _slice_conditioning(self, positive, negative, batch_idx): + """Helper to slice conditioning for a specific batch index.""" + # Clone and slice conditioning properly, including ControlNet hints + positive_sliced = clone_conditioning(positive) + negative_sliced = clone_conditioning(negative) + + for cond_list in [positive_sliced, negative_sliced]: + for i in range(len(cond_list)): + emb, cond_dict = cond_list[i] + if emb.shape[0] > 1: + cond_list[i][0] = emb[batch_idx:batch_idx+1] + if 'control' in cond_dict: + control = cond_dict['control'] + while control is not None: + hint = control.cond_hint_original + if hint.shape[0] > 1: + control.cond_hint_original = hint[batch_idx:batch_idx+1] + control = control.previous_controlnet + if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: + cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1] + + return positive_sliced, negative_sliced + + async def _get_all_completed_tasks(self, multi_job_id): + """Helper to retrieve all completed tasks from the job data.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and JOB_COMPLETED_TASKS in job_data: + return dict(job_data[JOB_COMPLETED_TASKS]) # Return a copy + return {} + + async def _process_worker_static_async(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_workers): + """Enhanced static mode worker with health monitoring and retry logic.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Calculate all tile positions + batch_size = upscaled_image.shape[0] + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + # Calculate total tiles across all images + num_tiles_per_image = len(all_tiles) + total_tiles = batch_size * num_tiles_per_image + + # Get initial assigned global tile indices for this worker + assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) + + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} initially assigned {len(assigned_global_indices)} tiles") + + # Process initially assigned tiles + processed_tiles = [] + sliced_conditioning_cache = {} + + # Process assigned tiles with heartbeat after each + for global_idx in assigned_global_indices: + # Process the tile + processed_tile = await self._process_single_tile( + global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat after each tile + try: + await _send_heartbeat_to_master(multi_job_id, master_url, worker_id) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send initial batch of processed tiles + if processed_tiles: + await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id) + processed_tiles = [] + + # Check for requeued tiles from failed workers + debug_log(f"Worker {worker_id} checking for requeued tiles...") + while True: + # Request next task from pending queue + task_id = await _get_next_task(multi_job_id) + if task_id is None: + break + + debug_log(f"Worker {worker_id} processing requeued tile {task_id}") + + # Process the requeued tile + processed_tile = await self._process_single_tile( + task_id, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + await _send_heartbeat_to_master(multi_job_id, master_url, worker_id) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id) + processed_tiles = [] + + # Send any remaining tiles + if processed_tiles: + await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id) + + debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") + return (upscaled_image,) + + def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, master_url, + worker_id, enabled_workers): + """Worker static mode processing with sync tensor operations and async communication.""" + # Round tile dimensions + tile_width = self.round_to_multiple(tile_width) + tile_height = self.round_to_multiple(tile_height) + + # Get dimensions and calculate tiles + _, height, width, _ = upscaled_image.shape + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + num_tiles_per_image = len(all_tiles) + + # Calculate worker assignments + num_workers = len(enabled_workers) + batch_size = upscaled_image.shape[0] + total_tiles = batch_size * num_tiles_per_image + worker_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) + + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(worker_global_indices)} tiles") + + # Process assigned tiles synchronously + processed_tiles = [] + sliced_conditioning_cache = {} + + for global_idx in worker_global_indices: + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Process tile synchronously + processed_tile = self._process_single_tile( + global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat (async operation) + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Send initial batch of processed tiles + if processed_tiles: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Check for requeued tiles from failed workers + debug_log(f"Worker {worker_id} checking for requeued tiles...") + while True: + # Request next task from pending queue (async operation) + task_id = run_async_in_server_loop( + _get_next_task(multi_job_id), + timeout=5.0 + ) + if task_id is None: + break + + debug_log(f"Worker {worker_id} processing requeued tile {task_id}") + + # Process the requeued tile synchronously + batch_idx = task_id // num_tiles_per_image + tile_idx = task_id % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + processed_tile = self._process_single_tile( + task_id, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Send any remaining tiles + if processed_tiles: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + + debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") + return (upscaled_image,) + + async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total): + """Async helper for collection and monitoring in static mode. + Returns collected tasks dict. Caller should check if all tasks are complete.""" + last_progress_log = time.time() + progress_interval = 5.0 + last_heartbeat_check = time.time() + stall_count = 0 + max_stall_count = 10 # Allow 10 iterations with no progress before returning + last_completed_count = 0 + + while True: + # Check for user interruption + if comfy.model_management.processing_interrupted(): + log("Processing interrupted by user") + raise comfy.model_management.InterruptProcessingException() + + # Drain any pending results + collected_count = await _drain_results_queue(multi_job_id) + + # Check and requeue timed-out workers periodically + current_time = time.time() + if current_time - last_heartbeat_check >= 10.0: + requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total) + if requeued_count > 0: + log(f"Requeued {requeued_count} tasks from timed-out workers") + stall_count = 0 # Reset stall count when we requeue + last_heartbeat_check = current_time + + # Get current completion count + completed_count = await _get_completed_count(multi_job_id) + + # Progress logging + if current_time - last_progress_log >= progress_interval: + log(f"Progress: {completed_count}/{expected_total} tasks completed") + last_progress_log = current_time + + # Check if all tasks are completed + if completed_count >= expected_total: + debug_log(f"All {expected_total} tasks completed") + break + + # Check if we're making progress + if completed_count > last_completed_count: + stall_count = 0 # Reset stall count on progress + last_completed_count = completed_count + else: + stall_count += 1 + + # If no progress for too long, check if we need to process locally + if stall_count >= max_stall_count: + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data: + pending_queue = job_data.get(JOB_PENDING_TASKS) + # If there are pending tasks but no active workers, or we're stalled + if pending_queue and not pending_queue.empty(): + log(f"Collection stalled with {expected_total - completed_count} tasks remaining. Returning for local processing.") + break + + # Wait a bit before next check + await asyncio.sleep(0.5) + + # Get all completed tasks for return + return await self._get_all_completed_tasks(multi_job_id) + + def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, + tile_width, tile_height, padding, mask_blur, + force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, + all_tiles, num_tiles_per_image): + """Static mode master processing with sync tensor operations and async coordination.""" + batch_size = upscaled_image.shape[0] + _, height, width, _ = upscaled_image.shape + total_tiles = batch_size * num_tiles_per_image + + # Distribute tasks across master and workers + all_indices = list(range(total_tiles)) + num_participants = len(enabled_workers) + 1 + task_assignments = _distribute_tasks(all_indices, num_participants) + master_indices = task_assignments[0] + + # Initialize job queue with task assignments (async operation) + run_async_in_server_loop( + self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, task_assignments), + timeout=10.0 + ) + + debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles") + + # Convert batch to PIL list for processing + result_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + result_images.append(image_pil.copy()) + + debug_log(f"UltimateSDUpscale Master - Processing {len(master_indices)} tiles locally") + + # Process master's tiles synchronously + sliced_conditioning_cache = {} + for global_idx in master_indices: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously - no async context for tensor operations + result_images[batch_idx] = self._process_and_blend_tile( + tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed (async operation) + run_async_in_server_loop( + _mark_task_completed(multi_job_id, global_idx, {'batch_idx': batch_idx, 'tile_idx': tile_idx}), + timeout=5.0 + ) + + # Collect worker results using async operations + try: + collected_tasks = run_async_in_server_loop( + self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), + timeout=300.0 # 5 minutes to allow for retries and requeuing + ) + except comfy.model_management.InterruptProcessingException: + # Clean up job on interruption + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) + raise + + # Check if we need to process any remaining tasks locally + completed_count = len(collected_tasks) + if completed_count < total_tiles: + log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") + + # Process any pending tasks synchronously + while True: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + + # Get next task from pending queue + task_id = run_async_in_server_loop( + _get_next_task(multi_job_id), + timeout=5.0 + ) + + if task_id is None: + break + + batch_idx = task_id // num_tiles_per_image + tile_idx = task_id % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously + result_images[batch_idx] = self._process_and_blend_tile( + tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed and store in collected_tasks + run_async_in_server_loop( + _mark_task_completed(multi_job_id, task_id, {'batch_idx': batch_idx, 'tile_idx': tile_idx}), + timeout=5.0 + ) + # Add to collected_tasks so we skip it in the blending loop + collected_tasks[task_id] = {'batch_idx': batch_idx, 'tile_idx': tile_idx} + + # Blend worker tiles synchronously + for global_idx, tile_data in collected_tasks.items(): + if global_idx in master_indices: + continue # Skip master's tiles + + if 'tensor' in tile_data: + batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) + tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) + + if batch_idx >= batch_size: + continue + + # Blend tile synchronously + x = tile_data.get('x', 0) + y = tile_data.get('y', 0) + tile_tensor = tile_data['tensor'] + tile_pil = tensor_to_pil(tile_tensor, 0) + orig_x, orig_y = all_tiles[tile_idx] + tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) + extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) + extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) + result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, + x, y, (extracted_width, extracted_height), tile_mask, padding) + + try: + # Convert back to tensor + if batch_size == 1: + result_tensor = pil_to_tensor(result_images[0]) + else: + result_tensors = [pil_to_tensor(img) for img in result_images] + result_tensor = torch.cat(result_tensors, dim=0) + + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") + return (result_tensor,) + finally: + # Cleanup (async operation) - always execute + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) + + def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive, negative, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache): + """Process a single tile.""" + # Calculate which image and tile this corresponds to + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image + + # Skip if batch_idx is out of range + if batch_idx >= upscaled_image.shape[0]: + debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}") + return None + + # Get or create sliced conditioning for this batch index + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + x, y = all_tiles[tile_idx] + + # Extract tile from the specific image in the batch + tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( + upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding + ) + + # Process tile through SD with unique seed + image_seed = seed + batch_idx * 1000 + processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, + image_seed + tile_idx, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, batch_idx=batch_idx, + region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height)) + + return { + 'tile': processed_tile, + 'global_idx': global_idx, + 'batch_idx': batch_idx, + 'tile_idx': tile_idx, + 'x': x1, + 'y': y1, + 'extracted_width': ew, + 'extracted_height': eh + } async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, padding, worker_id): @@ -993,7 +1330,7 @@ class UltimateSDUpscaleDistributed: for attempt in range(max_retries): try: session = await get_client_session() - url = f"{master_url}/distributed/tile_complete" + url = f"{master_url}/distributed/submit_tiles" async with session.post(url, data=data) as response: response.raise_for_status() @@ -1297,17 +1634,6 @@ class UltimateSDUpscaleDistributed: except asyncio.TimeoutError: return None - async def _get_completed_count(self, multi_job_id): - """Get count of completed images from workers.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if job_data and 'completed_images' in job_data: - # Count only worker-completed images (not master's) - completed = job_data['completed_images'] - worker_count = len([idx for idx in completed if isinstance(completed[idx], Image.Image)]) - return worker_count - return 0 async def _get_total_completed_count(self, multi_job_id): """Get total count of all completed images (master + workers).""" @@ -1374,33 +1700,8 @@ class UltimateSDUpscaleDistributed: async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): """Check for timed out workers and requeue their assigned images. Returns count of requeued images.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if not job_data: - return 0 - - current_time = time.time() - requeued_count = 0 - completed_images = job_data.get('completed_images', {}) - - # Check for timed out workers (60 second timeout) - for worker, last_heartbeat in list(job_data.get('worker_status', {}).items()): - if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: - log(f"Worker {worker} timed out") - # Requeue assigned images from this worker - for idx in job_data.get('assigned_to_workers', {}).get(worker, []): - if idx not in completed_images: - await job_data['pending_images'].put(idx) - requeued_count += 1 - debug_log(f"Requeued image {idx} from timed out worker {worker}") - # Clean up worker tracking - if 'worker_status' in job_data: - del job_data['worker_status'][worker] - if 'assigned_to_workers' in job_data: - job_data['assigned_to_workers'][worker] = [] - - return requeued_count + # Use the original function from usdu_managment.py + return await _check_and_requeue_timed_out_workers(multi_job_id, batch_size) def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae, @@ -1490,7 +1791,7 @@ class UltimateSDUpscaleDistributed: padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) run_async_in_server_loop( - self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), timeout=5.0 ) @@ -1505,7 +1806,7 @@ class UltimateSDUpscaleDistributed: ) # Send heartbeat after processing run_async_in_server_loop( - self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), timeout=5.0 ) if is_last: @@ -1611,7 +1912,7 @@ class UltimateSDUpscaleDistributed: for attempt in range(max_retries): try: session = await get_client_session() - url = f"{master_url}/distributed/tile_complete" + url = f"{master_url}/distributed/submit_image" async with session.post(url, data=data) as response: response.raise_for_status() @@ -1627,17 +1928,7 @@ class UltimateSDUpscaleDistributed: log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}") raise - async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id): - """Send heartbeat signal to master.""" - try: - data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)} - session = await get_client_session() - url = f"{master_url}/distributed/heartbeat" - - async with session.post(url, json=data) as response: - response.raise_for_status() - except Exception as e: - pass + # Note: Using _send_heartbeat_to_master from usdu_managment.py instead of duplicate async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): """Send completion signal to master in dynamic mode.""" @@ -1646,11 +1937,10 @@ class UltimateSDUpscaleDistributed: data.add_field('multi_job_id', multi_job_id) data.add_field('worker_id', str(worker_id)) data.add_field('is_last', 'true') - # Send empty tile data to signal completion - data.add_field('batch_size', '0') + # No image data - just completion signal session = await get_client_session() - url = f"{master_url}/distributed/tile_complete" + url = f"{master_url}/distributed/submit_image" async with session.post(url, data=data) as response: response.raise_for_status() @@ -1664,23 +1954,7 @@ class UltimateSDUpscaleDistributed: return "static" return "dynamic" - def _distribute_items(self, items: list, num_participants: int) -> List[List[any]]: - """Distributes a list of items among N participants.""" - if num_participants == 0: - return [items] # All items for the single participant (master) - - items_per_participant = len(items) // num_participants - remainder = len(items) % num_participants - - assignments = [] - start_idx = 0 - for i in range(num_participants): - count = items_per_participant + (1 if i < remainder else 0) - end_idx = start_idx + count - assignments.append(items[start_idx:end_idx]) - start_idx = end_idx - - return assignments + # Note: Using _distribute_tasks from usdu_managment.py instead of duplicate _distribute_items # Ensure initialization before registering routes ensure_tile_jobs_initialized() diff --git a/utils/constants.py b/utils/constants.py index 2ce53c4..a4ca5bd 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -1,6 +1,7 @@ """ Shared constants for ComfyUI-Distributed. """ +import os # Timeouts (in seconds) WORKER_JOB_TIMEOUT = 30.0 @@ -34,4 +35,7 @@ TILE_SEND_TIMEOUT = 60.0 MEMORY_CLEAR_DELAY = 0.5 # Batch processing -MAX_BATCH = 20 # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs) \ No newline at end of file +MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs) + +# Heartbeat monitoring +HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds \ No newline at end of file diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py new file mode 100644 index 0000000..a7cfb38 --- /dev/null +++ b/utils/usdu_managment.py @@ -0,0 +1,803 @@ +import asyncio +import time +import json +import copy +import os +import io +from aiohttp import web +import server +from PIL import Image +import numpy as np +import torch + +# Import from other utilities +from .logging import debug_log, log +from .network import handle_api_error, get_client_session +from .image import tensor_to_pil + +# Configure maximum payload size (50MB default, configurable via environment variable) +MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) + +# Import HEARTBEAT_TIMEOUT from constants +from .constants import HEARTBEAT_TIMEOUT + + +# Unified Job Data Structure Keys +JOB_QUEUE = 'queue' +JOB_MODE = 'mode' +JOB_COMPLETED_TASKS = 'completed_tasks' +JOB_WORKER_STATUS = 'worker_status' +JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers' +JOB_PENDING_TASKS = 'pending_tasks' +JOB_BATCH_SIZE = 'batch_size' # For dynamic +JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static + +# Task Types +TASK_TYPE_TILE = 'tile' +TASK_TYPE_IMAGE = 'image' + +async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None): + """Unified initialization for job queues in static and dynamic modes.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + debug_log(f"Queue already exists for {multi_job_id}") + return + + job_data = { + JOB_QUEUE: asyncio.Queue(), + JOB_MODE: mode, + JOB_COMPLETED_TASKS: {}, + JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []}, + JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []}, + JOB_PENDING_TASKS: asyncio.Queue(), + } + + if mode == 'dynamic': + job_data[JOB_BATCH_SIZE] = batch_size + pending_queue = job_data[JOB_PENDING_TASKS] + for i in all_indices or range(batch_size): + await pending_queue.put(i) + debug_log(f"Initialized dynamic queue with {batch_size} pending images") + elif mode == 'static': + job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image + # For static, pending_tasks starts empty; requeues will be added later + debug_log(f"Initialized static queue for {batch_size * num_tiles_per_image} tiles") + + # If task assignments provided, populate JOB_ASSIGNED_TO_WORKERS + if task_assignments and enabled_workers: + # task_assignments[0] is for master, 1+ are for workers + for i, worker_id in enumerate(enabled_workers): + if i + 1 < len(task_assignments): + job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1] + debug_log(f"Worker {worker_id} assigned {len(task_assignments[i + 1])} tasks") + else: + raise ValueError(f"Unknown mode: {mode}") + + prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data + +def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]: + """Distribute a list of items among N participants (master + workers).""" + if num_participants <= 1: + return [items] + + items_per_participant = len(items) // num_participants + remainder = len(items) % num_participants + assignments = [] + start_idx = 0 + for i in range(num_participants): + count = items_per_participant + (1 if i < remainder else 0) + end_idx = start_idx + count + assignments.append(items[start_idx:end_idx]) + start_idx = end_idx + return assignments + +async def _get_next_task(multi_job_id): + """Get next task from pending queue (generalized for tiles/images).""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or JOB_PENDING_TASKS not in job_data: + return None + try: + task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0) + return task_id + except asyncio.TimeoutError: + return None + +async def _drain_results_queue(multi_job_id): + """Drain pending results from queue and update completed_tasks. Returns count drained.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or JOB_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data: + return 0 + q = job_data[JOB_QUEUE] + completed_tasks = job_data[JOB_COMPLETED_TASKS] + + collected = 0 + while not q.empty(): + try: + result = await asyncio.wait_for(q.get(), timeout=0.1) + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + task_id = result['image_idx'] + if task_id not in completed_tasks: + completed_tasks[task_id] = result['image'] + collected += 1 + elif 'tiles' in result: + for tile_data in result['tiles']: + task_id = tile_data.get('global_idx', tile_data['tile_idx']) + if task_id not in completed_tasks: + completed_tasks[task_id] = tile_data + collected += 1 + elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat + task_id = result.get('global_idx', result['tile_idx']) + if task_id not in completed_tasks: + completed_tasks[task_id] = { + 'tensor': result['tensor'], + 'tile_idx': result['tile_idx'], + 'x': result['x'], + 'y': result['y'], + 'extracted_width': result['extracted_width'], + 'extracted_height': result['extracted_height'], + 'padding': result['padding'], + 'batch_idx': result.get('batch_idx', 0), + 'global_idx': task_id + } + collected += 1 + + if is_last: + # Track worker completion + if worker_id in job_data[JOB_WORKER_STATUS]: + del job_data[JOB_WORKER_STATUS][worker_id] + except asyncio.TimeoutError: + break + + return collected + +async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): + """Check timed out workers and requeue their tasks. Returns requeued count.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data: + return 0 + + current_time = time.time() + requeued_count = 0 + completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {}) + + for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()): + if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: + log(f"Worker {worker} timed out") + for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []): + if task_id not in completed_tasks: + await job_data[JOB_PENDING_TASKS].put(task_id) + requeued_count += 1 + if JOB_WORKER_STATUS in job_data: + del job_data[JOB_WORKER_STATUS][worker] + if JOB_ASSIGNED_TO_WORKERS in job_data: + job_data[JOB_ASSIGNED_TO_WORKERS][worker] = [] + + return requeued_count + +async def _get_completed_count(multi_job_id): + """Get count of completed tasks.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and JOB_COMPLETED_TASKS in job_data: + return len(job_data[JOB_COMPLETED_TASKS]) + return 0 + +async def _mark_task_completed(multi_job_id, task_id, result): + """Mark a task as completed.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if job_data and JOB_COMPLETED_TASKS in job_data: + job_data[JOB_COMPLETED_TASKS][task_id] = result + +async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id): + """Send heartbeat to master.""" + try: + data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)} + session = await get_client_session() + url = f"{master_url}/distributed/heartbeat" + async with session.post(url, json=data) as response: + response.raise_for_status() + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + +async def _cleanup_job(multi_job_id): + """Cleanup the job data.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + del prompt_server.distributed_pending_tile_jobs[multi_job_id] + debug_log(f"Cleaned up job {multi_job_id}") + +# API Endpoints (generalized) + +@server.PromptServer.instance.routes.post("/distributed/request_task") +async def request_task_endpoint(request): + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + mode = job_data.get(JOB_MODE) + pending_queue = job_data.get(JOB_PENDING_TASKS) + if pending_queue: + try: + task_id = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + if JOB_ASSIGNED_TO_WORKERS in job_data and worker_id in job_data[JOB_ASSIGNED_TO_WORKERS]: + job_data[JOB_ASSIGNED_TO_WORKERS][worker_id].append(task_id) + if JOB_WORKER_STATUS in job_data: + job_data[JOB_WORKER_STATUS][worker_id] = time.time() + remaining = pending_queue.qsize() + debug_log(f"Assigned task {task_id} to worker {worker_id} in {mode} mode") + return web.json_response({"task_id": task_id, "estimated_remaining": remaining, "mode": mode}) + except asyncio.TimeoutError: + return web.json_response({"task_id": None}) + else: + return await handle_api_error(request, "No pending tasks", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.post("/distributed/heartbeat") +async def heartbeat_endpoint(request): + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_WORKER_STATUS in job_data: + job_data[JOB_WORKER_STATUS][worker_id] = time.time() + debug_log(f"Heartbeat from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Worker status tracking not available", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.post("/distributed/submit_tiles") +async def submit_tiles_endpoint(request): + """Endpoint for workers to submit processed tiles in static mode.""" + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + batch_size = int(data.get('batch_size', 0)) + tiles = [] + + # Handle completion signal + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + return await handle_api_error(request, "Mode mismatch: expected static mode", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + + # Handle batch tiles with metadata + if batch_size > 0: + padding = int(data.get('padding', 32)) + metadata_field = data.get('tiles_metadata') + if metadata_field: + if hasattr(metadata_field, 'file'): + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + metadata_str = metadata_field.decode('utf-8') + else: + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + + tile_data_list = [] + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + if i < len(metadata): + tile_meta = metadata[i] + tile_idx = tile_meta.get('tile_idx', i) + tile_info = { + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': tile_meta['x'], + 'y': tile_meta['y'], + 'extracted_width': tile_meta['extracted_width'], + 'extracted_height': tile_meta['extracted_height'], + 'padding': padding, + 'batch_idx': tile_meta.get('batch_idx', 0), + 'global_idx': tile_meta.get('global_idx', tile_idx) + } + tile_data_list.append(tile_info) + + tile_data_list.sort(key=lambda x: x['tile_idx']) + tiles.extend(tile_data_list) + else: + # Legacy format + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + else: + # Single tile legacy + image_file = data.get('image') + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + tile_idx = int(data.get('tile_idx', 0)) + x = int(data.get('x', 0)) + y = int(data.get('y', 0)) + extracted_width = int(data.get('extracted_width', 512)) + extracted_height = int(data.get('extracted_height', 512)) + padding = int(data.get('padding', 32)) + + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles = [{ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding, + 'batch_idx': 0, + 'global_idx': tile_idx + }] + + # Submit tiles to queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + return await handle_api_error(request, "Mode mismatch: expected static mode", 400) + + q = job_data[JOB_QUEUE] + if batch_size > 0 or len(tiles) > 0: + await q.put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last + }) + else: + await q.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.post("/distributed/submit_image") +async def submit_image_endpoint(request): + """Endpoint for workers to submit processed images in dynamic mode.""" + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + # Handle image submission + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"Received full image {image_idx} from worker {worker_id}") + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) + + # Handle completion signal (no image data) + elif is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] # For compatibility + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Missing image data or invalid request", 400) + + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +# Keep legacy endpoint for backward compatibility +@server.PromptServer.instance.routes.post("/distributed/tile_complete") +async def tile_complete_endpoint(request): + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"Received full image {image_idx} from worker {worker_id}") + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic': + return await handle_api_error(request, "Mode mismatch for image submission", 400) + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last + }) + return web.json_response({"status": "success"}) + + batch_size = int(data.get('batch_size', 0)) + tiles = [] + + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_QUEUE in job_data: + await job_data[JOB_QUEUE].put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + + if batch_size > 0: + padding = int(data.get('padding', 32)) + metadata_field = data.get('tiles_metadata') + if metadata_field: + if hasattr(metadata_field, 'file'): + metadata_str = metadata_field.file.read().decode('utf-8') + elif isinstance(metadata_field, (bytes, bytearray)): + metadata_str = metadata_field.decode('utf-8') + else: + metadata_str = str(metadata_field) + + metadata = json.loads(metadata_str) + if len(metadata) != batch_size: + return await handle_api_error(request, "Metadata length mismatch", 400) + + tile_data_list = [] + + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + if i < len(metadata): + tile_meta = metadata[i] + tile_idx = tile_meta.get('tile_idx', i) + tile_info = { + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': tile_meta['x'], + 'y': tile_meta['y'], + 'extracted_width': tile_meta['extracted_width'], + 'extracted_height': tile_meta['extracted_height'], + 'padding': padding, + 'batch_idx': tile_meta.get('batch_idx', 0), + 'global_idx': tile_meta.get('global_idx', tile_idx) + } + tile_data_list.append(tile_info) + + tile_data_list.sort(key=lambda x: x['tile_idx']) + tiles.extend(tile_data_list) + else: + # Legacy + for i in range(batch_size): + tile_field = data.get(f'tile_{i}') + if tile_field is None: + continue + + tile_idx = int(data.get(f'tile_{i}_idx', i)) + x = int(data.get(f'tile_{i}_x', 0)) + y = int(data.get(f'tile_{i}_y', 0)) + extracted_width = int(data.get(f'tile_{i}_width', 512)) + extracted_height = int(data.get(f'tile_{i}_height', 512)) + + img_data = tile_field.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles.append({ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding + }) + else: + # Single tile legacy + image_file = data.get('image') + if not image_file: + return await handle_api_error(request, "Missing image data", 400) + + tile_idx = int(data.get('tile_idx', 0)) + x = int(data.get('x', 0)) + y = int(data.get('y', 0)) + extracted_width = int(data.get('extracted_width', 512)) + extracted_height = int(data.get('extracted_height', 512)) + padding = int(data.get('padding', 32)) + + img_data = image_file.file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + img_np = np.array(img).astype(np.float32) / 255.0 + tensor = torch.from_numpy(img_np)[None,] + + tiles = [{ + 'tensor': tensor, + 'tile_idx': tile_idx, + 'x': x, + 'y': y, + 'extracted_width': extracted_width, + 'extracted_height': extracted_height, + 'padding': padding, + 'batch_idx': 0, + 'global_idx': tile_idx + }] + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if JOB_MODE in job_data and job_data[JOB_MODE] != 'static': + return await handle_api_error(request, "Mode mismatch for tile submission", 400) + + q = job_data[JOB_QUEUE] + if batch_size > 0 or len(tiles) > 0: + await q.put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last + }) + else: + await q.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [] + }) + + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + + +# Helper functions for shallow copying conditioning without duplicating models +def clone_control_chain(control, clone_hint=True): + """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" + if control is None: + return None + new_control = copy.copy(control) # Shallow copy (shares model) + if clone_hint and hasattr(control, 'cond_hint_original'): + new_control.cond_hint_original = control.cond_hint_original.clone() + if hasattr(control, 'previous_controlnet'): + new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) + return new_control + +def clone_conditioning(cond_list, clone_hints=True): + """Clone conditioning without duplicating ControlNet models.""" + new_cond = [] + for emb, cond_dict in cond_list: + new_emb = emb.clone() if emb is not None else None + new_dict = cond_dict.copy() + if 'control' in new_dict: + new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) + if 'mask' in new_dict: + new_dict['mask'] = new_dict['mask'].clone() + # Handle other potential fields if needed + if 'pooled_output' in new_dict: + new_dict['pooled_output'] = new_dict['pooled_output'].clone() + if 'area' in new_dict: + new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple + new_cond.append([new_emb, new_dict]) + return new_cond + +# Helper function to ensure persistent state is initialized +def ensure_tile_jobs_initialized(): + """Ensure tile job storage is initialized on the server instance.""" + prompt_server = server.PromptServer.instance + if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): + debug_log("Initializing persistent tile job queue on server instance.") + prompt_server.distributed_pending_tile_jobs = {} + prompt_server.distributed_tile_jobs_lock = asyncio.Lock() + else: + # Clean up any legacy queue structures that don't have the 'mode' field + # (Should be rare after fixes, but keep for safety) + to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() + if not isinstance(job_data, dict) or 'mode' not in job_data] + for job_id in to_remove: + debug_log(f"Removing legacy queue structure for job {job_id}") + del prompt_server.distributed_pending_tile_jobs[job_id] + return prompt_server + +# API Endpoint for tile completion +@server.PromptServer.instance.routes.post("/distributed/request_image") +async def request_image_endpoint(request): + """Endpoint for workers to request images in dynamic mode.""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, dict) or 'mode' not in job_data: + return await handle_api_error(request, "Invalid job data structure", 500) + if job_data['mode'] != 'dynamic': + return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) + if 'pending_images' in job_data: + pending_queue = job_data['pending_images'] + try: + image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + # Track assigned image + if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: + job_data['assigned_to_workers'][worker_id].append(image_idx) + # Update worker heartbeat + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + # Get estimated remaining count + remaining = pending_queue.qsize() # Approximate + debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") + return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) + except asyncio.TimeoutError: + return web.json_response({"image_idx": None}) + else: + return await handle_api_error(request, "Invalid dynamic mode configuration", 400) + else: + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.get("/distributed/job_status") +async def job_status_endpoint(request): + """Endpoint to check if a job is ready.""" + multi_job_id = request.query.get('multi_job_id') + if not multi_job_id: + return web.json_response({"ready": False}) + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) + return web.json_response({"ready": ready}) + + diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py index 23693fd..21214a7 100644 --- a/utils/usdu_utils.py +++ b/utils/usdu_utils.py @@ -4,27 +4,13 @@ import torch import torch.nn.functional as F from torchvision.transforms import GaussianBlur import math -import copy -import asyncio -import json -import io -import time import os -import server -from aiohttp import web - -# Import other utilities from the same package -from .logging import debug_log, log -from .network import handle_api_error if (not hasattr(Image, 'Resampling')): # For older versions of Pillow Image.Resampling = Image BLUR_KERNEL_SIZE = 15 -# Configure maximum payload size (50MB default, configurable via environment variable) -MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) - def tensor_to_pil(img_tensor, batch_index=0): # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels] @@ -467,400 +453,4 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0) crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) cropped.append(n) - return cropped - - -# Helper functions for shallow copying conditioning without duplicating models -def clone_control_chain(control, clone_hint=True): - """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" - if control is None: - return None - new_control = copy.copy(control) # Shallow copy (shares model) - if clone_hint and hasattr(control, 'cond_hint_original'): - new_control.cond_hint_original = control.cond_hint_original.clone() - if hasattr(control, 'previous_controlnet'): - new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) - return new_control - -def clone_conditioning(cond_list, clone_hints=True): - """Clone conditioning without duplicating ControlNet models.""" - new_cond = [] - for emb, cond_dict in cond_list: - new_emb = emb.clone() if emb is not None else None - new_dict = cond_dict.copy() - if 'control' in new_dict: - new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) - if 'mask' in new_dict: - new_dict['mask'] = new_dict['mask'].clone() - # Handle other potential fields if needed - if 'pooled_output' in new_dict: - new_dict['pooled_output'] = new_dict['pooled_output'].clone() - if 'area' in new_dict: - new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple - new_cond.append([new_emb, new_dict]) - return new_cond - -# Helper function to ensure persistent state is initialized -def ensure_tile_jobs_initialized(): - """Ensure tile job storage is initialized on the server instance.""" - prompt_server = server.PromptServer.instance - if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): - debug_log("Initializing persistent tile job queue on server instance.") - prompt_server.distributed_pending_tile_jobs = {} - prompt_server.distributed_tile_jobs_lock = asyncio.Lock() - else: - # Clean up any legacy queue structures that don't have the 'mode' field - # (Should be rare after fixes, but keep for safety) - to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() - if not isinstance(job_data, dict) or 'mode' not in job_data] - for job_id in to_remove: - debug_log(f"Removing legacy queue structure for job {job_id}") - del prompt_server.distributed_pending_tile_jobs[job_id] - return prompt_server - -# API Endpoint for tile completion -@server.PromptServer.instance.routes.post("/distributed/request_image") -async def request_image_endpoint(request): - """Endpoint for workers to request images in dynamic mode.""" - try: - data = await request.json() - worker_id = data.get('worker_id') - multi_job_id = data.get('multi_job_id') - - if not worker_id or not multi_job_id: - return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'dynamic': - return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) - if 'pending_images' in job_data: - pending_queue = job_data['pending_images'] - try: - image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) - # Track assigned image - if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: - job_data['assigned_to_workers'][worker_id].append(image_idx) - # Update worker heartbeat - if 'worker_status' in job_data: - job_data['worker_status'][worker_id] = time.time() - # Get estimated remaining count - remaining = pending_queue.qsize() # Approximate - debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") - return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) - except asyncio.TimeoutError: - return web.json_response({"image_idx": None}) - else: - return await handle_api_error(request, "Invalid dynamic mode configuration", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.get("/distributed/job_status") -async def job_status_endpoint(request): - """Endpoint to check if a job is ready.""" - multi_job_id = request.query.get('multi_job_id') - if not multi_job_id: - return web.json_response({"ready": False}) - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data) - return web.json_response({"ready": ready}) - -@server.PromptServer.instance.routes.post("/distributed/heartbeat") -async def heartbeat_endpoint(request): - """Endpoint for workers to send heartbeat signals.""" - try: - data = await request.json() - worker_id = data.get('worker_id') - multi_job_id = data.get('multi_job_id') - - if not worker_id or not multi_job_id: - return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if 'worker_status' in job_data: - job_data['worker_status'][worker_id] = time.time() - debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}") - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Worker status tracking not available", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.post("/distributed/tile_complete") -async def tile_complete_endpoint(request): - """Endpoint for receiving completed tiles from workers.""" - try: - # Check payload size before processing - content_length = request.headers.get('content-length') - if content_length and int(content_length) > MAX_PAYLOAD_SIZE: - return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413) - - data = await request.post() - multi_job_id = data.get('multi_job_id') - worker_id = data.get('worker_id') - is_last = data.get('is_last', 'False').lower() == 'true' - - if multi_job_id is None or worker_id is None: - return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - - # Check if this is a full image submission (dynamic mode) - if 'full_image' in data and 'image_idx' in data: - image_idx = int(data.get('image_idx')) - is_last = data.get('is_last', 'False').lower() == 'true' - - try: - # Process full image - img_data = data['full_image'].file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - - debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}") - - # Put into dynamic mode queue - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'dynamic': - return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) - # Check for duplicate image submission - if 'completed_images' in job_data and image_idx in job_data['completed_images']: - log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring") - return web.json_response({"status": "duplicate"}) - if 'queue' in job_data: - await job_data['queue'].put({ - 'worker_id': worker_id, - 'image_idx': image_idx, - 'image': img, - 'is_last': is_last - }) - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Invalid dynamic mode configuration", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - log(f"Error processing full image from worker {worker_id}: {e}") - return await handle_api_error(request, f"Image processing error: {e}", 400) - - # Check for batch mode - batch_size = int(data.get('batch_size', 0)) - tiles = [] - - # Handle completion signals (batch_size=0 with is_last=True) - if batch_size == 0 and is_last: - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if isinstance(job_data, dict) and 'queue' in job_data: - await job_data['queue'].put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Invalid job structure", 500) - else: - return await handle_api_error(request, "Job not found", 404) - - if batch_size > 0: - # Batch mode: Extract multiple tiles - padding = int(data.get('padding', 32)) - debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}") - - # Check for JSON metadata - metadata_field = data.get('tiles_metadata') - if metadata_field: - # New JSON metadata format - try: - # Handle different types of metadata field - if hasattr(metadata_field, 'file'): - # File-like object - metadata_str = metadata_field.file.read().decode('utf-8') - elif isinstance(metadata_field, (bytes, bytearray)): - # Direct bytes/bytearray - metadata_str = metadata_field.decode('utf-8') - else: - # String - metadata_str = str(metadata_field) - - metadata = json.loads(metadata_str) - if len(metadata) != batch_size: - return await handle_api_error(request, "Metadata length mismatch", 400) - - tile_data_list = [] # Temporary list to collect and sort tiles - - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - log(f"Missing tile_{i} from worker {worker_id}, skipping") - continue - - try: - # Process image - img_data = tile_field.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] - - # Get metadata for this tile - if i < len(metadata): - tile_data = metadata[i] - tile_idx = tile_data.get('tile_idx', i) - - tile_info = { - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - 'padding': padding, - 'batch_idx': tile_data.get('batch_idx', 0), - 'global_idx': tile_data.get('global_idx', tile_idx) - } - tile_data_list.append(tile_info) - else: - log(f"Missing metadata for tile {i} from worker {worker_id}, skipping") - continue - - except Exception as e: - log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") - continue - - # Sort tiles by tile_idx to ensure correct order - tile_data_list.sort(key=lambda x: x['tile_idx']) - tiles.extend(tile_data_list) - - if tile_data_list: - debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx") - except Exception as e: - log(f"Error processing JSON metadata from worker {worker_id}: {e}") - return await handle_api_error(request, f"Metadata processing error: {e}", 400) - else: - # Legacy format: individual fields per tile (backward compatibility) - debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.") - - for i in range(batch_size): - tile_field = data.get(f'tile_{i}') - if tile_field is None: - log(f"Missing tile_{i} from worker {worker_id}, skipping") - continue - - # Get metadata for this tile - tile_idx = int(data.get(f'tile_{i}_idx', i)) - x = int(data.get(f'tile_{i}_x', 0)) - y = int(data.get(f'tile_{i}_y', 0)) - extracted_width = int(data.get(f'tile_{i}_width', 512)) - extracted_height = int(data.get(f'tile_{i}_height', 512)) - - try: - # Process image - img_data = tile_field.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] - - tiles.append({ - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': x, - 'y': y, - 'extracted_width': extracted_width, - 'extracted_height': extracted_height, - 'padding': padding - }) - except Exception as e: - log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile") - continue - else: - # Single tile mode (backward compatibility) - image_file = data.get('image') - if not image_file: - return await handle_api_error(request, "Missing image data", 400) - - tile_idx = int(data.get('tile_idx', 0)) - x = int(data.get('x', 0)) - y = int(data.get('y', 0)) - extracted_width = int(data.get('extracted_width', 512)) - extracted_height = int(data.get('extracted_height', 512)) - padding = int(data.get('padding', 32)) - - debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}") - - try: - # Process image - img_data = image_file.file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - img_np = np.array(img).astype(np.float32) / 255.0 - tensor = torch.from_numpy(img_np)[None,] - - tiles = [{ - 'tensor': tensor, - 'tile_idx': tile_idx, - 'x': x, - 'y': y, - 'extracted_width': extracted_width, - 'extracted_height': extracted_height, - 'padding': padding, - 'batch_idx': 0, # Default for legacy single-tile - 'global_idx': tile_idx # Default for legacy single-tile - }] - except Exception as e: - log(f"Error processing tile from worker {worker_id}: {e}") - return await handle_api_error(request, f"Tile processing error: {e}", 400) - - # Put tiles into queue - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, dict) or 'mode' not in job_data: - return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'static': - return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400) - - q = job_data['queue'] - if batch_size > 0: - # Put batch as single item - await q.put({ - 'worker_id': worker_id, - 'tiles': tiles, - 'is_last': is_last - }) - else: - # Put single tile (backward compat) - tile_data = tiles[0] - await q.put({ - 'tensor': tile_data['tensor'], - 'worker_id': worker_id, - 'tile_idx': tile_data['tile_idx'], - 'x': tile_data['x'], - 'y': tile_data['y'], - 'extracted_width': tile_data['extracted_width'], - 'extracted_height': tile_data['extracted_height'], - 'padding': tile_data['padding'], - 'is_last': is_last - }) - - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Job not found or already complete", 404) - except Exception as e: - return await handle_api_error(request, e, 500) \ No newline at end of file + return cropped \ No newline at end of file From a461c795318ce6cb599b6ba0d1c2ece52c4fd960 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 22 Jul 2025 15:52:12 +1000 Subject: [PATCH 43/93] static/dynamic switch --- distributed_upscale.py | 664 +++++++++++++++++++++++++++------------- utils/usdu_managment.py | 64 ++-- 2 files changed, 492 insertions(+), 236 deletions(-) diff --git a/distributed_upscale.py b/distributed_upscale.py index 32e6546..75a0eb9 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -110,6 +110,7 @@ class UltimateSDUpscaleDistributed: "mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}), "force_uniform_tiles": ("BOOLEAN", {"default": True}), "tiled_decode": ("BOOLEAN", {"default": False}), + "static_distribution": ("BOOLEAN", {"default": False, "label": "Use Static Distribution (legacy)"}), }, "hidden": { "multi_job_id": ("STRING", {"default": ""}), @@ -134,8 +135,9 @@ class UltimateSDUpscaleDistributed: def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, - mask_blur, force_uniform_tiles, tiled_decode, multi_job_id="", is_worker=False, - master_url="", enabled_worker_ids="[]", worker_id="", tile_indices="", dynamic_threshold=8): + mask_blur, force_uniform_tiles, tiled_decode, static_distribution=False, + multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", + worker_id="", tile_indices="", dynamic_threshold=8): """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler.""" if not multi_job_id: # No distributed processing, run single GPU version @@ -149,19 +151,20 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_worker_ids, dynamic_threshold) + worker_id, enabled_worker_ids, dynamic_threshold, static_distribution) else: # Master mode: distribute and collect synchronously return self.process_master(upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold) + force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, + dynamic_threshold, static_distribution) def process_worker(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_worker_ids, dynamic_threshold): + worker_id, enabled_worker_ids, dynamic_threshold, static_distribution=False): """Unified worker processing - handles both static and dynamic modes.""" # Get batch size to determine mode batch_size = upscaled_image.shape[0] @@ -187,12 +190,13 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_workers) + worker_id, enabled_workers, static_distribution) def process_master(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, - force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, dynamic_threshold): + force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, + dynamic_threshold, static_distribution=False): """Unified master processing with enhanced monitoring and failure handling.""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) @@ -234,7 +238,7 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, - all_tiles, num_tiles_per_image) + all_tiles, num_tiles_per_image, static_distribution) def _get_worker_global_indices(self, total_tiles, enabled_workers, worker_id): """Calculate which global tile indices are assigned to a specific worker in flattened mode.""" @@ -851,8 +855,8 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_workers): - """Worker static mode processing with sync tensor operations and async communication.""" + worker_id, enabled_workers, static_distribution=False): + """Worker static mode processing with optional dynamic queue pulling.""" # Round tile dimensions tile_width = self.round_to_multiple(tile_width) tile_height = self.round_to_multiple(tile_height) @@ -861,122 +865,199 @@ class UltimateSDUpscaleDistributed: _, height, width, _ = upscaled_image.shape all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) num_tiles_per_image = len(all_tiles) - - # Calculate worker assignments - num_workers = len(enabled_workers) batch_size = upscaled_image.shape[0] total_tiles = batch_size * num_tiles_per_image - worker_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) - debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(worker_global_indices)} tiles") - - # Process assigned tiles synchronously processed_tiles = [] sliced_conditioning_cache = {} - for global_idx in worker_global_indices: - batch_idx = global_idx // num_tiles_per_image - tile_idx = global_idx % num_tiles_per_image + if static_distribution: + # Legacy mode: Use pre-assigned tiles + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} using static tile assignment") + worker_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id) + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(worker_global_indices)} pre-assigned tiles") - if batch_idx >= batch_size: - continue - - # Get or create sliced conditioning - if batch_idx not in sliced_conditioning_cache: - positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) - sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) - else: - positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] - - # Process tile synchronously - processed_tile = self._process_single_tile( - global_idx, num_tiles_per_image, upscaled_image, all_tiles, - model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode, tile_width, tile_height, padding, - width, height, sliced_conditioning_cache - ) - - if processed_tile: - processed_tiles.append(processed_tile) + # Process assigned tiles synchronously + for global_idx in worker_global_indices: + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image - # Send heartbeat (async operation) - try: - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) - except Exception as e: - debug_log(f"Heartbeat failed: {e}") + if batch_idx >= batch_size: + continue - # Send tiles in batches - if len(processed_tiles) >= MAX_BATCH: - run_async_in_server_loop( - self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), - timeout=TILE_SEND_TIMEOUT - ) - processed_tiles = [] + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Process tile synchronously + processed_tile = self._process_single_tile( + global_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat (async operation) + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Send remaining tiles + if processed_tiles: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] + + # Check for requeued tiles from failed workers + debug_log(f"Worker {worker_id} checking for requeued tiles...") + while True: + # Request next task from pending queue (async operation) + task_id = run_async_in_server_loop( + _get_next_task(multi_job_id), + timeout=5.0 + ) + if task_id is None: + break + + debug_log(f"Worker {worker_id} processing requeued tile {task_id}") + + # Process the requeued tile synchronously + batch_idx = task_id // num_tiles_per_image + tile_idx = task_id % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + processed_tile = self._process_single_tile( + task_id, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] - # Send initial batch of processed tiles - if processed_tiles: - run_async_in_server_loop( - self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), - timeout=TILE_SEND_TIMEOUT - ) - processed_tiles = [] - - # Check for requeued tiles from failed workers - debug_log(f"Worker {worker_id} checking for requeued tiles...") - while True: - # Request next task from pending queue (async operation) - task_id = run_async_in_server_loop( - _get_next_task(multi_job_id), - timeout=5.0 - ) - if task_id is None: - break - - debug_log(f"Worker {worker_id} processing requeued tile {task_id}") + else: + # Dynamic queue mode + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} ready to process tiles dynamically") + processed_count = 0 - # Process the requeued tile synchronously - batch_idx = task_id // num_tiles_per_image - tile_idx = task_id % num_tiles_per_image - - if batch_idx >= batch_size: - continue - - # Get or create sliced conditioning - if batch_idx not in sliced_conditioning_cache: - positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) - sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + # Poll for job readiness + max_poll_attempts = 20 + for attempt in range(max_poll_attempts): + ready = run_async_in_server_loop( + self._check_job_status(multi_job_id, master_url), + timeout=5.0 + ) + if ready: + debug_log(f"Job {multi_job_id} ready after {attempt} attempts") + break + time.sleep(1.0) else: - positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") + return (upscaled_image,) - processed_tile = self._process_single_tile( - task_id, num_tiles_per_image, upscaled_image, all_tiles, - model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, - scheduler, denoise, tiled_decode, tile_width, tile_height, padding, - width, height, sliced_conditioning_cache - ) - - if processed_tile: - processed_tiles.append(processed_tile) + # Main processing loop - pull tiles from queue + while True: + # Request a tile to process + tile_idx, estimated_remaining = run_async_in_server_loop( + self._request_tile_from_master(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT + ) - # Send heartbeat - try: - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) - except Exception as e: - debug_log(f"Heartbeat failed: {e}") + if tile_idx is None: + debug_log(f"UltimateSDUpscale Worker - No more tiles to process") + break - # Send tiles in batches - if len(processed_tiles) >= MAX_BATCH: - run_async_in_server_loop( - self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), - timeout=TILE_SEND_TIMEOUT - ) - processed_tiles = [] + debug_log(f"UltimateSDUpscale Worker - Assigned tile {tile_idx} to worker {worker_id}") + processed_count += 1 + + # Calculate batch and local tile indices + batch_idx = tile_idx // num_tiles_per_image + local_tile_idx = tile_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Process tile synchronously + processed_tile = self._process_single_tile( + tile_idx, num_tiles_per_image, upscaled_image, all_tiles, + model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name, + scheduler, denoise, tiled_decode, tile_width, tile_height, padding, + width, height, sliced_conditioning_cache + ) + + if processed_tile: + processed_tiles.append(processed_tile) + + # Send heartbeat + try: + run_async_in_server_loop( + _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + except Exception as e: + debug_log(f"Heartbeat failed: {e}") + + # Send tiles in batches + if len(processed_tiles) >= MAX_BATCH: + run_async_in_server_loop( + self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id), + timeout=TILE_SEND_TIMEOUT + ) + processed_tiles = [] # Send any remaining tiles if processed_tiles: @@ -1058,101 +1139,48 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, - all_tiles, num_tiles_per_image): - """Static mode master processing with sync tensor operations and async coordination.""" + all_tiles, num_tiles_per_image, static_distribution=False): + """Static mode master processing with optional dynamic queue pulling.""" batch_size = upscaled_image.shape[0] _, height, width, _ = upscaled_image.shape total_tiles = batch_size * num_tiles_per_image - # Distribute tasks across master and workers - all_indices = list(range(total_tiles)) - num_participants = len(enabled_workers) + 1 - task_assignments = _distribute_tasks(all_indices, num_participants) - master_indices = task_assignments[0] - - # Initialize job queue with task assignments (async operation) - run_async_in_server_loop( - self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, task_assignments), - timeout=10.0 - ) - - debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles") - # Convert batch to PIL list for processing result_images = [] for b in range(batch_size): image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) result_images.append(image_pil.copy()) - debug_log(f"UltimateSDUpscale Master - Processing {len(master_indices)} tiles locally") - - # Process master's tiles synchronously sliced_conditioning_cache = {} - for global_idx in master_indices: - # Check for user interruption - comfy.model_management.throw_exception_if_processing_interrupted() - batch_idx = global_idx // num_tiles_per_image - tile_idx = global_idx % num_tiles_per_image + master_indices = set() # Track which tiles master processed + master_processed_count = 0 # Track how many tiles master processed + + if static_distribution: + # Legacy mode: Pre-assign tiles to master and workers + debug_log(f"UltimateSDUpscale Master - Using static tile distribution") - if batch_idx >= batch_size: - continue + # Distribute tasks across master and workers + all_indices = list(range(total_tiles)) + num_participants = len(enabled_workers) + 1 + task_assignments = _distribute_tasks(all_indices, num_participants) + master_indices = set(task_assignments[0]) # Convert to set for efficient lookup - # Get or create sliced conditioning - if batch_idx not in sliced_conditioning_cache: - positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) - sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) - else: - positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] - - # Use unique seed per image - image_seed = seed + batch_idx * 1000 - - # Process tile synchronously - no async context for tensor operations - result_images[batch_idx] = self._process_and_blend_tile( - tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], - model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx - ) - - # Mark as completed (async operation) + # Initialize job queue with task assignments (don't populate pending queue) run_async_in_server_loop( - _mark_task_completed(multi_job_id, global_idx, {'batch_idx': batch_idx, 'tile_idx': tile_idx}), - timeout=5.0 + self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, task_assignments), + timeout=10.0 ) - - # Collect worker results using async operations - try: - collected_tasks = run_async_in_server_loop( - self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), - timeout=300.0 # 5 minutes to allow for retries and requeuing - ) - except comfy.model_management.InterruptProcessingException: - # Clean up job on interruption - run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) - raise - - # Check if we need to process any remaining tasks locally - completed_count = len(collected_tasks) - if completed_count < total_tiles: - log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") - # Process any pending tasks synchronously - while True: + debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles") + debug_log(f"UltimateSDUpscale Master - Processing {len(master_indices)} tiles locally") + master_processed_count = len(master_indices) + + # Process master's tiles synchronously + for global_idx in master_indices: # Check for user interruption comfy.model_management.throw_exception_if_processing_interrupted() - - # Get next task from pending queue - task_id = run_async_in_server_loop( - _get_next_task(multi_job_id), - timeout=5.0 - ) - - if task_id is None: - break - - batch_idx = task_id // num_tiles_per_image - tile_idx = task_id % num_tiles_per_image + batch_idx = global_idx // num_tiles_per_image + tile_idx = global_idx % num_tiles_per_image if batch_idx >= batch_size: continue @@ -1167,7 +1195,7 @@ class UltimateSDUpscaleDistributed: # Use unique seed per image image_seed = seed + batch_idx * 1000 - # Process tile synchronously + # Process tile synchronously - no async context for tensor operations result_images[batch_idx] = self._process_and_blend_tile( tile_idx, all_tiles[tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, @@ -1175,37 +1203,187 @@ class UltimateSDUpscaleDistributed: padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx ) - # Mark as completed and store in collected_tasks + # Mark as completed (async operation) run_async_in_server_loop( - _mark_task_completed(multi_job_id, task_id, {'batch_idx': batch_idx, 'tile_idx': tile_idx}), + _mark_task_completed(multi_job_id, global_idx, {'batch_idx': batch_idx, 'tile_idx': tile_idx}), timeout=5.0 ) - # Add to collected_tasks so we skip it in the blending loop - collected_tasks[task_id] = {'batch_idx': batch_idx, 'tile_idx': tile_idx} + else: + # Dynamic queue mode: All tiles go to pending queue + debug_log(f"UltimateSDUpscale Master - Using dynamic tile distribution") + + # Initialize job queue with all tiles in pending queue for dynamic distribution + run_async_in_server_loop( + self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image), + timeout=10.0 + ) + + debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles with dynamic distribution") + debug_log(f"UltimateSDUpscale Master - Processing tiles dynamically from queue") + + # Process tiles by pulling from queue (similar to dynamic mode) + processed_count = 0 + consecutive_no_tile = 0 + max_consecutive_no_tile = 5 + + while processed_count < total_tiles: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + + # Try to get a tile to process + tile_idx = run_async_in_server_loop( + self._get_next_tile_index(multi_job_id), + timeout=5.0 + ) + + if tile_idx is not None: + consecutive_no_tile = 0 + processed_count += 1 + + # Calculate batch and local tile indices + batch_idx = tile_idx // num_tiles_per_image + local_tile_idx = tile_idx % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + debug_log(f"Master processing tile {tile_idx} (batch {batch_idx}, tile {local_tile_idx})") + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously + result_images[batch_idx] = self._process_and_blend_tile( + local_tile_idx, all_tiles[local_tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed + run_async_in_server_loop( + _mark_task_completed(multi_job_id, tile_idx, {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}), + timeout=5.0 + ) + else: + consecutive_no_tile += 1 + if consecutive_no_tile >= max_consecutive_no_tile: + # No more tiles available, break to start collecting + debug_log(f"Master processed {processed_count} tiles, moving to collection phase") + break + # Small wait before retry + time.sleep(0.1) + + # Set master_processed_count for the collection phase + master_processed_count = processed_count + + # Continue processing any remaining tiles while collecting worker results + remaining_tiles = total_tiles - master_processed_count + if remaining_tiles > 0: + debug_log(f"Master waiting for {remaining_tiles} tiles from workers") + + # Collect worker results using async operations + try: + collected_tasks = run_async_in_server_loop( + self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), + timeout=300.0 # 5 minutes to allow for retries and requeuing + ) + except comfy.model_management.InterruptProcessingException: + # Clean up job on interruption + run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0) + raise + + # Check if we need to process any remaining tasks locally after collection + completed_count = len(collected_tasks) + if completed_count < total_tiles: + log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") + + # Process any remaining pending tasks + while True: + # Check for user interruption + comfy.model_management.throw_exception_if_processing_interrupted() + + # Get next task from pending queue + task_id = run_async_in_server_loop( + self._get_next_tile_index(multi_job_id), + timeout=5.0 + ) + + if task_id is None: + break + + batch_idx = task_id // num_tiles_per_image + local_tile_idx = task_id % num_tiles_per_image + + if batch_idx >= batch_size: + continue + + # Get or create sliced conditioning + if batch_idx not in sliced_conditioning_cache: + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx) + sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced) + else: + positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx] + + # Use unique seed per image + image_seed = seed + batch_idx * 1000 + + # Process tile synchronously + result_images[batch_idx] = self._process_and_blend_tile( + local_tile_idx, all_tiles[local_tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx], + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx + ) + + # Mark as completed and store in collected_tasks + run_async_in_server_loop( + _mark_task_completed(multi_job_id, task_id, {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}), + timeout=5.0 + ) + # Add to collected_tasks so we skip it in the blending loop + collected_tasks[task_id] = {'batch_idx': batch_idx, 'tile_idx': local_tile_idx} + else: + # Master processed all tiles + collected_tasks = run_async_in_server_loop( + self._get_all_completed_tasks(multi_job_id), + timeout=5.0 + ) # Blend worker tiles synchronously for global_idx, tile_data in collected_tasks.items(): - if global_idx in master_indices: - continue # Skip master's tiles + # Skip tiles already processed by master + if static_distribution and global_idx in master_indices: + continue # Skip master's tiles in static mode - if 'tensor' in tile_data: - batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) - tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) - - if batch_idx >= batch_size: - continue - - # Blend tile synchronously - x = tile_data.get('x', 0) - y = tile_data.get('y', 0) - tile_tensor = tile_data['tensor'] - tile_pil = tensor_to_pil(tile_tensor, 0) - orig_x, orig_y = all_tiles[tile_idx] - tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) - extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) - extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) - result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, - x, y, (extracted_width, extracted_height), tile_mask, padding) + # Skip tiles that don't have tensor data (already processed) + if 'tensor' not in tile_data: + continue + + batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) + tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) + + if batch_idx >= batch_size: + continue + + # Blend tile synchronously + x = tile_data.get('x', 0) + y = tile_data.get('y', 0) + tile_tensor = tile_data['tensor'] + tile_pil = tensor_to_pil(tile_tensor, 0) + orig_x, orig_y = all_tiles[tile_idx] + tile_mask = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur) + extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) + extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) + result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil, + x, y, (extracted_width, extracted_height), tile_mask, padding) try: # Convert back to tensor @@ -1634,6 +1812,19 @@ class UltimateSDUpscaleDistributed: except asyncio.TimeoutError: return None + async def _get_next_tile_index(self, multi_job_id): + """Get next tile index from pending queue for master in static mode.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not job_data or JOB_PENDING_TASKS not in job_data: + return None + try: + tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0) + return tile_idx + except asyncio.TimeoutError: + return None + async def _get_total_completed_count(self, multi_job_id): """Get total count of all completed images (master + workers).""" @@ -1874,6 +2065,53 @@ class UltimateSDUpscaleDistributed: return None, 0 + async def _request_tile_from_master(self, multi_job_id, master_url, worker_id): + """Request a tile index to process from master in static mode (reusing dynamic infrastructure).""" + # Reuse the same retry logic as dynamic mode + max_retries = 10 + retry_delay = 0.5 + start_time = time.time() + + for attempt in range(max_retries): + # Check total timeout + if time.time() - start_time > 30: + log(f"Total request timeout after 30s for worker {worker_id}") + return None, 0 + + try: + session = await get_client_session() + url = f"{master_url}/distributed/request_image" # Same endpoint + + async with session.post(url, json={ + 'worker_id': str(worker_id), + 'multi_job_id': multi_job_id + }) as response: + if response.status == 200: + data = await response.json() + tile_idx = data.get('tile_idx') + estimated_remaining = data.get('estimated_remaining', 0) + return tile_idx, estimated_remaining + elif response.status == 404: + # Special handling for 404 - job not found yet + text = await response.text() + debug_log(f"Job not found (404), will retry: {text}") + await asyncio.sleep(1.0) + else: + text = await response.text() + debug_log(f"Request tile failed: {response.status} - {text}") + + except Exception as e: + if attempt < max_retries - 1: + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + await asyncio.sleep(retry_delay) + retry_delay *= 2 + retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + else: + log(f"Failed to request tile after {max_retries} attempts: {e}") + raise + + return None, 0 + async def _check_job_status(self, multi_job_id, master_url): """Check if job is ready on the master.""" try: diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py index a7cfb38..28fe055 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_managment.py @@ -61,16 +61,20 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima debug_log(f"Initialized dynamic queue with {batch_size} pending images") elif mode == 'static': job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image - # For static, pending_tasks starts empty; requeues will be added later - debug_log(f"Initialized static queue for {batch_size * num_tiles_per_image} tiles") + # For static with dynamic distribution, populate all tile indices in pending queue + pending_queue = job_data[JOB_PENDING_TASKS] + total_tiles = batch_size * num_tiles_per_image + for i in range(total_tiles): + await pending_queue.put(i) + debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution") - # If task assignments provided, populate JOB_ASSIGNED_TO_WORKERS + # Keep backward compatibility - if task assignments provided, still track them if task_assignments and enabled_workers: # task_assignments[0] is for master, 1+ are for workers for i, worker_id in enumerate(enabled_workers): if i + 1 < len(task_assignments): job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1] - debug_log(f"Worker {worker_id} assigned {len(task_assignments[i + 1])} tasks") + debug_log(f"Worker {worker_id} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)") else: raise ValueError(f"Unknown mode: {mode}") @@ -748,7 +752,7 @@ def ensure_tile_jobs_initialized(): # API Endpoint for tile completion @server.PromptServer.instance.routes.post("/distributed/request_image") async def request_image_endpoint(request): - """Endpoint for workers to request images in dynamic mode.""" + """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode).""" try: data = await request.json() worker_id = data.get('worker_id') @@ -763,26 +767,40 @@ async def request_image_endpoint(request): job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: return await handle_api_error(request, "Invalid job data structure", 500) - if job_data['mode'] != 'dynamic': - return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400) - if 'pending_images' in job_data: + + mode = job_data['mode'] + + # Handle both dynamic and static modes + if mode == 'dynamic' and 'pending_images' in job_data: pending_queue = job_data['pending_images'] - try: - image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) - # Track assigned image - if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: - job_data['assigned_to_workers'][worker_id].append(image_idx) - # Update worker heartbeat - if 'worker_status' in job_data: - job_data['worker_status'][worker_id] = time.time() - # Get estimated remaining count - remaining = pending_queue.qsize() # Approximate - debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}") - return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining}) - except asyncio.TimeoutError: - return web.json_response({"image_idx": None}) + elif mode == 'static' and JOB_PENDING_TASKS in job_data: + pending_queue = job_data[JOB_PENDING_TASKS] else: - return await handle_api_error(request, "Invalid dynamic mode configuration", 400) + return await handle_api_error(request, f"Invalid {mode} mode configuration", 400) + + try: + task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + # Track assigned task + if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']: + job_data['assigned_to_workers'][worker_id].append(task_idx) + # Update worker heartbeat + if 'worker_status' in job_data: + job_data['worker_status'][worker_id] = time.time() + # Get estimated remaining count + remaining = pending_queue.qsize() # Approximate + + # Return appropriate response based on mode + if mode == 'dynamic': + debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}") + return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining}) + else: # static + debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}") + return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining}) + except asyncio.TimeoutError: + if mode == 'dynamic': + return web.json_response({"image_idx": None}) + else: + return web.json_response({"tile_idx": None}) else: return await handle_api_error(request, "Job not found", 404) except Exception as e: From 809b840c6e7d81434140c054c3252c2f728fc1c4 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 12:07:23 +1000 Subject: [PATCH 44/93] Update README.md --- README.md | 137 +++++++++--------------------------------------------- 1 file changed, 22 insertions(+), 115 deletions(-) diff --git a/README.md b/README.md index adb7d15..fc12b58 100644 --- a/README.md +++ b/README.md @@ -1,38 +1,44 @@ -# ComfyUI-Distributed + -![Clipboard Image (3)](https://github.com/user-attachments/assets/19fdd1be-8f6e-4df5-bcd9-538ef566fa82) - -> **Supercharge your ComfyUI workflows with multi-GPU processing** - -A powerful extension for ComfyUI that enables parallel and distributed processing across multiple GPUs and machines. Speed up your image generation and upscaling workflows by leveraging all available GPU resources in your network. - -📺 [Watch the Tutorial](https://www.youtube.com/watch?v=p6eE3IlAbOs) +> **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network.** --- ## Key Features #### Parallel Workflow Processing -- **Parallel Generation** - Run the same workflow on multiple GPUs simultaneously with different seeds -- **Automatic Load Balancing** - Distribute workflow execution across available workers -- **Batch Acceleration** - Generate multiple variations faster by using all your GPUs +- Run the same workflow on multiple GPUs simultaneously with different seeds +- Get multiple outputs in the same time it takes to generate one +- The more workers you have, the more output you get +- Supports both video and images #### Distributed Upscaling -- **True Distributed Processing** - Split large upscaling tasks into tiles processed across multiple GPUs -- **Tile-based Upscaling** - Intelligent work distribution for Ultimate SD Upscale +- Supercharge Ultimate SD Upscale by having tiles processed across multiple GPUs +- Intelligent work distribution for asymmetrical GPU setups +- Static distribution for similar/same GPUs +- Supports both single images and image batches #### Management & Monitoring -- **Automatic Worker Management** - Launch and monitor workers from the UI +- **Worker Management** - Launch and monitor workers from the UI - **Network Support** - Use GPUs across different machines on your network +- **Cloud Support** - Use GPUs from cloud providers - **Real-time Monitoring** - Track worker status and performance from the UI - **Easy Configuration** - JSON-based configuration with UI controls -- **Memory Management** - Built-in VRAM clearing --- ## Requirements - ComfyUI installation + - Desktop app not supported currently - Multiple GPUs - That's it @@ -45,61 +51,11 @@ A powerful extension for ComfyUI that enables parallel and distributed processin git clone https://github.com/robertvoy/ComfyUI-Distributed.git ``` -2. **Restart ComfyUI** - If you'll be using remote workers, add `--enable-cors-header` to your launch arguments +2. **Restart ComfyUI** - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master --- -## Quick Start -### Adding Local Workers -![Distributed GPU Panel](https://github.com/user-attachments/assets/9c1d6d0e-3fd1-43e3-97c4-7c6bf2952b19) -> Local Workers: Additional ComfyUI instances running on the same computer (with multi-GPUs) as your main ComfyUI installation. - -1. **Open** the Distributed GPU panel. -2. **Click** "Add Worker" in the UI. -3. **Configure** your local worker: - - **Name**: A descriptive name for the worker (e.g., "My Gaming PC GPU 0") - - **Port**: A unique port number for this worker (e.g., 8189, 8190...). - - **CUDA Device**: The GPU index from `nvidia-smi` (e.g., 0, 1). - - **Extra Args**: Optional ComfyUI arguments for this specific worker. -4. **Save** and optionally launch the local worker. - -### Adding Remote Workers -📺 [Watch the Tutorial](https://www.youtube.com/watch?v=p6eE3IlAbOs) - -> Remote Workers: ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address. - -1. **On the Remote Worker Machine:** - - **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** - - **Add** workers in the UI panel if the remote machine has more than one GPU. - - Make sure that they also have `--listen` set in `Extra Args`. - - Then launch them. - - **Open** the configured worker port(s) (e.g., 8189, 8190) in the remote worker's firewall. - -2. **On the Main Machine:** - - **Open** the Distributed GPU panel (sidebar on the left). - - **Click** "Add Worker." - - **Enable** "Remote Worker" checkbox. - - **Configure** your remote worker: - - **Name**: A descriptive name for the worker (e.g., "Server Rack GPU 0") - - **Host**: The remote worker's IP address. - - **Port**: The port number used when launching ComfyUI on the remote worker (e.g., 8189). - - **Save** the remote worker configuration. - -### Configuration Tips - -| Setting | Description | Example | -|---------|-------------|---------| -| **CUDA Devices** | Use `nvidia-smi` to see GPU indices | 0, 1, 2... | -| **Ports** | Each worker needs a unique port | 8189, 8190... | -| **Extra Args** | Additional ComfyUI arguments | See below | - -**Common Extra Args:** -- `--listen` - **Required** for remote workers -- `--enable-cors-header` - **Required** if using remoter workers -- `--lowvram` - For GPUs with less memory -- `--highvram` - For high-end GPUs -- `--reserve-vram 2` - Reserves 2GB of VRAM. Recommended for your primary/display GPU --- @@ -166,47 +122,6 @@ The control centre for your distributed setup: --- -## Troubleshooting - -### Common Issues - -
-Workers won't start - -- Check ports are not in use: `netstat -an | grep 8189` -- Verify CUDA device exists: `nvidia-smi` -- Check ComfyUI path in worker logs -
- -
-"Worker not managed by UI" message - -- Worker was started outside the UI -- Stop the worker manually and use the UI to relaunch -
- -
-Images not combining properly - -- Ensure all remote workers have the same models available -- Check that custom nodes are installed on all remote workers -
- -
-Network connection issues - -- Check firewall settings for required ports -- Verify master IP is accessible: `ping 192.168.1.100` -- Ensure same ComfyUI version on all machines -- Ensure ComfyUI-Distributed is installed on remote workers -
- -
-Custom validation failed for node: image - Invalid image file - -- Add `--enable-cors-header` to your launch argument, on both master and remote worker -
- ## UI Panel Settings The Distributed GPU panel includes several configuration options in the Settings section: @@ -231,11 +146,3 @@ The Distributed GPU panel includes several configuration options in the Settings ## Development This project is under active development. Contributions are welcome! - -### Planned Features - -- [x] Support for Cloud workers (release coming soon) -- [ ] Remote worker control via SSH -- [ ] View remote worker logs in UI -- [ ] Improve worker timeout logic -- [ ] ComfyUI Desktop app compatibility From 19d2b442db09cfeea56cbe16b95fa1ec059494ad Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 12:34:07 +1000 Subject: [PATCH 45/93] Update README.md --- README.md | 59 +++++++++++++++++++++++-------------------------------- 1 file changed, 25 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index fc12b58..77ee34f 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,11 @@ - Static distribution for similar/same GPUs - Supports both single images and image batches +#### Convenience +- Automatically set up your local workers and easily add remote and cloud workers +- Manage the workers through the UI panel +- Easily convert any workflow to a distributed one, using no more than 2 nodes + #### Management & Monitoring - **Worker Management** - Launch and monitor workers from the UI - **Network Support** - Use GPUs across different machines on your network @@ -53,10 +58,6 @@ 2. **Restart ComfyUI** - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master ---- - - - --- ## Nodes @@ -71,6 +72,11 @@ **Usage**: Connect to any seed input. Each worker automatically receives an offset seed to ensure randomisation. Alternatively, you can connect a seed node directly to this node's seed input, which will automatically handle seed offsetting across all workers. +### Image Batch Divider +> Divides image batches, used video + +**Usage**: Place after the Distributed Collector and set the divide_by to the number of GPUs you are using (including the master) + ### Ultimate SD Upscale Distributed > Distributed version of Ultimate SD Upscale that processes tiles across multiple GPUs, making upscaling much faster @@ -82,22 +88,6 @@ --- -## UI Features - -### Distributed GPU Panel - -The control centre for your distributed setup: - -| Feature | Description | -|---------|-------------| -| **Worker Status** | Real-time status indicators | -| **Launch/Stop** | Control individual local workers | -| **Clear Memory** | Free VRAM on all workers | -| **Interrupt** | Stop current processing | -| **Worker Logs** | View real-time logs | - ---- - ## Workflow Examples ### Basic Distributed Generation @@ -122,24 +112,25 @@ The control centre for your distributed setup: --- -## UI Panel Settings +## FAQ -The Distributed GPU panel includes several configuration options in the Settings section: +Does it combine VRAM of multiple GPUs? +No. -### Debug Mode -- **Description**: Enable detailed logging for troubleshooting and monitoring distributed operations -- **Default**: Disabled -- **Usage**: When enabled, detailed debug information is logged to the browser console, including worker status updates, job distribution details, and network communication logs +Does it speed up the generation of a single image/video? +No, it gives you more images/videos, rather than a faster single one. However, it does speed up the upscaling of a single image using the Ultimate SD Upscale Distributed. -### Auto-launch Workers on Startup -- **Description**: Automatically launch enabled local workers when ComfyUI starts -- **Default**: Disabled -- **Usage**: When enabled, any local workers that are marked as "enabled" will be automatically launched in the background when the ComfyUI server starts, eliminating the need to launch each worker manually +Does it work with ComfyUI desktop app? +Not currently. -### Stop Workers on Master Exit -- **Description**: Automatically stop all managed local workers when the master ComfyUI instance shuts down -- **Default**: Enabled -- **Usage**: When enabled, ensures clean shutdown by stopping all UI-managed worker processes when the main ComfyUI server exits, preventing orphaned background processes +Can I combine my RTX 5090 with a GTX 980 to get faster results? +You can, but this works best with cards that are similar. If you have a large imbalance between GPUs, you will run into bottlenecks. Although you could still benefit from Ultimate SD Upscale Distributed with static_distribution set to false. This will allow the faster card to process more. Note this only works for upscaling. + +Does this work with cloud providers? +Yes, see setup guides. + +Can I make this work with my Docker setup? +Yes, it does work, but you need to know how to set up your Docker environment. I won't be able to assist you with that. --- From 98b4026c79c9d1cf68c48d0839cf181ba80b1b1c Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 12:46:00 +1000 Subject: [PATCH 46/93] Update README.md --- README.md | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 77ee34f..e19cb77 100644 --- a/README.md +++ b/README.md @@ -134,6 +134,15 @@ Yes, it does work, but you need to know how to set up your Docker environment. I --- -## Development +## Disclaimer -This project is under active development. Contributions are welcome! +This software is provided “as is” without any warranties, express or implied, including but not limited to merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities—whether in contract, tort, or otherwise—arising from the use, modification, distribution, or other dealings with the software. + +Users are solely responsible for ensuring their use of the software, including any content it generates, complies with all applicable laws and regulations in their jurisdiction. The developers are not responsible for any legal violations resulting from user actions. + +Additionally, users are responsible for the security and integrity of their networks, including protecting against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages, losses, or incidents arising from improper configuration, misuse, or external threats to user environments. + +--- + +## Support the Project +If my custom nodes have added value to your workflow, consider fueling future development with a coffee! Your support helps keep this project thriving. Buy me a coffee at: https://buymeacoffee.com/robertvoy From d60d93c29dec62345d1de682066bca8f0a4c4dc7 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 12:55:18 +1000 Subject: [PATCH 47/93] Update README.md --- README.md | 50 +++++++++++++++++++++++++------------------------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index e19cb77..2174752 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,30 @@ --- +## Workflow Examples + +### Basic Distributed Generation + +![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) + +1. Create your normal ComfyUI workflow +2. Add **Distributed Seed** → connect to sampler's seed +3. Add **Distributed Collector** → after VAE Decode +4. Enable workers in the UI +5. Run the workflow! + +### Distributed Upscaling + +![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) + +1. Load your image +2. Upscale with ESRGAN or similar +3. Connect to **Ultimate SD Upscale Distributed** +4. Configure tile settings +5. Enable workers for faster processing + +--- + ## Nodes ### Distributed Collector @@ -88,30 +112,6 @@ --- -## Workflow Examples - -### Basic Distributed Generation - -![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) - -1. Create your normal ComfyUI workflow -2. Add **Distributed Seed** → connect to sampler's seed -3. Add **Distributed Collector** → after VAE Decode -4. Enable workers in the UI -5. Run the workflow! - -### Distributed Upscaling - -![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) - -1. Load your image -2. Upscale with ESRGAN or similar -3. Connect to **Ultimate SD Upscale Distributed** -4. Configure tile settings -5. Enable workers for faster processing - ---- - ## FAQ Does it combine VRAM of multiple GPUs? @@ -136,7 +136,7 @@ Yes, it does work, but you need to know how to set up your Docker environment. I ## Disclaimer -This software is provided “as is” without any warranties, express or implied, including but not limited to merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities—whether in contract, tort, or otherwise—arising from the use, modification, distribution, or other dealings with the software. +By using this software, you agree to be bound by the terms of this Disclaimer. The software is provided “as is” without any warranties, express or implied, including but not limited to merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities—whether in contract, tort, or otherwise—arising from the use, modification, distribution, or other dealings with the software. Users are solely responsible for ensuring their use of the software, including any content it generates, complies with all applicable laws and regulations in their jurisdiction. The developers are not responsible for any legal violations resulting from user actions. From bbcc2b7789f381fbfe24a69e78ab1fe1b24eb6a0 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 13:07:46 +1000 Subject: [PATCH 48/93] Update README.md --- README.md | 66 ++++++++++++++++++++++++++++++------------------------- 1 file changed, 36 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 2174752..d456a49 100644 --- a/README.md +++ b/README.md @@ -15,28 +15,22 @@ ## Key Features #### Parallel Workflow Processing -- Run the same workflow on multiple GPUs simultaneously with different seeds -- Get multiple outputs in the same time it takes to generate one -- The more workers you have, the more output you get -- Supports both video and images +- Run workflows on multiple GPUs simultaneously with varied seeds +- Generate multiple outputs in the time of one +- Scale output with more workers +- Supports images and videos #### Distributed Upscaling -- Supercharge Ultimate SD Upscale by having tiles processed across multiple GPUs -- Intelligent work distribution for asymmetrical GPU setups -- Static distribution for similar/same GPUs -- Supports both single images and image batches +- Accelerate Ultimate SD Upscale by distributing tiles across GPUs +- Intelligent distribution for asymmetric setups +- Static distribution for similar GPUs +- Handles single images and batches #### Convenience -- Automatically set up your local workers and easily add remote and cloud workers -- Manage the workers through the UI panel -- Easily convert any workflow to a distributed one, using no more than 2 nodes - -#### Management & Monitoring -- **Worker Management** - Launch and monitor workers from the UI -- **Network Support** - Use GPUs across different machines on your network -- **Cloud Support** - Use GPUs from cloud providers -- **Real-time Monitoring** - Track worker status and performance from the UI -- **Easy Configuration** - JSON-based configuration with UI controls +- Auto-setup local workers; easily add remote/cloud ones +- Manage workers via UI panel +- Convert workflows to distributed with less than 2 nodes +- JSON configuration with UI controls --- @@ -114,23 +108,35 @@ ## FAQ -Does it combine VRAM of multiple GPUs? -No. +
+Does it combine VRAM of multiple GPUs? +No, it does not combine VRAM of multiple GPUs. +
-Does it speed up the generation of a single image/video? -No, it gives you more images/videos, rather than a faster single one. However, it does speed up the upscaling of a single image using the Ultimate SD Upscale Distributed. +
+Does it speed up the generation of a single image or video? +No, it does not speed up the generation of a single image or video. Instead, it enables the generation of more images or videos simultaneously. However, it can speed up the upscaling of a single image when using the Ultimate SD Upscale Distributed feature. +
-Does it work with ComfyUI desktop app? -Not currently. +
+Does it work with the ComfyUI desktop app? +Currently, it is not compatible with the ComfyUI desktop app. +
-Can I combine my RTX 5090 with a GTX 980 to get faster results? -You can, but this works best with cards that are similar. If you have a large imbalance between GPUs, you will run into bottlenecks. Although you could still benefit from Ultimate SD Upscale Distributed with static_distribution set to false. This will allow the faster card to process more. Note this only works for upscaling. +
+Can I combine my RTX 5090 with a GTX 980 to get faster results? +Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks. For upscaling, setting `static_distribution` to `false` allows the faster GPU to handle more processing, which can mitigate some bottlenecks. Note that this setting only applies to upscaling tasks. +
-Does this work with cloud providers? -Yes, see setup guides. +
+Does this work with cloud providers? +Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions. +
-Can I make this work with my Docker setup? -Yes, it does work, but you need to know how to set up your Docker environment. I won't be able to assist you with that. +
+Can I make this work with my Docker setup? +Yes, it is compatible with Docker setups, but you will need to configure your Docker environment yourself. Unfortunately, assistance with Docker configuration is not provided. +
--- From 054800ade3935cc01367e003c8811d7e3f1fca1f Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 14:16:35 +1000 Subject: [PATCH 49/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index d456a49..7e00abd 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@
-image +distributed-logo

📺Video Tutorial | 📝Setup Guides | From 9386f302aad2ea4caa437d8b807d8a4ba78c65bf Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 14:20:22 +1000 Subject: [PATCH 50/93] Update README.md --- README.md | 34 ---------------------------------- 1 file changed, 34 deletions(-) diff --git a/README.md b/README.md index 7e00abd..f317b6e 100644 --- a/README.md +++ b/README.md @@ -106,40 +106,6 @@ --- -## FAQ - -
-Does it combine VRAM of multiple GPUs? -No, it does not combine VRAM of multiple GPUs. -
- -
-Does it speed up the generation of a single image or video? -No, it does not speed up the generation of a single image or video. Instead, it enables the generation of more images or videos simultaneously. However, it can speed up the upscaling of a single image when using the Ultimate SD Upscale Distributed feature. -
- -
-Does it work with the ComfyUI desktop app? -Currently, it is not compatible with the ComfyUI desktop app. -
- -
-Can I combine my RTX 5090 with a GTX 980 to get faster results? -Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks. For upscaling, setting `static_distribution` to `false` allows the faster GPU to handle more processing, which can mitigate some bottlenecks. Note that this setting only applies to upscaling tasks. -
- -
-Does this work with cloud providers? -Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions. -
- -
-Can I make this work with my Docker setup? -Yes, it is compatible with Docker setups, but you will need to configure your Docker environment yourself. Unfortunately, assistance with Docker configuration is not provided. -
- ---- - ## Disclaimer By using this software, you agree to be bound by the terms of this Disclaimer. The software is provided “as is” without any warranties, express or implied, including but not limited to merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities—whether in contract, tort, or otherwise—arising from the use, modification, distribution, or other dealings with the software. From 52fdc096b04bc99d72725dde0c85ef80d25fadb8 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 15:01:59 +1000 Subject: [PATCH 51/93] Update README.md --- README.md | 44 +++++++++++++++++++++++++------------------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index f317b6e..8f76d24 100644 --- a/README.md +++ b/README.md @@ -78,31 +78,37 @@ --- -## Nodes +## FAQ -### Distributed Collector -> Collects and combines results from distributed processing +
+Does it combine VRAM of multiple GPUs? +No, it does not combine VRAM of multiple GPUs. +
-**Usage**: Place after any image generation node to enable distributed processing. Works automatically when workers are enabled. +
+Does it speed up the generation of a single image or video? +No, it does not speed up the generation of a single image or video. Instead, it enables the generation of more images or videos simultaneously. However, it can speed up the upscaling of a single image when using the Ultimate SD Upscale Distributed feature. +
-### Distributed Seed -> Ensures unique seeds across distributed workers for varied generations +
+Does it work with the ComfyUI desktop app? +Currently, it is not compatible with the ComfyUI desktop app. +
-**Usage**: Connect to any seed input. Each worker automatically receives an offset seed to ensure randomisation. Alternatively, you can connect a seed node directly to this node's seed input, which will automatically handle seed offsetting across all workers. +
+Can I combine my RTX 5090 with a GTX 980 to get faster results? +Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks. For upscaling, setting `static_distribution` to `false` allows the faster GPU to handle more processing, which can mitigate some bottlenecks. Note that this setting only applies to upscaling tasks. +
-### Image Batch Divider -> Divides image batches, used video +
+Does this work with cloud providers? +Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions. +
-**Usage**: Place after the Distributed Collector and set the divide_by to the number of GPUs you are using (including the master) - -### Ultimate SD Upscale Distributed -> Distributed version of Ultimate SD Upscale that processes tiles across multiple GPUs, making upscaling much faster - -**Usage**: -1. Upscale your image with a regular upscale model (ESRGAN, etc.) -2. Feed the upscaled image into this node -3. Configure tile settings -4. Enable workers for fast processing +
+Can I make this work with my Docker setup? +Yes, it is compatible with Docker setups, but you will need to configure your Docker environment yourself. Unfortunately, assistance with Docker configuration is not provided. +
--- From 726a6ffb50a0d17801e1fc1889056a673a708b6a Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 15:28:15 +1000 Subject: [PATCH 52/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8f76d24..5f92834 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@
-distributed-logo +

📺Video Tutorial | 📝Setup Guides | From 58485632c25d6d0aa51da514b5793329521d2b88 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 26 Jul 2025 15:36:03 +1000 Subject: [PATCH 53/93] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 5f92834..2a7eed5 100644 --- a/README.md +++ b/README.md @@ -36,8 +36,8 @@ ## Requirements -- ComfyUI installation - - Desktop app not supported currently +- ComfyUI +> Note: Desktop app not supported currently - Multiple GPUs - That's it From f2b66e5bc61e1d6a4e43158685d28e0e214e16ca Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 27 Jul 2025 09:28:25 +1000 Subject: [PATCH 54/93] Update README.md --- README.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 2a7eed5..ed2400b 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,8 @@ > **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network.** +![Clipboard Image (3)](https://github.com/user-attachments/assets/19fdd1be-8f6e-4df5-bcd9-538ef566fa82) + --- ## Key Features @@ -29,7 +31,7 @@ #### Convenience - Auto-setup local workers; easily add remote/cloud ones - Manage workers via UI panel -- Convert workflows to distributed with less than 2 nodes +- Convert any workflow to distributed with 2 nodes - JSON configuration with UI controls --- @@ -56,7 +58,7 @@ ## Workflow Examples -### Basic Distributed Generation +### Basic Parallel Generation ![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) @@ -66,6 +68,8 @@ 4. Enable workers in the UI 5. Run the workflow! +### Parallel WAN Generation + ### Distributed Upscaling ![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) From e4096b74a0e07c10efdb926146c8840b4d0455dc Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sun, 27 Jul 2025 14:39:23 +1000 Subject: [PATCH 55/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index ed2400b..f1df4a7 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@
- +

📺Video Tutorial | 📝Setup Guides | From 209fe5388b020f874f1fe5fa705793aaa7b9b3dd Mon Sep 17 00:00:00 2001 From: robertvoy Date: Mon, 28 Jul 2025 12:50:25 +1000 Subject: [PATCH 56/93] added docs, updated workflow, updated url for Runpod in UI --- docs/model-download-script.md | 43 + docs/worker-setup-guides.md | 88 ++ web/main.js | 10 + web/ui.js | 4 +- workflows/distributed-upscale-batch.json | 1277 ++++++++++++++++ workflows/distributed-upscale-video.json | 1414 ++++++++++++++++++ workflows/distributed-upscale.json | 1235 ++++++++++++---- workflows/distributed-wan.json | 1688 ++++++++++++++++++++++ 8 files changed, 5473 insertions(+), 286 deletions(-) create mode 100644 docs/model-download-script.md create mode 100644 docs/worker-setup-guides.md create mode 100644 workflows/distributed-upscale-batch.json create mode 100644 workflows/distributed-upscale-video.json create mode 100644 workflows/distributed-wan.json diff --git a/docs/model-download-script.md b/docs/model-download-script.md new file mode 100644 index 0000000..4c41511 --- /dev/null +++ b/docs/model-download-script.md @@ -0,0 +1,43 @@ +## Prompt for LLM - Automating ComfyUI Model Downloads + +This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM). + +1. Export your workflow as an API workflow +2. Copy the below prompt and upload the API workflow to a LLM that has access to the internet (like Gemini Flash) +3. Review (make sure all downloads are correct) and save the output from the LLM as a .sh file +4. Upload it onto your Runpod instance +5. Run the file. For example: `./download_comfyui_models.sh` + +``` +Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: +comfyui: +base_path: /workspace/ComfyUI +BiRefNet: models/BiRefNet/ +checkpoints: models/checkpoints/ +clip: models/clip/ +clip_vision: models/clip_vision/ +configs: models/configs/ +controlnet: models/controlnet/ +diffusers: models/diffusers/ +diffusion_models: models/diffusion_models/ +embeddings: models/embeddings/ +florence2: models/florence2/ +facerestore_models: models/facerestore_models/ +gligen: models/gligen/ +grounding-dino: models/grounding-dino/ +hypernetworks: models/hypernetworks/ +ipadapter: models/ipadapter/ +lama: models/lama/ +loras: models/loras/ +onnx: models/onnx/ +photomaker: models/photomaker/ +RMBG: models/RMBG/ +sams: models/sams/ +style_models: models/style_models/ +text_encoders: models/text_encoders/ +unet: models/unet/ +upscale_models: models/upscale_models/ +vae: models/vae/ +vae_approx: models/vae_approx/ +vitmatte: models/vitmatte/ +``` \ No newline at end of file diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md new file mode 100644 index 0000000..2dbde7f --- /dev/null +++ b/docs/worker-setup-guides.md @@ -0,0 +1,88 @@ +## Local workers +> Local Workers: These are added automatically on first launch, but you can add them manually if you need. + +1. **Open** the Distributed GPU panel. +2. **Click** "Add Worker" in the UI. +3. **Configure** your local worker: + - **Name**: A descriptive name for the worker (e.g., "Studio PC 1") + - **Port**: A unique port number for this worker (e.g., 8189, 8190...). + - **CUDA Device**: The GPU index from `nvidia-smi` (e.g., 0, 1). + - **Extra Args**: Optional ComfyUI arguments for this specific worker. +4. **Save** and launch the local worker. + +## Remote workers +> Remote Workers: ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address. + +**On the Remote Worker Machine:** +1. **Launch** the master (main ComfyUI instance) with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** +2. **Add** workers in the UI panel if the remote machine has more than one GPU. + - Make sure that they also have `--listen` set in `Extra Args`. + - Then launch them. +3. **Open** the configured worker port(s) (e.g., 8189, 8190) in the remote worker's firewall. + +**On the Main Machine:** +1. **Open** the Distributed GPU panel (sidebar on the left). +2. **Click** "Add Worker." +3. **Choose** "Remote". +4. **Configure** your remote worker: + - **Name**: A descriptive name for the worker (e.g., "Server Rack GPU 0") + - **Host**: The remote worker's IP address. + - **Port**: The port number used when launching ComfyUI on the remote worker (e.g., 8189). +5. **Save** the remote worker configuration. + +## Cloud workers +> Cloud Workers: ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p). + +### Deploy Cloud Worker on Runpod + +**On Runpod:** +1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account. +2. On Runpod, go to Storage > New Network Volume and create a volume which will store the models you need. +3. Now go to Pods and find a suitable GPU for your workflows. +> To use the ComfyUI-Distributed-Pod, you will need to filter instances by CUDA 12.8. +4. Choose the `ComfyUI-Distributed-Pod` template and make sure your network drive is mounted. +5. Launch your pod. +6. Access your pod using JupyterLab. +7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). +> You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow. +8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!** +9. Download any additional custom nodes using ComfyUI manager + +**On the Main Machine:** +1. Launch a Cloudflare tunnel + - Download from here `https://github.com/cloudflare/cloudflared/releases/tag/2025.7.0` + - Then run for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` +2. Copy the Cloudflare address +3. **Open** the Distributed GPU panel (sidebar on the left). +4. **Edit** the Master's host address and replace it with the Cloudflare address. + - **Click** "Add Worker." + - **Choose** "Cloud". + - **Configure** your cloud worker: + - **Host**: The Runpod address + - **Port**: 443 + - **Save** the remote worker configuration. + +### Deploy Cloud Worker on Other Platforms + +**On the Cloud Worker machine:** + - Your cloud worker container needs to have the same models and custom nodes as the workflow you want to run on your local machine. + - If your cloud platform doesn't provide a secure connection, use Cloudflare to create a tunnel for the worker. Each GPU needs their own tunnel for their respective port. + - For example: `./cloudflared tunnel --url http://localhost:8188` +1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** +2. **Add** workers in the UI panel if the remote machine has more than one GPU. + - Make sure that they also have `--listen` set in `Extra Args`. + - Then launch them. + +**On the Main Machine:** +1. **Launch** a Cloudflare tunnel on your local machine. + - Download from here `https://github.com/cloudflare/cloudflared/releases/tag/2025.7.0` + - For example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` +2. **Copy** the Cloudflare address +3. **Open** the Distributed GPU panel (sidebar on the left). +4. **Edit** the Master's host address and replace it with the Cloudflare address. +5. **Click** "Add Worker." +6. **Choose** "Cloud". +7. **Configure** your cloud worker: + - **Host**: The remote worker's IP address/domain + - **Port**: 443 +8. **Save** the remote worker configuration. \ No newline at end of file diff --git a/web/main.js b/web/main.js index 4ae885c..4830a30 100644 --- a/web/main.js +++ b/web/main.js @@ -945,6 +945,16 @@ class DistributedExtension { this.log(`Auto-populated ${newWorkers.length} workers and saved config`, "info"); + // Show success notification + if (app.extensionManager?.toast) { + app.extensionManager.toast.add({ + severity: "success", + summary: "Workers Auto-populated", + detail: `Automatically created ${newWorkers.length} workers based on detected CUDA devices`, + life: 5000 + }); + } + // Reload the config to include the new workers await this.loadConfig(); } else { diff --git a/web/ui.js b/web/ui.js index 93b6274..e3681ea 100644 --- a/web/ui.js +++ b/web/ui.js @@ -697,9 +697,9 @@ export class DistributedUI { // Create powered by Runpod text (initially hidden) const runpodText = document.createElement("a"); runpodText.id = `runpod-text-${worker.id}`; - runpodText.href = "https://www.runpod.io"; + runpodText.href = "https://get.runpod.io/0bw29uf3ug0p"; runpodText.target = "_blank"; - runpodText.textContent = "Deploy cloud worker with Runpod"; + runpodText.textContent = "Deploy Cloud Worker with Runpod"; runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;"; // Store the onchange function to be assigned later diff --git a/workflows/distributed-upscale-batch.json b/workflows/distributed-upscale-batch.json new file mode 100644 index 0000000..89f6871 --- /dev/null +++ b/workflows/distributed-upscale-batch.json @@ -0,0 +1,1277 @@ +{ + "id": "6b169846-6204-4336-bda1-f9dc43cd52e7", + "revision": 0, + "last_node_id": 26, + "last_link_id": 26, + "nodes": [ + { + "id": 1, + "type": "Reroute", + "pos": [ + 3355.76220703125, + 1302.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 1 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 6 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 2, + "type": "Reroute", + "pos": [ + 3355.76220703125, + 1252.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 2 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 4 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 3, + "type": "Reroute", + "pos": [ + 4035.76220703125, + 1302.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 3 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 23 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 4, + "type": "Reroute", + "pos": [ + 4035.76220703125, + 1252.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 4 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 20 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 5, + "type": "Reroute", + "pos": [ + 4035.76220703125, + 1222.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 5 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 19 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 6, + "type": "CheckpointLoaderSimple", + "pos": [ + 2925.76220703125, + 1522.14697265625 + ], + "size": [ + 370, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 2 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [ + 7, + 18 + ] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 1 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CheckpointLoaderSimple", + "models": [ + { + "name": "v1-5-pruned-emaonly-fp16.safetensors", + "url": "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true", + "directory": "checkpoints" + } + ] + }, + "widgets_values": [ + "SDXL\\juggernautXL_ragnarokBy.safetensors" + ] + }, + { + "id": 7, + "type": "Reroute", + "pos": [ + 3685.76220703125, + 1302.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 6 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 3, + 15 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 8, + "type": "ControlNetLoader", + "pos": [ + 3135.76220703125, + 1762.14697265625 + ], + "size": [ + 390, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 8 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetLoader" + }, + "widgets_values": [ + "SDXL\\xinsir-controlnet-union-sdxl-1.0-promax.safetensors" + ] + }, + { + "id": 9, + "type": "CLIPTextEncode", + "pos": [ + 3355.76220703125, + 1622.14697265625 + ], + "size": [ + 420, + 88 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 7 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 12 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "text, watermark" + ] + }, + { + "id": 10, + "type": "SetUnionControlNetType", + "pos": [ + 3565.76220703125, + 1762.14697265625 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 8 + } + ], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 13 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SetUnionControlNetType" + }, + "widgets_values": [ + "tile" + ] + }, + { + "id": 11, + "type": "Reroute", + "pos": [ + 3675.76220703125, + 1132.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 9 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 14 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 12, + "type": "ImageResize+", + "pos": [ + 2925.76220703125, + 1222.14697265625 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 10 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 5 + ] + }, + { + "name": "width", + "type": "INT", + "links": null + }, + { + "name": "height", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 2048, + 2048, + "lanczos", + "keep proportion", + "always", + 8 + ] + }, + { + "id": 13, + "type": "ControlNetApplyAdvanced", + "pos": [ + 3865.76220703125, + 1402.14697265625 + ], + "size": [ + 270, + 186 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 11 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 12 + }, + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 13 + }, + { + "name": "image", + "type": "IMAGE", + "link": 14 + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": 15 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 21 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 22 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetApplyAdvanced" + }, + "widgets_values": [ + 1.0000000000000002, + 0, + 0.9000000000000002 + ] + }, + { + "id": 14, + "type": "ImageUpscaleWithModel", + "pos": [ + 2635.76220703125, + 1222.14697265625 + ], + "size": [ + 222.75416564941406, + 46 + ], + "flags": {}, + "order": 17, + "mode": 4, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 16 + }, + { + "name": "image", + "type": "IMAGE", + "link": 17 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 10 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ImageUpscaleWithModel" + }, + "widgets_values": [] + }, + { + "id": 15, + "type": "MarkdownNote", + "pos": [ + 3145.76220703125, + 1862.14697265625 + ], + "size": [ + 370, + 88 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "[https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main)" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 17, + "type": "Note", + "pos": [ + 2685.76220703125, + 1532.14697265625 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Choose an SDXL model" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 18, + "type": "CLIPTextEncode", + "pos": [ + 3355.76220703125, + 1412.14697265625 + ], + "size": [ + 420, + 160 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 18 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 11 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 19, + "type": "UltimateSDUpscaleDistributed", + "pos": [ + 4225.76220703125, + 1362.14697265625 + ], + "size": [ + 326.691650390625, + 450 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "upscaled_image", + "type": "IMAGE", + "link": 19 + }, + { + "name": "model", + "type": "MODEL", + "link": 20 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 21 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 22 + }, + { + "name": "vae", + "type": "VAE", + "link": 23 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 25 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "dd23503883fdf319e8beb6e7a190445ecf89973c", + "Node name for S&R": "UltimateSDUpscaleDistributed" + }, + "widgets_values": [ + 320425997971252, + "randomize", + 20, + 7, + "dpmpp_2m_sde", + "karras", + 0.7000000000000002, + 1024, + 1024, + 32, + 16, + true, + false, + true + ] + }, + { + "id": 20, + "type": "Reroute", + "pos": [ + 2455.76220703125, + 1132.14697265625 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 26 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 9, + 17 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 21, + "type": "UpscaleModelLoader", + "pos": [ + 2355.76220703125, + 1322.14697265625 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 16 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "4xNomos8kDAT.pth" + ] + }, + { + "id": 23, + "type": "PreviewImage", + "pos": [ + 4605.76220703125, + 1372.14697265625 + ], + "size": [ + 490, + 550 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 25 + } + ], + "outputs": [], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 25, + "type": "VHS_LoadVideoPath", + "pos": [ + 1950, + 1130 + ], + "size": [ + 330, + 286 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 26 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_LoadVideoPath" + }, + "widgets_values": { + "video": "input/test/", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "None", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "input/test/", + "type": "path", + "format": "video/input/test/", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1 + } + } + } + }, + { + "id": 22, + "type": "Note", + "pos": [ + 2645.76220703125, + 1322.14697265625 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Optional" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 26, + "type": "Note", + "pos": [ + 1950, + 980 + ], + "size": [ + 320, + 88 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Use this node to load a batch of images from a folder.\n\nNOTE: images have to be the same dimensions." + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [ + [ + 1, + 6, + 2, + 1, + 0, + "*" + ], + [ + 2, + 6, + 0, + 2, + 0, + "*" + ], + [ + 3, + 7, + 0, + 3, + 0, + "*" + ], + [ + 4, + 2, + 0, + 4, + 0, + "*" + ], + [ + 5, + 12, + 0, + 5, + 0, + "*" + ], + [ + 6, + 1, + 0, + 7, + 0, + "*" + ], + [ + 7, + 6, + 1, + 9, + 0, + "CLIP" + ], + [ + 8, + 8, + 0, + 10, + 0, + "CONTROL_NET" + ], + [ + 9, + 20, + 0, + 11, + 0, + "*" + ], + [ + 10, + 14, + 0, + 12, + 0, + "IMAGE" + ], + [ + 11, + 18, + 0, + 13, + 0, + "CONDITIONING" + ], + [ + 12, + 9, + 0, + 13, + 1, + "CONDITIONING" + ], + [ + 13, + 10, + 0, + 13, + 2, + "CONTROL_NET" + ], + [ + 14, + 11, + 0, + 13, + 3, + "IMAGE" + ], + [ + 15, + 7, + 0, + 13, + 4, + "VAE" + ], + [ + 16, + 21, + 0, + 14, + 0, + "UPSCALE_MODEL" + ], + [ + 17, + 20, + 0, + 14, + 1, + "IMAGE" + ], + [ + 18, + 6, + 1, + 18, + 0, + "CLIP" + ], + [ + 19, + 5, + 0, + 19, + 0, + "IMAGE" + ], + [ + 20, + 4, + 0, + 19, + 1, + "MODEL" + ], + [ + 21, + 13, + 0, + 19, + 2, + "CONDITIONING" + ], + [ + 22, + 13, + 1, + 19, + 3, + "CONDITIONING" + ], + [ + 23, + 3, + 0, + 19, + 4, + "VAE" + ], + [ + 25, + 19, + 0, + 23, + 0, + "IMAGE" + ], + [ + 26, + 25, + 0, + 20, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9229599817706583, + "offset": [ + -1960.8728618221426, + -538.777843033621 + ] + }, + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/distributed-upscale-video.json b/workflows/distributed-upscale-video.json new file mode 100644 index 0000000..fc82bd5 --- /dev/null +++ b/workflows/distributed-upscale-video.json @@ -0,0 +1,1414 @@ +{ + "id": "a778fe2e-2eb6-417f-8b97-86177243bce8", + "revision": 0, + "last_node_id": 140, + "last_link_id": 210, + "nodes": [ + { + "id": 86, + "type": "Reroute", + "pos": [ + 2130, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 121 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 127 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 107, + "type": "Reroute", + "pos": [ + 2810, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 171 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 172 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 110, + "type": "Reroute", + "pos": [ + 2810, + 970 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 210 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 176 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 7, + "type": "CheckpointLoaderSimple", + "pos": [ + 1700, + 1270 + ], + "size": [ + 370, + 98 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "slot_index": 0, + "links": [ + 124 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "slot_index": 1, + "links": [ + 5, + 6 + ] + }, + { + "name": "VAE", + "type": "VAE", + "slot_index": 2, + "links": [ + 121 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CheckpointLoaderSimple", + "models": [ + { + "name": "v1-5-pruned-emaonly-fp16.safetensors", + "url": "https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors?download=true", + "directory": "checkpoints" + } + ] + }, + "widgets_values": [ + "SDXL\\juggernautXL_ragnarokBy.safetensors" + ] + }, + { + "id": 90, + "type": "Reroute", + "pos": [ + 2460, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 127 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 129, + 171 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 44, + "type": "ControlNetLoader", + "pos": [ + 1910, + 1510 + ], + "size": [ + 390, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 81 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetLoader" + }, + "widgets_values": [ + "SDXL\\xinsir-controlnet-union-sdxl-1.0-promax.safetensors" + ] + }, + { + "id": 52, + "type": "SetUnionControlNetType", + "pos": [ + 2340, + 1510 + ], + "size": [ + 210, + 58 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 81 + } + ], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 192 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SetUnionControlNetType" + }, + "widgets_values": [ + "tile" + ] + }, + { + "id": 125, + "type": "Reroute", + "pos": [ + 2450, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 200 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 206 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 71, + "type": "RIFE VFI", + "pos": [ + 3367.979248046875, + 1110 + ], + "size": [ + 290.73333740234375, + 198 + ], + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "frames", + "type": "IMAGE", + "link": 108 + }, + { + "name": "optional_interpolation_states", + "shape": 7, + "type": "INTERPOLATION_STATES", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 109 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-frame-interpolation", + "ver": "a969c01dbccd9e5510641be04eb51fe93f6bfc3d", + "Node name for S&R": "RIFE VFI" + }, + "widgets_values": [ + "rife47.pth", + 10, + 2, + true, + true, + 1 + ] + }, + { + "id": 34, + "type": "ImageResize+", + "pos": [ + 1700, + 970 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 197 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 210 + ] + }, + { + "name": "width", + "type": "INT", + "links": null + }, + { + "name": "height", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 2048, + 2048, + "lanczos", + "keep proportion", + "always", + 8 + ] + }, + { + "id": 43, + "type": "ControlNetApplyAdvanced", + "pos": [ + 2640, + 1150 + ], + "size": [ + 270, + 186 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 63 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 64 + }, + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 192 + }, + { + "name": "image", + "type": "IMAGE", + "link": 206 + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": 129 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 190 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 191 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetApplyAdvanced" + }, + "widgets_values": [ + 1.0000000000000002, + 0, + 0.9000000000000002 + ] + }, + { + "id": 123, + "type": "ImageUpscaleWithModel", + "pos": [ + 1410, + 970 + ], + "size": [ + 222.75416564941406, + 46 + ], + "flags": {}, + "order": 18, + "mode": 4, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 195 + }, + { + "name": "image", + "type": "IMAGE", + "link": 209 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 197 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ImageUpscaleWithModel" + }, + "widgets_values": [] + }, + { + "id": 131, + "type": "Note", + "pos": [ + 1430, + 1070 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Optional" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 88, + "type": "Reroute", + "pos": [ + 2130, + 1020 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 124 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 137 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 89, + "type": "Reroute", + "pos": [ + 2810, + 1020 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 137 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 170 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 124, + "type": "Reroute", + "pos": [ + 1230, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 208 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 200, + 209 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 132, + "type": "Note", + "pos": [ + 1460, + 1270 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Choose an SDXL model" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 2140, + 1360 + ], + "size": [ + 420, + 88 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 6 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 64 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "text, watermark" + ] + }, + { + "id": 41, + "type": "VHS_VideoCombine", + "pos": [ + 3700, + 1110 + ], + "size": [ + 450, + 594.9140625 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 109 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 30, + "loop_count": 0, + "filename_prefix": "AnimateDiff", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": false, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00231.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 30, + "workflow": "AnimateDiff_00231.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00231.mp4" + } + } + } + }, + { + "id": 134, + "type": "MarkdownNote", + "pos": [ + 1910, + 1610 + ], + "size": [ + 380, + 88 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "[https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main)" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 122, + "type": "UpscaleModelLoader", + "pos": [ + 1120, + 1040 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 195 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "4xNomos8kDAT.pth" + ] + }, + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + 2140, + 1150 + ], + "size": [ + 420, + 160 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 63 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 30, + "type": "UltimateSDUpscaleDistributed", + "pos": [ + 3000, + 1110 + ], + "size": [ + 326.691650390625, + 450 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "upscaled_image", + "type": "IMAGE", + "link": 176 + }, + { + "name": "model", + "type": "MODEL", + "link": 170 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 190 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 191 + }, + { + "name": "vae", + "type": "VAE", + "link": 172 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 108 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "dd23503883fdf319e8beb6e7a190445ecf89973c", + "Node name for S&R": "UltimateSDUpscaleDistributed" + }, + "widgets_values": [ + 320425997971252, + "randomize", + 20, + 7, + "euler", + "normal", + 0.7000000000000002, + 1024, + 1024, + 32, + 16, + true, + false, + true + ] + }, + { + "id": 139, + "type": "Note", + "pos": [ + 3000, + 1610 + ], + "size": [ + 330, + 88 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "If all your GPUs are the same/similar, set static_distribution to true\n" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 138, + "type": "VHS_LoadVideo", + "pos": [ + 670, + 880 + ], + "size": [ + 340, + 509 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 208 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "AnimateDiff_00155.mp4", + "force_rate": 16, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "Wan", + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00155.mp4", + "type": "input", + "format": "video/mp4", + "force_rate": 16, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1 + } + } + } + }, + { + "id": 140, + "type": "MarkdownNote", + "pos": [ + 680, + 680 + ], + "size": [ + 390, + 140 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "# WARNING: Experimental workflow\n\nThe tiled upscaling works with video, but it is experimental.\n\nThe results can be flickery." + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [ + [ + 5, + 7, + 1, + 5, + 0, + "CLIP" + ], + [ + 6, + 7, + 1, + 6, + 0, + "CLIP" + ], + [ + 63, + 5, + 0, + 43, + 0, + "CONDITIONING" + ], + [ + 64, + 6, + 0, + 43, + 1, + "CONDITIONING" + ], + [ + 81, + 44, + 0, + 52, + 0, + "CONTROL_NET" + ], + [ + 108, + 30, + 0, + 71, + 0, + "IMAGE" + ], + [ + 109, + 71, + 0, + 41, + 0, + "IMAGE" + ], + [ + 121, + 7, + 2, + 86, + 0, + "*" + ], + [ + 124, + 7, + 0, + 88, + 0, + "*" + ], + [ + 127, + 86, + 0, + 90, + 0, + "*" + ], + [ + 129, + 90, + 0, + 43, + 4, + "VAE" + ], + [ + 137, + 88, + 0, + 89, + 0, + "*" + ], + [ + 170, + 89, + 0, + 30, + 1, + "MODEL" + ], + [ + 171, + 90, + 0, + 107, + 0, + "*" + ], + [ + 172, + 107, + 0, + 30, + 4, + "VAE" + ], + [ + 176, + 110, + 0, + 30, + 0, + "IMAGE" + ], + [ + 190, + 43, + 0, + 30, + 2, + "CONDITIONING" + ], + [ + 191, + 43, + 1, + 30, + 3, + "CONDITIONING" + ], + [ + 192, + 52, + 0, + 43, + 2, + "CONTROL_NET" + ], + [ + 195, + 122, + 0, + 123, + 0, + "UPSCALE_MODEL" + ], + [ + 197, + 123, + 0, + 34, + 0, + "IMAGE" + ], + [ + 200, + 124, + 0, + 125, + 0, + "*" + ], + [ + 206, + 125, + 0, + 43, + 3, + "IMAGE" + ], + [ + 208, + 138, + 0, + 124, + 0, + "*" + ], + [ + 209, + 124, + 0, + 123, + 1, + "IMAGE" + ], + [ + 210, + 34, + 0, + 110, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.6934334949441466, + "offset": [ + -745.9669584283469, + -262.3877185095017 + ] + }, + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/distributed-upscale.json b/workflows/distributed-upscale.json index 339a7df..044a9f9 100644 --- a/workflows/distributed-upscale.json +++ b/workflows/distributed-upscale.json @@ -1,317 +1,197 @@ { - "id": "0a68c65f-df29-492e-8525-7621fada0ddd", + "id": "817bbfe2-06b8-44c8-8c14-b82b63b335d5", "revision": 0, - "last_node_id": 17, - "last_link_id": 23, + "last_node_id": 137, + "last_link_id": 211, "nodes": [ { - "id": 17, - "type": "ImageScale", + "id": 86, + "type": "Reroute", "pos": [ - 1540, + 2130, 1050 ], "size": [ - 270, - 130 + 75, + 26 ], "flags": {}, - "order": 6, + "order": 11, "mode": 0, "inputs": [ { - "name": "image", - "type": "IMAGE", - "link": 22 + "name": "", + "type": "*", + "link": 121 } ], "outputs": [ { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 23 - ] - } - ], - "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "ImageScale" - }, - "widgets_values": [ - "lanczos", - 2048, - 2048, - "disabled" - ] - }, - { - "id": 13, - "type": "UltimateSDUpscaleDistributed", - "pos": [ - 1890, - 660 - ], - "size": [ - 326.691650390625, - 402 - ], - "flags": {}, - "order": 7, - "mode": 0, - "inputs": [ - { - "name": "upscaled_image", - "type": "IMAGE", - "link": 23 - }, - { - "name": "model", - "type": "MODEL", - "link": 15 - }, - { - "name": "positive", - "type": "CONDITIONING", - "link": 16 - }, - { - "name": "negative", - "type": "CONDITIONING", - "link": 17 - }, - { - "name": "vae", + "name": "", "type": "VAE", - "link": 18 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", "links": [ - 12 + 127 ] } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "aux_id": "robertvoy/ComfyUI-Distributed", - "ver": "99021363d65cc2b2f0f3a0f12a76a358f0fb330f", - "Node name for S&R": "UltimateSDUpscaleDistributed" - }, - "widgets_values": [ - 402634549992051, - "randomize", - 10, - 8, - "euler", - "normal", - 0.24000000000000005, - 512, - 512, - 32, - 16, - true - ] + "showOutputText": false, + "horizontal": false + } }, { - "id": 10, - "type": "PreviewImage", + "id": 88, + "type": "Reroute", "pos": [ - 2260, - 680 + 2130, + 1000 ], "size": [ - 410, - 410 + 75, + 26 ], "flags": {}, "order": 8, "mode": 0, "inputs": [ { - "name": "images", - "type": "IMAGE", - "link": 12 + "name": "", + "type": "*", + "link": 124 + } + ], + "outputs": [ + { + "name": "", + "type": "MODEL", + "links": [ + 137 + ] } ], - "outputs": [], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "PreviewImage" - }, - "widgets_values": [] + "showOutputText": false, + "horizontal": false + } }, { - "id": 5, - "type": "CLIPTextEncode", + "id": 107, + "type": "Reroute", "pos": [ - 1390, - 540 + 2810, + 1050 ], "size": [ - 422.84503173828125, - 164.31304931640625 + 75, + 26 ], "flags": {}, - "order": 4, + "order": 18, "mode": 0, "inputs": [ { - "name": "clip", - "type": "CLIP", - "link": 5 + "name": "", + "type": "*", + "link": 171 } ], "outputs": [ { - "name": "CONDITIONING", - "type": "CONDITIONING", - "slot_index": 0, + "name": "", + "type": "VAE", "links": [ - 16 + 172 ] } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.41", - "Node name for S&R": "CLIPTextEncode" - }, - "widgets_values": [ - "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," - ] + "showOutputText": false, + "horizontal": false + } }, { - "id": 14, - "type": "UpscaleModelLoader", + "id": 89, + "type": "Reroute", "pos": [ - 970, - 950 + 2810, + 1000 ], "size": [ - 270, - 58 + 75, + 26 ], "flags": {}, - "order": 0, + "order": 14, "mode": 0, - "inputs": [], + "inputs": [ + { + "name": "", + "type": "*", + "link": 137 + } + ], "outputs": [ { - "name": "UPSCALE_MODEL", - "type": "UPSCALE_MODEL", + "name": "", + "type": "MODEL", "links": [ - 19 + 170 ] } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "UpscaleModelLoader" - }, - "widgets_values": [ - "4xNomos8kDAT.pth" - ] + "showOutputText": false, + "horizontal": false + } }, { - "id": 12, - "type": "LoadImage", + "id": 110, + "type": "Reroute", "pos": [ - 960, - 1080 + 2810, + 970 ], "size": [ - 274.375, - 314.00006103515625 + 75, + 26 ], "flags": {}, - "order": 1, + "order": 21, "mode": 0, - "inputs": [], + "inputs": [ + { + "name": "", + "type": "*", + "link": 211 + } + ], "outputs": [ { - "name": "IMAGE", + "name": "", "type": "IMAGE", "links": [ - 20 + 176 ] - }, - { - "name": "MASK", - "type": "MASK", - "links": null } ], "properties": { - "enableTabs": false, - "tabWidth": 65, - "tabXOffset": 10, - "hasSecondTab": false, - "secondTabText": "Send Back", - "secondTabOffset": 80, - "secondTabWidth": 65, - "cnr_id": "comfy-core", - "ver": "0.3.43", - "Node name for S&R": "LoadImage" - }, - "widgets_values": [ - "example.png", - "image" - ] + "showOutputText": false, + "horizontal": false + } }, { "id": 7, "type": "CheckpointLoaderSimple", "pos": [ - 960, - 700 + 1700, + 1270 ], "size": [ - 315, + 370, 98 ], "flags": {}, - "order": 2, + "order": 0, "mode": 0, "inputs": [], "outputs": [ @@ -320,7 +200,7 @@ "type": "MODEL", "slot_index": 0, "links": [ - 15 + 124 ] }, { @@ -337,7 +217,7 @@ "type": "VAE", "slot_index": 2, "links": [ - 18 + 121 ] } ], @@ -364,19 +244,95 @@ "SDXL\\juggernautXL_ragnarokBy.safetensors" ] }, + { + "id": 90, + "type": "Reroute", + "pos": [ + 2460, + 1050 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 127 + } + ], + "outputs": [ + { + "name": "", + "type": "VAE", + "links": [ + 129, + 171 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 44, + "type": "ControlNetLoader", + "pos": [ + 1910, + 1510 + ], + "size": [ + 390, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 81 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetLoader" + }, + "widgets_values": [ + "SDXL\\xinsir-controlnet-union-sdxl-1.0-promax.safetensors" + ] + }, { "id": 6, "type": "CLIPTextEncode", "pos": [ - 1390, - 760 + 2130, + 1370 ], "size": [ - 425.27801513671875, - 180.6060791015625 + 420, + 88 ], "flags": {}, - "order": 5, + "order": 10, "mode": 0, "inputs": [ { @@ -391,7 +347,7 @@ "type": "CONDITIONING", "slot_index": 0, "links": [ - 17 + 64 ] } ], @@ -412,29 +368,105 @@ ] }, { - "id": 15, - "type": "ImageUpscaleWithModel", + "id": 52, + "type": "SetUnionControlNetType", "pos": [ - 1280, - 1050 + 2340, + 1510 ], "size": [ - 221.96665954589844, - 46 + 210, + 58 ], "flags": {}, - "order": 3, + "order": 12, "mode": 0, "inputs": [ { - "name": "upscale_model", - "type": "UPSCALE_MODEL", - "link": 19 - }, + "name": "control_net", + "type": "CONTROL_NET", + "link": 81 + } + ], + "outputs": [ + { + "name": "CONTROL_NET", + "type": "CONTROL_NET", + "links": [ + 192 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SetUnionControlNetType" + }, + "widgets_values": [ + "tile" + ] + }, + { + "id": 125, + "type": "Reroute", + "pos": [ + 2450, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 200 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 206 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 34, + "type": "ImageResize+", + "pos": [ + 1700, + 970 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ { "name": "image", "type": "IMAGE", - "link": 20 + "link": 197 } ], "outputs": [ @@ -442,7 +474,73 @@ "name": "IMAGE", "type": "IMAGE", "links": [ - 22 + 211 + ] + }, + { + "name": "width", + "type": "INT", + "links": null + }, + { + "name": "height", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 2048, + 2048, + "lanczos", + "keep proportion", + "always", + 8 + ] + }, + { + "id": 123, + "type": "ImageUpscaleWithModel", + "pos": [ + 1410, + 970 + ], + "size": [ + 222.75416564941406, + 46 + ], + "flags": {}, + "order": 17, + "mode": 4, + "inputs": [ + { + "name": "upscale_model", + "type": "UPSCALE_MODEL", + "link": 195 + }, + { + "name": "image", + "type": "IMAGE", + "link": 210 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 197 ] } ], @@ -457,7 +555,460 @@ "cnr_id": "comfy-core", "ver": "0.3.43", "Node name for S&R": "ImageUpscaleWithModel" + }, + "widgets_values": [] + }, + { + "id": 134, + "type": "MarkdownNote", + "pos": [ + 1920, + 1610 + ], + "size": [ + 370, + 88 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "[https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main](https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main)" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 135, + "type": "LoadImage", + "pos": [ + 780, + 880 + ], + "size": [ + 274.375, + 314.00006103515625 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 207 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example.png", + "image" + ] + }, + { + "id": 132, + "type": "Note", + "pos": [ + 1460, + 1280 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Choose an SDXL model" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + 2130, + 1160 + ], + "size": [ + 420, + 160 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "slot_index": 0, + "links": [ + 63 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.41", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 124, + "type": "Reroute", + "pos": [ + 1230, + 880 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 207 + } + ], + "outputs": [ + { + "name": "", + "type": "IMAGE", + "links": [ + 200, + 210 + ] + } + ], + "properties": { + "showOutputText": false, + "horizontal": false } + }, + { + "id": 122, + "type": "UpscaleModelLoader", + "pos": [ + 1130, + 1070 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "UPSCALE_MODEL", + "type": "UPSCALE_MODEL", + "links": [ + 195 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "UpscaleModelLoader" + }, + "widgets_values": [ + "4xNomos8kDAT.pth" + ] + }, + { + "id": 131, + "type": "Note", + "pos": [ + 1420, + 1070 + ], + "size": [ + 210, + 88 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Optional" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 137, + "type": "Note", + "pos": [ + 3000, + 1600 + ], + "size": [ + 330, + 90 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "If all your GPUs are the same/similar, set static_distribution to true\n" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 136, + "type": "PreviewImage", + "pos": [ + 3380, + 1100 + ], + "size": [ + 490, + 550 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 209 + } + ], + "outputs": [], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + }, + { + "id": 43, + "type": "ControlNetApplyAdvanced", + "pos": [ + 2640, + 1150 + ], + "size": [ + 270, + 186 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 63 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 64 + }, + { + "name": "control_net", + "type": "CONTROL_NET", + "link": 192 + }, + { + "name": "image", + "type": "IMAGE", + "link": 206 + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": 129 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 190 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": [ + 191 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ControlNetApplyAdvanced" + }, + "widgets_values": [ + 1.0000000000000002, + 0, + 0.8000000000000002 + ] + }, + { + "id": 30, + "type": "UltimateSDUpscaleDistributed", + "pos": [ + 3000, + 1110 + ], + "size": [ + 326.691650390625, + 450 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "upscaled_image", + "type": "IMAGE", + "link": 176 + }, + { + "name": "model", + "type": "MODEL", + "link": 170 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 190 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 191 + }, + { + "name": "vae", + "type": "VAE", + "link": 172 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 209 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "dd23503883fdf319e8beb6e7a190445ecf89973c", + "Node name for S&R": "UltimateSDUpscaleDistributed" + }, + "widgets_values": [ + 269777990474642, + "randomize", + 20, + 7, + "dpmpp_2m_sde", + "karras", + 0.6000000000000001, + 1024, + 1024, + 32, + 16, + true, + false, + true + ] } ], "links": [ @@ -478,89 +1029,205 @@ "CLIP" ], [ - 12, - 13, - 0, - 10, - 0, - "IMAGE" - ], - [ - 15, - 7, - 0, - 13, - 1, - "MODEL" - ], - [ - 16, + 63, 5, 0, - 13, - 2, + 43, + 0, "CONDITIONING" ], [ - 17, + 64, 6, 0, - 13, - 3, + 43, + 1, "CONDITIONING" ], [ - 18, + 81, + 44, + 0, + 52, + 0, + "CONTROL_NET" + ], + [ + 121, 7, 2, - 13, + 86, + 0, + "*" + ], + [ + 124, + 7, + 0, + 88, + 0, + "*" + ], + [ + 127, + 86, + 0, + 90, + 0, + "*" + ], + [ + 129, + 90, + 0, + 43, 4, "VAE" ], [ - 19, - 14, + 137, + 88, 0, - 15, + 89, + 0, + "*" + ], + [ + 170, + 89, + 0, + 30, + 1, + "MODEL" + ], + [ + 171, + 90, + 0, + 107, + 0, + "*" + ], + [ + 172, + 107, + 0, + 30, + 4, + "VAE" + ], + [ + 176, + 110, + 0, + 30, + 0, + "IMAGE" + ], + [ + 190, + 43, + 0, + 30, + 2, + "CONDITIONING" + ], + [ + 191, + 43, + 1, + 30, + 3, + "CONDITIONING" + ], + [ + 192, + 52, + 0, + 43, + 2, + "CONTROL_NET" + ], + [ + 195, + 122, + 0, + 123, 0, "UPSCALE_MODEL" ], [ - 20, - 12, + 197, + 123, 0, - 15, + 34, + 0, + "IMAGE" + ], + [ + 200, + 124, + 0, + 125, + 0, + "*" + ], + [ + 206, + 125, + 0, + 43, + 3, + "IMAGE" + ], + [ + 207, + 135, + 0, + 124, + 0, + "*" + ], + [ + 209, + 30, + 0, + 136, + 0, + "IMAGE" + ], + [ + 210, + 124, + 0, + 123, 1, "IMAGE" ], [ - 22, - 15, + 211, + 34, 0, - 17, + 110, 0, - "IMAGE" - ], - [ - 23, - 17, - 0, - 13, - 0, - "IMAGE" + "*" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.6649272177973091, + "scale": 1.0152559799477252, "offset": [ - -510.50184649377775, - -337.73126769987704 + -2260.53316345765, + -499.7179536588252 ] }, - "frontendVersion": "1.23.4" + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true }, "version": 0.4 } \ No newline at end of file diff --git a/workflows/distributed-wan.json b/workflows/distributed-wan.json new file mode 100644 index 0000000..3e2faf0 --- /dev/null +++ b/workflows/distributed-wan.json @@ -0,0 +1,1688 @@ +{ + "id": "00000000-0000-0000-0000-000000000000", + "revision": 0, + "last_node_id": 234, + "last_link_id": 79, + "nodes": [ + { + "id": 67, + "type": "ModelSamplingSD3", + "pos": [ + 1700, + 550 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 23 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 17, + 26 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "ModelSamplingSD3" + }, + "widgets_values": [ + 8.000000000000002 + ] + }, + { + "id": 50, + "type": "WanImageToVideo", + "pos": [ + 1590, + 230 + ], + "size": [ + 270, + 210 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "link": 7 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 8 + }, + { + "name": "vae", + "type": "VAE", + "link": 9 + }, + { + "name": "clip_vision_output", + "shape": 7, + "type": "CLIP_VISION_OUTPUT", + "link": 10 + }, + { + "name": "start_image", + "shape": 7, + "type": "IMAGE", + "link": 11 + }, + { + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": 5 + }, + { + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": 6 + } + ], + "outputs": [ + { + "name": "positive", + "type": "CONDITIONING", + "links": [ + 27 + ] + }, + { + "name": "negative", + "type": "CONDITIONING", + "links": null + }, + { + "name": "latent", + "type": "LATENT", + "links": [ + 22 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "WanImageToVideo" + }, + "widgets_values": [ + 832, + 480, + 17, + 1 + ] + }, + { + "id": 60, + "type": "RandomNoise", + "pos": [ + 1850, + 10 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "noise_seed", + "type": "INT", + "widget": { + "name": "noise_seed" + }, + "link": 16 + } + ], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 18 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 200526686850175, + "randomize" + ] + }, + { + "id": 150, + "type": "DistributedSeed", + "pos": [ + 1540, + 0 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "seed", + "type": "INT", + "links": [ + 16 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "9650280c7f50898720bf9773e8874d75e7c20846", + "Node name for S&R": "DistributedSeed" + }, + "widgets_values": [ + 177896884005433, + "randomize" + ] + }, + { + "id": 61, + "type": "KSamplerSelect", + "pos": [ + 2140, + 350 + ], + "size": [ + 260, + 60 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 20 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "lcm" + ] + }, + { + "id": 221, + "type": "DistributedCollector", + "pos": [ + 2660, + 200 + ], + "size": [ + 166.50416564941406, + 26 + ], + "flags": { + "collapsed": false + }, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 67 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 78 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "b4ed92a363fab14d944b87c610aa9a0b4c87c085", + "Node name for S&R": "DistributedCollector" + }, + "widgets_values": [] + }, + { + "id": 231, + "type": "ImageBatchDivider", + "pos": [ + 2850, + 200 + ], + "size": [ + 210, + 118 + ], + "flags": {}, + "order": 23, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 78 + } + ], + "outputs": [ + { + "name": "batch_1", + "type": "IMAGE", + "links": [ + 74 + ] + }, + { + "name": "batch_2", + "type": "IMAGE", + "links": [ + 75 + ] + }, + { + "name": "batch_3", + "type": "IMAGE", + "links": [ + 76 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "ComfyUI-Distributed", + "ver": "fa09a2da10dfc63cecfbec1f8d4a0516e923b911", + "Node name for S&R": "ImageBatchDivider" + }, + "widgets_values": [ + 3 + ] + }, + { + "id": 139, + "type": "CLIPLoaderGGUF", + "pos": [ + 490, + 60 + ], + "size": [ + 270, + 82 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 1, + 2 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-gguf", + "ver": "1.1.1", + "Node name for S&R": "CLIPLoaderGGUF" + }, + "widgets_values": [ + "umt5-xxl-encoder-Q8_0.gguf", + "wan" + ] + }, + { + "id": 142, + "type": "UnetLoaderGGUF", + "pos": [ + 930, + 850 + ], + "size": [ + 340, + 60 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 24 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-gguf", + "ver": "1.1.1", + "Node name for S&R": "UnetLoaderGGUF" + }, + "widgets_values": [ + "wan2.1-i2v-14b-720p-Q8_0.gguf" + ] + }, + { + "id": 118, + "type": "LoraLoaderModelOnly", + "pos": [ + 1290, + 840 + ], + "size": [ + 310, + 82 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 24 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 23 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "LoraLoaderModelOnly" + }, + "widgets_values": [ + "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors", + 0.8100000000000002 + ] + }, + { + "id": 135, + "type": "ImageNoiseAugmentation", + "pos": [ + 1250, + 310 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 25 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 11 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-kjnodes", + "ver": "1.1.2", + "Node name for S&R": "ImageNoiseAugmentation" + }, + "widgets_values": [ + 0.10000000000000002, + 1100164865582526, + "randomize" + ] + }, + { + "id": 49, + "type": "CLIPVisionLoader", + "pos": [ + 1260, + 610 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP_VISION", + "type": "CLIP_VISION", + "links": [ + 12 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPVisionLoader" + }, + "widgets_values": [ + "clip_vision_h.safetensors" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": [ + 790, + 120 + ], + "size": [ + 380, + 110 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 2 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 8 + ] + } + ], + "title": "CLIP Text Encode (Negative Prompt)", + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" + ] + }, + { + "id": 64, + "type": "SamplerCustomAdvanced", + "pos": [ + 2140, + 190 + ], + "size": [ + 260, + 110 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 18 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 19 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 20 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 21 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 22 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 3 + ] + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "SamplerCustomAdvanced" + }, + "widgets_values": [] + }, + { + "id": 233, + "type": "Note", + "pos": [ + 2840, + 40 + ], + "size": [ + 220, + 90 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Use this node to set how many videos to output.\n\nExample: if you have 1x master and 2x workers, set it to 3." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 234, + "type": "Note", + "pos": [ + 1540, + -140 + ], + "size": [ + 270, + 90 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "This node will give you a different variation from each worker. Delete this node if you want them to be all the same." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 227, + "type": "VHS_VideoCombine", + "pos": [ + 3090, + 200 + ], + "size": [ + 380, + 555.6923217773438 + ], + "flags": {}, + "order": 24, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 74 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WAN", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00163.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8, + "workflow": "AnimateDiff_00163.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00163.mp4" + } + } + } + }, + { + "id": 228, + "type": "VHS_VideoCombine", + "pos": [ + 3490, + 190 + ], + "size": [ + 380, + 334 + ], + "flags": {}, + "order": 25, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 75 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WAN", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00165.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8, + "workflow": "AnimateDiff_00165.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00165.mp4" + } + } + } + }, + { + "id": 229, + "type": "VHS_VideoCombine", + "pos": [ + 3880, + 190 + ], + "size": [ + 390, + 334 + ], + "flags": {}, + "order": 26, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 76 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui-videohelpersuite", + "ver": "a7ce59e381934733bfae03b1be029756d6ce936d", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "WAN", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00164.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 8, + "workflow": "AnimateDiff_00164.png", + "fullpath": "C:\\venvs\\ComfyUI\\ComfyUI\\output\\AnimateDiff_00164.mp4" + } + } + } + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 2450, + 190 + ], + "size": [ + 140, + 46 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 3 + }, + { + "name": "vae", + "type": "VAE", + "link": 4 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 67 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 51, + "type": "CLIPVisionEncode", + "pos": [ + 1260, + 490 + ], + "size": [ + 290.97918701171875, + 78 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "clip_vision", + "type": "CLIP_VISION", + "link": 12 + }, + { + "name": "image", + "type": "IMAGE", + "link": 13 + } + ], + "outputs": [ + { + "name": "CLIP_VISION_OUTPUT", + "type": "CLIP_VISION_OUTPUT", + "links": [ + 10 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPVisionEncode" + }, + "widgets_values": [ + "none" + ] + }, + { + "id": 192, + "type": "LoadImage", + "pos": [ + 530, + 310 + ], + "size": [ + 310, + 370 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 79 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example.png", + "image" + ] + }, + { + "id": 55, + "type": "ImageResize+", + "pos": [ + 890, + 320 + ], + "size": [ + 270, + 218 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 79 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 13, + 25 + ] + }, + { + "name": "width", + "type": "INT", + "links": [ + 5 + ] + }, + { + "name": "height", + "type": "INT", + "links": [ + 6 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfyui_essentials", + "ver": "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9", + "Node name for S&R": "ImageResize+" + }, + "widgets_values": [ + 832, + 480, + "lanczos", + "fill / crop", + "always", + 0 + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 790, + -10 + ], + "size": [ + 380, + 88 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 1 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 7 + ] + } + ], + "title": "CLIP Text Encode (Positive Prompt)", + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "beautiful scenery nature glass bottle landscape, , purple galaxy bottle," + ] + }, + { + "id": 39, + "type": "VAELoader", + "pos": [ + 1250, + 200 + ], + "size": [ + 270, + 58 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 4, + 9 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "wan_2.1_vae.safetensors" + ] + }, + { + "id": 148, + "type": "BasicGuider", + "pos": [ + 1980, + 260 + ], + "size": [ + 156.0208282470703, + 46 + ], + "flags": { + "collapsed": true + }, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 26 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 27 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 19 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "BasicGuider" + }, + "widgets_values": [] + }, + { + "id": 62, + "type": "BasicScheduler", + "pos": [ + 2130, + 460 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 17 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 21 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "cnr_id": "comfy-core", + "ver": "0.3.43", + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 4, + 1 + ] + } + ], + "links": [ + [ + 1, + 139, + 0, + 6, + 0, + "CLIP" + ], + [ + 2, + 139, + 0, + 7, + 0, + "CLIP" + ], + [ + 3, + 64, + 0, + 8, + 0, + "LATENT" + ], + [ + 4, + 39, + 0, + 8, + 1, + "VAE" + ], + [ + 5, + 55, + 1, + 50, + 5, + "INT" + ], + [ + 6, + 55, + 2, + 50, + 6, + "INT" + ], + [ + 7, + 6, + 0, + 50, + 0, + "CONDITIONING" + ], + [ + 8, + 7, + 0, + 50, + 1, + "CONDITIONING" + ], + [ + 9, + 39, + 0, + 50, + 2, + "VAE" + ], + [ + 10, + 51, + 0, + 50, + 3, + "CLIP_VISION_OUTPUT" + ], + [ + 11, + 135, + 0, + 50, + 4, + "IMAGE" + ], + [ + 12, + 49, + 0, + 51, + 0, + "CLIP_VISION" + ], + [ + 13, + 55, + 0, + 51, + 1, + "IMAGE" + ], + [ + 16, + 150, + 0, + 60, + 0, + "INT" + ], + [ + 17, + 67, + 0, + 62, + 0, + "MODEL" + ], + [ + 18, + 60, + 0, + 64, + 0, + "NOISE" + ], + [ + 19, + 148, + 0, + 64, + 1, + "GUIDER" + ], + [ + 20, + 61, + 0, + 64, + 2, + "SAMPLER" + ], + [ + 21, + 62, + 0, + 64, + 3, + "SIGMAS" + ], + [ + 22, + 50, + 2, + 64, + 4, + "LATENT" + ], + [ + 23, + 118, + 0, + 67, + 0, + "MODEL" + ], + [ + 24, + 142, + 0, + 118, + 0, + "MODEL" + ], + [ + 25, + 55, + 0, + 135, + 0, + "IMAGE" + ], + [ + 26, + 67, + 0, + 148, + 0, + "MODEL" + ], + [ + 27, + 50, + 0, + 148, + 1, + "CONDITIONING" + ], + [ + 67, + 8, + 0, + 221, + 0, + "IMAGE" + ], + [ + 74, + 231, + 0, + 227, + 0, + "IMAGE" + ], + [ + 75, + 231, + 1, + 228, + 0, + "IMAGE" + ], + [ + 76, + 231, + 2, + 229, + 0, + "IMAGE" + ], + [ + 78, + 221, + 0, + 231, + 0, + "IMAGE" + ], + [ + 79, + 192, + 0, + 55, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.9090909090909091, + "offset": [ + -1587.6131275248988, + 280.3225523752761 + ] + }, + "frontendVersion": "1.23.4", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file From a045beb71ae63b49e7e29b07ed72332274f45ae4 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 12:55:04 +1000 Subject: [PATCH 57/93] Update README.md --- README.md | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index f1df4a7..3a8cf99 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@

📺Video Tutorial | -📝Setup Guides | +📝Setup Guides | ⚙️Workflows | 🎁Donation

@@ -18,19 +18,16 @@ #### Parallel Workflow Processing - Run workflows on multiple GPUs simultaneously with varied seeds -- Generate multiple outputs in the time of one - Scale output with more workers - Supports images and videos #### Distributed Upscaling - Accelerate Ultimate SD Upscale by distributing tiles across GPUs -- Intelligent distribution for asymmetric setups -- Static distribution for similar GPUs +- Intelligent distribution - Handles single images and batches #### Convenience - Auto-setup local workers; easily add remote/cloud ones -- Manage workers via UI panel - Convert any workflow to distributed with 2 nodes - JSON configuration with UI controls @@ -54,6 +51,8 @@ 2. **Restart ComfyUI** - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master +3. Read the [setup guide](https://github.com/robertvoy/ComfyUI-Distributed/blob/feature/cloud-workers/docs/worker-setup-guides.md) for adding workers + --- ## Workflow Examples From 84d42fd2d33291a7456ad5eb0a7cc5b3affc4a4e Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 12:58:41 +1000 Subject: [PATCH 58/93] Update README.md --- README.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 3a8cf99..1aacc0b 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@

📺Video Tutorial | 📝Setup Guides | -⚙️Workflows | +⚙️Workflows | 🎁Donation

@@ -26,7 +26,7 @@ - Intelligent distribution - Handles single images and batches -#### Convenience +#### Ease of Use - Auto-setup local workers; easily add remote/cloud ones - Convert any workflow to distributed with 2 nodes - JSON configuration with UI controls @@ -60,6 +60,7 @@ ### Basic Parallel Generation ![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) +> [Download workflow](/workflows/distributed-txt2img.json) 1. Create your normal ComfyUI workflow 2. Add **Distributed Seed** → connect to sampler's seed @@ -69,10 +70,14 @@ ### Parallel WAN Generation +> [Download workflow](/workflows/distributed-wan.json) +> ### Distributed Upscaling ![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) +> [Download workflow](/workflows/distributed-txt2img.json) + 1. Load your image 2. Upscale with ESRGAN or similar 3. Connect to **Ultimate SD Upscale Distributed** From 24115742d4c2b9916a1b7e457418852d163244fe Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 12:59:25 +1000 Subject: [PATCH 59/93] Update README.md --- README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index 1aacc0b..8a404bf 100644 --- a/README.md +++ b/README.md @@ -58,9 +58,8 @@ ## Workflow Examples ### Basic Parallel Generation - -![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) > [Download workflow](/workflows/distributed-txt2img.json) +![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) 1. Create your normal ComfyUI workflow 2. Add **Distributed Seed** → connect to sampler's seed From b1d619f38b6d3c0d3774ffca5d5788aff94d9592 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 13:01:14 +1000 Subject: [PATCH 60/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 8a404bf..afdbe26 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ ## Requirements - ComfyUI -> Note: Desktop app not supported currently +> Note: Desktop app not currently supported - Multiple GPUs - That's it From d55e9c50f9c0ff93f88b25e9454b573b1f3b4a0f Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 13:02:28 +1000 Subject: [PATCH 61/93] Update README.md --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index afdbe26..479ebf7 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ ```bash git clone https://github.com/robertvoy/ComfyUI-Distributed.git ``` - + 2. **Restart ComfyUI** - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master 3. Read the [setup guide](https://github.com/robertvoy/ComfyUI-Distributed/blob/feature/cloud-workers/docs/worker-setup-guides.md) for adding workers @@ -58,8 +58,9 @@ ## Workflow Examples ### Basic Parallel Generation -> [Download workflow](/workflows/distributed-txt2img.json) + ![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) +> [Download workflow](/workflows/distributed-txt2img.json) 1. Create your normal ComfyUI workflow 2. Add **Distributed Seed** → connect to sampler's seed From 01296fe78aaea209bd1d09181faf4bc072613f8f Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 13:04:00 +1000 Subject: [PATCH 62/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 479ebf7..eb6edf6 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,7 @@ Yes, you can combine different GPUs, but performance is optimized when using sim
Does this work with cloud providers? -Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions. +Yes, it is compatible with cloud providers. Refer to the [setup guides](https://github.com/robertvoy/ComfyUI-Distributed/blob/feature/cloud-workers/docs/worker-setup-guides.md) for detailed instructions.
From e6c0927e2caa0920abaef9da4c995208e4616a33 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 14:28:32 +1000 Subject: [PATCH 63/93] Update model-download-script.md --- docs/model-download-script.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index 4c41511..883b04c 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -1,4 +1,4 @@ -## Prompt for LLM - Automating ComfyUI Model Downloads +## Automating ComfyUI Model Downloads This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM). @@ -40,4 +40,4 @@ upscale_models: models/upscale_models/ vae: models/vae/ vae_approx: models/vae_approx/ vitmatte: models/vitmatte/ -``` \ No newline at end of file +``` From 4e7736594617a4c08ca4c03ef094513e43ef49ed Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 15:43:19 +1000 Subject: [PATCH 64/93] Update README.md --- README.md | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/README.md b/README.md index eb6edf6..e215fba 100644 --- a/README.md +++ b/README.md @@ -122,11 +122,7 @@ Yes, it is compatible with Docker setups, but you will need to configure your Do ## Disclaimer -By using this software, you agree to be bound by the terms of this Disclaimer. The software is provided “as is” without any warranties, express or implied, including but not limited to merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities—whether in contract, tort, or otherwise—arising from the use, modification, distribution, or other dealings with the software. - -Users are solely responsible for ensuring their use of the software, including any content it generates, complies with all applicable laws and regulations in their jurisdiction. The developers are not responsible for any legal violations resulting from user actions. - -Additionally, users are responsible for the security and integrity of their networks, including protecting against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages, losses, or incidents arising from improper configuration, misuse, or external threats to user environments. +This software is provided "as is" without any warranties, express or implied, including merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities arising from the use, modification, or distribution of the software. Users are solely responsible for ensuring compliance with applicable laws and regulations and for securing their networks against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages or incidents resulting from misuse, improper configuration, or external threats. --- From 28db418465f8e936e80eed84a2b1386010da4538 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Mon, 28 Jul 2025 15:48:21 +1000 Subject: [PATCH 65/93] docs fix --- docs/worker-setup-guides.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index 2dbde7f..eb95cc7 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -50,7 +50,7 @@ **On the Main Machine:** 1. Launch a Cloudflare tunnel - - Download from here `https://github.com/cloudflare/cloudflared/releases/tag/2025.7.0` + - Download from here `https://github.com/cloudflare/cloudflared/releases` - Then run for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` 2. Copy the Cloudflare address 3. **Open** the Distributed GPU panel (sidebar on the left). @@ -75,8 +75,8 @@ **On the Main Machine:** 1. **Launch** a Cloudflare tunnel on your local machine. - - Download from here `https://github.com/cloudflare/cloudflared/releases/tag/2025.7.0` - - For example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` + - Download from here `https://github.com/cloudflare/cloudflared/releases` + - For example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` 2. **Copy** the Cloudflare address 3. **Open** the Distributed GPU panel (sidebar on the left). 4. **Edit** the Master's host address and replace it with the Cloudflare address. From 6349d0b7f68373968214856e9cc4c14e838222b5 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:01:32 +1000 Subject: [PATCH 66/93] Update worker-setup-guides.md --- docs/worker-setup-guides.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index eb95cc7..b4633e3 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -62,6 +62,8 @@ - **Port**: 443 - **Save** the remote worker configuration. +--- + ### Deploy Cloud Worker on Other Platforms **On the Cloud Worker machine:** @@ -85,4 +87,4 @@ 7. **Configure** your cloud worker: - **Host**: The remote worker's IP address/domain - **Port**: 443 -8. **Save** the remote worker configuration. \ No newline at end of file +8. **Save** the remote worker configuration. From 1458d24cc324b9bff8f5effda20d467b816e2f97 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:03:37 +1000 Subject: [PATCH 67/93] Update worker-setup-guides.md --- docs/worker-setup-guides.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index b4633e3..4f37f46 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -39,10 +39,10 @@ 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account. 2. On Runpod, go to Storage > New Network Volume and create a volume which will store the models you need. 3. Now go to Pods and find a suitable GPU for your workflows. -> To use the ComfyUI-Distributed-Pod, you will need to filter instances by CUDA 12.8. +> To use the ComfyUI-Distributed-Pod template, you will need to filter instances by CUDA 12.8. 4. Choose the `ComfyUI-Distributed-Pod` template and make sure your network drive is mounted. 5. Launch your pod. -6. Access your pod using JupyterLab. +6. Access your pod using JupyterLabs. 7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). > You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow. 8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!** From c22e843493a8fb55de50c2aff7e5ab62469280bc Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:24:28 +1000 Subject: [PATCH 68/93] Update README.md --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index e215fba..3867312 100644 --- a/README.md +++ b/README.md @@ -70,8 +70,9 @@ ### Parallel WAN Generation +![Clipboard Image](https://github.com/user-attachments/assets/794c3928-12c1-4522-82fb-f6a022b9e9ab) > [Download workflow](/workflows/distributed-wan.json) -> + ### Distributed Upscaling ![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) From 4a6f5298884e437dc990e6fbca83cdac82770ee5 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:26:35 +1000 Subject: [PATCH 69/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 3867312..3500880 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ ### Parallel WAN Generation -![Clipboard Image](https://github.com/user-attachments/assets/794c3928-12c1-4522-82fb-f6a022b9e9ab) +![Clipboard Image (1)](https://github.com/user-attachments/assets/de2ee023-4650-4164-885e-404b3cedd67e) > [Download workflow](/workflows/distributed-wan.json) ### Distributed Upscaling From fb4c70dfebce90201ea8fd25971f53b027f156d8 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:29:50 +1000 Subject: [PATCH 70/93] Update README.md --- README.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 3500880..fad4a21 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,7 @@ ![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) > [Download workflow](/workflows/distributed-txt2img.json) -1. Create your normal ComfyUI workflow +1. Open your ComfyUI workflow 2. Add **Distributed Seed** → connect to sampler's seed 3. Add **Distributed Collector** → after VAE Decode 4. Enable workers in the UI @@ -73,6 +73,15 @@ ![Clipboard Image (1)](https://github.com/user-attachments/assets/de2ee023-4650-4164-885e-404b3cedd67e) > [Download workflow](/workflows/distributed-wan.json) +1. Open your WAN ComfyUI workflow +2. Add **Distributed Seed** → connect to sampler's seed +3. Add **Distributed Collector** → after VAE Decode +4. Add **Image Batch Divider** → after Distributed Collector +5. Set the `divide_by` to the number of GPUs you have available +> For example: if you have a master and 2x workers, set it to 3 +7. Enable workers in the UI +8. Run the workflow! + ### Distributed Upscaling ![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) From 74d5daf353d3ec835730c552aa3335a11b414349 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:31:19 +1000 Subject: [PATCH 71/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index fad4a21..f565326 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ ### Parallel WAN Generation -![Clipboard Image (1)](https://github.com/user-attachments/assets/de2ee023-4650-4164-885e-404b3cedd67e) +![Clipboard Image (2)](https://github.com/user-attachments/assets/ba7536e7-675c-44d7-9f3f-bf11a2216cd6) > [Download workflow](/workflows/distributed-wan.json) 1. Open your WAN ComfyUI workflow From 94e2419decb2560b1b0aa16a273758bd48c7e610 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:44:42 +1000 Subject: [PATCH 72/93] Update README.md --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index f565326..a3b5183 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,8 @@ ### Parallel WAN Generation -![Clipboard Image (2)](https://github.com/user-attachments/assets/ba7536e7-675c-44d7-9f3f-bf11a2216cd6) +![Clipboard Image (5)](https://github.com/user-attachments/assets/5382b845-833b-43b7-b238-a91c5579581a) + > [Download workflow](/workflows/distributed-wan.json) 1. Open your WAN ComfyUI workflow @@ -84,7 +85,7 @@ ### Distributed Upscaling -![Clipboard Image (2)](https://github.com/user-attachments/assets/ec2548d0-1fc7-4705-801f-3270d720cfce) +![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a) > [Download workflow](/workflows/distributed-txt2img.json) From bcf992acebb23a9b31f2f0c9ef1a7cfd970e34df Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:46:53 +1000 Subject: [PATCH 73/93] Update README.md --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a3b5183..726cf01 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ > **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network.** -![Clipboard Image (3)](https://github.com/user-attachments/assets/19fdd1be-8f6e-4df5-bcd9-538ef566fa82) +![Clipboard Image (7)](https://github.com/user-attachments/assets/66aaadef-f195-48a1-a368-17dd0dae477d) --- @@ -59,7 +59,8 @@ ### Basic Parallel Generation -![Clipboard Image (1)](https://github.com/user-attachments/assets/e8e46d97-d698-4c18-b4e5-1e1a2f4f7da3) +![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb) + > [Download workflow](/workflows/distributed-txt2img.json) 1. Open your ComfyUI workflow From 7d123fd924ce7c70a4de84fd55477b9785a209a1 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:52:31 +1000 Subject: [PATCH 74/93] Update README.md --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 726cf01..7427a27 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ ## Workflow Examples ### Basic Parallel Generation - +> Generate multiple images in the time it takes to generate one. Each worker uses a different seed. ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb) > [Download workflow](/workflows/distributed-txt2img.json) @@ -70,7 +70,7 @@ 5. Run the workflow! ### Parallel WAN Generation - +> Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. ![Clipboard Image (5)](https://github.com/user-attachments/assets/5382b845-833b-43b7-b238-a91c5579581a) > [Download workflow](/workflows/distributed-wan.json) @@ -85,7 +85,7 @@ 8. Run the workflow! ### Distributed Upscaling - +> Distributed tiled upscaling: Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs. ![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a) > [Download workflow](/workflows/distributed-txt2img.json) From 56f7275096490e395d62321c50390c203a9d893d Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:52:58 +1000 Subject: [PATCH 75/93] Update README.md --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 7427a27..16bee48 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ ## Workflow Examples ### Basic Parallel Generation -> Generate multiple images in the time it takes to generate one. Each worker uses a different seed. +Generate multiple images in the time it takes to generate one. Each worker uses a different seed. ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb) > [Download workflow](/workflows/distributed-txt2img.json) @@ -70,7 +70,7 @@ 5. Run the workflow! ### Parallel WAN Generation -> Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. +Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. ![Clipboard Image (5)](https://github.com/user-attachments/assets/5382b845-833b-43b7-b238-a91c5579581a) > [Download workflow](/workflows/distributed-wan.json) @@ -85,7 +85,7 @@ 8. Run the workflow! ### Distributed Upscaling -> Distributed tiled upscaling: Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs. +Distributed tiled upscaling: Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs. ![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a) > [Download workflow](/workflows/distributed-txt2img.json) From 923234d069cb87d136c73248827885f693a24752 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:53:11 +1000 Subject: [PATCH 76/93] Update README.md --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 16bee48..2e14370 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,7 @@ ### Basic Parallel Generation Generate multiple images in the time it takes to generate one. Each worker uses a different seed. + ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb) > [Download workflow](/workflows/distributed-txt2img.json) From 1b74f854e1b762253f454b2d2e5f5403a98bb64c Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:53:24 +1000 Subject: [PATCH 77/93] Update README.md --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 2e14370..0a8aad9 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,7 @@ Generate multiple images in the time it takes to generate one. Each worker uses ### Parallel WAN Generation Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. + ![Clipboard Image (5)](https://github.com/user-attachments/assets/5382b845-833b-43b7-b238-a91c5579581a) > [Download workflow](/workflows/distributed-wan.json) @@ -87,6 +88,7 @@ Generate multiple videos in the time it takes to generate one. Each worker uses ### Distributed Upscaling Distributed tiled upscaling: Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs. + ![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a) > [Download workflow](/workflows/distributed-txt2img.json) From 75e04028e02f40c47351500405c9754d9c51d431 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:53:46 +1000 Subject: [PATCH 78/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 0a8aad9..37d0a86 100644 --- a/README.md +++ b/README.md @@ -87,7 +87,7 @@ Generate multiple videos in the time it takes to generate one. Each worker uses 8. Run the workflow! ### Distributed Upscaling -Distributed tiled upscaling: Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs. +Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs. ![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a) From 1263ee13e2a97111729d4d57667da5f30072c2ab Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:55:05 +1000 Subject: [PATCH 79/93] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 37d0a86..d8c434e 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ ## Workflow Examples ### Basic Parallel Generation -Generate multiple images in the time it takes to generate one. Each worker uses a different seed. +Generate multiple images in the time it takes to generate one. Each worker uses a different seed. Works with any workflow that outputs images. ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb) @@ -71,7 +71,7 @@ Generate multiple images in the time it takes to generate one. Each worker uses 5. Run the workflow! ### Parallel WAN Generation -Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. +Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. Works with any workflow that outputs images, including video frames. ![Clipboard Image (5)](https://github.com/user-attachments/assets/5382b845-833b-43b7-b238-a91c5579581a) From 21a79df092c268e856e38c66a585a944a952d378 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 16:55:21 +1000 Subject: [PATCH 80/93] Update README.md --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index d8c434e..37d0a86 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ ## Workflow Examples ### Basic Parallel Generation -Generate multiple images in the time it takes to generate one. Each worker uses a different seed. Works with any workflow that outputs images. +Generate multiple images in the time it takes to generate one. Each worker uses a different seed. ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb) @@ -71,7 +71,7 @@ Generate multiple images in the time it takes to generate one. Each worker uses 5. Run the workflow! ### Parallel WAN Generation -Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. Works with any workflow that outputs images, including video frames. +Generate multiple videos in the time it takes to generate one. Each worker uses a different seed. ![Clipboard Image (5)](https://github.com/user-attachments/assets/5382b845-833b-43b7-b238-a91c5579581a) From e80a8b56d9328ac6c74c54b39c1efaee1733ccf5 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 17:47:46 +1000 Subject: [PATCH 81/93] Update README.md --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 37d0a86..854b9cd 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,10 @@ From 425329ba0923bea4ff1eb2e0e80f892086d905f7 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 28 Jul 2025 17:56:49 +1000 Subject: [PATCH 82/93] Update README.md --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 854b9cd..68914a6 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ Setup Guides Workflows Donation +Twitter

From 56f4f2abc5ad30c4b778ea60d307934cbd47e2b3 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 07:52:04 +1000 Subject: [PATCH 83/93] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 68914a6..774d3f5 100644 --- a/README.md +++ b/README.md @@ -126,7 +126,7 @@ Yes, you can combine different GPUs, but performance is optimized when using sim
Does this work with cloud providers? -Yes, it is compatible with cloud providers. Refer to the [setup guides](https://github.com/robertvoy/ComfyUI-Distributed/blob/feature/cloud-workers/docs/worker-setup-guides.md) for detailed instructions. +Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions.
From 4396591d131cdbe21bbd197ade9eaafb74f8427c Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 09:07:14 +1000 Subject: [PATCH 84/93] Update model-download-script.md --- docs/model-download-script.md | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index 883b04c..6893a57 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -3,7 +3,7 @@ This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM). 1. Export your workflow as an API workflow -2. Copy the below prompt and upload the API workflow to a LLM that has access to the internet (like Gemini Flash) +2. Copy the below prompt and upload the API workflow to a LLM that has access to the internet 3. Review (make sure all downloads are correct) and save the output from the LLM as a .sh file 4. Upload it onto your Runpod instance 5. Run the file. For example: `./download_comfyui_models.sh` @@ -12,32 +12,19 @@ This guide will walk you through creating a shell script to automatically downlo Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: comfyui: base_path: /workspace/ComfyUI -BiRefNet: models/BiRefNet/ checkpoints: models/checkpoints/ clip: models/clip/ clip_vision: models/clip_vision/ -configs: models/configs/ controlnet: models/controlnet/ -diffusers: models/diffusers/ diffusion_models: models/diffusion_models/ embeddings: models/embeddings/ florence2: models/florence2/ -facerestore_models: models/facerestore_models/ -gligen: models/gligen/ -grounding-dino: models/grounding-dino/ -hypernetworks: models/hypernetworks/ ipadapter: models/ipadapter/ -lama: models/lama/ loras: models/loras/ -onnx: models/onnx/ -photomaker: models/photomaker/ -RMBG: models/RMBG/ -sams: models/sams/ style_models: models/style_models/ text_encoders: models/text_encoders/ unet: models/unet/ upscale_models: models/upscale_models/ vae: models/vae/ -vae_approx: models/vae_approx/ -vitmatte: models/vitmatte/ +Make sure you find the correct URLs for the models online. ``` From ec08dfb1996a87a8ea6ffbf828bce5efa60481b2 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 09:14:46 +1000 Subject: [PATCH 85/93] Update model-download-script.md --- docs/model-download-script.md | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index 6893a57..63c4bc4 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -2,11 +2,12 @@ This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM). -1. Export your workflow as an API workflow -2. Copy the below prompt and upload the API workflow to a LLM that has access to the internet -3. Review (make sure all downloads are correct) and save the output from the LLM as a .sh file -4. Upload it onto your Runpod instance -5. Run the file. For example: `./download_comfyui_models.sh` +1. Make sure you launched the Pod with these Environment Variables: CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account)) and HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens)) +2. Export your workflow as an API workflow +3. Copy the below prompt and upload the API workflow to a LLM that has access to the internet +4. Review (make sure all downloads are correct) and save the output from the LLM as a .sh file +5. Upload it onto your Runpod instance +6. Run the file. For example: `./download_comfyui_models.sh` ``` Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: @@ -27,4 +28,5 @@ unet: models/unet/ upscale_models: models/upscale_models/ vae: models/vae/ Make sure you find the correct URLs for the models online. +Use comfy cli to download the models: comfy model download --url ?[--relative-path ] ``` From e793269dd4029c21166ef6e762abfd3518feae77 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 09:17:20 +1000 Subject: [PATCH 86/93] Update model-download-script.md --- docs/model-download-script.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index 63c4bc4..22ce85e 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -2,12 +2,14 @@ This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM). -1. Make sure you launched the Pod with these Environment Variables: CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account)) and HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens)) -2. Export your workflow as an API workflow -3. Copy the below prompt and upload the API workflow to a LLM that has access to the internet -4. Review (make sure all downloads are correct) and save the output from the LLM as a .sh file -5. Upload it onto your Runpod instance -6. Run the file. For example: `./download_comfyui_models.sh` +1. Launched the Pod with these Environment Variables: + - CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account)) + - HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens)) +3. In ComfyUI, export your workflow as an API workflow +4. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet** +5. Review the LLMs output to make sure all download links are correct and save it as a .sh file +6. Upload it onto your Runpod instance +7. Run the file. For example: `./download_comfyui_models.sh` ``` Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: From 9d69aef597f5343d256939d03e938bacfc0df4f7 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 09:31:29 +1000 Subject: [PATCH 87/93] Update model-download-script.md --- docs/model-download-script.md | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index 22ce85e..b0959ac 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -7,10 +7,6 @@ This guide will walk you through creating a shell script to automatically downlo - HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens)) 3. In ComfyUI, export your workflow as an API workflow 4. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet** -5. Review the LLMs output to make sure all download links are correct and save it as a .sh file -6. Upload it onto your Runpod instance -7. Run the file. For example: `./download_comfyui_models.sh` - ``` Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: comfyui: @@ -32,3 +28,6 @@ vae: models/vae/ Make sure you find the correct URLs for the models online. Use comfy cli to download the models: comfy model download --url ?[--relative-path ] ``` +5. Review the LLMs output to make sure all download links are correct and save it as a .sh file +6. Upload it onto your Runpod instance +7. Run the file. For example: `./download_comfyui_models.sh` From 175afc966600d026fc38bbc34ac20be99309922d Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 10:18:19 +1000 Subject: [PATCH 88/93] Update model-download-script.md --- docs/model-download-script.md | 37 ++++++++++++++++++++++++++++------- 1 file changed, 30 insertions(+), 7 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index b0959ac..cd96c79 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -5,11 +5,11 @@ This guide will walk you through creating a shell script to automatically downlo 1. Launched the Pod with these Environment Variables: - CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account)) - HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens)) -3. In ComfyUI, export your workflow as an API workflow -4. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet** +2. In ComfyUI, export your workflow as an API workflow +3. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet** ``` Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths: -comfyui: + base_path: /workspace/ComfyUI checkpoints: models/checkpoints/ clip: models/clip/ @@ -25,9 +25,32 @@ text_encoders: models/text_encoders/ unet: models/unet/ upscale_models: models/upscale_models/ vae: models/vae/ + +--- + +Important: Make sure you find the correct URLs for the models online. -Use comfy cli to download the models: comfy model download --url ?[--relative-path ] +Use comfy cli to download the models: `comfy model download --url [--relative-path ] [--set-civitai-api-token ] [--set-hf-api-token ]` +Make sure you add `--set-civitai-api-token $CIVITAI_API_TOKEN` for CivitAI download and `--set-hf-api-token $HF_API_TOKEN` for Hugging Face downloads. + +--- + +Example: + +#!/bin/bash + +# Download from CivitAI +comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN + +# Download model from Hugging Face +comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN + +# If a model in the workflow was in a subfolder +comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints/SDXL --set-civitai-api-token $CIVITAI_API_TOKEN + ``` -5. Review the LLMs output to make sure all download links are correct and save it as a .sh file -6. Upload it onto your Runpod instance -7. Run the file. For example: `./download_comfyui_models.sh` +4. Review the LLMs output to make sure all download links are correct and save it as a .sh file +5. Upload it onto your Runpod instance, into `/workspace` +6. Then run these commands: + - `chmod 755 /workspace/download_comfyui_models.sh` + - `/workspace/download_comfyui_models.sh` From 6fb1cdaea4b72c4878d677f52b44ef60480e8d49 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 10:23:10 +1000 Subject: [PATCH 89/93] Update model-download-script.md --- docs/model-download-script.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/model-download-script.md b/docs/model-download-script.md index cd96c79..0ec346f 100644 --- a/docs/model-download-script.md +++ b/docs/model-download-script.md @@ -49,8 +49,8 @@ comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/r comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints/SDXL --set-civitai-api-token $CIVITAI_API_TOKEN ``` -4. Review the LLMs output to make sure all download links are correct and save it as a .sh file +4. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh` 5. Upload it onto your Runpod instance, into `/workspace` 6. Then run these commands: - - `chmod 755 /workspace/download_comfyui_models.sh` - - `/workspace/download_comfyui_models.sh` + - `chmod 755 /workspace/download_models.sh` + - `/workspace/download_models.sh` From 1df61c9b8e550c014c0ea583eab30ffabf5bed4d Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 29 Jul 2025 12:21:12 +1000 Subject: [PATCH 90/93] fix runpod worker auto host name construction --- web/main.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/web/main.js b/web/main.js index 4830a30..51c2f07 100644 --- a/web/main.js +++ b/web/main.js @@ -906,7 +906,7 @@ class DistributedExtension { const worker = { id: crypto.randomUUID(), name: `Worker ${workerNum}`, - host: "localhost", + host: isRunpod ? null : "localhost", port: 8189 + portOffset, cuda_device: i, enabled: true, From 8322b253d19c16b4bbedc0f8337a752f4f63f208 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 12:53:30 +1000 Subject: [PATCH 91/93] Update worker-setup-guides.md --- docs/worker-setup-guides.md | 60 ++++++++++++++++++++++++------------- 1 file changed, 39 insertions(+), 21 deletions(-) diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index 4f37f46..7f2afee 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -1,5 +1,19 @@ +## Overview + +ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs, whether they're in the same machine, on your local network, or in the cloud. + +**Master**: The main ComfyUI instance that coordinates and distributes work. This is where you load workflows, manage the queue, and view results. + +**Worker**: A ComfyUI instance that receives and processes tasks from the master. Workers handle just the GPU computation and send results back to the master. You can have multiple workers connected to a single master, each utilizing their own GPU. + +### Types of Workers + +- **Local workers**: Additional GPUs on the same machine as the master +- **Remote workers**: GPUs on different computers within your network +- **Cloud workers**: GPUs hosted on cloud services like Runpod + ## Local workers -> Local Workers: These are added automatically on first launch, but you can add them manually if you need. +> These are added automatically on first launch, but you can add them manually if you need to. 1. **Open** the Distributed GPU panel. 2. **Click** "Add Worker" in the UI. @@ -11,14 +25,17 @@ 4. **Save** and launch the local worker. ## Remote workers -> Remote Workers: ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address. +> ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address. **On the Remote Worker Machine:** -1. **Launch** the master (main ComfyUI instance) with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** -2. **Add** workers in the UI panel if the remote machine has more than one GPU. - - Make sure that they also have `--listen` set in `Extra Args`. - - Then launch them. -3. **Open** the configured worker port(s) (e.g., 8189, 8190) in the remote worker's firewall. +1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** + - This ComfyUI instance will serve as a worker for your main master. +2. *Optionally* add additional local workers on this machine if it has multiple GPUs: + - Access the Distributed GPU panel in this ComfyUI instance + - Add workers for any additional GPUs (if they haven't been added automatically) + - Make sure they have `--listen` set in `Extra Args` + - Launch them +3. **Open** the ComfyUI port (e.g., 8188) and any additional worker ports (e.g., 8189, 8190) in the firewall. **On the Main Machine:** 1. **Open** the Distributed GPU panel (sidebar on the left). @@ -27,31 +44,32 @@ 4. **Configure** your remote worker: - **Name**: A descriptive name for the worker (e.g., "Server Rack GPU 0") - **Host**: The remote worker's IP address. - - **Port**: The port number used when launching ComfyUI on the remote worker (e.g., 8189). + - **Port**: The port number used when launching ComfyUI on the remote master/worker (e.g., 8188). 5. **Save** the remote worker configuration. ## Cloud workers -> Cloud Workers: ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p). +> ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p). ### Deploy Cloud Worker on Runpod **On Runpod:** 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account. -2. On Runpod, go to Storage > New Network Volume and create a volume which will store the models you need. +2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. 3. Now go to Pods and find a suitable GPU for your workflows. -> To use the ComfyUI-Distributed-Pod template, you will need to filter instances by CUDA 12.8. -4. Choose the `ComfyUI-Distributed-Pod` template and make sure your network drive is mounted. -5. Launch your pod. -6. Access your pod using JupyterLabs. -7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). +4. Choose the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template and make sure your network drive is mounted. +> To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8. +6. Launch your pod. +7. Access your pod using JupyterLabs. +8. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). > You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow. 8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!** -9. Download any additional custom nodes using ComfyUI manager +9. Download any additional custom nodes using the ComfyUI Manager. **On the Main Machine:** 1. Launch a Cloudflare tunnel - - Download from here `https://github.com/cloudflare/cloudflared/releases` - - Then run for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` + - Download from here: `https://github.com/cloudflare/cloudflared/releases` + - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` +> Cloudflare tunnels create secure connections without exposing ports directly to the internet 2. Copy the Cloudflare address 3. **Open** the Distributed GPU panel (sidebar on the left). 4. **Edit** the Master's host address and replace it with the Cloudflare address. @@ -69,7 +87,7 @@ **On the Cloud Worker machine:** - Your cloud worker container needs to have the same models and custom nodes as the workflow you want to run on your local machine. - If your cloud platform doesn't provide a secure connection, use Cloudflare to create a tunnel for the worker. Each GPU needs their own tunnel for their respective port. - - For example: `./cloudflared tunnel --url http://localhost:8188` + - For example: `./cloudflared tunnel --url http://localhost:8188` 1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!** 2. **Add** workers in the UI panel if the remote machine has more than one GPU. - Make sure that they also have `--listen` set in `Extra Args`. @@ -77,8 +95,8 @@ **On the Main Machine:** 1. **Launch** a Cloudflare tunnel on your local machine. - - Download from here `https://github.com/cloudflare/cloudflared/releases` - - For example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` + - Download from here: `https://github.com/cloudflare/cloudflared/releases` + - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` 2. **Copy** the Cloudflare address 3. **Open** the Distributed GPU panel (sidebar on the left). 4. **Edit** the Master's host address and replace it with the Cloudflare address. From 19d8e8f731d99730ec08b72135cb890f08b759fd Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 12:58:45 +1000 Subject: [PATCH 92/93] Update worker-setup-guides.md --- docs/worker-setup-guides.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index 7f2afee..6d66784 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -67,7 +67,7 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs **On the Main Machine:** 1. Launch a Cloudflare tunnel - - Download from here: `https://github.com/cloudflare/cloudflared/releases` + - Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases) - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` > Cloudflare tunnels create secure connections without exposing ports directly to the internet 2. Copy the Cloudflare address @@ -95,7 +95,7 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs **On the Main Machine:** 1. **Launch** a Cloudflare tunnel on your local machine. - - Download from here: `https://github.com/cloudflare/cloudflared/releases` + - Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases) - Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188` 2. **Copy** the Cloudflare address 3. **Open** the Distributed GPU panel (sidebar on the left). From b9d06703f4980e9539ee24eb7449776aefde7739 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 29 Jul 2025 12:59:17 +1000 Subject: [PATCH 93/93] Update worker-setup-guides.md --- docs/worker-setup-guides.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/worker-setup-guides.md b/docs/worker-setup-guides.md index 6d66784..0b2f242 100644 --- a/docs/worker-setup-guides.md +++ b/docs/worker-setup-guides.md @@ -1,4 +1,4 @@ -## Overview +## ComfyUI Distributed Worker Setup Guide ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs, whether they're in the same machine, on your local network, or in the cloud.