diff --git a/__init__.py b/__init__.py index 5c8a98d..e669a16 100644 --- a/__init__.py +++ b/__init__.py @@ -50,19 +50,13 @@ from .distributed_upscale import ( NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS ) -# Import WAN helpers -from .wan_helpers import ( - NODE_CLASS_MAPPINGS as WAN_HELPERS_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as WAN_HELPERS_DISPLAY_NAME_MAPPINGS, -) - WEB_DIRECTORY = "./web" ensure_config_exists() # Merge node mappings -NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS, **WAN_HELPERS_CLASS_MAPPINGS} -NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS, **WAN_HELPERS_DISPLAY_NAME_MAPPINGS} +NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS} +NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS} __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/distributed_upscale.py b/distributed_upscale.py index 8f9c6a8..e3c4c27 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -837,6 +837,7 @@ class UltimateSDUpscaleDistributed: # Always batched-per-tile in static mode debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}") + processed_count += batch_size tile_id = tile_idx tx, ty = all_tiles[tile_id] # Extract tile for entire batch @@ -850,9 +851,8 @@ class UltimateSDUpscaleDistributed: seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, region, (width, height) ) - # Queue results (clamped to actual outputs) - out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size - for b in range(min(batch_size, out_bs)): + # Queue results + for b in range(batch_size): processed_tiles.append({ 'tile': processed_batch[b:b+1], 'tile_idx': tile_id, @@ -864,7 +864,6 @@ class UltimateSDUpscaleDistributed: 'batch_idx': b, 'global_idx': b * num_tiles_per_image + tile_id }) - processed_count += min(batch_size, out_bs) # Send heartbeat try: @@ -876,12 +875,12 @@ class UltimateSDUpscaleDistributed: debug_log(f"Worker[{worker_id[:8]}] 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 = [] + 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: @@ -968,7 +967,6 @@ class UltimateSDUpscaleDistributed: batch_size = upscaled_image.shape[0] _, height, width, _ = upscaled_image.shape total_tiles = batch_size * num_tiles_per_image - effective_bs = batch_size # Adjusted if model returns fewer outputs per batch # Convert batch to PIL list for processing result_images = [] @@ -1004,6 +1002,7 @@ class UltimateSDUpscaleDistributed: ) if tile_idx is not None: consecutive_no_tile = 0 + processed_count += batch_size tile_id = tile_idx tx, ty = all_tiles[tile_id] tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( @@ -1016,9 +1015,7 @@ class UltimateSDUpscaleDistributed: region, (width, height) ) tile_mask = tile_masks[tile_id] - out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size - effective_bs = min(effective_bs, out_bs) - for b in range(min(batch_size, out_bs)): + for b in range(batch_size): tile_pil = tensor_to_pil(processed_batch, b) if tile_pil.size != (ew, eh): tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) @@ -1028,7 +1025,6 @@ class UltimateSDUpscaleDistributed: _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}), timeout=5.0 ) - processed_count += min(batch_size, out_bs) log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})") else: consecutive_no_tile += 1 @@ -1045,9 +1041,8 @@ class UltimateSDUpscaleDistributed: # Collect worker results using async operations try: - expected_total = num_tiles_per_image * max(1, effective_bs) collected_tasks = run_async_in_server_loop( - self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=expected_total), + 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: @@ -1060,51 +1055,44 @@ class UltimateSDUpscaleDistributed: if completed_count < total_tiles: log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") - # Process any remaining pending tasks + # Process any remaining pending tasks (batched-per-tile) 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 tile_id from pending queue + tile_id = run_async_in_server_loop( self._get_next_tile_index(multi_job_id), timeout=5.0 ) - - if task_id is None: + + if tile_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, force_uniform_tiles, tiled_decode, batch_idx=batch_idx + + # Extract batched tile and process across available batch + tx, ty = all_tiles[tile_id] + tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( + upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles ) - - # 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 + region = (x1, y1, x1 + ew, y1 + eh) + processed_batch = self.process_tiles_batch( + tile_batch, model, positive, negative, vae, + seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, + region, (width, height) ) - # 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} + tile_mask = tile_masks[tile_id] + out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size + for b in range(min(batch_size, out_bs)): + tile_pil = tensor_to_pil(processed_batch, b) + if tile_pil.size != (ew, eh): + tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) + result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding) + global_idx = b * num_tiles_per_image + tile_id + # Mark as completed so the collector state is consistent + run_async_in_server_loop( + _mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}), + timeout=5.0 + ) else: # Master processed all tiles collected_tasks = run_async_in_server_loop( @@ -1209,18 +1197,13 @@ class UltimateSDUpscaleDistributed: if not processed_tiles: return # Early exit if empty + total_tiles = len(processed_tiles) + debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)") + # Prepare encoded images and sizes to enable size-aware chunking encoded = [] for idx, tile_data in enumerate(processed_tiles): - tile = tile_data.get('tile') - # Skip empty or invalid tiles - try: - if tile is None or (hasattr(tile, 'shape') and tile.shape[0] == 0): - continue - img = tensor_to_pil(tile, 0) - except Exception: - # Any issue converting this tile -> skip it - continue + img = tensor_to_pil(tile_data['tile'], 0) bio = io.BytesIO() # Keep compression low to balance speed and size; adjust if needed img.save(bio, format='PNG', compress_level=0) @@ -1237,11 +1220,6 @@ class UltimateSDUpscaleDistributed: **({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}), } }) - if not encoded: - debug_log(f"Worker[{worker_id[:8]}] - No valid tiles to send (all empty)") - return - total_tiles = len(encoded) - debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)") # Size-aware chunking max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024) # 1MB headroom @@ -1351,8 +1329,7 @@ class UltimateSDUpscaleDistributed: # Blend results back into each image using cached mask tile_mask = tile_masks[tile_idx] - out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size - for b in range(min(batch_size, out_bs)): + for b in range(batch_size): tile_pil = tensor_to_pil(processed_batch, b) # Resize back to extracted size if tile_pil.size != (ew, eh): diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py index c4cf544..6a1843d 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_managment.py @@ -18,6 +18,7 @@ MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 # Import HEARTBEAT_TIMEOUT from constants from .constants import HEARTBEAT_TIMEOUT +from .config import load_config def _parse_tiles_from_form(data): @@ -271,13 +272,35 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): requeued_count = 0 completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {}) + # Allow override via config setting 'worker_timeout_seconds' + cfg = load_config() + hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT)) + for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()): - if current_time - last_heartbeat > HEARTBEAT_TIMEOUT: + if current_time - last_heartbeat > hb_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 batched_static, task_id is a tile_idx; consider it complete only if + # all corresponding global_idx entries are present in completed_tasks. + batched_static = bool(job_data.get('batched_static', False)) + if batched_static: + num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1) + batch_size = job_data.get(JOB_BATCH_SIZE, 1) + # Check all global indices for this tile across the batch + all_done = True + for b in range(batch_size): + gidx = b * num_tiles_per_image + task_id + if gidx not in completed_tasks: + all_done = False + break + if not all_done: + await job_data[JOB_PENDING_TASKS].put(task_id) + requeued_count += 1 + else: + # Legacy/global-idx mode: task_id is a global index key + 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: diff --git a/web/main.js b/web/main.js index 51c2f07..e1026d2 100644 --- a/web/main.js +++ b/web/main.js @@ -217,22 +217,7 @@ class DistributedExtension { this.panelElement = null; } - 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."; - } - } - } + // updateSummary removed // --- Core Logic & Execution --- @@ -1333,4 +1318,4 @@ app.registerExtension({ async setup() { new DistributedExtension(); } -}); \ No newline at end of file +}); diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js index 9fc95a0..a8b42a3 100644 --- a/web/sidebarRenderer.js +++ b/web/sidebarRenderer.js @@ -148,57 +148,67 @@ export async function renderSidebarContent(extension, el) { // Settings section const settingsSection = document.createElement("div"); - settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;"; + // Top separator only; spacing handled by the clickable toggle area for equal top/bottom spacing + settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;"; - // Settings header with toggle + // Settings header with toggle (full-area clickable between separators) + const settingsToggleArea = document.createElement("div"); + // Equal spacing above header (to top separator) and below header (to bottom separator) + settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;"; + const settingsHeader = document.createElement("div"); - settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;"; - + settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;"; 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); + settingsToggleArea.appendChild(settingsHeader); + // Hover effect for toggle area + settingsToggleArea.onmouseover = () => { workerSettingsToggle.style.color = "#fff"; }; + settingsToggleArea.onmouseout = () => { workerSettingsToggle.style.color = "#888"; }; - // Hover effect for header - settingsHeader.onmouseover = () => { - workerSettingsToggle.style.color = "#fff"; - }; - settingsHeader.onmouseout = () => { - workerSettingsToggle.style.color = "#888"; - }; - + // A small separator shown only when collapsed (to make the section boundary obvious) + const settingsSeparator = document.createElement("div"); + // No margin so the bottom spacing is controlled by settingsToggleArea padding-bottom + settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;"; + // Collapsible settings content const settingsContent = document.createElement("div"); 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;"; + settingsDiv.style.cssText = "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;"; // Toggle functionality let settingsExpanded = false; - settingsHeader.onclick = () => { + settingsToggleArea.onclick = () => { settingsExpanded = !settingsExpanded; if (settingsExpanded) { settingsContent.style.maxHeight = "200px"; settingsContent.style.opacity = "1"; workerSettingsToggle.style.transform = "rotate(90deg)"; + settingsSeparator.style.display = "none"; } else { settingsContent.style.maxHeight = "0"; settingsContent.style.opacity = "0"; workerSettingsToggle.style.transform = "rotate(0deg)"; + settingsSeparator.style.display = "block"; } }; // Debug mode setting + // Section: General + const generalLabel = document.createElement("div"); + generalLabel.textContent = "GENERAL"; + generalLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;"; + const debugGroup = document.createElement("div"); - debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; - + debugGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;"; + const debugCheckbox = document.createElement("input"); debugCheckbox.type = "checkbox"; debugCheckbox.id = "setting-debug"; @@ -209,13 +219,14 @@ export async function renderSidebarContent(extension, el) { debugLabel.htmlFor = "setting-debug"; debugLabel.textContent = "Debug Mode"; debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + debugLabel.title = "Enable verbose logging in the browser console."; 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;"; + autoLaunchGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;"; const autoLaunchCheckbox = document.createElement("input"); autoLaunchCheckbox.type = "checkbox"; @@ -227,13 +238,14 @@ export async function renderSidebarContent(extension, el) { autoLaunchLabel.htmlFor = "setting-auto-launch"; autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup"; autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + autoLaunchLabel.title = "Start local worker processes automatically when the master starts."; autoLaunchGroup.appendChild(autoLaunchCheckbox); autoLaunchGroup.appendChild(autoLaunchLabel); - // Stop workers on exit setting + // Stop workers on exit setting (under General) const stopOnExitGroup = document.createElement("div"); - stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;"; + stopOnExitGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;"; const stopOnExitCheckbox = document.createElement("input"); stopOnExitCheckbox.type = "checkbox"; @@ -245,28 +257,56 @@ export async function renderSidebarContent(extension, el) { stopOnExitLabel.htmlFor = "setting-stop-on-exit"; stopOnExitLabel.textContent = "Stop Local Workers on Master Exit"; stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;"; + stopOnExitLabel.title = "Stop local worker processes automatically when the master exits."; stopOnExitGroup.appendChild(stopOnExitCheckbox); stopOnExitGroup.appendChild(stopOnExitLabel); + settingsDiv.appendChild(generalLabel); settingsDiv.appendChild(debugGroup); settingsDiv.appendChild(autoLaunchGroup); settingsDiv.appendChild(stopOnExitGroup); + + // Worker Timeout setting (seconds) + // Section: Timeouts + const timeoutsLabel = document.createElement("div"); + timeoutsLabel.textContent = "TIMEOUTS"; + timeoutsLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;"; + + const timeoutGroup = document.createElement("div"); + timeoutGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 6px;"; + + const timeoutLabel = document.createElement("label"); + timeoutLabel.htmlFor = "setting-worker-timeout"; + timeoutLabel.textContent = "Worker Timeout"; + timeoutLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: default;"; + timeoutLabel.title = "Seconds without a heartbeat before a worker is considered timed out and its tasks are requeued. Default 60. For WAN, consider 300–600."; + + const timeoutInput = document.createElement("input"); + timeoutInput.type = "number"; + timeoutInput.id = "setting-worker-timeout"; + timeoutInput.min = "10"; + timeoutInput.step = "1"; + timeoutInput.style.cssText = "width: 80px; padding: 2px 6px; background: #222; color: #ddd; border: 1px solid #333; border-radius: 3px;"; + timeoutInput.value = (extension.config?.settings?.worker_timeout_seconds ?? 60); + timeoutInput.onchange = (e) => { + const v = parseInt(e.target.value, 10); + if (!Number.isFinite(v) || v <= 0) return; + extension._updateSetting('worker_timeout_seconds', v); + }; + + timeoutGroup.appendChild(timeoutLabel); + timeoutGroup.appendChild(timeoutInput); + settingsDiv.appendChild(timeoutsLabel); + settingsDiv.appendChild(timeoutGroup); settingsContent.appendChild(settingsDiv); - settingsSection.appendChild(settingsHeader); + settingsSection.appendChild(settingsToggleArea); + settingsSection.appendChild(settingsSeparator); 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); @@ -274,4 +314,4 @@ export async function renderSidebarContent(extension, el) { // Always reset the rendering flag extension._isRendering = false; } -} \ No newline at end of file +}