commit d4384781dc518c439db343466e0e37865104be4e Author: Robert Date: Wed Jun 25 11:59:35 2025 +1000 Initial commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..7b78a9e --- /dev/null +++ b/README.md @@ -0,0 +1,165 @@ +# ComfyUI-Distributed + +> ⚠️ **DEVELOPMENT STATUS**: This project is currently in active development and is not yet production-ready. Features may change, and bugs may exist. Use at your own risk. + +A custom node extension for ComfyUI that enables distributed image generation across multiple GPUs through a master-worker architecture. + +## Overview + +ComfyUI-Distributed extends ComfyUI with the ability to distribute batch processing across multiple GPUs, allowing for linear scaling of image generation performance. It achieves this by coordinating multiple ComfyUI instances, each running on a different GPU, with one instance acting as the master coordinator. + +## Features + +- **Multi-GPU Distributed Processing**: Leverage all available GPUs simultaneously +- **Master-Worker Architecture**: Automatic coordination between GPU instances +- **Dynamic Worker Management**: Enable/disable GPUs on-the-fly through the UI +- **Execution Modes**: + - **Parallel Mode**: All GPUs process simultaneously for maximum speed + - **Staggered Mode**: Sequential GPU startup with configurable delays to manage VRAM usage +- **VRAM Management**: Built-in tools to clear GPU memory across all instances +- **Web UI Integration**: Seamless control panel integrated into ComfyUI's interface + +## How It Works + +1. Launch multiple ComfyUI instances, each assigned to a different GPU +2. One instance acts as the master (typically on CUDA device 0) +3. Add the "Multi-GPU Collector" node to your workflow +4. The extension automatically distributes batches across available GPUs +5. Results are collected and combined by the master instance + +For example, with 4 GPUs and a batch size of 2, you'll generate 8 images in parallel. + +## Requirements + +- ComfyUI installation +- Multiple NVIDIA GPUs +- Python 3.8+ +- Windows (for the provided launcher script) or Linux/Mac with manual setup + +## Installation + +1. Clone this repository into your ComfyUI custom nodes directory: + ```bash + cd ComfyUI/custom_nodes + git clone https://github.com/yourusername/ComfyUI-Distributed.git + ``` + +2. No additional dependencies required - uses ComfyUI's existing libraries + +## Usage + +### Quick Start (Windows) + +1. Copy the `launch_distributed.bat` file to your ComfyUI root directory (alongside other .bat files like `run_nvidia_gpu.bat`) + +2. Run the launcher script from the ComfyUI root directory: + ```bash + launch_distributed.bat + ``` + This will automatically start ComfyUI instances for each GPU + +3. Open the master instance (usually http://localhost:8188) + +4. Add the "Multi-GPU Collector" node to your workflow + +5. Use the GPU control panel in the sidebar to: + - Select which GPUs to use + - Choose between parallel or staggered execution + - Clear VRAM when needed + +### Manual Setup (All Platforms) + +Start each ComfyUI instance with the appropriate CUDA device: + +```bash +# Master instance (GPU 0) +CUDA_VISIBLE_DEVICES=0 python main.py --port 8188 + +# Worker instance (GPU 1) +CUDA_VISIBLE_DEVICES=1 python main.py --port 8189 --enable-cors-header + +# Worker instance (GPU 2) +CUDA_VISIBLE_DEVICES=2 python main.py --port 8190 --enable-cors-header + +# Continue for additional GPUs... +``` + +### Configuration + +The system uses a `gpu_config.json` file for configuration. Default configuration: + +```json +{ + "master": { + "port": 8188, + "device": 0, + "enabled": true + }, + "workers": [ + { + "port": 8189, + "device": 1, + "enabled": true, + "extra_args": "--enable-cors-header" + }, + { + "port": 8190, + "device": 2, + "enabled": true, + "extra_args": "--enable-cors-header" + } + ], + "settings": { + "stagger_delay": 2.0, + "job_timeout": 120 + } +} +``` + +## Workflow Integration + +The Multi-GPU Collector node integrates seamlessly with existing ComfyUI workflows: + +1. Connect the node after your image generation pipeline +2. The node will automatically detect batch sizes and distribute work +3. All GPUs will process their assigned batches +4. Results are collected and passed to the next node as a single batch + +## Performance Considerations + +- **Parallel Mode**: Best for workflows with consistent VRAM usage +- **Staggered Mode**: Ideal for memory-intensive workflows or when GPUs have different VRAM capacities +- **Batch Size**: Total generation = batch_size × number_of_active_GPUs + +## Troubleshooting + +- **Workers not connecting**: Ensure CORS headers are enabled (`--enable-cors-header`) +- **VRAM errors**: Use staggered mode with appropriate delays +- **Timeout errors**: Increase `job_timeout` in configuration +- **Port conflicts**: Ensure each instance uses a unique port + +## Development + +This project is under active development. Contributions are welcome! + +### Known Issues + +- Windows Terminal required for the launcher script +- Manual configuration needed for more than 4 GPUs +- Some custom nodes may not be compatible with distributed processing + +### Planned Features + +- Linux/Mac launcher scripts +- Automatic GPU detection and configuration +- Load balancing based on GPU capabilities +- Support for remote workers over network +- Queue management improvements + +## License + +Same as ComfyUI - see ComfyUI repository for details. + +## Acknowledgments + +Built as an extension for the amazing [ComfyUI](https://github.com/comfyanonymous/ComfyUI) project. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2e88060 --- /dev/null +++ b/__init__.py @@ -0,0 +1,23 @@ +import os +import sys + +# Add the directory to Python path +sys.path.append(os.path.dirname(__file__)) + +# Import everything needed from the main module +from .multigpu import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + ensure_config_exists, + CONFIG_FILE +) + +WEB_DIRECTORY = "./web" + +# Use the centralized config management from multigpu.py +ensure_config_exists() + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +print("[MultiGPU] Loaded Multi-GPU nodes.") +print(f"[MultiGPU] Config file: {CONFIG_FILE}") \ No newline at end of file diff --git a/gpu_config.json b/gpu_config.json new file mode 100644 index 0000000..ccfc23a --- /dev/null +++ b/gpu_config.json @@ -0,0 +1,35 @@ +{ + "master": { + "port": 8480, + "cuda_device": 0 + }, + "workers": [ + { + "id": 1, + "name": "GPU 1", + "cuda_device": 1, + "port": 8180, + "enabled": false, + "extra_args": "" + }, + { + "id": 2, + "name": "GPU 2", + "cuda_device": 2, + "port": 8280, + "enabled": false, + "extra_args": "" + }, + { + "id": 3, + "name": "GPU 3", + "cuda_device": 3, + "port": 8380, + "enabled": false, + "extra_args": "" + } + ], + "settings": { + "stagger_delay_ms": 15000 + } +} \ No newline at end of file diff --git a/launch_distributed.bat b/launch_distributed.bat new file mode 100644 index 0000000..77781bd --- /dev/null +++ b/launch_distributed.bat @@ -0,0 +1,119 @@ +@echo off +title ComfyUI Multi-GPU CLI Launcher +cd /d "%~dp0" +color 0B +cls + +echo ComfyUI Multi-GPU CLI Launcher +echo =============================== + +:: Check if gpu_config.json exists (relative path) +set "CONFIG_PATH=ComfyUI\custom_nodes\ComfyUI-Distributed\gpu_config.json" +if not exist "%CONFIG_PATH%" ( + color 0C + echo Error: Configuration file not found! + echo Expected location: %CONFIG_PATH% + pause + exit /b 1 +) + +:: Check if venv exists (relative path) +set "VENV_PATH=venv" +if not exist "%VENV_PATH%\Scripts\activate.bat" ( + color 0C + echo Error: Virtual environment not found! + echo Expected location: %CD%\%VENV_PATH% + echo Make sure this batch file is in the ComfyUI_Dev directory with venv folder + pause + exit /b 1 +) + +:: Check if ComfyUI CLI is available in venv +call "%VENV_PATH%\Scripts\activate.bat" && comfy --help >nul 2>&1 +if %ERRORLEVEL% NEQ 0 ( + color 0C + echo Error: ComfyUI CLI not found in virtual environment! + echo Make sure ComfyUI CLI is installed in: %VENV_PATH% + pause + exit /b 1 +) + +:: Read configuration with PowerShell and store in environment variables +echo Loading configuration... + +powershell -NoProfile -Command "try { $json = Get-Content '%CONFIG_PATH%' -Raw | ConvertFrom-Json; $master = $json.master; $settings = $json.settings; $retryDelay = if ($settings.retry_delay_ms) { $settings.retry_delay_ms } else { 500 }; Write-Output \"RETRY_DELAY=$retryDelay\"; Write-Output \"MASTER_PORT=$($master.port)\"; Write-Output \"MASTER_CUDA=$($master.cuda_device)\"; foreach ($worker in $json.workers) { if ($worker.enabled) { Write-Output \"WORKER_$($worker.id)_NAME=$($worker.name)\"; Write-Output \"WORKER_$($worker.id)_PORT=$($worker.port)\"; Write-Output \"WORKER_$($worker.id)_CUDA=$($worker.cuda_device)\"; Write-Output \"WORKER_$($worker.id)_EXTRA=$($worker.extra_args)\" } } } catch { Write-Output \"PowerShell Error: $_\"; pause; exit 1 }" > %TEMP%\config_vars.txt + +if %ERRORLEVEL% NEQ 0 ( + color 0C + echo Error running PowerShell configuration reader! + pause + exit /b 1 +) + +for /f "tokens=*" %%i in (%TEMP%\config_vars.txt) do ( + set %%i +) + +del %TEMP%\config_vars.txt >nul 2>&1 + +:: Check Master +echo Starting Master instance... +:MASTER +powershell -NoProfile -Command "$ErrorActionPreference = 'SilentlyContinue'; $tcpClient = New-Object System.Net.Sockets.TcpClient; $connect = $tcpClient.BeginConnect('127.0.0.1', %MASTER_PORT%, $null, $null); $wait = $connect.AsyncWaitHandle.WaitOne(200, $false); if (!$wait) { exit 1 } else { $tcpClient.EndConnect($connect); $tcpClient.Close(); exit 0 }" +if %ERRORLEVEL% == 1 ( + if not defined MASTER_STARTED ( + start wt -w 0 new-tab --title "Master - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title Master - ComfyUI CLI && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && comfy --here launch -- --port %MASTER_PORT%" + set MASTER_STARTED=1 + ) + %TEMP%\worker_ids.txt + +for /f %%w in (%TEMP%\worker_ids.txt) do ( + call :PROCESS_WORKER %%w +) + +del %TEMP%\worker_ids.txt >nul 2>&1 +goto :END_SCRIPT + +:PROCESS_WORKER +set "WORKER_ID=%1" +call set "WORKER_NAME=%%WORKER_%WORKER_ID%_NAME%%" +call set "WORKER_PORT=%%WORKER_%WORKER_ID%_PORT%%" +call set "WORKER_CUDA=%%WORKER_%WORKER_ID%_CUDA%%" +call set "WORKER_EXTRA=%%WORKER_%WORKER_ID%_EXTRA%%" + +echo Starting %WORKER_NAME%... + +:WORKER_CHECK_LOOP +powershell -NoProfile -Command "$ErrorActionPreference = 'SilentlyContinue'; $tcpClient = New-Object System.Net.Sockets.TcpClient; $connect = $tcpClient.BeginConnect('127.0.0.1', %WORKER_PORT%, $null, $null); $wait = $connect.AsyncWaitHandle.WaitOne(200, $false); if (!$wait) { exit 1 } else { $tcpClient.EndConnect($connect); $tcpClient.Close(); exit 0 }" +if %ERRORLEVEL% == 1 ( + if not defined WORKER%WORKER_ID%_STARTED ( + start wt -w 0 new-tab --title "%WORKER_NAME% - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title %WORKER_NAME% - ComfyUI CLI && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%WORKER_CUDA% && comfy --here launch -- --port %WORKER_PORT% %WORKER_EXTRA% --enable-cors-header" + set WORKER%WORKER_ID%_STARTED=1 + ) + nul +exit \ No newline at end of file diff --git a/multigpu.py b/multigpu.py new file mode 100644 index 0000000..34d2094 --- /dev/null +++ b/multigpu.py @@ -0,0 +1,286 @@ +import torch +import numpy as np +from PIL import Image +import folder_paths +import os +import json +import asyncio +import aiohttp +from aiohttp import web +import io +import server +from concurrent.futures import Future +import comfy.model_management as mm +import gc + +# --- Config Management --- +CONFIG_FILE = os.path.join(os.path.dirname(__file__), "gpu_config.json") + +def get_default_config(): + """Returns the default configuration dictionary. Single source of truth.""" + return { + "master": {"port": 8188, "cuda_device": 0}, + "workers": [], + "settings": { + "stagger_delay_ms": 15000 + } + } + +def load_config(): + """Loads the config, falling back to defaults if the file is missing or invalid.""" + if os.path.exists(CONFIG_FILE): + try: + with open(CONFIG_FILE, 'r') as f: + return json.load(f) + except Exception as e: + print(f"[MultiGPU Config] Error loading config, using defaults: {e}") + return get_default_config() + +def save_config(config): + """Saves the configuration to file.""" + try: + with open(CONFIG_FILE, 'w') as f: + json.dump(config, f, indent=2) + return True + except Exception as e: + print(f"[MultiGPU Config] Error saving config: {e}") + return False + +def ensure_config_exists(): + """Creates default config file if it doesn't exist. Used by __init__.py""" + if not os.path.exists(CONFIG_FILE): + default_config = get_default_config() + if save_config(default_config): + print("[MultiGPU] Created default config file") + else: + print("[MultiGPU] Could not create default config file") + +# --- Helper Functions --- +async def handle_api_error(request, error, status=500): + """Standardized error response handler""" + print(f"[MultiGPU] API Error: {error}") + return web.json_response({"status": "error", "message": str(error)}, status=status) + +# --- API Endpoints --- +@server.PromptServer.instance.routes.get("/multigpu/config") +async def get_config_endpoint(request): + config = load_config() + return web.json_response(config) + +@server.PromptServer.instance.routes.post("/multigpu/config/update_worker") +async def update_worker_endpoint(request): + try: + data = await request.json() + worker_id = data.get("worker_id") + enabled = data.get("enabled", False) + + if worker_id is None: + return await handle_api_error(request, "Missing worker_id", 400) + + config = load_config() + worker_found = False + + for worker in config.get("workers", []): + if worker["id"] == worker_id: + worker["enabled"] = enabled + worker_found = True + break + + if not worker_found: + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + + if save_config(config): + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Failed to save config") + except Exception as e: + return await handle_api_error(request, e, 400) + +@server.PromptServer.instance.routes.post("/multigpu/config/update_setting") +async def update_setting_endpoint(request): + """Updates a specific key in the settings object.""" + try: + data = await request.json() + key = data.get("key") + value = data.get("value") + + if not key or value is None: + return await handle_api_error(request, "Missing 'key' or 'value' in request", 400) + + config = load_config() + if 'settings' not in config: + config['settings'] = {} + + config['settings'][key] = value + + if save_config(config): + return web.json_response({"status": "success", "message": f"Setting '{key}' updated."}) + else: + return await handle_api_error(request, "Failed to save config") + except Exception as e: + return await handle_api_error(request, e, 400) + + +@server.PromptServer.instance.routes.post("/multigpu/prepare_job") +async def prepare_job_endpoint(request): + try: + data = await request.json() + multi_job_id = data.get('multi_job_id') + if not multi_job_id: + return await handle_api_error(request, "Missing multi_job_id", 400) + + async with PENDING_JOBS_LOCK: + if multi_job_id not in PENDING_JOBS: + PENDING_JOBS[multi_job_id] = asyncio.Queue() + + print(f"[MultiGPU] Prepared queue for job {multi_job_id}") + return web.json_response({"status": "success"}) + except Exception as e: + return await handle_api_error(request, e) + +@server.PromptServer.instance.routes.post("/multigpu/clear_memory") +async def clear_memory_endpoint(request): + print("[MultiGPU] Received request to clear VRAM.") + try: + gc.collect() + mm.unload_all_models() + mm.soft_empty_cache() + print("[MultiGPU] VRAM cleared successfully.") + return web.json_response({"status": "success", "message": "GPU memory cleared."}) + except Exception as e: + return await handle_api_error(request, f"Error clearing VRAM: {e}") + +# --- Global State & Callback Endpoint --- +PENDING_JOBS = {} +PENDING_JOBS_LOCK = asyncio.Lock() + +@server.PromptServer.instance.routes.post("/multigpu/job_complete") +async def job_complete_endpoint(request): + try: + data = await request.post() + multi_job_id = data.get('multi_job_id') + image_file = data.get('image') + + 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,] + + async with PENDING_JOBS_LOCK: + if multi_job_id in PENDING_JOBS: + await PENDING_JOBS[multi_job_id].put(tensor) + print(f"[MultiGPU] Received result for job {multi_job_id}") + 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) + +# --- Collector Node --- +class MultiGPUCollectorNode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { "images": ("IMAGE",) }, + "hidden": { + "multi_job_id": ("STRING", {"default": ""}), + "is_worker": ("BOOLEAN", {"default": False}), + "master_url": ("STRING", {"default": ""}), + "enabled_worker_ids": ("STRING", {"default": "[]"}), + # --- FIX [7 of 7]: Type Mismatch --- + # Changed from FLOAT to INT for type safety and clarity. + "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}), + # Added a dynamic timeout for staggered mode + "stagger_timeout": ("FLOAT", {"default": 300.0, "min": 0.0}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "run" + CATEGORY = "image" + + def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, stagger_timeout=300.0): + if not multi_job_id: + return (images,) + + main_loop = server.PromptServer.instance.loop + future = Future() + + async def run_and_set_future(): + try: + result = await self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, stagger_timeout) + future.set_result(result) + except Exception as e: + future.set_exception(e) + + asyncio.run_coroutine_threadsafe(run_and_set_future(), main_loop) + return future.result() + + async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index): + """Helper method to send a single image to the master""" + img_np = (image_tensor.cpu().numpy() * 255.).astype(np.uint8) + img = Image.fromarray(img_np) + byte_io = io.BytesIO() + img.save(byte_io, format='PNG') + byte_io.seek(0) + + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('image', byte_io, filename=f'image_{image_index}.png', content_type='image/png') + + try: + async with aiohttp.ClientSession() as session: + async with session.post(f"{master_url}/multigpu/job_complete", data=data) as response: + response.raise_for_status() + except Exception as e: + print(f"[MultiGPU Worker] Failed to send image {image_index+1} to master: {e}") + + async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, stagger_timeout=300.0): + if is_worker: + # Worker mode: send images to master + print(f"[MultiGPU Worker] Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master at {master_url}") + + tasks = [self.send_image_to_master(images[i], multi_job_id, master_url, i) for i in range(images.shape[0])] + await asyncio.gather(*tasks) + return (images,) + else: + # Master mode: collect images from workers + num_workers = len(json.loads(enabled_worker_ids)) + if num_workers == 0: + return (images,) + + images_on_cpu = images.cpu() + total_images_to_collect = int(num_workers * worker_batch_size) + print(f"[MultiGPU Master] Job {multi_job_id}: Waiting for {total_images_to_collect} image(s)...") + + all_tensors = [images_on_cpu] + + async with PENDING_JOBS_LOCK: + if multi_job_id not in PENDING_JOBS: + PENDING_JOBS[multi_job_id] = asyncio.Queue() + q = PENDING_JOBS[multi_job_id] + + try: + # Collect all images with a potentially dynamic timeout + for i in range(total_images_to_collect): + worker_tensor = await asyncio.wait_for(q.get(), timeout=stagger_timeout) + all_tensors.append(worker_tensor) + print(f"[MultiGPU Master] Collected image {i+1}/{total_images_to_collect}") + except asyncio.TimeoutError: + print(f"[MultiGPU Master] Timeout waiting for worker images (got {len(all_tensors)-1}/{total_images_to_collect})") + finally: + # Clean up job queue + async with PENDING_JOBS_LOCK: + if multi_job_id in PENDING_JOBS: + del PENDING_JOBS[multi_job_id] + + combined = torch.cat(all_tensors, dim=0) + print(f"[MultiGPU Master] Job {multi_job_id} complete. Combined {combined.shape[0]} images total.") + return (combined,) + +NODE_CLASS_MAPPINGS = { "MultiGPUCollector": MultiGPUCollectorNode } +NODE_DISPLAY_NAME_MAPPINGS = { "MultiGPUCollector": "Multi-GPU Collector" } \ No newline at end of file diff --git a/web/gpupanel.js b/web/gpupanel.js new file mode 100644 index 0000000..76e16ac --- /dev/null +++ b/web/gpupanel.js @@ -0,0 +1,452 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +class MultiGPUExtension { + // 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;", + stagger: "background-color: #224422; margin-top: 8px;", + success: "background-color: #3ca03c;", + error: "background-color: #c04c4c;" + }; + + constructor() { + this.config = null; + this.originalQueuePrompt = api.queuePrompt.bind(api); + this.isStaggeredRun = false; + + // --- FIX [1 of 7]: Constructor Race Condition --- + // The interceptor is now set up *after* the config is guaranteed to be loaded. + this.loadConfig().then(() => { + this.registerSidebarTab(); + this.setupInterceptor(); + }); + } + + // --- 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 { + const response = await fetch(`${window.location.origin}/multigpu/config`); + if (response.ok) { + this.config = await response.json(); + console.log("[MultiGPU] Loaded config:", this.config); + } else { + console.error("[MultiGPU] Failed to load config"); + this.config = { workers: [], settings: {} }; + } + } catch (error) { + console.error("[MultiGPU] Error loading config:", error); + 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 fetch(`${window.location.origin}/multigpu/config/update_worker`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ worker_id: workerId, enabled: enabled }) + }); + } catch (error) { + console.error("[MultiGPU] Error updating worker:", error); + } + } + + async _updateSetting(key, value) { + try { + await fetch(`${window.location.origin}/multigpu/config/update_setting`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ key, value }) + }); + } catch (error) { + console.error(`[MultiGPU] Error updating setting '${key}':`, error); + } + } + + // --- UI Rendering --- + + registerSidebarTab() { + app.extensionManager.registerSidebarTab({ + id: "multiGPU", + icon: "pi pi-server", + title: "Multi-GPU", + tooltip: "Multi-GPU Control Panel", + type: "custom", + render: (el) => this.renderSidebarContent(el) + }); + } + + _createButton(text, onClick, style) { + const button = document.createElement("button"); + button.textContent = text; + button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + style; + button.onclick = onClick; + return button; + } + + async _handleClearMemory(button) { + const originalText = button.textContent; + button.textContent = "Clearing..."; + button.disabled = true; + + try { + const urlsToClear = [{ name: 'Master', url: window.location.origin }, ...this.enabledWorkers.map(w => ({ name: w.name, url: `http://localhost:${w.port}` }))]; + + const promises = urlsToClear.map(target => + fetch(`${target.url}/multigpu/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!"; + button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + } else { + button.textContent = "Error! See Console"; + button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + } + + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory; + }, 3000); + } finally { + // --- FIX [2 of 7]: Button State --- + // Ensure the button is re-enabled even if an error occurs. + button.disabled = false; + } + } + + renderSidebarContent(el) { + el.innerHTML = ''; + const container = document.createElement("div"); + container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: 100%;"; + + const title = document.createElement("h3"); + title.textContent = "Multi-GPU Control"; + title.style.cssText = "margin: 0 0 15px 0; padding-bottom: 10px; border-bottom: 1px solid #444;"; + container.appendChild(title); + + 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"); + (this.config?.workers || []).forEach(worker => { + const gpuDiv = document.createElement("div"); + gpuDiv.style.cssText = "margin-bottom: 8px; padding: 8px; background: #2a2a2a; border-radius: 4px;"; + const checkbox = document.createElement("input"); + checkbox.type = "checkbox"; + checkbox.id = `gpu-${worker.id}`; + checkbox.checked = worker.enabled; + checkbox.onchange = async (e) => { + await this.updateWorkerEnabled(worker.id, e.target.checked); + this.updateSummary(); + }; + const label = document.createElement("label"); + label.htmlFor = `gpu-${worker.id}`; + label.style.cssText = "cursor: pointer; display: flex; align-items: center; gap: 8px;"; + const gpuInfo = document.createElement("span"); + gpuInfo.innerHTML = `${worker.name}
CUDA ${worker.cuda_device} • Port ${worker.port}`; + label.appendChild(checkbox); + label.appendChild(gpuInfo); + gpuDiv.appendChild(label); + gpuList.appendChild(gpuDiv); + }); + gpuSection.appendChild(gpuList); + container.appendChild(gpuSection); + + const actionsSection = document.createElement("div"); + actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; + + const clearMemButton = this._createButton("Clear VRAM", (e) => this._handleClearMemory(e.target), MultiGPUExtension.BUTTON_STYLES.clearMemory); + actionsSection.appendChild(clearMemButton); + + const staggerButton = this._createButton("Staggered Start", (e) => { + // --- FIX [3 of 7]: Button State & Race Condition --- + const button = e.target; + button.disabled = true; + try { + this.isStaggeredRun = true; + app.queuePrompt(0, 1); + } finally { + // Re-enable the button after a short delay to allow the prompt to be processed. + setTimeout(() => { button.disabled = false; }, 1000); + } + }, MultiGPUExtension.BUTTON_STYLES.stagger); + staggerButton.id = "stagger-queue-button"; + staggerButton.title = "Run on Master immediately, then on each Worker with a delay."; + actionsSection.appendChild(staggerButton); + + const settingsSection = document.createElement("div"); + settingsSection.style.cssText = "margin-top: 10px;"; + + const delayDiv = document.createElement("div"); + delayDiv.style.cssText = "display: flex; align-items: center; justify-content: space-between;"; + const delayLabel = document.createElement("label"); + delayLabel.textContent = "Stagger Delay (s)"; + delayLabel.htmlFor = "stagger-delay-input"; + delayLabel.style.cssText = "font-size: 12px; color: #888;"; + const delayInput = document.createElement("input"); + delayInput.id = "stagger-delay-input"; + delayInput.type = "number"; + delayInput.min = "0"; + delayInput.step = "0.1"; + delayInput.style.cssText = "width: 60px; background-color: #222; border: 1px solid #444; border-radius: 4px; padding: 4px; color: white; font-size: 12px;"; + delayInput.value = (this.config?.settings?.stagger_delay_ms || 15000) / 1000; + delayInput.onchange = async (e) => { + const valueInSeconds = parseFloat(e.target.value); + if(isNaN(valueInSeconds) || valueInSeconds < 0) return; + const valueInMs = Math.round(valueInSeconds * 1000); + this.config.settings.stagger_delay_ms = valueInMs; + await this._updateSetting('stagger_delay_ms', valueInMs); + }; + delayDiv.appendChild(delayLabel); + delayDiv.appendChild(delayInput); + settingsSection.appendChild(delayDiv); + actionsSection.appendChild(settingsSection); + + container.appendChild(actionsSection); + + const summarySection = document.createElement("div"); + summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;"; + const summary = document.createElement("div"); + summary.id = "multigpu-summary"; + summary.style.cssText = "font-size: 11px; color: #888;"; + summarySection.appendChild(summary); + container.appendChild(summarySection); + el.appendChild(container); + this.updateSummary(); + } + + updateSummary() { + const summaryEl = document.getElementById('multigpu-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."; + } + } + } + + // --- Core Logic & Execution --- + + setupInterceptor() { + api.queuePrompt = async (number, prompt) => { + // --- FIX [4 of 7]: isStaggeredRun Race Condition --- + // The check for the flag is now safer. It's read and reset in a way that + // prevents the race condition caused by the original 'finally' block. + if (this.isEnabled && this.findNodesByClass(prompt.output, "MultiGPUCollector").length > 0) { + const handleStaggered = this.isStaggeredRun; + this.isStaggeredRun = false; // Consume the flag immediately + + console.log(`[MultiGPU] Intercepting prompt. Staggered mode: ${handleStaggered}`); + if (handleStaggered) { + return await this.executeStaggeredMultiGPU(prompt); + } else { + return await this.executeParallelMultiGPU(prompt); + } + } + return this.originalQueuePrompt(number, prompt); + }; + } + + async executeParallelMultiGPU(promptWrapper) { + try { + console.log("[MultiGPU] Collector node found. Activating standard Multi-GPU execution."); + const multi_job_id = "mgpu_" + Date.now(); + const enabledWorkers = this.enabledWorkers; + const batchSize = this._getBatchSizeFromPrompt(promptWrapper); + await this._prepareMultiGpuJob(multi_job_id); + + const jobs = []; + const participants = ['master', ...enabledWorkers.map(w => w.id)]; + + for (const participantId of participants) { + const jobApiPrompt = this._prepareApiPromptForParticipant( + promptWrapper.output, multi_job_id, participantId, + { enabled_worker_ids: enabledWorkers.map(w => w.id), worker_batch_size: batchSize } + ); + if (participantId === 'master') { + jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); + } else { + const worker = this.config.workers.find(w => w.id === participantId); + if (worker) jobs.push({ type: 'worker', worker, prompt: jobApiPrompt, workflow: promptWrapper.workflow }); + } + } + await this._executeJobs(jobs); + return { "prompt_id": "multi-gpu-parallel-job-dispatched" }; + } catch (error) { + console.error("[MultiGPU] Parallel execution failed:", error); + alert(`[MultiGPU] Parallel execution failed: ${error.message}`); + throw error; + } + } + + async executeStaggeredMultiGPU(promptWrapper) { + try { + console.log("[MultiGPU] Staggered Start activated."); + const multi_job_id = "mgpu_" + Date.now(); + const enabledWorkers = this.enabledWorkers; + const batchSize = this._getBatchSizeFromPrompt(promptWrapper); + await this._prepareMultiGpuJob(multi_job_id); + + // --- FIX [5 of 7]: Synchronization --- + // Calculate a dynamic timeout for the master node. + // Baseline of 300s + stagger delay for each worker. + const staggerDelay = this.config.settings?.stagger_delay_ms || 15000; + const requiredTimeout = ((enabledWorkers.length * staggerDelay) / 1000) + 300; + console.log(`[MultiGPU] Master collector timeout set to ${requiredTimeout}s to accommodate staggered workers.`); + + const masterPrompt = this._prepareApiPromptForParticipant( + promptWrapper.output, multi_job_id, 'master', + { + enabled_worker_ids: enabledWorkers.map(w => w.id), + worker_batch_size: batchSize, + stagger_timeout: requiredTimeout // Pass the new timeout to the master + } + ); + this.originalQueuePrompt(0, { ...promptWrapper, output: masterPrompt }); + console.log("[MultiGPU] Dispatched to Master immediately."); + + console.log(`[MultiGPU] Beginning staggered dispatch to workers with a ${staggerDelay / 1000}s delay.`); + + for (const worker of enabledWorkers) { + await new Promise(resolve => setTimeout(resolve, staggerDelay)); + // --- FIX [6 of 7]: Missing Parameters --- + // The options object was missing here, which prevented the 'master_url' + // from being set correctly for the worker. + const workerPrompt = this._prepareApiPromptForParticipant( + promptWrapper.output, + multi_job_id, + worker.id, + { enabled_worker_ids: [], worker_batch_size: 0 } // This ensures master_url is set + ); + await this._dispatchToWorker(worker, workerPrompt, promptWrapper.workflow); + } + return { "prompt_id": "multi-gpu-staggered-job-dispatched" }; + } catch (error) { + console.error("[MultiGPU] Staggered execution failed:", error); + alert(`[MultiGPU] Staggered execution failed: ${error.message}`); + throw error; + } + } + + async _executeJobs(jobs) { + const promises = jobs.map(job => + job.type === 'master' + ? this.originalQueuePrompt(0, job.promptWrapper) + : this._dispatchToWorker(job.worker, job.prompt, job.workflow) + ); + await Promise.all(promises); + } + + // --- Helper Methods --- + + findNodesByClass(apiPrompt, className) { + return Object.entries(apiPrompt) + .filter(([, nodeData]) => nodeData.class_type === className) + .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); + } + + findBatchNode(workflow) { + return workflow.nodes.find(node => { + const title = node.title?.trim().toUpperCase(); + const displayName = node.properties?.["Node name for S&R"]?.trim().toUpperCase(); + return title === "BATCH" || title === "BATCH SIZE" || displayName === "BATCH" || displayName === "BATCH SIZE"; + }); + } + + _getBatchSizeFromPrompt(promptWrapper) { + const batchNode = this.findBatchNode(promptWrapper.workflow); + if (!batchNode) throw new Error("MultiGPU: Could not find a 'Batch' node."); + return promptWrapper.output[String(batchNode.id)].inputs.batch_size || promptWrapper.output[String(batchNode.id)].inputs.value || 1; + } + + _prepareApiPromptForParticipant(baseApiPrompt, multi_job_id, participantId, options = {}) { + const jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); + const collectorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUCollector"); + const isMaster = participantId === 'master'; + + for (const collector of collectorNodes) { + const { inputs } = jobApiPrompt[collector.id]; + inputs.multi_job_id = multi_job_id; + inputs.is_worker = !isMaster; + if (isMaster) { + inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); + inputs.worker_batch_size = options.worker_batch_size || 1; + // Pass the dynamic timeout for staggered mode + if (options.stagger_timeout) { + inputs.stagger_timeout = options.stagger_timeout; + } + } else { + inputs.master_url = window.location.origin; + } + } + return jobApiPrompt; + } + + async _prepareMultiGpuJob(multi_job_id) { + try { + await fetch(`${window.location.origin}/multigpu/prepare_job`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ multi_job_id }) + }); + } catch (error) { + console.error("[MultiGPU] Error preparing job:", error); + throw error; + } + } + + async _dispatchToWorker(worker, prompt, workflow) { + const workerUrl = `http://localhost:${worker.port}`; + console.log(`[MultiGPU] Dispatching to worker: ${worker.name} (${workerUrl})`); + const promptToSend = { + prompt, + extra_data: { extra_pnginfo: { workflow } }, + client_id: api.clientId + }; + try { + await fetch(`${workerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify(promptToSend) + }); + } catch (e) { + console.error(`[MultiGPU] Failed to connect to worker ${worker.name} at ${workerUrl}`, e); + } + } +} + +app.registerExtension({ + name: "MultiGPU.Panel", + async setup() { + new MultiGPUExtension(); + } +}); \ No newline at end of file