Main release
This commit is contained in:
robertvoy
@@ -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.
|
||||
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
<details>
|
||||
<summary><b>Workers won't start</b></summary>
|
||||
|
||||
- Check ports are not in use: `netstat -an | grep 8189`
|
||||
- Verify CUDA device exists: `nvidia-smi`
|
||||
- Check ComfyUI path in worker logs
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>"Worker not managed by UI" message</b></summary>
|
||||
|
||||
- Worker was started outside the UI
|
||||
- Stop the worker manually and use the UI to relaunch
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Images not combining properly</b></summary>
|
||||
|
||||
- Ensure all workers have the same models loaded
|
||||
- Check that custom nodes are installed on all workers
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Network connection issues</b></summary>
|
||||
|
||||
- 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
|
||||
</details>
|
||||
|
||||
### 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.
|
||||
<div align="center">
|
||||
|
||||
## 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)
|
||||
|
||||
</div>
|
||||
+20
-8
@@ -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}")
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
+1483
File diff suppressed because it is too large
Load Diff
@@ -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)",
|
||||
}
|
||||
@@ -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": {}
|
||||
}
|
||||
@@ -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
|
||||
)
|
||||
<nul set /p ="."
|
||||
powershell -NoProfile -Command "Start-Sleep -m %RETRY_DELAY%"
|
||||
goto MASTER_SLOW
|
||||
)
|
||||
echo Master ready!
|
||||
|
||||
:: Process workers sequentially - wait for each to be ready before starting next
|
||||
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 :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
|
||||
)
|
||||
<nul set /p ="."
|
||||
powershell -NoProfile -Command "Start-Sleep -m %RETRY_DELAY%"
|
||||
goto WORKER_CHECK_LOOP_SLOW
|
||||
)
|
||||
echo %WORKER_NAME% ready!
|
||||
goto :eof
|
||||
|
||||
::=============================================================================
|
||||
:: ALL-FAST MODE - Launch all instances with 2-second delays
|
||||
::=============================================================================
|
||||
:MODE_ALL_FAST
|
||||
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; 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
|
||||
|
||||
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
|
||||
-438
@@ -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"
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
"""
|
||||
Utility modules for ComfyUI-Distributed extension.
|
||||
"""
|
||||
|
||||
# Make utils importable as a package
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+1980
-89
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
Reference in New Issue
Block a user