diff --git a/README.md b/README.md index 3ce1585..992052a 100644 --- a/README.md +++ b/README.md @@ -1,256 +1,212 @@ # 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. +> **Supercharge your ComfyUI workflows with distributed GPU processing** -A custom node extension for ComfyUI that enables distributed image generation across multiple GPUs through a master-worker architecture. +A powerful extension for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that enables distributed processing across multiple GPUs and machines. Speed up your image generation and upscaling workflows by leveraging all available GPU resources in your network. -![image](https://github.com/user-attachments/assets/1e95e9e6-5237-4b24-8500-ac3b2703a1ea) - -## 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 -- **Parallel Execution**: All enabled GPUs process batches simultaneously for maximum speed -- **Automatic Seed Distribution**: Smart seed offsetting across workers -- **Web UI Integration**: Seamless control panel integrated into ComfyUI's sidebar -- **Smart Launcher**: Unified launcher with multiple modes and instance detection -- **VRAM Management**: Clear worker VRAM remotely through the UI -- **Flexible Configuration**: JSON-based configuration with runtime updates +- **Distributed Processing** - Split workloads across multiple GPUs automatically +- **Distributed Upscaling** - Tile-based upscaling with intelligent work distribution +- **Automatic Worker Management** - Launch and monitor workers from the UI +- **Network Support** - Use GPUs across different machines on your network +- **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 and optimization -## How It Works - -1. Launch multiple ComfyUI instances using the unified launcher -2. One instance acts as the master (typically on CUDA device 0) -3. Add the "Multi-GPU Collector" node to your workflow -4. Optionally add "Multi-GPU Distributor" nodes for automatic seed distribution -5. The extension automatically distributes batches across available GPUs -6. 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. - -## Nodes - -### Multi-GPU Distributor -Handles automatic seed distribution for diverse results: -- Master uses the original seed value -- Workers automatically get offset seeds (seed + worker_index + 1) -- Can connect to any node that accepts seed/batch_size inputs -- Ensures each GPU generates unique images - -### Multi-GPU Collector -The core node that handles job result collection: -- Automatically detects enabled workers -- Collects and combines results in a predictable order +--- ## Requirements -- Virtual environment with ComfyUI installed -- ComfyUI installation with ComfyUI CLI (`comfy` command) -- Multiple NVIDIA GPUs +- ComfyUI installation +- PyTorch with CUDA support - Python 3.8+ +--- + ## Installation -1. Clone this repository into your ComfyUI custom nodes directory: +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 +2. **Restart ComfyUI** - the extension will create a default configuration file on first load -## Usage +--- -### Quick Start (Windows) +## Quick Start -1. Copy the `launch_distributed.bat` file to your ComfyUI root directory (alongside your `venv` folder) +### Adding Local Workers -2. Edit the `gpu_config.json` file in `ComfyUI/custom_nodes/ComfyUI-Distributed/` to match your GPU setup: - - Adjust the number of workers based on your available GPUs - - Update CUDA device numbers if needed - - Set appropriate ports (ensure they're not in use) - - Add custom launch arguments in `extra_args` +1. **Open** the Distributed GPU panel (sidebar on the left) +2. **Click** "Add Worker" in the UI +3. **Configure** your worker: + - **Name**: Descriptive name for the worker + - **Host**: `localhost` for same machine + - **Port**: Unique port number (e.g., 8189, 8190...) + - **CUDA Device**: GPU index from `nvidia-smi` + - **Extra Args**: Optional ComfyUI arguments +4. **Save** and optionally launch the worker -3. Run the launcher from the ComfyUI root directory: - ```bash - launch_distributed.bat - ``` - - Or use command-line options: - ```bash - launch_distributed.bat master # Launch master only - launch_distributed.bat all-fast # Launch all with 2-second delays - launch_distributed.bat all-slow # Launch all, wait for each to be ready - launch_distributed.bat venv # Open command prompt with venv activated - launch_distributed.bat update # Update ComfyUI - launch_distributed.bat update-nodes # Update all custom nodes - ``` +### Adding Remote Workers -4. Open the master instance (check the launcher output for the URL) +1. **Follow** the same steps as above, but: +2. **Add** `--listen` to the 'Extra Args' field ⚠️ **Required!** +3. **Open** the worker port in your firewall +4. **Launch** ComfyUI on the remote machine +5. **Add** the worker using the remote machine's IP address -5. Add the "Multi-GPU Collector" node to your workflow +> **Tip**: If the remote machine has multiple GPUs, create multiple workers with different CUDA devices and ports -6. Use the Multi-GPU sidebar panel to: - - Enable/disable worker GPUs - - View worker status - - Clear worker VRAM +### Configuration Tips -### Configuration +| 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 | -The system uses a `gpu_config.json` file for configuration. Example configuration: +**Common Extra Args:** +- `--lowvram` - For GPUs with less memory +- `--listen` - **Required** for remote workers +- `--highvram` - For GPUs with 12GB+ VRAM -```json -{ - "master": { - "port": 8080, - "cuda_device": 0, - "extra_args": "--mmap-torch-files --listen" - }, - "workers": [ - { - "id": 1, - "name": "GPU 1", - "cuda_device": 1, - "port": 8180, - "enabled": true, - "extra_args": "--mmap-torch-files --listen" - }, - { - "id": 2, - "name": "GPU 2", - "cuda_device": 2, - "port": 8280, - "enabled": true, - "extra_args": "--mmap-torch-files --listen" - }, - { - "id": 3, - "name": "GPU 3", - "cuda_device": 3, - "port": 8380, - "enabled": false, - "extra_args": "--mmap-torch-files --listen --reserve-vram 2" - } - ], - "settings": {} -} +--- + +## Nodes + +### Distributed Collector +> Collects and combines results from distributed processing + +**Usage**: Place after any image generation node to enable distributed processing. Works automatically when workers are enabled. + +### Distributed Seed +> Ensures unique seeds across distributed workers for varied generation + +**Usage**: Connect to any seed input. Each worker automatically receives an offset seed for variety. + +### Ultimate SD Upscale Distributed +> Distributed version of Ultimate SD Upscale that processes tiles across multiple GPUs + +**Usage**: +1. Upscale your image with a regular upscale model (ESRGAN, etc.) +2. Feed the upscaled image into this node +3. Configure tile settings based on your VRAM +4. Enable workers for faster processing + +--- + +## UI Features + +### Distributed GPU Panel + +The control center for your distributed setup: + +| Feature | Description | +|---------|-------------| +| **Worker Status** | Real-time status indicators | +| **Launch/Stop** | Control individual workers | +| **Clear Memory** | Free VRAM on all workers | +| **Interrupt** | Stop current processing | +| **Auto-detect IP** | Find best network interface | +| **Worker Logs** | View real-time logs | + +--- + +## Workflow Examples + +### Basic Distributed Generation + +``` +Your Workflow → Distributed Seed → KSampler → VAE Decode → Distributed Collector → Output ``` -**Configuration Options:** -- `master`: Master instance configuration - - `port`: Port number for the master instance - - `cuda_device`: CUDA device index - - `extra_args`: Custom launch arguments -- `workers`: Array of worker configurations - - `id`: Unique worker identifier - - `name`: Display name in UI - - `cuda_device`: CUDA device index - - `port`: Port number for the worker - - `enabled`: Whether the worker is active - - `extra_args`: Custom launch arguments - - `host` (optional): Remote host address -- `settings`: Additional settings (reserved for future use) +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! -### Remote Workers +### Distributed Upscaling -Remote workers are fully supported! Add a `host` field to any worker configuration: - -```json -{ - "id": 3, - "name": "Remote GPU", - "host": "192.168.1.100", - "port": 8190, - "enabled": true, - "extra_args": "--enable-cors-header --listen" -} +``` +Load Image → Upscale Model → Ultimate SD Upscale Distributed → Save Image ``` -Ensure remote workers: -- Are started with `--listen --enable-cors-header` -- Are accessible from the master instance -- Have the same models and custom nodes installed +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 -## Workflow Integration - -### Basic Setup -1. Add "Multi-GPU Collector" node after your image generation pipeline -2. The node automatically handles distribution and collection -3. Connect subsequent nodes normally - they'll receive the combined batch - -### Seed Distribution -This node is optional, it will help you control seeds across workers. Without it, you will generate exactly the same workflow as the master (same seed): -1. Add "Multi-GPU Seed" node(s) before your samplers -2. Connect the seed output to your sampling node(s) -3. Each GPU will automatically use different seeds: - - Master: original seed - - Worker 0: seed + 1 - - Worker 1: seed + 2 - - etc. - -### Result Ordering -Images are returned in a predictable order: -1. All images from master (in batch order) -2. All images from worker 1 (in batch order) -3. All images from worker 2 (in batch order) -4. etc. - -## API Endpoints - -The extension provides REST API endpoints for integration: - -- `GET /multigpu/config` - Retrieve current configuration -- `POST /multigpu/config/update_worker` - Enable/disable a worker -- `POST /multigpu/config/update_setting` - Update configuration settings -- `POST /multigpu/prepare_job` - Prepare a multi-GPU job -- `POST /multigpu/clear_memory` - Clear VRAM on workers -- `POST /multigpu/job_complete` - Worker callback for job completion - -## Performance Considerations - -- **Parallel Execution**: All enabled GPUs process simultaneously for maximum throughput -- **Batch Size**: Total generation = batch_size × (1 master + number_of_enabled_workers) -- **Smart Launching**: The launcher only starts instances that aren't already running -- **Memory Management**: Use appropriate `extra_args` in configuration for memory optimization -- **Network Overhead**: Minimal - only final images are transferred between instances +--- ## Troubleshooting -- **Workers not connecting**: Ensure CORS headers are enabled (`--enable-cors-header`) -- **VRAM errors**: Add memory management flags in `extra_args` (e.g., `--highvram`, `--disable-smart-memory`) -- **Port conflicts**: Ensure each instance uses a unique port -- **Launcher errors**: Ensure the batch file is in the ComfyUI root directory with the `venv` folder -- **Batch node not found**: Ensure your workflow has a node titled "BATCH" or "BATCH SIZE" -- **ComfyUI CLI not found**: Make sure ComfyUI CLI is installed (`pip install comfyui-cli`) +### 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 workers have the same models loaded +- Check that custom nodes are installed on all 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 +
+ +### Debug Mode + +Enable detailed logging for troubleshooting: + +1. **UI Method**: Settings → Debug Mode → Enable +2. **Check** console output for detailed logs + +--- ## Development This project is under active development. Contributions are welcome! -### Planned Features/Improvements +### Planned Features -- ~Custom launcher~ - - ~Launch local workers~ - - Launch remote workers (via SSH) -- ~Support for remote workers over network~ - - Worker status (ready/in progress/offline) - - Restart worker button -- Improve worker timeout logic -- ~Better VRAM management for multi-GPU setups~ -- ~Better Seed handling~ -- ~Handle worker failure~ -- ~Handle multiple collector nodes~ -- ~Compatibility with FLUX Continuum `CTRL+SHIFT+C` shortcut~ [Commit 027629c](https://github.com/robertvoy/ComfyUI-Flux-Continuum/commit/027629c753dd3aae1ceeff5214ceb701943dd3fe) +- [ ] Remote worker control via SSH +- [ ] View remote worker logs in UI +- [ ] Improve worker timeout logic +- [ ] Authentication for network security +- [ ] Automatic model synchronization -## License +--- -Same as ComfyUI - see ComfyUI repository for details. +
-## Acknowledgments +**Made with ❤️ for the ComfyUI community** -Built as an extension for the amazing [ComfyUI](https://github.com/comfyanonymous/ComfyUI) project. +[Report Issues](https://github.com/yourusername/ComfyUI-Distributed/issues) • [Request Features](https://github.com/yourusername/ComfyUI-Distributed/issues) + +
\ No newline at end of file diff --git a/__init__.py b/__init__.py index bc655ec..9bc2dde 100644 --- a/__init__.py +++ b/__init__.py @@ -5,19 +5,31 @@ import sys 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 +from .distributed import ( + NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS +) + +# Import utilities +from .utils.config import ensure_config_exists, CONFIG_FILE +from .utils.logging import debug_log + +# Import distributed upscale nodes +from .distributed_upscale import ( + NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS ) WEB_DIRECTORY = "./web" -# Use the centralized config management from multigpu.py ensure_config_exists() +# Merge node 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'] -print("[MultiGPU] Loaded Multi-GPU nodes.") -print(f"[MultiGPU] Config file: {CONFIG_FILE}") \ No newline at end of file +debug_log("Loaded Distributed nodes.") +debug_log(f"Config file: {CONFIG_FILE}") +debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}") \ No newline at end of file diff --git a/distributed.py b/distributed.py new file mode 100644 index 0000000..54e0f59 --- /dev/null +++ b/distributed.py @@ -0,0 +1,1483 @@ +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 +import subprocess +import platform +import time +import atexit +import signal + +# Import shared utilities +from .utils.logging import debug_log, log +from .utils.config import CONFIG_FILE, get_default_config, load_config, save_config, ensure_config_exists +from .utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous +from .utils.process import is_process_alive, terminate_process, get_python_executable +from .utils.network import handle_api_error, get_server_port, get_server_loop, get_client_session, cleanup_client_session +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 +) + +# Try to import psutil for better process management +try: + import psutil + PSUTIL_AVAILABLE = True +except ImportError: + log("psutil not available, using fallback process management") + PSUTIL_AVAILABLE = False + +# Register cleanup for aiohttp session +def cleanup(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(cleanup_client_session()) + loop.close() + +atexit.register(cleanup) + +# --- API Endpoints --- +@server.PromptServer.instance.routes.get("/distributed/config") +async def get_config_endpoint(request): + config = load_config() + return web.json_response(config) + +@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}") +async def queue_status_endpoint(request): + """Check if a job queue is initialized.""" + try: + job_id = request.match_info['job_id'] + + # Import to ensure initialization + from .distributed_upscale import ensure_tile_jobs_initialized + prompt_server = ensure_tile_jobs_initialized() + + async with prompt_server.distributed_tile_jobs_lock: + exists = job_id in prompt_server.distributed_pending_tile_jobs + + debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}") + return web.json_response({"exists": exists, "job_id": job_id}) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching") +async def clear_launching_state(request): + """Clear the launching flag when worker is confirmed running.""" + try: + data = await request.json() + worker_id = str(data.get('worker_id')) + + if not worker_id: + return await handle_api_error(request, "worker_id is required", 400) + + # Clear launching flag in managed processes + if worker_id in worker_manager.processes: + if 'launching' in worker_manager.processes[worker_id]: + del worker_manager.processes[worker_id]['launching'] + worker_manager.save_processes() + debug_log(f"Cleared launching state for worker {worker_id}") + + return web.json_response({"status": "success"}) + except Exception as e: + return await handle_api_error(request, e, 500) + +@server.PromptServer.instance.routes.get("/distributed/network_info") +async def get_network_info_endpoint(request): + """Get network interfaces and recommend best IP for master.""" + import socket + + def get_network_ips(): + """Get all network IPs, trying multiple methods.""" + ips = [] + hostname = socket.gethostname() + + # Method 1: Try socket.getaddrinfo + try: + addr_info = socket.getaddrinfo(hostname, None) + for info in addr_info: + ip = info[4][0] + if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now + ips.append(ip) + except: + pass + + # Method 2: Try to connect to external server and get local IP + try: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s.connect(("8.8.8.8", 80)) # Google DNS + local_ip = s.getsockname()[0] + s.close() + if local_ip not in ips: + ips.append(local_ip) + except: + pass + + # Method 3: Platform-specific commands + + try: + if platform.system() == "Windows": + # Windows ipconfig + result = subprocess.run(["ipconfig"], capture_output=True, text=True) + lines = result.stdout.split('\n') + for i, line in enumerate(lines): + if 'IPv4' in line and i + 1 < len(lines): + ip = lines[i].split(':')[-1].strip() + if ip and ip not in ips: + ips.append(ip) + else: + # Unix/Linux/Mac ifconfig or ip addr + try: + result = subprocess.run(["ip", "addr"], capture_output=True, text=True) + except: + result = subprocess.run(["ifconfig"], capture_output=True, text=True) + + import re + ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)') + for match in ip_pattern.finditer(result.stdout): + ip = match.group(1) + if ip and ip not in ips: + ips.append(ip) + except: + pass + + return ips + + def get_recommended_ip(ips): + """Choose the best IP for master-worker communication.""" + # Priority order: + # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x) + # 2. Other non-localhost IPs + # 3. Localhost as last resort + + private_ips = [] + public_ips = [] + + for ip in ips: + if ip.startswith('127.') or ip == 'localhost': + continue + elif (ip.startswith('192.168.') or + ip.startswith('10.') or + (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)): + private_ips.append(ip) + else: + public_ips.append(ip) + + # Prefer private IPs + if private_ips: + # Prefer 192.168 range as it's most common + for ip in private_ips: + if ip.startswith('192.168.'): + return ip + return private_ips[0] + elif public_ips: + return public_ips[0] + elif ips: + return ips[0] + else: + return None + + try: + hostname = socket.gethostname() + all_ips = get_network_ips() + recommended_ip = get_recommended_ip(all_ips) + + return web.json_response({ + "status": "success", + "hostname": hostname, + "all_ips": all_ips, + "recommended_ip": recommended_ip, + "message": "Auto-detected network configuration" + }) + except Exception as e: + return web.json_response({ + "status": "error", + "message": str(e), + "hostname": "unknown", + "all_ips": [], + "recommended_ip": None + }) + +@server.PromptServer.instance.routes.post("/distributed/config/update_worker") +async def update_worker_endpoint(request): + try: + data = await request.json() + worker_id = data.get("worker_id") + + 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: + # Update all provided fields + if "enabled" in data: + worker["enabled"] = data["enabled"] + if "name" in data: + worker["name"] = data["name"] + if "port" in data: + worker["port"] = data["port"] + + # Handle host field - remove it if None + if "host" in data: + if data["host"] is None: + worker.pop("host", None) + else: + worker["host"] = data["host"] + + # Handle cuda_device field - remove it if None + if "cuda_device" in data: + if data["cuda_device"] is None: + worker.pop("cuda_device", None) + else: + worker["cuda_device"] = data["cuda_device"] + + # Handle extra_args field - remove it if None + if "extra_args" in data: + if data["extra_args"] is None: + worker.pop("extra_args", None) + else: + worker["extra_args"] = data["extra_args"] + worker_found = True + break + + if not worker_found: + # If worker not found and all required fields are provided, create new worker + if all(key in data for key in ["name", "port", "cuda_device"]): + new_worker = { + "id": worker_id, + "name": data["name"], + "host": data.get("host", "localhost"), + "port": data["port"], + "cuda_device": data["cuda_device"], + "enabled": data.get("enabled", False), + "extra_args": data.get("extra_args", "") + } + if "workers" not in config: + config["workers"] = [] + config["workers"].append(new_worker) + worker_found = True + else: + return await handle_api_error(request, f"Worker {worker_id} not found and missing required fields for creation", 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("/distributed/config/delete_worker") +async def delete_worker_endpoint(request): + try: + data = await request.json() + worker_id = data.get("worker_id") + + if worker_id is None: + return await handle_api_error(request, "Missing worker_id", 400) + + config = load_config() + workers = config.get("workers", []) + + # Find and remove the worker + worker_index = -1 + for i, worker in enumerate(workers): + if worker["id"] == worker_id: + worker_index = i + break + + if worker_index == -1: + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + + # Remove the worker + removed_worker = workers.pop(worker_index) + + if save_config(config): + return web.json_response({ + "status": "success", + "message": f"Worker {removed_worker.get('name', worker_id)} deleted" + }) + 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("/distributed/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("/distributed/config/update_master") +async def update_master_endpoint(request): + """Updates master configuration.""" + try: + data = await request.json() + + config = load_config() + if 'master' not in config: + config['master'] = {} + + # Update all provided fields + if "host" in data: + config['master']['host'] = data['host'] + if "port" in data: + config['master']['port'] = data['port'] + if "cuda_device" in data: + config['master']['cuda_device'] = data['cuda_device'] + if "extra_args" in data: + config['master']['extra_args'] = data['extra_args'] + + if save_config(config): + return web.json_response({"status": "success", "message": "Master configuration 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("/distributed/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 prompt_server.distributed_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_jobs: + prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() + + debug_log(f"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("/distributed/clear_memory") +async def clear_memory_endpoint(request): + debug_log("Received request to clear VRAM.") + try: + # Use ComfyUI's prompt server queue system like the /free endpoint does + if hasattr(server.PromptServer.instance, 'prompt_queue'): + server.PromptServer.instance.prompt_queue.set_flag("unload_models", True) + server.PromptServer.instance.prompt_queue.set_flag("free_memory", True) + debug_log("Set queue flags for memory clearing.") + + # Wait a bit for the queue to process + await asyncio.sleep(MEMORY_CLEAR_DELAY) + + # Also do direct cleanup as backup, but with error handling + import gc + import comfy.model_management as mm + + try: + mm.unload_all_models() + except AttributeError as e: + debug_log(f"Warning during model unload: {e}") + + try: + mm.soft_empty_cache() + except Exception as e: + debug_log(f"Warning during cache clear: {e}") + + for _ in range(3): + gc.collect() + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + debug_log("VRAM cleared successfully.") + return web.json_response({"status": "success", "message": "GPU memory cleared."}) + except Exception as e: + # Even if there's an error, try to do basic cleanup + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + debug_log(f"Partial VRAM clear completed with warning: {e}") + return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"}) + + +@server.PromptServer.instance.routes.post("/distributed/launch_worker") +async def launch_worker_endpoint(request): + """Launch a worker process from the UI.""" + try: + data = await request.json() + worker_id = data.get("worker_id") + + if not worker_id: + return await handle_api_error(request, "Missing worker_id", 400) + + # Find worker config + config = load_config() + worker = next((w for w in config.get("workers", []) if w["id"] == worker_id), None) + if not worker: + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + + # Ensure consistent string ID + worker_id_str = str(worker_id) + + # Check if already running (managed by this instance) + if worker_id_str in worker_manager.processes: + proc_info = worker_manager.processes[worker_id_str] + process = proc_info.get('process') + + # Check if still running + is_running = False + if process: + is_running = process.poll() is None + else: + # Restored process without subprocess object + is_running = worker_manager._is_process_running(proc_info['pid']) + + if is_running: + return web.json_response({ + "status": "error", + "message": "Worker already running (managed by UI)", + "pid": proc_info['pid'], + "log_file": proc_info.get('log_file') + }, status=409) + else: + # Process is dead, remove it + del worker_manager.processes[worker_id_str] + worker_manager.save_processes() + + # Launch the worker + try: + pid = worker_manager.launch_worker(worker) + log_file = worker_manager.processes[worker_id_str].get('log_file') + return web.json_response({ + "status": "success", + "pid": pid, + "message": f"Worker {worker['name']} launched", + "log_file": log_file + }) + except Exception as e: + return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500) + + except Exception as e: + return await handle_api_error(request, e, 400) + + +@server.PromptServer.instance.routes.post("/distributed/stop_worker") +async def stop_worker_endpoint(request): + """Stop a worker process that was launched from the UI.""" + try: + data = await request.json() + worker_id = data.get("worker_id") + + if not worker_id: + return await handle_api_error(request, "Missing worker_id", 400) + + success, message = worker_manager.stop_worker(worker_id) + + if success: + return web.json_response({"status": "success", "message": message}) + else: + return web.json_response({"status": "error", "message": message}, + status=404 if "not managed" in message else 409) + + except Exception as e: + return await handle_api_error(request, e, 400) + + +@server.PromptServer.instance.routes.get("/distributed/managed_workers") +async def get_managed_workers_endpoint(request): + """Get list of workers managed by this UI instance.""" + try: + managed = worker_manager.get_managed_workers() + return web.json_response({ + "status": "success", + "managed_workers": managed + }) + except Exception as 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.""" + try: + worker_id = request.match_info['worker_id'] + + # Ensure worker_id is string + worker_id = str(worker_id) + + # Check if we manage this worker + if worker_id not in worker_manager.processes: + return await handle_api_error(request, f"Worker {worker_id} not managed by UI", 404) + + proc_info = worker_manager.processes[worker_id] + log_file = proc_info.get('log_file') + + if not log_file or not os.path.exists(log_file): + return await handle_api_error(request, "Log file not found", 404) + + # Read last N lines (or full file if small) + lines_to_read = int(request.query.get('lines', 1000)) + + try: + # Get file size + file_size = os.path.getsize(log_file) + + with open(log_file, 'r', encoding='utf-8', errors='replace') as f: + if lines_to_read > 0 and file_size > 1024 * 1024: # If file > 1MB and limited lines requested + # Read last N lines efficiently + lines = [] + # Start from end and work backwards + f.seek(0, 2) # Go to end + file_length = f.tell() + + # Read chunks from end + chunk_size = min(CHUNK_SIZE, file_length) + while len(lines) < lines_to_read and f.tell() > 0: + # Move back and read chunk + current_pos = max(0, f.tell() - chunk_size) + f.seek(current_pos) + chunk = f.read(chunk_size) + + # Process chunk + chunk_lines = chunk.splitlines() + if current_pos > 0: + # Partial line at beginning, combine with next chunk + chunk_lines = chunk_lines[1:] + + lines = chunk_lines + lines + + # Move back for next chunk + f.seek(current_pos) + + # Take only last N lines + content = '\n'.join(lines[-lines_to_read:]) + truncated = len(lines) > lines_to_read + else: + # Read entire file + content = f.read() + truncated = False + + return web.json_response({ + "status": "success", + "content": content, + "log_file": log_file, + "file_size": file_size, + "truncated": truncated, + "lines_shown": lines_to_read if truncated else content.count('\n') + 1 + }) + + except Exception as e: + return await handle_api_error(request, f"Error reading log file: {str(e)}", 500) + + except Exception as e: + return await handle_api_error(request, e, 500) + + +# --- Worker Process Management --- +class WorkerProcessManager: + def __init__(self): + self.processes = {} # worker_id -> process info + self.load_processes() + + def find_comfy_root(self): + """Find the ComfyUI root directory.""" + # Start from current file location and go up + current_dir = os.path.dirname(os.path.abspath(__file__)) + # This file is in ComfyUI/custom_nodes/ComfyUI-Distributed/ + # So go up two levels to get to ComfyUI root + comfy_root = os.path.dirname(os.path.dirname(current_dir)) + return comfy_root + + def _find_windows_terminal(self): + """Find Windows Terminal executable.""" + # Common locations for Windows Terminal + possible_paths = [ + os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"), + os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"), + "wt.exe" # Try PATH + ] + + for path in possible_paths: + if os.path.exists(path): + return path + # Handle wildcard for WindowsApps + if '*' in path: + import glob + matches = glob.glob(path) + if matches: + return matches[0] + + # Try to find it in PATH + import shutil + wt_path = shutil.which("wt") + if wt_path: + return wt_path + + return None + + def build_launch_command(self, worker_config, comfy_root): + """Build the command to launch a worker.""" + # Use main.py directly - it's the most reliable method + main_py = os.path.join(comfy_root, "main.py") + + if os.path.exists(main_py): + cmd = [ + get_python_executable(), + main_py, + "--port", str(worker_config['port']), + "--enable-cors-header" + ] + debug_log(f"Using main.py: {main_py}") + else: + # Fallback error + raise RuntimeError(f"Could not find main.py in {comfy_root}") + + # Add any extra arguments + if worker_config.get('extra_args'): + cmd.extend(worker_config['extra_args'].split()) + + return cmd + + def launch_worker(self, worker_config, show_window=False): + """Launch a worker process with logging.""" + comfy_root = self.find_comfy_root() + + # Set up environment + env = os.environ.copy() + env['CUDA_VISIBLE_DEVICES'] = str(worker_config.get('cuda_device', 0)) + env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' + + # Pass master PID to worker so it can monitor if master is still alive + env['COMFYUI_MASTER_PID'] = str(os.getpid()) + + cmd = self.build_launch_command(worker_config, comfy_root) + + # Change to ComfyUI root directory for the process + cwd = comfy_root + + # Create log directory and file + log_dir = os.path.join(comfy_root, "logs", "workers") + os.makedirs(log_dir, exist_ok=True) + + # Use daily log files instead of timestamp + date_stamp = time.strftime("%Y%m%d") + worker_name = worker_config.get('name', f'Worker{worker_config["id"]}') + # Clean worker name for filename + safe_name = "".join(c if c.isalnum() or c in ('-', '_') else '_' for c in worker_name) + log_file = os.path.join(log_dir, f"{safe_name}_{date_stamp}.log") + + # Launch process with logging (append mode for daily logs) + with open(log_file, 'a') as log_handle: + # Write startup info to log with timestamp + log_handle.write(f"\n\n{'='*50}\n") + log_handle.write(f"=== ComfyUI Worker Session Started ===\n") + log_handle.write(f"Worker: {worker_name}\n") + log_handle.write(f"Port: {worker_config['port']}\n") + log_handle.write(f"CUDA Device: {worker_config.get('cuda_device', 0)}\n") + log_handle.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n") + log_handle.write(f"Command: {' '.join(cmd)}\n") + + # Note about worker behavior + config = load_config() + stop_on_master_exit = config.get('settings', {}).get('stop_workers_on_master_exit', True) + + if stop_on_master_exit: + log_handle.write("Note: Worker will stop when master shuts down\n") + else: + log_handle.write("Note: Worker will continue running after master shuts down\n") + + log_handle.write("=" * 30 + "\n\n") + log_handle.flush() + + # Wrap command with monitor if needed + if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'): + # Use the monitor wrapper + monitor_script = os.path.join(os.path.dirname(__file__), 'worker_monitor.py') + monitored_cmd = [get_python_executable(), monitor_script] + cmd + log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n") + log_handle.flush() + else: + monitored_cmd = cmd + + # Platform-specific process creation - always hidden with logging + if platform.system() == "Windows": + CREATE_NO_WINDOW = 0x08000000 + process = subprocess.Popen( + monitored_cmd, env=env, cwd=cwd, + stdout=log_handle, + stderr=subprocess.STDOUT, + creationflags=CREATE_NO_WINDOW + ) + else: + # Unix-like systems + process = subprocess.Popen( + monitored_cmd, env=env, cwd=cwd, + stdout=log_handle, + stderr=subprocess.STDOUT, + start_new_session=True # Detach from parent + ) + + # Track the process with log file info - use string ID for consistency + worker_id = str(worker_config['id']) + self.processes[worker_id] = { + 'pid': process.pid, + 'process': process, + 'started_at': time.time(), + 'config': worker_config, + 'log_file': log_file, + 'is_monitor': stop_on_master_exit and env.get('COMFYUI_MASTER_PID'), # Track if using monitor + 'launching': True # Mark as launching until confirmed running + } + + # Save process info for persistence + self.save_processes() + + if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'): + debug_log(f"Launched worker {worker_name} via monitor (Monitor PID: {process.pid})") + else: + log(f"Launched worker {worker_name} directly (PID: {process.pid})") + debug_log(f"Log file: {log_file}") + return process.pid + + def stop_worker(self, worker_id): + """Stop a worker process.""" + # Ensure worker_id is string + worker_id = str(worker_id) + if worker_id not in self.processes: + return False, "Worker not managed by UI" + + proc_info = self.processes[worker_id] + process = proc_info.get('process') + pid = proc_info['pid'] + + debug_log(f"Attempting to stop worker {worker_id} (PID: {pid})") + + # For restored processes without subprocess object + if not process: + try: + print(f"[Distributed] Stopping restored process (no subprocess object)") + if self._kill_process_tree(pid): + del self.processes[worker_id] + self.save_processes() + debug_log(f"Successfully stopped worker {worker_id} and all child processes") + return True, "Worker stopped" + else: + return False, "Failed to stop worker process" + except Exception as e: + print(f"[MultiGPU] Exception during stop: {e}") + return False, f"Error stopping worker: {str(e)}" + + # Normal case with subprocess object + # Check if still running + if process.poll() is not None: + # Already stopped + print(f"[Distributed] Worker {worker_id} already stopped") + del self.processes[worker_id] + self.save_processes() + return False, "Worker already stopped" + + # Try to kill the entire process tree + try: + debug_log(f"Using process tree kill for worker {worker_id}") + if self._kill_process_tree(pid): + # Clean up tracking + del self.processes[worker_id] + self.save_processes() + debug_log(f"Successfully stopped worker {worker_id} and all child processes") + return True, "Worker stopped" + else: + # Fallback to normal termination + print(f"[Distributed] Process tree kill failed, trying normal termination") + if process: + terminate_process(process, timeout=PROCESS_TERMINATION_TIMEOUT) + + del self.processes[worker_id] + self.save_processes() + return True, "Worker stopped (fallback)" + + except Exception as e: + print(f"[Distributed] Exception during stop: {e}") + return False, f"Error stopping worker: {str(e)}" + + def get_managed_workers(self): + """Get list of workers managed by this process.""" + managed = {} + for worker_id, proc_info in list(self.processes.items()): + # Check if process is still running + is_running, _ = self._check_worker_process(worker_id, proc_info) + + if is_running: + managed[worker_id] = { + 'pid': proc_info['pid'], + 'started_at': proc_info['started_at'], + 'log_file': proc_info.get('log_file'), + 'launching': proc_info.get('launching', False) + } + else: + # Process has stopped, remove from tracking + del self.processes[worker_id] + + return managed + + def cleanup_all(self): + """Stop all managed workers (called on shutdown).""" + for worker_id in list(self.processes.keys()): + try: + self.stop_worker(worker_id) + except Exception as e: + print(f"[Distributed] Error stopping worker {worker_id}: {e}") + + # Clear all managed processes from config + config = load_config() + config['managed_processes'] = {} + save_config(config) + + def load_processes(self): + """Load persisted process information from config.""" + config = load_config() + managed_processes = config.get('managed_processes', {}) + + # Verify each saved process is still running + for worker_id, proc_info in managed_processes.items(): + pid = proc_info.get('pid') + if pid and self._is_process_running(pid): + # Reconstruct process info + self.processes[worker_id] = { + 'pid': pid, + 'process': None, # Can't reconstruct subprocess object + 'started_at': proc_info.get('started_at'), + 'config': proc_info.get('config'), + 'log_file': proc_info.get('log_file') + } + print(f"[Distributed] Restored worker {worker_id} (PID: {pid})") + else: + if pid: + print(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running") + + def save_processes(self): + """Save process information to config.""" + config = load_config() + + # Create serializable version of process info + managed_processes = {} + for worker_id, proc_info in self.processes.items(): + # Only save if process is running + is_running, _ = self._check_worker_process(worker_id, proc_info) + + if is_running: + managed_processes[worker_id] = { + 'pid': proc_info['pid'], + 'started_at': proc_info['started_at'], + 'config': proc_info['config'], + 'log_file': proc_info.get('log_file'), + 'launching': proc_info.get('launching', False) + } + + # Update config with managed processes + config['managed_processes'] = managed_processes + save_config(config) + + def _is_process_running(self, pid): + """Check if a process with given PID is running.""" + return is_process_alive(pid) + + def _check_worker_process(self, worker_id, proc_info): + """Check if a worker process is still running and return status. + + Returns: + tuple: (is_running, has_subprocess_object) + """ + process = proc_info.get('process') + pid = proc_info.get('pid') + + if process: + # Normal case with subprocess object + return process.poll() is None, True + elif pid: + # Restored process without subprocess object + return self._is_process_running(pid), False + else: + # No process or PID + return False, False + + def _kill_process_tree(self, pid): + """Kill a process and all its children.""" + if PSUTIL_AVAILABLE: + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + + # Log what we're about to kill + debug_log(f"Killing process tree for PID {pid} ({parent.name()})") + for child in children: + debug_log(f" - Child PID {child.pid} ({child.name()})") + + # Kill children first + for child in children: + try: + debug_log(f"Terminating child {child.pid}") + child.terminate() + except psutil.NoSuchProcess: + pass + + # Wait a bit for graceful termination + gone, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT) + + # Force kill any remaining + for child in alive: + try: + debug_log(f"Force killing child {child.pid}") + child.kill() + except psutil.NoSuchProcess: + pass + + # Finally kill the parent + try: + debug_log(f"Terminating parent {pid}") + parent.terminate() + parent.wait(timeout=PROCESS_WAIT_TIMEOUT) + except psutil.TimeoutExpired: + debug_log(f"Force killing parent {pid}") + parent.kill() + except psutil.NoSuchProcess: + debug_log(f"Parent process {pid} already gone") + + return True + + except psutil.NoSuchProcess: + debug_log(f"Process {pid} does not exist") + return False + except Exception as e: + debug_log(f"Error killing process tree: {e}") + # Fall through to OS commands + + # Fallback to OS-specific commands + print(f"[Distributed] Using OS commands to kill process tree") + if platform.system() == "Windows": + try: + # Use wmic to find child processes + result = subprocess.run(['wmic', 'process', 'where', f'ParentProcessId={pid}', 'get', 'ProcessId'], + capture_output=True, text=True) + if result.returncode == 0: + lines = result.stdout.strip().split('\n')[1:] # Skip header + child_pids = [line.strip() for line in lines if line.strip() and line.strip().isdigit()] + + print(f"[Distributed] Found child processes: {child_pids}") + + # Kill each child + for child_pid in child_pids: + try: + subprocess.run(['taskkill', '/F', '/PID', child_pid], + capture_output=True, check=False) + except: + pass + + # Kill the parent with tree flag + result = subprocess.run(['taskkill', '/F', '/PID', str(pid), '/T'], + capture_output=True, text=True) + print(f"[Distributed] Taskkill result: {result.stdout.strip()}") + return result.returncode == 0 + except Exception as e: + print(f"[Distributed] Error with taskkill: {e}") + return False + else: + # Unix: use pkill + try: + subprocess.run(['pkill', '-TERM', '-P', str(pid)], check=False) + time.sleep(WORKER_CHECK_INTERVAL) + subprocess.run(['pkill', '-KILL', '-P', str(pid)], check=False) + os.kill(pid, signal.SIGKILL) + return True + except: + return False + +# Create global instance +worker_manager = WorkerProcessManager() + +# Auto-launch workers if enabled +def auto_launch_workers(): + """Launch enabled workers if auto_launch_workers is set to true.""" + try: + config = load_config() + if config.get('settings', {}).get('auto_launch_workers', False): + log("Auto-launch workers is enabled, checking for workers to start...") + + # Clear managed_processes before launching new workers + # This handles cases where the master was killed without proper cleanup + if config.get('managed_processes'): + log("Clearing old managed_processes before auto-launch...") + config['managed_processes'] = {} + save_config(config) + + workers = config.get('workers', []) + launched_count = 0 + + for worker in workers: + if worker.get('enabled', False): + worker_id = worker.get('id') + worker_name = worker.get('name', f'Worker {worker_id}') + + # Skip remote workers + host = worker.get('host', 'localhost').lower() + if host not in ['localhost', '127.0.0.1', '', None]: + debug_log(f"Skipping remote worker {worker_name} (host: {host})") + continue + + # Check if already running + if str(worker_id) in worker_manager.processes: + proc_info = worker_manager.processes[str(worker_id)] + if worker_manager._is_process_running(proc_info['pid']): + debug_log(f"Worker {worker_name} already running, skipping") + continue + + # Launch the worker + try: + pid = worker_manager.launch_worker(worker) + log(f"Auto-launched worker {worker_name} (PID: {pid})") + + # Mark as launching in managed processes + if str(worker_id) in worker_manager.processes: + worker_manager.processes[str(worker_id)]['launching'] = True + worker_manager.save_processes() + + launched_count += 1 + except Exception as e: + log(f"Failed to auto-launch worker {worker_name}: {e}") + + if launched_count > 0: + log(f"Auto-launched {launched_count} worker(s)") + else: + debug_log("No workers to auto-launch") + else: + debug_log("Auto-launch workers is disabled") + except Exception as e: + log(f"Error during auto-launch: {e}") + +# Schedule auto-launch after a short delay to ensure server is ready +def delayed_auto_launch(): + """Delay auto-launch to ensure server is fully initialized.""" + import threading + timer = threading.Timer(WORKER_STARTUP_DELAY, auto_launch_workers) + 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""" + 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") + # Still save the current state + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + +# Register cleanup handlers +atexit.register(cleanup_on_exit) + +# 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}") + +# --- Persistent State Storage --- +# Store job queue on the persistent server instance to survive script reloads +prompt_server = server.PromptServer.instance + +# Initialize persistent state if not already present +if not hasattr(prompt_server, 'distributed_pending_jobs'): + debug_log("Initializing persistent job queue on server instance.") + prompt_server.distributed_pending_jobs = {} + prompt_server.distributed_jobs_lock = asyncio.Lock() + +@server.PromptServer.instance.routes.post("/distributed/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') + worker_id = data.get('worker_id') + image_index = data.get('image_index') + is_last = data.get('is_last', 'False').lower() == 'true' + + 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) + + 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})") + debug_log(f"Queue size after put: {prompt_server.distributed_pending_jobs[multi_job_id].qsize()}") + return web.json_response({"status": "success"}) + else: + log(f"ERROR: Job {multi_job_id} not found in distributed_pending_jobs") + 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 DistributedCollectorNode: + @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": "[]"}), + "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}), + "worker_id": ("STRING", {"default": ""}), + "pass_through": ("BOOLEAN", {"default": False}), + }, + } + + 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, worker_id="", pass_through=False): + if not multi_job_id or pass_through: + if pass_through: + print(f"[Distributed Collector] Pass-through mode enabled, returning images unchanged") + return (images,) + + # Use async helper to run in server loop + result = run_async_in_server_loop( + self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id) + ) + return result + + 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) + + 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') + + try: + session = await get_client_session() + async with session.post(f"{master_url}/distributed/job_complete", 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}") + + 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 + 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) + + return (images,) + else: + # Master mode: collect images from workers + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) + if num_workers == 0: + return (images,) + + images_on_cpu = images.cpu() + master_batch_size = images.shape[0] + debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...") + + # 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 + + # Get the existing queue - it should already exist from prepare_job + async with prompt_server.distributed_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_jobs: + log(f"Master - WARNING: Queue doesn't exist for job {multi_job_id}, creating one") + prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() + else: + existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize() + debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})") + + # Collect images until all workers report they're done + collected_count = 0 + workers_done = set() + + # 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}") + + while len(workers_done) < num_workers: + try: + # Get the queue again each time to ensure we have the right reference + async with prompt_server.distributed_jobs_lock: + 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'] + 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 + + # Once we start receiving images, use shorter timeout + timeout = WORKER_JOB_TIMEOUT + + if is_last: + workers_done.add(worker_id) + 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 + log(f"Master - Timeout. Still waiting for workers: {list(missing_workers)}") + + # Check queue size again with lock + async with prompt_server.distributed_jobs_lock: + 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 = [] + while not final_q.empty(): + try: + item = final_q.get_nowait() + remaining_items.append(item) + except asyncio.QueueEmpty: + 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'] + 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 + + if is_last: + workers_done.add(worker_id) + 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: + if multi_job_id in prompt_server.distributed_pending_jobs: + del prompt_server.distributed_pending_jobs[multi_job_id] + + # Reorder images according to seed distribution pattern + # Pattern: master img 1, master img 2, worker 1 img 1, worker 1 img 2, worker 2 img 1, worker 2 img 2, etc. + ordered_tensors = [] + + # Add master images first + for i in range(master_batch_size): + ordered_tensors.append(images_on_cpu[i:i+1]) + + # Add worker images in order + # The worker IDs in worker_images are already strings (e.g., "1", "2") + # Just iterate through what we actually received + for worker_id_str in sorted(worker_images.keys()): + # Sort by image index for each worker + for idx in sorted(worker_images[worker_id_str].keys()): + ordered_tensors.append(worker_images[worker_id_str][idx]) + + # Ensure all tensors are on CPU and properly formatted before concatenation + cpu_tensors = [] + for t in ordered_tensors: + if t.is_cuda: + t = t.cpu() + # Ensure tensor is contiguous in memory + t = ensure_contiguous(t) + cpu_tensors.append(t) + + try: + combined = torch.cat(cpu_tensors, dim=0) + # Ensure the combined tensor is contiguous and properly formatted + combined = ensure_contiguous(combined) + debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total (master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") + 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,) + +# --- Distributor Node --- +class DistributedSeed: + """ + Distributes seed values across multiple GPUs. + On master: passes through the original seed. + On workers: adds offset based on worker ID. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 1125899906842624, + "forceInput": False # Widget by default, can be converted to input + }), + }, + "hidden": { + "is_worker": ("BOOLEAN", {"default": False}), + "worker_id": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("seed",) + FUNCTION = "distribute" + CATEGORY = "utils" + + def distribute(self, seed, is_worker=False, worker_id=""): + if not is_worker: + # Master node: pass through original values + debug_log(f"Distributor - Master: seed={seed}") + return (seed,) + else: + # Worker node: apply offset based on worker index + # Find worker index from enabled_worker_ids + try: + # Worker IDs are passed as "worker_0", "worker_1", etc. + if worker_id.startswith("worker_"): + worker_index = int(worker_id.split("_")[1]) + else: + # Fallback: try to parse as direct index + worker_index = int(worker_id) + + offset = worker_index + 1 + new_seed = seed + offset + debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}") + return (new_seed,) + except (ValueError, IndexError) as e: + debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}") + # Fallback: return original seed + return (seed,) + +NODE_CLASS_MAPPINGS = { + "DistributedCollector": DistributedCollectorNode, + "DistributedSeed": DistributedSeed +} +NODE_DISPLAY_NAME_MAPPINGS = { + "DistributedCollector": "Distributed Collector", + "DistributedSeed": "Distributed Seed" +} diff --git a/distributed_upscale.py b/distributed_upscale.py new file mode 100644 index 0000000..7bf702b --- /dev/null +++ b/distributed_upscale.py @@ -0,0 +1,770 @@ +import torch +import numpy as np +from PIL import Image, ImageFilter, ImageDraw +import json +import asyncio +import aiohttp +from aiohttp import web +import io +import server +from typing import List, Tuple + +# Import ComfyUI modules +import comfy.samplers + +# Import shared utilities +from .utils.logging import debug_log, log +from .utils.config import CONFIG_FILE +from .utils.image import tensor_to_pil, pil_to_tensor +from .utils.network import get_server_port, get_server_loop, get_client_session, handle_api_error +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 +) + +# 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() + return prompt_server + +# Note: tensor_to_pil and pil_to_tensor are imported from utils.image + +class UltimateSDUpscaleDistributed: + """ + Distributed version of Ultimate SD Upscale (No Upscale). + Distributes tile processing across multiple workers. + """ + + def __init__(self): + """Initialize the node and ensure persistent storage exists.""" + # Pre-initialize the persistent storage on node creation + ensure_tile_jobs_initialized() + debug_log("UltimateSDUpscaleDistributed - Node initialized") + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "upscaled_image": ("IMAGE",), + "model": ("MODEL",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}), + "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), + "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), + "denoise": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "tile_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}), + "tile_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}), + "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}), + }, + "hidden": { + "multi_job_id": ("STRING", {"default": ""}), + "is_worker": ("BOOLEAN", {"default": False}), + "master_url": ("STRING", {"default": ""}), + "enabled_worker_ids": ("STRING", {"default": "[]"}), + "worker_id": ("STRING", {"default": ""}), + "tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "run" + CATEGORY = "image/upscaling" + + @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}") + + 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, multi_job_id="", is_worker=False, + master_url="", enabled_worker_ids="[]", worker_id="", tile_indices=""): + """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) + + 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) + 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) + + 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.""" + # 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) + + # Get assigned tiles for this worker + assigned_tiles = self._get_worker_tiles(all_tiles, enabled_worker_ids, worker_id) + if not assigned_tiles: + return (upscaled_image,) + + debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(assigned_tiles)} tiles") + + # Process tiles SYNCHRONOUSLY + processed_tiles = [] + for tile_idx in assigned_tiles: + x, y = all_tiles[tile_idx] + + # Extract tile + tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( + upscaled_image, x, y, tile_width, tile_height, padding + ) + + # Process tile through SD (SYNCHRONOUS - no async!) + processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, + seed + tile_idx, steps, cfg, sampler_name, + scheduler, denoise) + + processed_tiles.append({ + 'tile': processed_tile, + '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 + try: + run_async_in_server_loop( + self._async_send_tiles_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,) + + 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, + force_uniform_tiles, multi_job_id, enabled_worker_ids): + """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 + + # Calculate all tiles + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + total_tiles = len(all_tiles) + + debug_log(f"UltimateSDUpscale Master - Total tiles: {total_tiles}") + + # Parse enabled workers + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) + + if num_workers == 0: + # 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) + + # 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}") + # 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") + + # Don't try to prepare via API since we're using local queues + # 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() + + # 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) + + debug_log(f"UltimateSDUpscale Master - Processing {master_tiles} tiles locally") + + # 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, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height + ) + + # Collect worker tiles using async + worker_tiles_expected = total_tiles - 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) + + # 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 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) + + # 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_image = self.blend_tile(result_image, tile_pil, + x, y, (extracted_width, extracted_height), tile_mask, padding) + + # Convert back to tensor + result_tensor = pil_to_tensor(result_image) + 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_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)) + + 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)) + + 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): + """Process a single tile and blend it into the result image.""" + x, y = tile_pos + + # Extract and process tile + tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding( + upscaled_image, x, y, tile_width, tile_height, padding + ) + + processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae, + seed + tile_idx, steps, cfg, sampler_name, + scheduler, denoise) + + # Convert and blend + processed_pil = tensor_to_pil(processed_tile, 0) + # Create mask for this specific tile + tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur) + # Use extraction position and size for blending + result_image = self.blend_tile(result_image, processed_pil, + x1, y1, (ew, eh), tile_mask, padding) + + 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(self, multi_job_id): + """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] = asyncio.Queue() + 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.""" + # Get the already initialized queue + 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}") + q = prompt_server.distributed_pending_tile_jobs[multi_job_id] + + collected_tiles = {} + workers_done = set() + timeout = TILE_WAIT_TIMEOUT + + 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) + 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})") + + 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") + + # 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 + + 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]]: + """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 + + 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]: + """Extract a tile from the image with padding.""" + _, h, w, _ = image.shape + + # Calculate extraction bounds with padding + x1 = max(0, x - padding) + y1 = max(0, y - padding) + x2 = min(w, x + tile_width + padding) + y2 = min(h, y + tile_height + padding) + + # Store the actual extracted size before resizing + extracted_width = x2 - x1 + extracted_height = y2 - y1 + + # Debug logging (uncomment if needed) + # debug_log(f"[Extract] Tile at ({x}, {y}) -> Extract from ({x1}, {y1}) to ({x2}, {y2}), size: {extracted_width}x{extracted_height}") + + # Extract tile + tile = image[:, y1:y2, x1:x2, :] + + # Convert to PIL using utility function + tile_pil = tensor_to_pil(tile, 0) + + # Resize to target dimensions for processing + tile_pil = tile_pil.resize((tile_width, tile_height), Image.LANCZOS) + + # Convert back to tensor using utility function + tile_tensor = pil_to_tensor(tile_pil) + + # Move to same device as original image + if image.is_cuda: + tile_tensor = tile_tensor.cuda() + + 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) -> torch.Tensor: + """Process a single tile through SD sampling.""" + # Import here to avoid circular dependencies + from nodes import common_ksampler, VAEEncode, VAEDecode + + # 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) + + # Move to correct device + if tile_tensor.is_cuda: + clean_tensor = clean_tensor.cuda() + + # Encode to latent + 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] + + return image + + def create_tile_mask(self, image_width: int, image_height: int, + x: int, y: int, tile_width: int, tile_height: int, + mask_blur: int) -> Image.Image: + """Create a mask for blending tiles - matches Ultimate SD Upscale approach. + + Creates a black image with a white rectangle at the tile position, + then applies blur to create soft edges. + """ + # Create a full-size mask matching the image dimensions + mask = Image.new('L', (image_width, image_height), 0) # Black background + + # Draw white rectangle at tile position + draw = ImageDraw.Draw(mask) + draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) + + # Apply blur to soften edges + if mask_blur > 0: + mask = mask.filter(ImageFilter.GaussianBlur(mask_blur)) + + return mask + + 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: + """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 + """ + extracted_width, extracted_height = extracted_size + + # Debug logging (uncomment if needed) + # debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}") + + # Calculate the crop region that was used for extraction + crop_region = (x, y, x + extracted_width, y + extracted_height) + + # The mask is already full-size, no need to crop + + # Resize the processed tile back to the extracted size + if tile_image.size != (extracted_width, extracted_height): + tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS) + else: + tile_resized = tile_image + + # Follow Ultimate SD Upscale blending approach: + # Put the tile into position + image_tile_only = Image.new('RGBA', base_image.size) + image_tile_only.paste(tile_resized, crop_region[:2]) + + # Add the mask as an alpha channel + # Must make a copy due to the possibility of an edge becoming black + temp = image_tile_only.copy() + temp.putalpha(mask) # Use the full image mask + image_tile_only.paste(temp, image_tile_only) + + # Add back the tile to the initial image according to the mask in the alpha channel + result = base_image.convert('RGBA') + result.alpha_composite(image_tile_only) + + # Convert back to RGB + return result.convert('RGB') + + + 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...") + 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) + + + 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): + """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 + + # Calculate all tiles + all_tiles = self.calculate_tiles(width, height, tile_width, tile_height) + + debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally") + + # Convert to PIL + image_pil = tensor_to_pil(upscaled_image, 0) + result_image = 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 + ) + + # Convert back to tensor + result_tensor = pil_to_tensor(result_image) + if upscaled_image.is_cuda: + result_tensor = result_tensor.cuda() + + return (result_tensor,) + + +# Ensure initialization before registering routes +ensure_tile_jobs_initialized() + +# API Endpoint for tile completion +@server.PromptServer.instance.routes.post("/distributed/tile_complete") +async def tile_complete_endpoint(request): + """Endpoint for receiving completed tiles from workers.""" + 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,] + + 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({ + '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}") + 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) + + +# Node registration +NODE_CLASS_MAPPINGS = { + "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)", +} \ No newline at end of file diff --git a/gpu_config.json b/gpu_config.json deleted file mode 100644 index 66d1a49..0000000 --- a/gpu_config.json +++ /dev/null @@ -1,34 +0,0 @@ -{ - "master": { - "port": 8080, - "cuda_device": 0, - "extra_args": "--mmap-torch-files --listen" - }, - "workers": [ - { - "id": 1, - "name": "GPU 1", - "cuda_device": 1, - "port": 8180, - "enabled": false, - "extra_args": "--mmap-torch-files --listen" - }, - { - "id": 2, - "name": "GPU 2", - "cuda_device": 2, - "port": 8280, - "enabled": false, - "extra_args": "--mmap-torch-files --listen" - }, - { - "id": 3, - "name": "GPU 3", - "cuda_device": 3, - "port": 8380, - "enabled": false, - "extra_args": "--mmap-torch-files --listen --reserve-vram 2" - } - ], - "settings": {} -} \ No newline at end of file diff --git a/launch_distributed.bat b/launch_distributed.bat deleted file mode 100644 index bce8c97..0000000 --- a/launch_distributed.bat +++ /dev/null @@ -1,502 +0,0 @@ -@echo off -title ComfyUI Unified Launcher -cd /d "%~dp0" -color 0B -cls - -echo. -echo ================================================================ -echo ComfyUI Unified Launcher -echo ================================================================ -echo. - -:: Check command line parameter -set "MODE=%1" -if "%MODE%"=="" goto :SHOW_MENU -if /i "%MODE%"=="master" goto :SETUP -if /i "%MODE%"=="all-slow" goto :SETUP -if /i "%MODE%"=="all-fast" goto :SETUP -if /i "%MODE%"=="venv" goto :ACTIVATE_VENV -if /i "%MODE%"=="update" goto :UPDATE_COMFYUI -if /i "%MODE%"=="update-nodes" goto :UPDATE_NODES - -echo. -echo ERROR: Invalid parameter: %MODE% -echo. -timeout /t 2 /nobreak >nul -goto :SHOW_MENU - -:SHOW_MENU -color 0B -echo ---------------------------------------------------------------- -echo Launch Mode Selection -echo ---------------------------------------------------------------- -echo. -echo [1] Master only -echo ^> Launch only the master instance -echo. -echo [2] All instances (fast) -echo ^> Launch all instances with 2-second delays -echo. -echo [3] All instances (sequential) -echo ^> Launch all instances, wait for each to be ready -echo. -echo [4] Activate virtual environment -echo ^> Open command prompt with venv activated -echo. -echo [5] Update ComfyUI -echo ^> Update ComfyUI to latest version -echo. -echo [6] Update custom nodes -echo ^> Update all custom nodes to latest versions -echo. -echo [Q] Quit -echo. -echo ---------------------------------------------------------------- -echo TIP: You can also launch directly from command line: -echo launch_comfyui.bat master -echo launch_comfyui.bat all-slow -echo launch_comfyui.bat all-fast -echo launch_comfyui.bat venv -echo launch_comfyui.bat update -echo launch_comfyui.bat update-nodes -echo ---------------------------------------------------------------- -echo. - -:: Use CHOICE to wait for a single keypress without needing Enter -:: /C specifies the allowed characters. /N hides the default "[1,2,3,4,5,6,Q]?" prompt. /M sets a custom message. -choice /c 123456Q /n /m " ^> Enter your choice: " -echo. - -:: Check the ERRORLEVEL. IMPORTANT: Must check from highest to lowest! -:: CHOICE sets ERRORLEVEL to the index of the character in /C -:: 1=1, 2=2, 3=3, 4=4, 5=5, 6=6, Q=7 - -if errorlevel 7 goto :QUIT_CHOICE -if errorlevel 6 goto :UPDATE_NODES -if errorlevel 5 goto :UPDATE_COMFYUI -if errorlevel 4 goto :ACTIVATE_VENV -if errorlevel 3 ( - set "MODE=all-slow" - goto :SETUP -) -if errorlevel 2 ( - set "MODE=all-fast" - goto :SETUP -) -if errorlevel 1 ( - set "MODE=master" - goto :SETUP -) - -:: This part of the script should now be unreachable, but it's good practice to handle it. -echo. -echo ERROR: Invalid choice. -echo. -timeout /t 2 /nobreak >nul -goto :SHOW_MENU - -:QUIT_CHOICE -echo Exiting... -timeout /t 1 /nobreak >nul -exit /b 0 - -:SETUP -echo Mode: %MODE% -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% - echo. - echo Make sure ComfyUI-Distributed custom node is installed. - 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 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 -) - -:: Jump to appropriate mode -if /i "%MODE%"=="master" goto :MODE_MASTER -if /i "%MODE%"=="all-slow" goto :MODE_ALL_SLOW -if /i "%MODE%"=="all-fast" goto :MODE_ALL_FAST - -::============================================================================= -:: MASTER MODE - Launch only the master instance -::============================================================================= -:MODE_MASTER -echo Configuring memory management for large models... -set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -echo Memory management configured. - -echo Loading master configuration... - -:: Read master configuration -powershell -NoProfile -Command "try { $json = Get-Content '%CONFIG_PATH%' -Raw | ConvertFrom-Json; $master = $json.master; Write-Output \"MASTER_PORT=$($master.port)\"; Write-Output \"MASTER_CUDA=$($master.cuda_device)\"; Write-Output \"MASTER_EXTRA=$($master.extra_args)\" } catch { Write-Output \"PowerShell Error: $_\"; pause; exit 1 }" > %TEMP%\master_config.txt - -if %ERRORLEVEL% NEQ 0 ( - color 0C - echo Error reading master configuration! - pause - exit /b 1 -) - -for /f "tokens=*" %%i in (%TEMP%\master_config.txt) do ( - set %%i -) - -del %TEMP%\master_config.txt >nul 2>&1 - -:: Display master configuration -echo. -echo Master Configuration: -echo Port: %MASTER_PORT% -echo CUDA Device: %MASTER_CUDA% -echo Extra Args: %MASTER_EXTRA% -echo. - -:: Check if master is already running -echo Checking if master is already running... -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% == 0 ( - color 0E - echo Master is already running on port %MASTER_PORT%! - echo URL: http://localhost:%MASTER_PORT% - echo. - echo Press any key to exit... - pause >nul - exit /b 0 -) - -:: Launch master -echo Launching Master ComfyUI instance... - -:: Try Windows Terminal first, fall back to cmd -wt -w 0 new-tab --title "Master - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title Master - ComfyUI CLI && echo CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && echo Starting ComfyUI Master... && comfy --here launch -- --port %MASTER_PORT% %MASTER_EXTRA%" 2>nul - -if %ERRORLEVEL% NEQ 0 ( - echo Windows Terminal not available, using regular command prompt... - start "Master - ComfyUI CLI" cmd /k "title Master - ComfyUI CLI && cd /d "%CD%" && echo CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && echo Starting ComfyUI Master... && comfy --here launch -- --port %MASTER_PORT% %MASTER_EXTRA%" -) - -echo. -echo ====================================== -echo Master ComfyUI launcher started! -echo ====================================== -echo. -echo Expected URL: http://localhost:%MASTER_PORT% -echo CUDA Device: %MASTER_CUDA% -echo. -echo The ComfyUI window will open shortly. -echo This launcher will close in 1 second... - -timeout /t 1 /nobreak >nul -exit - -::============================================================================= -:: ALL-SLOW MODE - Launch all instances sequentially -::============================================================================= -:MODE_ALL_SLOW -echo Configuring memory management for large models... -set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -echo Memory management configured. - -:: 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)\"; Write-Output \"MASTER_EXTRA=$($master.extra_args)\"; foreach ($worker in $json.workers) { if ($worker.enabled -and -not $worker.host) { 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_SLOW -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 ( - echo Starting Master with CUDA_VISIBLE_DEVICES=%MASTER_CUDA% - start wt -w 0 new-tab --title "Master - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title Master - ComfyUI CLI && echo CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && comfy --here launch -- --port %MASTER_PORT% %MASTER_EXTRA%" - set MASTER_STARTED=1 - ) - %TEMP%\worker_ids.txt - -for /f %%w in (%TEMP%\worker_ids.txt) do ( - call :PROCESS_WORKER_SLOW %%w -) - -del %TEMP%\worker_ids.txt >nul 2>&1 -goto :FINISH_ALL - -:PROCESS_WORKER_SLOW -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_SLOW -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 ( - echo Starting %WORKER_NAME% with CUDA_VISIBLE_DEVICES=%WORKER_CUDA% - start wt -w 0 new-tab --title "%WORKER_NAME% - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title %WORKER_NAME% - ComfyUI CLI && echo CUDA_VISIBLE_DEVICES=%WORKER_CUDA% && 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 - ) - %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 - -echo. -echo Checking running instances... -echo. - -set "INSTANCES_STARTED=0" - -:: Check and start Master if not running -echo Checking Master (port %MASTER_PORT%)... -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 ( - echo Master NOT running - Starting with CUDA_VISIBLE_DEVICES=%MASTER_CUDA% - start wt -w 0 new-tab --title "Master - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title Master - ComfyUI CLI && echo CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%MASTER_CUDA% && comfy --here launch -- --port %MASTER_PORT% %MASTER_EXTRA%" - set /a INSTANCES_STARTED+=1 - echo Waiting 2 seconds before checking next instance... - timeout /t 2 /nobreak >nul -) else ( - echo Master already running on port %MASTER_PORT% -) - -:: Check and start workers -powershell -NoProfile -Command "try { $json = Get-Content '%CONFIG_PATH%' -Raw | ConvertFrom-Json; foreach ($worker in $json.workers) { if ($worker.enabled -and -not $worker.host) { Write-Output $worker.id } } } catch { exit 1 }" > %TEMP%\worker_ids.txt - -for /f %%w in (%TEMP%\worker_ids.txt) do ( - call :CHECK_AND_START_WORKER_FAST %%w -) - -del %TEMP%\worker_ids.txt >nul 2>&1 - -echo. -echo ====================================== -if %INSTANCES_STARTED% == 0 ( - echo All instances were already running! -) else ( - echo Started %INSTANCES_STARTED% new instance^(s^)! -) -echo ====================================== -goto :FINISH_ALL - -:CHECK_AND_START_WORKER_FAST -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 Checking %WORKER_NAME% (port %WORKER_PORT%)... -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 ( - echo %WORKER_NAME% NOT running - Starting with CUDA_VISIBLE_DEVICES=%WORKER_CUDA% - start wt -w 0 new-tab --title "%WORKER_NAME% - ComfyUI CLI" --startingDirectory "%CD%" cmd /k "title %WORKER_NAME% - ComfyUI CLI && echo CUDA_VISIBLE_DEVICES=%WORKER_CUDA% && call venv\Scripts\activate.bat && set CUDA_VISIBLE_DEVICES=%WORKER_CUDA% && comfy --here launch -- --port %WORKER_PORT% %WORKER_EXTRA% --enable-cors-header" - set /a INSTANCES_STARTED+=1 - echo Waiting 2 seconds before checking next instance... - timeout /t 2 /nobreak >nul -) else ( - echo %WORKER_NAME% already running on port %WORKER_PORT% -) -goto :eof - -::============================================================================= -:: FINISH - Common ending for all modes -::============================================================================= -:FINISH_ALL -echo. -echo Master: http://localhost:%MASTER_PORT% -powershell -NoProfile -Command "try { $json = Get-Content '%CONFIG_PATH%' -Raw | ConvertFrom-Json; foreach ($worker in $json.workers) { if ($worker.enabled -and -not $worker.host) { Write-Output \"$($worker.name): http://localhost:$($worker.port)\" } } } catch { }" -echo. -echo Launcher will close in 5 seconds... -timeout /t 5 /nobreak >nul -exit - -::============================================================================= -:: ACTIVATE VIRTUAL ENVIRONMENT - Open command prompt with venv activated -::============================================================================= -:ACTIVATE_VENV -echo. -echo ================================================================ -echo Activating Virtual Environment -echo ================================================================ -echo. - -:: Check if venv exists -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% - pause - exit /b 1 -) - -echo Activating virtual environment... -echo Project folder: %CD% -echo. -echo Opening new command prompt with activated environment... -echo (You can close this window) -echo. - -:: Open new command prompt with venv activated -start "ComfyUI (Active venv)" cmd /k "call "%CD%\%VENV_PATH%\Scripts\activate.bat" & cd /d "%CD%"" -exit - -::============================================================================= -:: UPDATE COMFYUI - Update ComfyUI to latest version -::============================================================================= -:UPDATE_COMFYUI -echo. -echo ================================================================ -echo Updating ComfyUI -echo ================================================================ -echo. - -:: Check if venv exists -set "VENV_PATH=venv" -if not exist "%VENV_PATH%\Scripts\activate.bat" ( - color 0C - echo Error: Virtual environment not found! - pause - exit /b 1 -) - -:: Check if ComfyUI CLI is available -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! - pause - exit /b 1 -) - -echo Starting ComfyUI update... -echo This may take a few minutes... -echo. - -:: Activate venv and run update -call "%VENV_PATH%\Scripts\activate.bat" && comfy --here update - -echo. -echo Update completed! -echo. -pause -exit - -::============================================================================= -:: UPDATE CUSTOM NODES - Update all custom nodes -::============================================================================= -:UPDATE_NODES -echo. -echo ================================================================ -echo Updating Custom Nodes -echo ================================================================ -echo. - -:: Check if venv exists -set "VENV_PATH=venv" -if not exist "%VENV_PATH%\Scripts\activate.bat" ( - color 0C - echo Error: Virtual environment not found! - pause - exit /b 1 -) - -:: Check if ComfyUI CLI is available -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! - pause - exit /b 1 -) - -echo Starting custom nodes update... -echo This may take several minutes... -echo. - -:: Activate venv and run node update -call "%VENV_PATH%\Scripts\activate.bat" && comfy --here node update-all - -echo. -echo Custom nodes update completed! -echo. -pause -exit diff --git a/multigpu.py b/multigpu.py deleted file mode 100644 index c91f968..0000000 --- a/multigpu.py +++ /dev/null @@ -1,438 +0,0 @@ -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 - -# --- 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": {} - } - -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: - # Use ComfyUI's prompt server queue system like the /free endpoint does - if hasattr(server.PromptServer.instance, 'prompt_queue'): - server.PromptServer.instance.prompt_queue.set_flag("unload_models", True) - server.PromptServer.instance.prompt_queue.set_flag("free_memory", True) - print("[MultiGPU] Set queue flags for memory clearing.") - - # Wait a bit for the queue to process - await asyncio.sleep(0.5) - - # Also do direct cleanup as backup, but with error handling - import gc - import comfy.model_management as mm - - try: - mm.unload_all_models() - except AttributeError as e: - print(f"[MultiGPU] Warning during model unload: {e}") - - try: - mm.soft_empty_cache() - except Exception as e: - print(f"[MultiGPU] Warning during cache clear: {e}") - - for _ in range(3): - gc.collect() - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.ipc_collect() - - print("[MultiGPU] VRAM cleared successfully.") - return web.json_response({"status": "success", "message": "GPU memory cleared."}) - except Exception as e: - # Even if there's an error, try to do basic cleanup - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - print(f"[MultiGPU] Partial VRAM clear completed with warning: {e}") - return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"}) - - -# --- 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') - worker_id = data.get('worker_id') - image_index = data.get('image_index') - 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,] - - async with PENDING_JOBS_LOCK: - if multi_job_id in PENDING_JOBS: - await 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 - }) - print(f"[MultiGPU] Received result for job {multi_job_id} from worker {worker_id} (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) - -# --- 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": "[]"}), - "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}), - "worker_id": ("STRING", {"default": ""}), - }, - } - - 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, worker_id=""): - 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, worker_id) - 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, worker_id, is_last=False): - """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('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') - - 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, worker_id=""): - 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") - - # 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) - - return (images,) - else: - # Master mode: collect images from workers - enabled_workers = json.loads(enabled_worker_ids) - num_workers = len(enabled_workers) - if num_workers == 0: - return (images,) - - images_on_cpu = images.cpu() - print(f"[MultiGPU Master] Job {multi_job_id}: Collecting images from {num_workers} workers...") - - # Store master images first - master_batch_size = images.shape[0] - - # Initialize storage for collected images - worker_images = {} # Dict to store images by worker_id and index - - 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] - - # Collect images until all workers report they're done - collected_count = 0 - workers_done = set() - - # Use a reasonable timeout for the first image - timeout = 10.0 - - while len(workers_done) < num_workers: - try: - 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) - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - collected_count += 1 - - # Once we start receiving images, use shorter timeout - timeout = 10.0 - - if is_last: - workers_done.add(worker_id) - print(f"[MultiGPU Master] Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") - else: - print(f"[MultiGPU Master] Collected image {image_index + 1} from worker {worker_id}") - - except asyncio.TimeoutError: - missing_workers = set(enabled_workers) - workers_done - print(f"[MultiGPU Master] Timeout. Still waiting for workers: {list(missing_workers)}") - break - - total_collected = sum(len(imgs) for imgs in worker_images.values()) - print(f"[MultiGPU Master] Collection complete. Received {total_collected} images from {len(workers_done)} workers") - - # Clean up job queue - async with PENDING_JOBS_LOCK: - if multi_job_id in PENDING_JOBS: - del PENDING_JOBS[multi_job_id] - - # Reorder images according to seed distribution pattern - # Pattern: master img 1, master img 2, worker 1 img 1, worker 1 img 2, worker 2 img 1, worker 2 img 2, etc. - ordered_tensors = [] - - # Add master images first - for i in range(master_batch_size): - ordered_tensors.append(images_on_cpu[i:i+1]) - - # Add worker images in order - for worker_id in enabled_workers: - # Convert worker_id to string since it comes as string from form data - worker_id_str = str(worker_id) - if worker_id_str in worker_images: - # Sort by image index for each worker - for idx in sorted(worker_images[worker_id_str].keys()): - ordered_tensors.append(worker_images[worker_id_str][idx]) - - combined = torch.cat(ordered_tensors, dim=0) - print(f"[MultiGPU Master] Job {multi_job_id} complete. Combined {combined.shape[0]} images total.") - return (combined,) - -# --- Distributor Node --- -class MultiGPUDistributor: - """ - Distributes seed values across multiple GPUs. - On master: passes through the original seed. - On workers: adds offset based on worker ID. - """ - - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "seed": ("INT", { - "default": 0, - "min": 0, - "max": 1125899906842624, - "forceInput": False # Widget by default, can be converted to input - }), - }, - "hidden": { - "is_worker": ("BOOLEAN", {"default": False}), - "worker_id": ("STRING", {"default": ""}), - }, - } - - RETURN_TYPES = ("INT",) - RETURN_NAMES = ("seed",) - FUNCTION = "distribute" - CATEGORY = "utils" - - def distribute(self, seed, is_worker=False, worker_id=""): - if not is_worker: - # Master node: pass through original values - print(f"[MultiGPU Distributor] Master: seed={seed}") - return (seed,) - else: - # Worker node: apply offset based on worker index - # Find worker index from enabled_worker_ids - try: - # Worker IDs are passed as "worker_0", "worker_1", etc. - if worker_id.startswith("worker_"): - worker_index = int(worker_id.split("_")[1]) - else: - # Fallback: try to parse as direct index - worker_index = int(worker_id) - - offset = worker_index + 1 - new_seed = seed + offset - print(f"[MultiGPU Distributor] Worker {worker_index}: seed={seed} → {new_seed}") - return (new_seed,) - except (ValueError, IndexError) as e: - print(f"[MultiGPU Distributor] Error parsing worker_id '{worker_id}': {e}") - # Fallback: return original seed - return (seed,) - -NODE_CLASS_MAPPINGS = { - "MultiGPUCollector": MultiGPUCollectorNode, - "MultiGPUDistributor": MultiGPUDistributor -} -NODE_DISPLAY_NAME_MAPPINGS = { - "MultiGPUCollector": "Multi-GPU Collector", - "MultiGPUDistributor": "Multi-GPU Seed" -} diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..51b62aa --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,5 @@ +""" +Utility modules for ComfyUI-Distributed extension. +""" + +# Make utils importable as a package \ No newline at end of file diff --git a/utils/async_helpers.py b/utils/async_helpers.py new file mode 100644 index 0000000..8fff8ba --- /dev/null +++ b/utils/async_helpers.py @@ -0,0 +1,50 @@ +""" +Async helper utilities for ComfyUI-Distributed. +""" +import asyncio +import threading +from typing import Optional, Any, Coroutine +from .network import get_server_loop + +def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -> Any: + """ + Run async coroutine in server's event loop and wait for result. + + This is useful when you need to run async code from a synchronous context + but want to use the server's existing event loop instead of creating a new one. + + Args: + coro: The coroutine to run + timeout: Optional timeout in seconds + + Returns: + The result of the coroutine + + Raises: + TimeoutError: If the operation times out + Exception: Any exception raised by the coroutine + """ + event = threading.Event() + result = None + error = None + + async def wrapper(): + nonlocal result, error + try: + result = await coro + except Exception as e: + error = e + finally: + event.set() + + # Schedule on server's event loop + loop = get_server_loop() + asyncio.run_coroutine_threadsafe(wrapper(), loop) + + # Wait for completion + if not event.wait(timeout): + raise TimeoutError(f"Async operation timed out after {timeout} seconds") + + if error: + raise error + return result \ No newline at end of file diff --git a/utils/config.py b/utils/config.py new file mode 100644 index 0000000..f1e8040 --- /dev/null +++ b/utils/config.py @@ -0,0 +1,50 @@ +""" +Configuration management for ComfyUI-Distributed. +""" +import os +import json +from .logging import log + +CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json") + +def get_default_config(): + """Returns the default configuration dictionary. Single source of truth.""" + return { + "master": {"host": ""}, + "workers": [], + "settings": { + "debug": False, + "auto_launch_workers": False, + "stop_workers_on_master_exit": True + } + } + +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: + log(f"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: + log(f"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): + from .logging import debug_log + debug_log("Created default config file") + else: + log("Could not create default config file") \ No newline at end of file diff --git a/utils/constants.py b/utils/constants.py new file mode 100644 index 0000000..a199f24 --- /dev/null +++ b/utils/constants.py @@ -0,0 +1,34 @@ +""" +Shared constants for ComfyUI-Distributed. +""" + +# Timeouts (in seconds) +WORKER_JOB_TIMEOUT = 10.0 +TILE_COLLECTION_TIMEOUT = 60.0 +TILE_WAIT_TIMEOUT = 30.0 +PROCESS_TERMINATION_TIMEOUT = 5.0 + +# Process monitoring +WORKER_CHECK_INTERVAL = 2.0 +STATUS_CHECK_INTERVAL = 5.0 + +# Network +CHUNK_SIZE = 8192 +LOG_TAIL_BYTES = 65536 # 64KB + +# File paths +WORKER_LOG_PATTERN = "distributed_worker_*.log" + +# Worker management +WORKER_STARTUP_DELAY = 2.0 + +# Tile transfer +TILE_TRANSFER_TIMEOUT = 30.0 + +# Process cleanup +PROCESS_WAIT_TIMEOUT = 3.0 +QUEUE_INIT_TIMEOUT = 5.0 +TILE_SEND_TIMEOUT = 60.0 + +# Memory operations +MEMORY_CLEAR_DELAY = 0.5 \ No newline at end of file diff --git a/utils/image.py b/utils/image.py new file mode 100644 index 0000000..c6786df --- /dev/null +++ b/utils/image.py @@ -0,0 +1,25 @@ +""" +Image and tensor conversion utilities for ComfyUI-Distributed. +""" +import torch +import numpy as np +from PIL import Image + +def tensor_to_pil(img_tensor, batch_index=0): + """Takes a batch of images in tensor form [B, H, W, C] and returns an RGB PIL Image.""" + 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, H, W, C].""" + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image).unsqueeze(0) + if len(image.shape) == 3: # If grayscale, add channel dimension + image = image.unsqueeze(-1) + return image + +def ensure_contiguous(tensor): + """Ensure tensor is contiguous in memory.""" + if not tensor.is_contiguous(): + return tensor.contiguous() + return tensor + diff --git a/utils/logging.py b/utils/logging.py new file mode 100644 index 0000000..15388a8 --- /dev/null +++ b/utils/logging.py @@ -0,0 +1,28 @@ +""" +Shared logging utilities for ComfyUI-Distributed. +""" +import os +import json + +# Config file is in parent directory +CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json") + +def is_debug_enabled(): + """Check if debug is enabled.""" + if os.path.exists(CONFIG_FILE): + try: + with open(CONFIG_FILE, 'r') as f: + config = json.load(f) + return config.get("settings", {}).get("debug", False) + except: + pass + return False + +def debug_log(message): + """Log debug messages only if debug is enabled in config.""" + if is_debug_enabled(): + print(f"[Distributed] {message}") + +def log(message): + """Always log important messages.""" + print(f"[Distributed] {message}") \ No newline at end of file diff --git a/utils/network.py b/utils/network.py new file mode 100644 index 0000000..d6199a3 --- /dev/null +++ b/utils/network.py @@ -0,0 +1,41 @@ +""" +Network and API utilities for ComfyUI-Distributed. +""" +import aiohttp +from aiohttp import web +from .logging import debug_log + +# Shared session for connection pooling +_client_session = None + +async def get_client_session(): + """Get or create a shared aiohttp client session.""" + global _client_session + if _client_session is None or _client_session.closed: + connector = aiohttp.TCPConnector(limit=100, limit_per_host=30) + # Don't set timeout here - set it per request + _client_session = aiohttp.ClientSession(connector=connector) + return _client_session + +async def cleanup_client_session(): + """Clean up the shared client session.""" + global _client_session + if _client_session and not _client_session.closed: + await _client_session.close() + _client_session = None + +async def handle_api_error(request, error, status=500): + """Standardized error response handler.""" + debug_log(f"API Error: {error}") + return web.json_response({"status": "error", "message": str(error)}, status=status) + +def get_server_port(): + """Get the ComfyUI server port.""" + import server + return server.PromptServer.instance.port + +def get_server_loop(): + """Get the ComfyUI server event loop.""" + import server + return server.PromptServer.instance.loop + diff --git a/utils/process.py b/utils/process.py new file mode 100644 index 0000000..85887b9 --- /dev/null +++ b/utils/process.py @@ -0,0 +1,38 @@ +""" +Process management utilities for ComfyUI-Distributed. +""" +import os +import subprocess +import platform +import signal + +def is_process_alive(pid): + """Check if a process with given PID is still alive.""" + try: + if platform.system() == "Windows": + # Windows: use tasklist + result = subprocess.run(['tasklist', '/FI', f'PID eq {pid}'], + capture_output=True, text=True) + return str(pid) in result.stdout + else: + # Unix: send signal 0 + os.kill(pid, 0) + return True + except (OSError, subprocess.SubprocessError): + return False + +def terminate_process(process, timeout=5): + """Gracefully terminate a process with timeout.""" + if process.poll() is None: # Still running + process.terminate() + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + +def get_python_executable(): + """Get the Python executable path.""" + import sys + return sys.executable + diff --git a/web/gpupanel.js b/web/gpupanel.js index 01765c2..d631dcd 100644 --- a/web/gpupanel.js +++ b/web/gpupanel.js @@ -1,24 +1,109 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; -class MultiGPUExtension { +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;" + 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.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 + + // Inject CSS for pulsing animation + this.injectStyles(); this.loadConfig().then(() => { this.registerSidebarTab(); this.setupInterceptor(); + this.startStatusChecking(); + this.loadManagedWorkers(); + this.detectMasterIP(); // Auto-detect master IP on startup }); } + + // Debug logging helpers + debugLog(message) { + if (this.config?.settings?.debug) { + console.log(`[Distributed] ${message}`); + } + } + + log(message) { + console.log(`[Distributed] ${message}`); + } + + 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) --- @@ -32,16 +117,16 @@ class MultiGPUExtension { async loadConfig() { try { - const response = await fetch(`${window.location.origin}/multigpu/config`); + const response = await fetch(`${window.location.origin}/distributed/config`); if (response.ok) { this.config = await response.json(); - console.log("[MultiGPU] Loaded config:", this.config); + this.debugLog("Loaded config: " + JSON.stringify(this.config)); } else { - console.error("[MultiGPU] Failed to load config"); + this.log("Failed to load config"); this.config = { workers: [], settings: {} }; } } catch (error) { - console.error("[MultiGPU] Error loading config:", error); + this.log("Error loading config: " + error.message); this.config = { workers: [], settings: {} }; } } @@ -53,25 +138,44 @@ class MultiGPUExtension { } try { - await fetch(`${window.location.origin}/multigpu/config/update_worker`, { + await fetch(`${window.location.origin}/distributed/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); + 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 fetch(`${window.location.origin}/multigpu/config/update_setting`, { + await fetch(`${window.location.origin}/distributed/config/update_setting`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ 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) { - console.error(`[MultiGPU] Error updating setting '${key}':`, error); + this.log(`Error updating setting '${key}': ${error.message}`); + app.extensionManager.toast.add({ + severity: "error", + summary: "Setting Update Failed", + detail: error.message, + life: 3000 + }); } } @@ -79,10 +183,10 @@ class MultiGPUExtension { registerSidebarTab() { app.extensionManager.registerSidebarTab({ - id: "multiGPU", + id: "distributed", icon: "pi pi-server", - title: "Multi-GPU", - tooltip: "Multi-GPU Control Panel", + title: "Distributed", + tooltip: "Distributed Control Panel", type: "custom", render: (el) => this.renderSidebarContent(el) }); @@ -91,11 +195,67 @@ class MultiGPUExtension { _createButton(text, onClick, style) { const button = document.createElement("button"); button.textContent = text; - button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + style; + button.style.cssText = DistributedExtension.BUTTON_STYLES.base + style; button.onclick = onClick; + button.className = "distributed-button"; return button; } + async _handleInterruptWorkers(button) { + const originalText = button.textContent; + const originalStyle = button.style.cssText; + button.textContent = "Interrupting..."; + button.disabled = true; + button.style.backgroundColor = "#c04c4c"; + + try { + const targetsToInterrupt = this.enabledWorkers.map(w => ({ + name: w.name, + url: `http://${w.host || "localhost"}:${w.port}` + })); + + if (targetsToInterrupt.length === 0) { + button.textContent = "No Workers"; + button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + }, 3000); + return; + } + + const promises = targetsToInterrupt.map(target => + fetch(`${target.url}/interrupt`, { + 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 = "Interrupted!"; + button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + // Update status after interrupt + setTimeout(() => this.checkAllWorkerStatuses(), 500); + } else { + button.textContent = "Error! See Console"; + button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + this.log("Failed to interrupt: " + failures.map(f => f.name).join(", ")); + } + + setTimeout(() => { + button.textContent = originalText; + button.style.cssText = originalStyle; + }, 3000); + } finally { + button.disabled = false; + } + } + async _handleClearMemory(button) { const originalText = button.textContent; button.textContent = "Clearing..."; @@ -109,16 +269,16 @@ class MultiGPUExtension { if (urlsToClear.length === 0) { button.textContent = "No Workers"; - button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); setTimeout(() => { button.textContent = originalText; - button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory; + button.style.cssText = DistributedExtension.BUTTON_STYLES.base + DistributedExtension.BUTTON_STYLES.clearMemory; }, 3000); return; } const promises = urlsToClear.map(target => - fetch(`${target.url}/multigpu/clear_memory`, { + fetch(`${target.url}/distributed/clear_memory`, { method: 'POST', mode: 'cors' }) @@ -131,16 +291,16 @@ class MultiGPUExtension { if (failures.length === 0) { button.textContent = "Success!"; - button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); } else { button.textContent = "Error! See Console"; - button.style.backgroundColor = MultiGPUExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); - console.error("[MultiGPU] Failed to clear memory on:", failures.map(f => f.name).join(", ")); + button.style.backgroundColor = DistributedExtension.BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + this.log("Failed to clear memory on: " + failures.map(f => f.name).join(", ")); } setTimeout(() => { button.textContent = originalText; - button.style.cssText = MultiGPUExtension.BUTTON_STYLES.base + MultiGPUExtension.BUTTON_STYLES.clearMemory; + button.style.cssText = DistributedExtension.BUTTON_STYLES.base + DistributedExtension.BUTTON_STYLES.clearMemory; }, 3000); } finally { button.disabled = false; @@ -149,12 +309,35 @@ class MultiGPUExtension { renderSidebarContent(el) { + // Panel is being opened/rendered + this.debugLog("Panel opened"); + + if (!el) { + this.debugLog("No element provided to renderSidebarContent"); + return; + } + + // Store reference to the panel element + this.panelElement = el; + + // 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 = ''; 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.textContent = "ComfyUI Distributed"; title.style.cssText = "margin: 0 0 15px 0; padding-bottom: 10px; border-bottom: 1px solid #444;"; container.appendChild(title); @@ -166,45 +349,434 @@ class MultiGPUExtension { gpuSection.appendChild(gpuTitle); const gpuList = document.createElement("div"); - (this.config?.workers || []).forEach(worker => { + const workers = this.config?.workers || []; + + // 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"; + + // 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 = document.createElement("div"); + blueprintContent.style.cssText = "flex: 1; display: flex; flex-direction: column;"; + + // Info row with placeholder content + const blueprintInfo = document.createElement("div"); + blueprintInfo.style.cssText = "display: flex; align-items: center; padding: 8px; min-height: 60px;"; + + const blueprintWorkerContent = document.createElement("div"); + blueprintWorkerContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + + const blueprintDot = document.createElement("span"); + blueprintDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; 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); + + // Click handler + blueprintDiv.onclick = () => this.addNewWorker(); + + // 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 = document.createElement("div"); - gpuDiv.style.cssText = "margin-bottom: 8px; padding: 8px; background: #2a2a2a; border-radius: 4px;"; + gpuDiv.style.cssText = "margin-bottom: 12px; background: #2a2a2a; border-radius: 4px; overflow: hidden; display: flex;"; + + // 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;"; + 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.onchange = async (e) => { - await this.updateWorkerEnabled(worker.id, e.target.checked); + 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(); }; - const label = document.createElement("label"); - label.htmlFor = `gpu-${worker.id}`; - label.style.cssText = "cursor: pointer; display: flex; align-items: center; gap: 8px;"; + + // 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: worker 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 = `status-${worker.id}`; + statusDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; background-color: #666;"; + statusDot.title = "Checking status..."; + const gpuInfo = document.createElement("span"); - const hostInfo = worker.host ? ` • ${worker.host}` : ''; - const cudaInfo = worker.cuda_device !== undefined ? `CUDA ${worker.cuda_device} • ` : ''; - gpuInfo.innerHTML = `${worker.name}
${cudaInfo}Port ${worker.port}${hostInfo}`; - label.appendChild(checkbox); - label.appendChild(gpuInfo); - gpuDiv.appendChild(label); + 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.expandedWorkers.has(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 launchBtn = this._createButton('Launch', + () => this.launchWorker(worker.id), + 'background-color: #555;'); // Neutral color initially + launchBtn.style.cssText += ' padding: 3px 8px; font-size: 11px;'; + 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.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"; + + controlsDiv.appendChild(launchBtn); + controlsDiv.appendChild(stopBtn); + controlsDiv.appendChild(logBtn); + } + + // 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)) { + settingsDiv.classList.add("expanded"); + } + settingsDiv.style.cssText = "margin: 0 8px; padding: 0 12px; background: #1e1e1e; border-radius: 4px;"; + if (this.expandedWorkers.has(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 = 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"; + + // 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 = document.createElement("div"); + addContentColumn.style.cssText = "flex: 1; display: flex; flex-direction: column;"; + + // 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) + + // Placeholder content + const addContent = document.createElement("div"); + addContent.style.cssText = "display: flex; align-items: center; gap: 8px; flex: 1;"; + + // Placeholder dot + const placeholderDot = document.createElement("span"); + placeholderDot.style.cssText = "display: inline-block; width: 8px; height: 8px; border-radius: 50%; 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"; + }; + + // Click handler + addWorkerDiv.onclick = () => this.addNewWorker(); + + 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;"; - const clearMemButton = this._createButton("Clear Worker VRAM", (e) => this._handleClearMemory(e.target), MultiGPUExtension.BUTTON_STYLES.clearMemory); + // 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)"; - actionsSection.appendChild(clearMemButton); + 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 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;"; + + settingsHeader.appendChild(settingsTitle); + settingsHeader.appendChild(settingsToggle); + + // Hover effect for header + settingsHeader.onmouseover = () => { + settingsToggle.style.color = "#fff"; + }; + settingsHeader.onmouseout = () => { + settingsToggle.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"; + settingsToggle.style.transform = "rotate(90deg)"; + } else { + settingsContent.style.maxHeight = "0"; + settingsToggle.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 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 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 = "multigpu-summary"; + summary.id = "distributed-summary"; summary.style.cssText = "font-size: 11px; color: #888;"; summarySection.appendChild(summary); container.appendChild(summarySection); @@ -213,7 +785,7 @@ class MultiGPUExtension { } updateSummary() { - const summaryEl = document.getElementById('multigpu-summary'); + const summaryEl = document.getElementById('distributed-summary'); if (summaryEl) { const totalGPUs = this.enabledWorkers.length + 1; if (this.isEnabled) { @@ -230,36 +802,58 @@ class MultiGPUExtension { setupInterceptor() { api.queuePrompt = async (number, prompt) => { - if (this.isEnabled && this.findNodesByClass(prompt.output, "MultiGPUCollector").length > 0) { - return await this.executeParallelMultiGPU(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 executeParallelMultiGPU(promptWrapper) { + async executeParallelDistributed(promptWrapper) { try { const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution const enabledWorkers = this.enabledWorkers; - // Find all collector nodes in the workflow - const collectorNodes = this.findNodesByClass(promptWrapper.output, "MultiGPUCollector"); + // 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(collectorNodes.map(node => [node.id, `${executionPrefix}_${node.id}`])); + 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._prepareMultiGpuJob(uniqueId)); + const preparePromises = Array.from(job_id_map.values()).map(uniqueId => this._prepareDistributedJob(uniqueId)); await Promise.all(preparePromises); const jobs = []; - const participants = ['master', ...enabledWorkers.map(w => w.id)]; + // Use only active workers + const participants = ['master', ...activeWorkers.map(w => w.id)]; for (const participantId of participants) { const jobApiPrompt = this._prepareApiPromptForParticipant( promptWrapper.output, participantId, { - enabled_worker_ids: enabledWorkers.map(w => w.id), + enabled_worker_ids: activeWorkers.map(w => w.id), workflow: promptWrapper.workflow, job_id_map: job_id_map // Pass the map of unique IDs } @@ -268,7 +862,7 @@ class MultiGPUExtension { if (participantId === 'master') { jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); } else { - const worker = this.config.workers.find(w => w.id === participantId); + const worker = activeWorkers.find(w => w.id === participantId); if (worker) jobs.push({ type: 'worker', worker, prompt: jobApiPrompt, workflow: promptWrapper.workflow }); } } @@ -276,8 +870,7 @@ class MultiGPUExtension { const result = await this._executeJobs(jobs); return result; } catch (error) { - console.error("[MultiGPU] Parallel execution failed:", error); - alert(`[MultiGPU] Parallel execution failed: ${error.message}`); + this.log("Parallel execution failed: " + error.message); throw error; } } @@ -296,7 +889,7 @@ class MultiGPUExtension { } }); await Promise.all(promises); - return masterPromptId || { "prompt_id": "multi-gpu-job-dispatched" }; + return masterPromptId || { "prompt_id": "distributed-job-dispatched" }; } // --- Helper Methods --- @@ -366,27 +959,29 @@ class MultiGPUExtension { } /** - * Prune workflow to only include nodes connected to collectors + * Prune workflow to only include nodes connected to distributed nodes * @param {Object} apiPrompt - The full workflow API prompt - * @param {Array} collectorNodes - Array of collector nodes (optional, will find if not provided) + * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided) * @returns {Object} Pruned API prompt with only required nodes */ - pruneWorkflowForWorker(apiPrompt, collectorNodes = null) { - // Find all collector nodes if not provided - if (!collectorNodes) { - collectorNodes = this.findNodesByClass(apiPrompt, "MultiGPUCollector"); + 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 (collectorNodes.length === 0) { - // No collectors, return full workflow + if (distributedNodes.length === 0) { + // No distributed nodes, return full workflow return apiPrompt; } - // Get all nodes connected to collectors (both upstream and downstream) - const collectorIds = collectorNodes.map(node => node.id); - const connectedNodes = this.findCollectorConnectedNodes(apiPrompt, collectorIds); + // Get all nodes connected to distributed nodes (both upstream and downstream) + const distributedIds = distributedNodes.map(node => node.id); + const connectedNodes = this.findCollectorConnectedNodes(apiPrompt, distributedIds); - console.log(`[MultiGPU] Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes`); + this.debugLog(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes`); // Create pruned prompt with only required nodes const prunedPrompt = {}; @@ -402,68 +997,226 @@ class MultiGPUExtension { let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); const isMaster = participantId === 'master'; - // Find collector nodes once (before pruning) - const collectorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUCollector"); + // 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 collector dependencies - if (!isMaster && collectorNodes.length > 0) { - jobApiPrompt = this.pruneWorkflowForWorker(jobApiPrompt, collectorNodes); + // For workers, prune the workflow to only include distributed node dependencies + if (!isMaster && allDistributedNodes.length > 0) { + jobApiPrompt = this.pruneWorkflowForWorker(jobApiPrompt, allDistributedNodes); } - // Handle MultiGPU distributor nodes - const distributorNodes = this.findNodesByClass(jobApiPrompt, "MultiGPUDistributor"); + // Handle Distributed seed nodes + const distributorNodes = this.findNodesByClass(jobApiPrompt, "DistributedSeed"); if (distributorNodes.length > 0) { - console.log(`[MultiGPU] Found ${distributorNodes.length} distributor node(s)`); + this.debugLog(`Found ${distributorNodes.length} seed node(s)`); } - for (const distributor of distributorNodes) { - const { inputs } = jobApiPrompt[distributor.id]; + 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}`; - console.log(`[MultiGPU] Set distributor ${distributor.id} for worker ${workerIndex}`); + this.debugLog(`Set seed node ${seedNode.id} for worker ${workerIndex}`); } } - // Handle MultiGPU collector nodes (already found above) + // Handle Distributed collector nodes (already found above) for (const collector of collectorNodes) { const { inputs } = jobApiPrompt[collector.id]; - // 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; + // 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; - // 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 = window.location.origin; - // Also make the worker_job_id unique to prevent potential caching issues - inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; + 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 _prepareMultiGpuJob(multi_job_id) { + async _prepareDistributedJob(multi_job_id) { try { - await fetch(`${window.location.origin}/multigpu/prepare_job`, { + await fetch(`${window.location.origin}/distributed/prepare_job`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ multi_job_id }) }); } catch (error) { - console.error("[MultiGPU] Error preparing job:", 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() { + if (!this.config || !this.config.workers) return; + + for (const worker of this.config.workers) { + this.checkWorkerStatus(worker); + } + } + + async checkWorkerStatus(worker) { + const host = worker.host || "localhost"; + const url = `http://${host}:${worker.port}/prompt`; + const statusDot = document.getElementById(`status-${worker.id}`); + + // Status dot will be updated when response comes back + + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: AbortSignal.timeout(300) // 300ms for faster response + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + const isProcessing = queueRemaining > 0; + + // Update status + this.workerStatuses.set(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.launchingWorkers.has(worker.id)) { + this.launchingWorkers.delete(worker.id); + + // Also clear launching flag in backend + this.clearLaunchingFlag(worker.id); + } + } else { + throw new Error(`HTTP ${response.status}`); + } + } catch (error) { + // Worker is offline or unreachable + this.workerStatuses.set(worker.id, { + online: false, + processing: false, + queueCount: 0 + }); + + // Check if worker is launching + if (this.launchingWorkers.has(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); + } + } + + // Update control buttons based on new status + this.updateWorkerControls(worker.id); + } + async _dispatchToWorker(worker, prompt, workflow) { - const workerUrl = `http://${worker.host || window.location.hostname}:${worker.port}`; + const host = worker.host || "localhost"; + const workerUrl = `http://${host}:${worker.port}`; const promptToSend = { prompt, @@ -480,14 +1233,1152 @@ class MultiGPUExtension { body: JSON.stringify(promptToSend) }); } catch (e) { - console.error(`[MultiGPU] Failed to connect to worker ${worker.name} at ${workerUrl}`, e); + this.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.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 host = worker.host || "localhost"; + const url = `http://${host}:${worker.port}/prompt`; + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + signal: AbortSignal.timeout(300) // 300ms timeout for pre-flight + }); + + 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`); + 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.launchingWorkers.add(workerId); + + setTimeout(() => { + this.launchingWorkers.delete(workerId); + }, 15000); + + if (!launchBtn) return; + + try { + // 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) { + this.log(`Launched ${worker.name} (PID: ${result.pid})`); + if (result.log_file) { + this.debugLog(`Log file: ${result.log_file}`); + } + + this.managedWorkers.set(String(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); + } 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; + } + } + } + + 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 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) { + this.log(`Stopped worker: ${result.message}`); + this.managedWorkers.delete(String(workerId)); + + // 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 fetch(`${window.location.origin}/distributed/worker/clear_launching`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ worker_id: workerId }) + }); + this.debugLog(`Cleared launching flag for worker ${workerId}`); + } catch (error) { + this.debugLog(`Error clearing launching flag: ${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 response = await fetch(`${window.location.origin}/distributed/managed_workers`); + if (response.ok) { + const result = await response.json(); + this.managedWorkers.clear(); + + // 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)); + } + } + } 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.managedWorkers.get(String(workerId)); + const status = this.workerStatuses.get(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.managedWorkers.get(String(workerId)); + 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 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(); + + // 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 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(); + + // 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 + }); + } + } + } + + 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'; + return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; + } + + isRemoteWorker(worker) { + const host = worker.host || "localhost"; + 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) { + return `http://${this.config.master.host}:${window.location.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 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(); + this.debugLog("Network info: " + JSON.stringify(data)); + + // 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.saveMasterIp(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 saveMasterIp(masterIp, 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() + }) + }); + + if (response.ok) { + // Update local config + if (!this.config.master) { + this.config.master = {}; + } + this.config.master.host = masterIp.trim(); + + // 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.expandedWorkers.has(workerId)) { + this.expandedWorkers.delete(workerId); + 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.expandedWorkers.add(workerId); + 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; + } + + 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; + } + + 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 : "localhost"; + 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 < 1024 || port > 65535) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail: "Port must be between 1024 and 65535", + life: 3000 + }); + return; + } + + // Check for port conflicts + const portConflict = this.config.workers.some(w => + w.id !== workerId && w.port === port + ); + + if (portConflict) { + app.extensionManager.toast.add({ + severity: "error", + summary: "Port Conflict", + detail: `Port ${port} is already in use by another worker`, + 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 + }); + } + } + + 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: false, + 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.expandedWorkers.add(newId); + 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: "MultiGPU.Panel", + name: "Distributed.Panel", async setup() { - new MultiGPUExtension(); + new DistributedExtension(); } }); \ No newline at end of file diff --git a/worker_monitor.py b/worker_monitor.py new file mode 100644 index 0000000..62fd856 --- /dev/null +++ b/worker_monitor.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +""" +Worker process monitor - monitors if the master process is still alive +and terminates the worker if the master dies. +""" +import os +import sys +import time +import subprocess +import platform +import signal + +# Add parent directory to path to import utils +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +try: + from ComfyUI_Distributed.utils.process import is_process_alive, terminate_process + from ComfyUI_Distributed.utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT +except ImportError: + # Fallback if running from different context + def is_process_alive(pid): + """Check if a process with given PID is still alive.""" + try: + if platform.system() == "Windows": + # Windows: use tasklist + result = subprocess.run(['tasklist', '/FI', f'PID eq {pid}'], + capture_output=True, text=True) + return str(pid) in result.stdout + else: + # Unix: send signal 0 + os.kill(pid, 0) + return True + except (OSError, subprocess.SubprocessError): + return False + + WORKER_CHECK_INTERVAL = 2.0 + PROCESS_TERMINATION_TIMEOUT = 5.0 + +def monitor_and_run(master_pid, command): + """Run command and monitor master process.""" + # Start the actual worker process + worker_process = subprocess.Popen(command) + + print(f"[Distributed] Started worker PID: {worker_process.pid}") + print(f"[Distributed] Monitoring master PID: {master_pid}") + + # Write worker PID to a file so parent can track it + monitor_pid = os.getpid() + pid_info_file = os.environ.get('WORKER_PID_FILE') + if pid_info_file: + try: + with open(pid_info_file, 'w') as f: + f.write(f"{monitor_pid},{worker_process.pid}") + print(f"[Distributed] Wrote PID info to {pid_info_file}") + except Exception as e: + print(f"[Distributed] Could not write PID file: {e}") + + # Define cleanup function + def cleanup_worker(signum=None, frame=None): + """Clean up worker process when monitor is terminated.""" + if signum: + print(f"\n[Distributed] Received signal {signum}, terminating worker...") + else: + print("\n[Distributed] Terminating worker...") + + if worker_process.poll() is None: # Still running + try: + terminate_process(worker_process, timeout=PROCESS_TERMINATION_TIMEOUT) + except NameError: + # Fallback if terminate_process wasn't imported + worker_process.terminate() + try: + worker_process.wait(timeout=PROCESS_TERMINATION_TIMEOUT) + except subprocess.TimeoutExpired: + print("[Distributed] Worker didn't terminate gracefully, forcing kill...") + worker_process.kill() + worker_process.wait() + + print("[Distributed] Worker terminated.") + sys.exit(0) + + # Register signal handlers for graceful shutdown + signal.signal(signal.SIGTERM, cleanup_worker) + signal.signal(signal.SIGINT, cleanup_worker) + if platform.system() != "Windows": + signal.signal(signal.SIGHUP, cleanup_worker) + + # Monitor loop + check_interval = WORKER_CHECK_INTERVAL + + try: + while True: + # Check if worker is still running + if worker_process.poll() is not None: + print(f"[Distributed] Worker process exited with code: {worker_process.returncode}") + sys.exit(worker_process.returncode) + + # Check if master is still alive + if not is_process_alive(master_pid): + print(f"[Distributed] Master process {master_pid} is no longer running. Terminating worker...") + cleanup_worker() + + time.sleep(check_interval) + + except KeyboardInterrupt: + cleanup_worker() + +if __name__ == "__main__": + # Get master PID from environment + master_pid = os.environ.get('COMFYUI_MASTER_PID') + if not master_pid: + print("[Distributed] Error: COMFYUI_MASTER_PID not set") + sys.exit(1) + + try: + master_pid = int(master_pid) + except ValueError: + print(f"[Distributed] Error: Invalid master PID: {master_pid}") + sys.exit(1) + + # Get the actual command to run (all remaining arguments) + if len(sys.argv) < 2: + print("[Distributed] Error: No command specified") + sys.exit(1) + + command = sys.argv[1:] + + # Start monitoring + monitor_and_run(master_pid, command) \ No newline at end of file