Merge branch 'feature/cloud-workers'

This commit is contained in:
robertvoy
2025-07-29 13:37:53 +10:00
22 changed files with 13629 additions and 3566 deletions
+30
View File
@@ -4,6 +4,36 @@ import sys
# Add the directory to Python path
sys.path.append(os.path.dirname(__file__))
# Patch ComfyUI execution validation for DistributedVideoCollector
try:
import execution
from .distributed import ImageBatchDivider
# Store original validate function if it exists
if hasattr(execution, 'validate_outputs'):
original_validate_outputs = execution.validate_outputs
def patched_validate_outputs(executor, node_id, result, node_class):
if node_class == ImageBatchDivider:
return # Skip validation for our dynamic output node
return original_validate_outputs(executor, node_id, result, node_class)
execution.validate_outputs = patched_validate_outputs
else:
# Fallback: patch the main execute method if validate_outputs doesn't exist
if hasattr(execution.PromptExecutor, 'execute'):
original_execute = execution.PromptExecutor.execute
def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
# This is a more complex patch - for now just call original
# The ByPassTypeTuple should handle most validation issues
return original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
execution.PromptExecutor.execute = patched_execute
except ImportError:
pass # ComfyUI execution module not available during import
# Import everything needed from the main module
from .distributed import (
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
+753 -133
View File
File diff suppressed because it is too large Load Diff
+1855 -418
View File
File diff suppressed because it is too large Load Diff
+56
View File
@@ -0,0 +1,56 @@
## Automating ComfyUI Model Downloads
This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM).
1. Launched the Pod with these Environment Variables:
- CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account))
- HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens))
2. In ComfyUI, export your workflow as an API workflow
3. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet**
```
Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths:
base_path: /workspace/ComfyUI
checkpoints: models/checkpoints/
clip: models/clip/
clip_vision: models/clip_vision/
controlnet: models/controlnet/
diffusion_models: models/diffusion_models/
embeddings: models/embeddings/
florence2: models/florence2/
ipadapter: models/ipadapter/
loras: models/loras/
style_models: models/style_models/
text_encoders: models/text_encoders/
unet: models/unet/
upscale_models: models/upscale_models/
vae: models/vae/
---
Important:
Make sure you find the correct URLs for the models online.
Use comfy cli to download the models: `comfy model download --url <URL> [--relative-path <PATH>] [--set-civitai-api-token <TOKEN>] [--set-hf-api-token <TOKEN>]`
Make sure you add `--set-civitai-api-token $CIVITAI_API_TOKEN` for CivitAI download and `--set-hf-api-token $HF_API_TOKEN` for Hugging Face downloads.
---
Example:
#!/bin/bash
# Download from CivitAI
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
# Download model from Hugging Face
comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN
# If a model in the workflow was in a subfolder
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints/SDXL --set-civitai-api-token $CIVITAI_API_TOKEN
```
4. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh`
5. Upload it onto your Runpod instance, into `/workspace`
6. Then run these commands:
- `chmod 755 /workspace/download_models.sh`
- `/workspace/download_models.sh`
+108
View File
@@ -0,0 +1,108 @@
## ComfyUI Distributed Worker Setup Guide
ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs, whether they're in the same machine, on your local network, or in the cloud.
**Master**: The main ComfyUI instance that coordinates and distributes work. This is where you load workflows, manage the queue, and view results.
**Worker**: A ComfyUI instance that receives and processes tasks from the master. Workers handle just the GPU computation and send results back to the master. You can have multiple workers connected to a single master, each utilizing their own GPU.
### Types of Workers
- **Local workers**: Additional GPUs on the same machine as the master
- **Remote workers**: GPUs on different computers within your network
- **Cloud workers**: GPUs hosted on cloud services like Runpod
## Local workers
> These are added automatically on first launch, but you can add them manually if you need to.
1. **Open** the Distributed GPU panel.
2. **Click** "Add Worker" in the UI.
3. **Configure** your local worker:
- **Name**: A descriptive name for the worker (e.g., "Studio PC 1")
- **Port**: A unique port number for this worker (e.g., 8189, 8190...).
- **CUDA Device**: The GPU index from `nvidia-smi` (e.g., 0, 1).
- **Extra Args**: Optional ComfyUI arguments for this specific worker.
4. **Save** and launch the local worker.
## Remote workers
> ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address.
**On the Remote Worker Machine:**
1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!**
- This ComfyUI instance will serve as a worker for your main master.
2. *Optionally* add additional local workers on this machine if it has multiple GPUs:
- Access the Distributed GPU panel in this ComfyUI instance
- Add workers for any additional GPUs (if they haven't been added automatically)
- Make sure they have `--listen` set in `Extra Args`
- Launch them
3. **Open** the ComfyUI port (e.g., 8188) and any additional worker ports (e.g., 8189, 8190) in the firewall.
**On the Main Machine:**
1. **Open** the Distributed GPU panel (sidebar on the left).
2. **Click** "Add Worker."
3. **Choose** "Remote".
4. **Configure** your remote worker:
- **Name**: A descriptive name for the worker (e.g., "Server Rack GPU 0")
- **Host**: The remote worker's IP address.
- **Port**: The port number used when launching ComfyUI on the remote master/worker (e.g., 8188).
5. **Save** the remote worker configuration.
## Cloud workers
> ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p).
### Deploy Cloud Worker on Runpod
**On Runpod:**
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later.
3. Now go to Pods and find a suitable GPU for your workflows.
4. Choose the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template and make sure your network drive is mounted.
> To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8.
6. Launch your pod.
7. Access your pod using JupyterLabs.
8. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod).
> You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow.
8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!**
9. Download any additional custom nodes using the ComfyUI Manager.
**On the Main Machine:**
1. Launch a Cloudflare tunnel
- Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases)
- Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188`
> Cloudflare tunnels create secure connections without exposing ports directly to the internet
2. Copy the Cloudflare address
3. **Open** the Distributed GPU panel (sidebar on the left).
4. **Edit** the Master's host address and replace it with the Cloudflare address.
- **Click** "Add Worker."
- **Choose** "Cloud".
- **Configure** your cloud worker:
- **Host**: The Runpod address
- **Port**: 443
- **Save** the remote worker configuration.
---
### Deploy Cloud Worker on Other Platforms
**On the Cloud Worker machine:**
- Your cloud worker container needs to have the same models and custom nodes as the workflow you want to run on your local machine.
- If your cloud platform doesn't provide a secure connection, use Cloudflare to create a tunnel for the worker. Each GPU needs their own tunnel for their respective port.
- For example: `./cloudflared tunnel --url http://localhost:8188`
1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!**
2. **Add** workers in the UI panel if the remote machine has more than one GPU.
- Make sure that they also have `--listen` set in `Extra Args`.
- Then launch them.
**On the Main Machine:**
1. **Launch** a Cloudflare tunnel on your local machine.
- Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases)
- Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188`
2. **Copy** the Cloudflare address
3. **Open** the Distributed GPU panel (sidebar on the left).
4. **Edit** the Master's host address and replace it with the Cloudflare address.
5. **Click** "Add Worker."
6. **Choose** "Cloud".
7. **Configure** your cloud worker:
- **Host**: The remote worker's IP address/domain
- **Port**: 443
8. **Save** the remote worker configuration.
+10 -3
View File
@@ -1,10 +1,11 @@
"""
Shared constants for ComfyUI-Distributed.
"""
import os
# Timeouts (in seconds)
WORKER_JOB_TIMEOUT = 10.0
TILE_COLLECTION_TIMEOUT = 60.0
WORKER_JOB_TIMEOUT = 30.0
TILE_COLLECTION_TIMEOUT = 30.0
TILE_WAIT_TIMEOUT = 30.0
PROCESS_TERMINATION_TIMEOUT = 5.0
@@ -31,4 +32,10 @@ QUEUE_INIT_TIMEOUT = 5.0
TILE_SEND_TIMEOUT = 60.0
# Memory operations
MEMORY_CLEAR_DELAY = 0.5
MEMORY_CLEAR_DELAY = 0.5
# Batch processing
MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
# Heartbeat monitoring
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds
+821
View File
@@ -0,0 +1,821 @@
import asyncio
import time
import json
import copy
import os
import io
from aiohttp import web
import server
from PIL import Image
import numpy as np
import torch
# Import from other utilities
from .logging import debug_log, log
from .network import handle_api_error, get_client_session
from .image import tensor_to_pil
# Configure maximum payload size (50MB default, configurable via environment variable)
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
# Import HEARTBEAT_TIMEOUT from constants
from .constants import HEARTBEAT_TIMEOUT
# Unified Job Data Structure Keys
JOB_QUEUE = 'queue'
JOB_MODE = 'mode'
JOB_COMPLETED_TASKS = 'completed_tasks'
JOB_WORKER_STATUS = 'worker_status'
JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers'
JOB_PENDING_TASKS = 'pending_tasks'
JOB_BATCH_SIZE = 'batch_size' # For dynamic
JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static
# Task Types
TASK_TYPE_TILE = 'tile'
TASK_TYPE_IMAGE = 'image'
async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None):
"""Unified initialization for job queues in static and dynamic modes."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
debug_log(f"Queue already exists for {multi_job_id}")
return
job_data = {
JOB_QUEUE: asyncio.Queue(),
JOB_MODE: mode,
JOB_COMPLETED_TASKS: {},
JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []},
JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []},
JOB_PENDING_TASKS: asyncio.Queue(),
}
if mode == 'dynamic':
job_data[JOB_BATCH_SIZE] = batch_size
pending_queue = job_data[JOB_PENDING_TASKS]
for i in all_indices or range(batch_size):
await pending_queue.put(i)
debug_log(f"Initialized dynamic queue with {batch_size} pending images")
elif mode == 'static':
job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
# For static with dynamic distribution, populate all tile indices in pending queue
pending_queue = job_data[JOB_PENDING_TASKS]
total_tiles = batch_size * num_tiles_per_image
for i in range(total_tiles):
await pending_queue.put(i)
debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution")
# Keep backward compatibility - if task assignments provided, still track them
if task_assignments and enabled_workers:
# task_assignments[0] is for master, 1+ are for workers
for i, worker_id in enumerate(enabled_workers):
if i + 1 < len(task_assignments):
job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1]
debug_log(f"Worker {worker_id} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)")
else:
raise ValueError(f"Unknown mode: {mode}")
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]:
"""Distribute a list of items among N participants (master + workers)."""
if num_participants <= 1:
return [items]
items_per_participant = len(items) // num_participants
remainder = len(items) % num_participants
assignments = []
start_idx = 0
for i in range(num_participants):
count = items_per_participant + (1 if i < remainder else 0)
end_idx = start_idx + count
assignments.append(items[start_idx:end_idx])
start_idx = end_idx
return assignments
async def _get_next_task(multi_job_id):
"""Get next task from pending queue (generalized for tiles/images)."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not job_data or JOB_PENDING_TASKS not in job_data:
return None
try:
task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0)
return task_id
except asyncio.TimeoutError:
return None
async def _drain_results_queue(multi_job_id):
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not job_data or JOB_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data:
return 0
q = job_data[JOB_QUEUE]
completed_tasks = job_data[JOB_COMPLETED_TASKS]
collected = 0
while not q.empty():
try:
result = await asyncio.wait_for(q.get(), timeout=0.1)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
if 'image_idx' in result and 'image' in result:
task_id = result['image_idx']
if task_id not in completed_tasks:
completed_tasks[task_id] = result['image']
collected += 1
elif 'tiles' in result:
for tile_data in result['tiles']:
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
if task_id not in completed_tasks:
completed_tasks[task_id] = tile_data
collected += 1
elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat
task_id = result.get('global_idx', result['tile_idx'])
if task_id not in completed_tasks:
completed_tasks[task_id] = {
'tensor': result['tensor'],
'tile_idx': result['tile_idx'],
'x': result['x'],
'y': result['y'],
'extracted_width': result['extracted_width'],
'extracted_height': result['extracted_height'],
'padding': result['padding'],
'batch_idx': result.get('batch_idx', 0),
'global_idx': task_id
}
collected += 1
if is_last:
# Track worker completion
if worker_id in job_data[JOB_WORKER_STATUS]:
del job_data[JOB_WORKER_STATUS][worker_id]
except asyncio.TimeoutError:
break
return collected
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
"""Check timed out workers and requeue their tasks. Returns requeued count."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not job_data:
return 0
current_time = time.time()
requeued_count = 0
completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {})
for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()):
if current_time - last_heartbeat > HEARTBEAT_TIMEOUT:
log(f"Worker {worker} timed out")
for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []):
if task_id not in completed_tasks:
await job_data[JOB_PENDING_TASKS].put(task_id)
requeued_count += 1
if JOB_WORKER_STATUS in job_data:
del job_data[JOB_WORKER_STATUS][worker]
if JOB_ASSIGNED_TO_WORKERS in job_data:
job_data[JOB_ASSIGNED_TO_WORKERS][worker] = []
return requeued_count
async def _get_completed_count(multi_job_id):
"""Get count of completed tasks."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and JOB_COMPLETED_TASKS in job_data:
return len(job_data[JOB_COMPLETED_TASKS])
return 0
async def _mark_task_completed(multi_job_id, task_id, result):
"""Mark a task as completed."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and JOB_COMPLETED_TASKS in job_data:
job_data[JOB_COMPLETED_TASKS][task_id] = result
async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
"""Send heartbeat to master."""
try:
data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
session = await get_client_session()
url = f"{master_url}/distributed/heartbeat"
async with session.post(url, json=data) as response:
response.raise_for_status()
except Exception as e:
debug_log(f"Heartbeat failed: {e}")
async def _cleanup_job(multi_job_id):
"""Cleanup the job data."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
debug_log(f"Cleaned up job {multi_job_id}")
# API Endpoints (generalized)
@server.PromptServer.instance.routes.post("/distributed/request_task")
async def request_task_endpoint(request):
try:
data = await request.json()
worker_id = data.get('worker_id')
multi_job_id = data.get('multi_job_id')
if not worker_id or not multi_job_id:
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
mode = job_data.get(JOB_MODE)
pending_queue = job_data.get(JOB_PENDING_TASKS)
if pending_queue:
try:
task_id = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
if JOB_ASSIGNED_TO_WORKERS in job_data and worker_id in job_data[JOB_ASSIGNED_TO_WORKERS]:
job_data[JOB_ASSIGNED_TO_WORKERS][worker_id].append(task_id)
if JOB_WORKER_STATUS in job_data:
job_data[JOB_WORKER_STATUS][worker_id] = time.time()
remaining = pending_queue.qsize()
debug_log(f"Assigned task {task_id} to worker {worker_id} in {mode} mode")
return web.json_response({"task_id": task_id, "estimated_remaining": remaining, "mode": mode})
except asyncio.TimeoutError:
return web.json_response({"task_id": None})
else:
return await handle_api_error(request, "No pending tasks", 400)
else:
return await handle_api_error(request, "Job not found", 404)
except Exception as e:
return await handle_api_error(request, e, 500)
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
async def heartbeat_endpoint(request):
try:
data = await request.json()
worker_id = data.get('worker_id')
multi_job_id = data.get('multi_job_id')
if not worker_id or not multi_job_id:
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_WORKER_STATUS in job_data:
job_data[JOB_WORKER_STATUS][worker_id] = time.time()
debug_log(f"Heartbeat from worker {worker_id}")
return web.json_response({"status": "success"})
else:
return await handle_api_error(request, "Worker status tracking not available", 400)
else:
return await handle_api_error(request, "Job not found", 404)
except Exception as e:
return await handle_api_error(request, e, 500)
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
async def submit_tiles_endpoint(request):
"""Endpoint for workers to submit processed tiles in static mode."""
try:
content_length = request.headers.get('content-length')
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
data = await request.post()
multi_job_id = data.get('multi_job_id')
worker_id = data.get('worker_id')
is_last = data.get('is_last', 'False').lower() == 'true'
if multi_job_id is None or worker_id is None:
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
prompt_server = ensure_tile_jobs_initialized()
batch_size = int(data.get('batch_size', 0))
tiles = []
# Handle completion signal
if batch_size == 0 and is_last:
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
'is_last': True,
'tiles': []
})
debug_log(f"Received completion signal from worker {worker_id}")
return web.json_response({"status": "success"})
# Handle batch tiles with metadata
if batch_size > 0:
padding = int(data.get('padding', 32))
metadata_field = data.get('tiles_metadata')
if metadata_field:
if hasattr(metadata_field, 'file'):
metadata_str = metadata_field.file.read().decode('utf-8')
elif isinstance(metadata_field, (bytes, bytearray)):
metadata_str = metadata_field.decode('utf-8')
else:
metadata_str = str(metadata_field)
metadata = json.loads(metadata_str)
if len(metadata) != batch_size:
return await handle_api_error(request, "Metadata length mismatch", 400)
tile_data_list = []
for i in range(batch_size):
tile_field = data.get(f'tile_{i}')
if tile_field is None:
continue
img_data = tile_field.file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np)[None,]
if i < len(metadata):
tile_meta = metadata[i]
tile_idx = tile_meta.get('tile_idx', i)
tile_info = {
'tensor': tensor,
'tile_idx': tile_idx,
'x': tile_meta['x'],
'y': tile_meta['y'],
'extracted_width': tile_meta['extracted_width'],
'extracted_height': tile_meta['extracted_height'],
'padding': padding,
'batch_idx': tile_meta.get('batch_idx', 0),
'global_idx': tile_meta.get('global_idx', tile_idx)
}
tile_data_list.append(tile_info)
tile_data_list.sort(key=lambda x: x['tile_idx'])
tiles.extend(tile_data_list)
else:
# Legacy format
for i in range(batch_size):
tile_field = data.get(f'tile_{i}')
if tile_field is None:
continue
tile_idx = int(data.get(f'tile_{i}_idx', i))
x = int(data.get(f'tile_{i}_x', 0))
y = int(data.get(f'tile_{i}_y', 0))
extracted_width = int(data.get(f'tile_{i}_width', 512))
extracted_height = int(data.get(f'tile_{i}_height', 512))
img_data = tile_field.file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np)[None,]
tiles.append({
'tensor': tensor,
'tile_idx': tile_idx,
'x': x,
'y': y,
'extracted_width': extracted_width,
'extracted_height': extracted_height,
'padding': padding
})
else:
# Single tile legacy
image_file = data.get('image')
if not image_file:
return await handle_api_error(request, "Missing image data", 400)
tile_idx = int(data.get('tile_idx', 0))
x = int(data.get('x', 0))
y = int(data.get('y', 0))
extracted_width = int(data.get('extracted_width', 512))
extracted_height = int(data.get('extracted_height', 512))
padding = int(data.get('padding', 32))
img_data = image_file.file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np)[None,]
tiles = [{
'tensor': tensor,
'tile_idx': tile_idx,
'x': x,
'y': y,
'extracted_width': extracted_width,
'extracted_height': extracted_height,
'padding': padding,
'batch_idx': 0,
'global_idx': tile_idx
}]
# Submit tiles to queue
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
q = job_data[JOB_QUEUE]
if batch_size > 0 or len(tiles) > 0:
await q.put({
'worker_id': worker_id,
'tiles': tiles,
'is_last': is_last
})
else:
await q.put({
'worker_id': worker_id,
'is_last': True,
'tiles': []
})
return web.json_response({"status": "success"})
else:
return await handle_api_error(request, "Job not found", 404)
except Exception as e:
return await handle_api_error(request, e, 500)
@server.PromptServer.instance.routes.post("/distributed/submit_image")
async def submit_image_endpoint(request):
"""Endpoint for workers to submit processed images in dynamic mode."""
try:
content_length = request.headers.get('content-length')
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
data = await request.post()
multi_job_id = data.get('multi_job_id')
worker_id = data.get('worker_id')
is_last = data.get('is_last', 'False').lower() == 'true'
if multi_job_id is None or worker_id is None:
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
prompt_server = ensure_tile_jobs_initialized()
# Handle image submission
if 'full_image' in data and 'image_idx' in data:
image_idx = int(data.get('image_idx'))
img_data = data['full_image'].file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
debug_log(f"Received full image {image_idx} from worker {worker_id}")
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
'image_idx': image_idx,
'image': img,
'is_last': is_last
})
return web.json_response({"status": "success"})
# Handle completion signal (no image data)
elif is_last:
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
'is_last': True,
'tiles': [] # For compatibility
})
debug_log(f"Received completion signal from worker {worker_id}")
return web.json_response({"status": "success"})
else:
return await handle_api_error(request, "Missing image data or invalid request", 400)
return await handle_api_error(request, "Job not found", 404)
except Exception as e:
return await handle_api_error(request, e, 500)
# Keep legacy endpoint for backward compatibility
@server.PromptServer.instance.routes.post("/distributed/tile_complete")
async def tile_complete_endpoint(request):
try:
content_length = request.headers.get('content-length')
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
data = await request.post()
multi_job_id = data.get('multi_job_id')
worker_id = data.get('worker_id')
is_last = data.get('is_last', 'False').lower() == 'true'
if multi_job_id is None or worker_id is None:
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
prompt_server = ensure_tile_jobs_initialized()
if 'full_image' in data and 'image_idx' in data:
image_idx = int(data.get('image_idx'))
img_data = data['full_image'].file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
debug_log(f"Received full image {image_idx} from worker {worker_id}")
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
return await handle_api_error(request, "Mode mismatch for image submission", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
'image_idx': image_idx,
'image': img,
'is_last': is_last
})
return web.json_response({"status": "success"})
batch_size = int(data.get('batch_size', 0))
tiles = []
if batch_size == 0 and is_last:
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
'is_last': True,
'tiles': []
})
debug_log(f"Received completion signal from worker {worker_id}")
return web.json_response({"status": "success"})
if batch_size > 0:
padding = int(data.get('padding', 32))
metadata_field = data.get('tiles_metadata')
if metadata_field:
if hasattr(metadata_field, 'file'):
metadata_str = metadata_field.file.read().decode('utf-8')
elif isinstance(metadata_field, (bytes, bytearray)):
metadata_str = metadata_field.decode('utf-8')
else:
metadata_str = str(metadata_field)
metadata = json.loads(metadata_str)
if len(metadata) != batch_size:
return await handle_api_error(request, "Metadata length mismatch", 400)
tile_data_list = []
for i in range(batch_size):
tile_field = data.get(f'tile_{i}')
if tile_field is None:
continue
img_data = tile_field.file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np)[None,]
if i < len(metadata):
tile_meta = metadata[i]
tile_idx = tile_meta.get('tile_idx', i)
tile_info = {
'tensor': tensor,
'tile_idx': tile_idx,
'x': tile_meta['x'],
'y': tile_meta['y'],
'extracted_width': tile_meta['extracted_width'],
'extracted_height': tile_meta['extracted_height'],
'padding': padding,
'batch_idx': tile_meta.get('batch_idx', 0),
'global_idx': tile_meta.get('global_idx', tile_idx)
}
tile_data_list.append(tile_info)
tile_data_list.sort(key=lambda x: x['tile_idx'])
tiles.extend(tile_data_list)
else:
# Legacy
for i in range(batch_size):
tile_field = data.get(f'tile_{i}')
if tile_field is None:
continue
tile_idx = int(data.get(f'tile_{i}_idx', i))
x = int(data.get(f'tile_{i}_x', 0))
y = int(data.get(f'tile_{i}_y', 0))
extracted_width = int(data.get(f'tile_{i}_width', 512))
extracted_height = int(data.get(f'tile_{i}_height', 512))
img_data = tile_field.file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np)[None,]
tiles.append({
'tensor': tensor,
'tile_idx': tile_idx,
'x': x,
'y': y,
'extracted_width': extracted_width,
'extracted_height': extracted_height,
'padding': padding
})
else:
# Single tile legacy
image_file = data.get('image')
if not image_file:
return await handle_api_error(request, "Missing image data", 400)
tile_idx = int(data.get('tile_idx', 0))
x = int(data.get('x', 0))
y = int(data.get('y', 0))
extracted_width = int(data.get('extracted_width', 512))
extracted_height = int(data.get('extracted_height', 512))
padding = int(data.get('padding', 32))
img_data = image_file.file.read()
img = Image.open(io.BytesIO(img_data)).convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
tensor = torch.from_numpy(img_np)[None,]
tiles = [{
'tensor': tensor,
'tile_idx': tile_idx,
'x': x,
'y': y,
'extracted_width': extracted_width,
'extracted_height': extracted_height,
'padding': padding,
'batch_idx': 0,
'global_idx': tile_idx
}]
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
return await handle_api_error(request, "Mode mismatch for tile submission", 400)
q = job_data[JOB_QUEUE]
if batch_size > 0 or len(tiles) > 0:
await q.put({
'worker_id': worker_id,
'tiles': tiles,
'is_last': is_last
})
else:
await q.put({
'worker_id': worker_id,
'is_last': True,
'tiles': []
})
return web.json_response({"status": "success"})
else:
return await handle_api_error(request, "Job not found", 404)
except Exception as e:
return await handle_api_error(request, e, 500)
# Helper functions for shallow copying conditioning without duplicating models
def clone_control_chain(control, clone_hint=True):
"""Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
if control is None:
return None
new_control = copy.copy(control) # Shallow copy (shares model)
if clone_hint and hasattr(control, 'cond_hint_original'):
new_control.cond_hint_original = control.cond_hint_original.clone()
if hasattr(control, 'previous_controlnet'):
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
return new_control
def clone_conditioning(cond_list, clone_hints=True):
"""Clone conditioning without duplicating ControlNet models."""
new_cond = []
for emb, cond_dict in cond_list:
new_emb = emb.clone() if emb is not None else None
new_dict = cond_dict.copy()
if 'control' in new_dict:
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
if 'mask' in new_dict:
new_dict['mask'] = new_dict['mask'].clone()
# Handle other potential fields if needed
if 'pooled_output' in new_dict:
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
if 'area' in new_dict:
new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple
new_cond.append([new_emb, new_dict])
return new_cond
# Helper function to ensure persistent state is initialized
def ensure_tile_jobs_initialized():
"""Ensure tile job storage is initialized on the server instance."""
prompt_server = server.PromptServer.instance
if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
debug_log("Initializing persistent tile job queue on server instance.")
prompt_server.distributed_pending_tile_jobs = {}
prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
else:
# Clean up any legacy queue structures that don't have the 'mode' field
# (Should be rare after fixes, but keep for safety)
to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
if not isinstance(job_data, dict) or 'mode' not in job_data]
for job_id in to_remove:
debug_log(f"Removing legacy queue structure for job {job_id}")
del prompt_server.distributed_pending_tile_jobs[job_id]
return prompt_server
# API Endpoint for tile completion
@server.PromptServer.instance.routes.post("/distributed/request_image")
async def request_image_endpoint(request):
"""Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
try:
data = await request.json()
worker_id = data.get('worker_id')
multi_job_id = data.get('multi_job_id')
if not worker_id or not multi_job_id:
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if not isinstance(job_data, dict) or 'mode' not in job_data:
return await handle_api_error(request, "Invalid job data structure", 500)
mode = job_data['mode']
# Handle both dynamic and static modes
if mode == 'dynamic' and 'pending_images' in job_data:
pending_queue = job_data['pending_images']
elif mode == 'static' and JOB_PENDING_TASKS in job_data:
pending_queue = job_data[JOB_PENDING_TASKS]
else:
return await handle_api_error(request, f"Invalid {mode} mode configuration", 400)
try:
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
# Track assigned task
if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
job_data['assigned_to_workers'][worker_id].append(task_idx)
# Update worker heartbeat
if 'worker_status' in job_data:
job_data['worker_status'][worker_id] = time.time()
# Get estimated remaining count
remaining = pending_queue.qsize() # Approximate
# Return appropriate response based on mode
if mode == 'dynamic':
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
else: # static
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining})
except asyncio.TimeoutError:
if mode == 'dynamic':
return web.json_response({"image_idx": None})
else:
return web.json_response({"tile_idx": None})
else:
return await handle_api_error(request, "Job not found", 404)
except Exception as e:
return await handle_api_error(request, e, 500)
@server.PromptServer.instance.routes.get("/distributed/job_status")
async def job_status_endpoint(request):
"""Endpoint to check if a job is ready."""
multi_job_id = request.query.get('multi_job_id')
if not multi_job_id:
return web.json_response({"ready": False})
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data)
return web.json_response({"ready": ready})
+456
View File
@@ -0,0 +1,456 @@
import numpy as np
from PIL import Image, ImageFilter
import torch
import torch.nn.functional as F
from torchvision.transforms import GaussianBlur
import math
import os
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
Image.Resampling = Image
BLUR_KERNEL_SIZE = 15
def tensor_to_pil(img_tensor, batch_index=0):
# Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels]
# and returns an RGB PIL Image. Assumes channels=3
return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
def pil_to_tensor(image):
# Takes a PIL image and returns a tensor of shape [1, height, width, channels]
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image).unsqueeze(0)
if len(image.shape) == 3: # If the image is grayscale, add a channel dimension
image = image.unsqueeze(-1)
return image
def controlnet_hint_to_pil(tensor, batch_index=0):
return tensor_to_pil(tensor.movedim(1, -1), batch_index)
def pil_to_controlnet_hint(img):
return pil_to_tensor(img).movedim(-1, 1)
def crop_tensor(tensor, region):
# Takes a tensor of shape [batch_size, height, width, channels] and crops it to the given region
x1, y1, x2, y2 = region
return tensor[:, y1:y2, x1:x2, :]
def resize_tensor(tensor, size, mode="nearest-exact"):
# Takes a tensor of shape [B, C, H, W] and resizes
# it to a shape of [B, C, size[0], size[1]] using the given mode
return torch.nn.functional.interpolate(tensor, size=size, mode=mode)
def get_crop_region(mask, pad=0):
# Takes a black and white PIL image in 'L' mode and returns the coordinates of the white rectangular mask region
# Should be equivalent to the get_crop_region function from https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/modules/masking.py
coordinates = mask.getbbox()
if coordinates is not None:
x1, y1, x2, y2 = coordinates
else:
x1, y1, x2, y2 = mask.width, mask.height, 0, 0
# Apply padding
x1 = max(x1 - pad, 0)
y1 = max(y1 - pad, 0)
x2 = min(x2 + pad, mask.width)
y2 = min(y2 + pad, mask.height)
return fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height))
def fix_crop_region(region, image_size):
# Remove the extra pixel added by the get_crop_region function
image_width, image_height = image_size
x1, y1, x2, y2 = region
if x2 < image_width:
x2 -= 1
if y2 < image_height:
y2 -= 1
return x1, y1, x2, y2
def expand_crop(region, width, height, target_width, target_height):
'''
Expands a crop region to a specified target size.
:param region: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points
of the rectangular region. Expected to have x2 > x1 and y2 > y1.
:param width: The width of the image the crop region is from.
:param height: The height of the image the crop region is from.
:param target_width: The desired width of the crop region.
:param target_height: The desired height of the crop region.
'''
x1, y1, x2, y2 = region
actual_width = x2 - x1
actual_height = y2 - y1
# target_width = math.ceil(actual_width / 8) * 8
# target_height = math.ceil(actual_height / 8) * 8
# Try to expand region to the right of half the difference
width_diff = target_width - actual_width
x2 = min(x2 + width_diff // 2, width)
# Expand region to the left of the difference including the pixels that could not be expanded to the right
width_diff = target_width - (x2 - x1)
x1 = max(x1 - width_diff, 0)
# Try the right again
width_diff = target_width - (x2 - x1)
x2 = min(x2 + width_diff, width)
# Try to expand region to the bottom of half the difference
height_diff = target_height - actual_height
y2 = min(y2 + height_diff // 2, height)
# Expand region to the top of the difference including the pixels that could not be expanded to the bottom
height_diff = target_height - (y2 - y1)
y1 = max(y1 - height_diff, 0)
# Try the bottom again
height_diff = target_height - (y2 - y1)
y2 = min(y2 + height_diff, height)
return (x1, y1, x2, y2), (target_width, target_height)
def resize_region(region, init_size, resize_size):
# Resize a crop so that it fits an image that was resized to the given width and height
x1, y1, x2, y2 = region
init_width, init_height = init_size
resize_width, resize_height = resize_size
x1 = math.floor(x1 * resize_width / init_width)
x2 = math.ceil(x2 * resize_width / init_width)
y1 = math.floor(y1 * resize_height / init_height)
y2 = math.ceil(y2 * resize_height / init_height)
return (x1, y1, x2, y2)
def pad_image(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
'''
Pads an image with the given number of pixels on each side and fills the padding with data from the edges.
:param image: A PIL image
:param left_pad: The number of pixels to pad on the left side
:param right_pad: The number of pixels to pad on the right side
:param top_pad: The number of pixels to pad on the top side
:param bottom_pad: The number of pixels to pad on the bottom side
:param blur: Whether to blur the padded edges
:return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad)
'''
left_edge = image.crop((0, 1, 1, image.height - 1))
right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1))
top_edge = image.crop((1, 0, image.width - 1, 1))
bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height))
new_width = image.width + left_pad + right_pad
new_height = image.height + top_pad + bottom_pad
padded_image = Image.new(image.mode, (new_width, new_height))
padded_image.paste(image, (left_pad, top_pad))
if fill:
for i in range(left_pad):
edge = left_edge.resize(
(1, new_height - i * (top_pad + bottom_pad) // left_pad), resample=Image.Resampling.NEAREST)
padded_image.paste(edge, (i, i * top_pad // left_pad))
for i in range(right_pad):
edge = right_edge.resize(
(1, new_height - i * (top_pad + bottom_pad) // right_pad), resample=Image.Resampling.NEAREST)
padded_image.paste(edge, (new_width - 1 - i, i * top_pad // right_pad))
for i in range(top_pad):
edge = top_edge.resize(
(new_width - i * (left_pad + right_pad) // top_pad, 1), resample=Image.Resampling.NEAREST)
padded_image.paste(edge, (i * left_pad // top_pad, i))
for i in range(bottom_pad):
edge = bottom_edge.resize(
(new_width - i * (left_pad + right_pad) // bottom_pad, 1), resample=Image.Resampling.NEAREST)
padded_image.paste(edge, (i * left_pad // bottom_pad, new_height - 1 - i))
if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0):
padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE))
padded_image.paste(image, (left_pad, top_pad))
return padded_image
def pad_image2(image, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
'''
Pads an image with the given number of pixels on each side and fills the padding with data from the edges.
Faster than pad_image, but only pads with edge data in straight lines.
:param image: A PIL image
:param left_pad: The number of pixels to pad on the left side
:param right_pad: The number of pixels to pad on the right side
:param top_pad: The number of pixels to pad on the top side
:param bottom_pad: The number of pixels to pad on the bottom side
:param blur: Whether to blur the padded edges
:return: A PIL image with size (image.width + left_pad + right_pad, image.height + top_pad + bottom_pad)
'''
left_edge = image.crop((0, 1, 1, image.height - 1))
right_edge = image.crop((image.width - 1, 1, image.width, image.height - 1))
top_edge = image.crop((1, 0, image.width - 1, 1))
bottom_edge = image.crop((1, image.height - 1, image.width - 1, image.height))
new_width = image.width + left_pad + right_pad
new_height = image.height + top_pad + bottom_pad
padded_image = Image.new(image.mode, (new_width, new_height))
padded_image.paste(image, (left_pad, top_pad))
if fill:
if left_pad > 0:
padded_image.paste(left_edge.resize((left_pad, new_height), resample=Image.Resampling.NEAREST), (0, 0))
if right_pad > 0:
padded_image.paste(right_edge.resize((right_pad, new_height),
resample=Image.Resampling.NEAREST), (new_width - right_pad, 0))
if top_pad > 0:
padded_image.paste(top_edge.resize((new_width, top_pad), resample=Image.Resampling.NEAREST), (0, 0))
if bottom_pad > 0:
padded_image.paste(bottom_edge.resize((new_width, bottom_pad),
resample=Image.Resampling.NEAREST), (0, new_height - bottom_pad))
if blur and not (left_pad == right_pad == top_pad == bottom_pad == 0):
padded_image = padded_image.filter(ImageFilter.GaussianBlur(BLUR_KERNEL_SIZE))
padded_image.paste(image, (left_pad, top_pad))
return padded_image
def pad_tensor(tensor, left_pad, right_pad, top_pad, bottom_pad, fill=False, blur=False):
'''
Pads an image tensor with the given number of pixels on each side and fills the padding with data from the edges.
:param tensor: A tensor of shape [B, H, W, C]
:param left_pad: The number of pixels to pad on the left side
:param right_pad: The number of pixels to pad on the right side
:param top_pad: The number of pixels to pad on the top side
:param bottom_pad: The number of pixels to pad on the bottom side
:param blur: Whether to blur the padded edges
:return: A tensor of shape [B, H + top_pad + bottom_pad, W + left_pad + right_pad, C]
'''
batch_size, channels, height, width = tensor.shape
h_pad = left_pad + right_pad
v_pad = top_pad + bottom_pad
new_width = width + h_pad
new_height = height + v_pad
# Create empty image
padded = torch.zeros((batch_size, channels, new_height, new_width), dtype=tensor.dtype)
# Copy the original image into the centor of the padded tensor
padded[:, :, top_pad:top_pad + height, left_pad:left_pad + width] = tensor
# Duplicate the edges of the original image into the padding
if top_pad > 0:
padded[:, :, :top_pad, :] = padded[:, :, top_pad:top_pad + 1, :] # Top edge
if bottom_pad > 0:
padded[:, :, -bottom_pad:, :] = padded[:, :, -bottom_pad - 1:-bottom_pad, :] # Bottom edge
if left_pad > 0:
padded[:, :, :, :left_pad] = padded[:, :, :, left_pad:left_pad + 1] # Left edge
if right_pad > 0:
padded[:, :, :, -right_pad:] = padded[:, :, :, -right_pad - 1:-right_pad] # Right edge
return padded
def resize_and_pad_image(image, width, height, fill=False, blur=False):
'''
Resizes an image to the given width and height and pads it to the given width and height.
:param image: A PIL image
:param width: The width of the resized image
:param height: The height of the resized image
:param fill: Whether to fill the padding with data from the edges
:param blur: Whether to blur the padded edges
:return: A PIL image of size (width, height)
'''
width_ratio = width / image.width
height_ratio = height / image.height
if height_ratio > width_ratio:
resize_ratio = width_ratio
else:
resize_ratio = height_ratio
resize_width = round(image.width * resize_ratio)
resize_height = round(image.height * resize_ratio)
resized = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS)
# Pad the sides of the image to get the image to the desired size that wasn't covered by the resize
horizontal_pad = (width - resize_width) // 2
vertical_pad = (height - resize_height) // 2
result = pad_image2(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur)
result = result.resize((width, height), resample=Image.Resampling.LANCZOS)
return result, (horizontal_pad, vertical_pad)
def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False):
'''
Resizes an image tensor to the given width and height and pads it to the given width and height.
:param tensor: A tensor of shape [B, H, W, C]
:param width: The width of the resized image
:param height: The height of the resized image
:param fill: Whether to fill the padding with data from the edges
:param blur: Whether to blur the padded edges
:return: A tensor of shape [B, height, width, C]
'''
# Resize the image to the closest size that maintains the aspect ratio
width_ratio = width / tensor.shape[3]
height_ratio = height / tensor.shape[2]
if height_ratio > width_ratio:
resize_ratio = width_ratio
else:
resize_ratio = height_ratio
resize_width = round(tensor.shape[3] * resize_ratio)
resize_height = round(tensor.shape[2] * resize_ratio)
resized = F.interpolate(tensor, size=(resize_height, resize_width), mode='nearest-exact')
# Pad the sides of the image to get the image to the desired size that wasn't covered by the resize
horizontal_pad = (width - resize_width) // 2
vertical_pad = (height - resize_height) // 2
result = pad_tensor(resized, horizontal_pad, horizontal_pad, vertical_pad, vertical_pad, fill, blur)
result = F.interpolate(result, size=(height, width), mode='nearest-exact')
return result
def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
if "control" not in cond_dict:
return
c = cond_dict["control"]
controlnet = c.copy()
cond_dict["control"] = controlnet
while c is not None:
# hint is shape (B, C, H, W)
hint = controlnet.cond_hint_original
resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1])
hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1)
hint = resize_tensor(hint, tile_size[::-1])
controlnet.cond_hint_original = hint
c = c.previous_controlnet
controlnet.set_previous_controlnet(c.copy() if c is not None else None)
controlnet = controlnet.previous_controlnet
def region_intersection(region1, region2):
"""
Returns the coordinates of the intersection of two rectangular regions.
:param region1: A tuple of the form (x1, y1, x2, y2) denoting the upper left and the lower right points
of the first rectangular region. Expected to have x2 > x1 and y2 > y1.
:param region2: The second rectangular region with the same format as the first.
:return: A tuple of the form (x1, y1, x2, y2) denoting the rectangular intersection.
None if there is no intersection.
"""
x1, y1, x2, y2 = region1
x1_, y1_, x2_, y2_ = region2
x1 = max(x1, x1_)
y1 = max(y1, y1_)
x2 = min(x2, x2_)
y2 = min(y2, y2_)
if x1 >= x2 or y1 >= y2:
return None
return (x1, y1, x2, y2)
def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
if "gligen" not in cond_dict:
return
type, model, cond = cond_dict["gligen"]
if type != "position":
from warnings import warn
warn(f"Unknown gligen type {type}")
return
cropped = []
for c in cond:
emb, h, w, y, x = c
# Get the coordinates of the box in the upscaled image
x1 = x * 8
y1 = y * 8
x2 = x1 + w * 8
y2 = y1 + h * 8
gligen_upscaled_box = resize_region((x1, y1, x2, y2), init_size, canvas_size)
# Calculate the intersection of the gligen box and the region
intersection = region_intersection(gligen_upscaled_box, region)
if intersection is None:
continue
x1, y1, x2, y2 = intersection
# Offset the gligen box so that the origin is at the top left of the tile region
x1 -= region[0]
y1 -= region[1]
x2 -= region[0]
y2 -= region[1]
# Add the padding
x1 += w_pad
y1 += h_pad
x2 += w_pad
y2 += h_pad
# Set the new position params
h = (y2 - y1) // 8
w = (x2 - x1) // 8
x = x1 // 8
y = y1 // 8
cropped.append((emb, h, w, y, x))
cond_dict["gligen"] = (type, model, cropped)
def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
if "area" not in cond_dict:
return
# Resize the area conditioning to the canvas size and confine it to the tile region
h, w, y, x = cond_dict["area"]
w, h, x, y = 8 * w, 8 * h, 8 * x, 8 * y
x1, y1, x2, y2 = resize_region((x, y, x + w, y + h), init_size, canvas_size)
intersection = region_intersection((x1, y1, x2, y2), region)
if intersection is None:
del cond_dict["area"]
del cond_dict["strength"]
return
x1, y1, x2, y2 = intersection
# Offset origin to the top left of the tile
x1 -= region[0]
y1 -= region[1]
x2 -= region[0]
y2 -= region[1]
# Add the padding
x1 += w_pad
y1 += h_pad
x2 += w_pad
y2 += h_pad
# Set the params for tile
w, h = (x2 - x1) // 8, (y2 - y1) // 8
x, y = x1 // 8, y1 // 8
cond_dict["area"] = (h, w, y, x)
def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
if "mask" not in cond_dict:
return
mask_tensor = cond_dict["mask"] # (B, H, W)
masks = []
for i in range(mask_tensor.shape[0]):
# Convert to PIL image
mask = tensor_to_pil(mask_tensor, i) # W x H
# Resize the mask to the canvas size
mask = mask.resize(canvas_size, Image.Resampling.BICUBIC)
# Crop the mask to the region
mask = mask.crop(region)
# Add padding
mask, _ = resize_and_pad_image(mask, tile_size[0], tile_size[1], fill=True)
# Resize the mask to the tile size
if tile_size != mask.size:
mask = mask.resize(tile_size, Image.Resampling.BICUBIC)
# Convert back to tensor
mask = pil_to_tensor(mask) # (1, H, W, 1)
mask = mask.squeeze(-1) # (1, H, W)
masks.append(mask)
cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W)
def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0):
cropped = []
for emb, x in cond:
cond_dict = x.copy()
n = [emb, cond_dict]
crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
cropped.append(n)
return cropped
+146
View File
@@ -0,0 +1,146 @@
import { TIMEOUTS } from './constants.js';
export function createApiClient(baseUrl) {
const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => {
let lastError;
let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff
for (let attempt = 0; attempt < retries; attempt++) {
try {
const response = await fetch(`${baseUrl}${endpoint}`, {
headers: { 'Content-Type': 'application/json' },
...options
});
if (!response.ok) {
const error = await response.json().catch(() => ({ message: 'Request failed' }));
throw new Error(error.message || `HTTP ${response.status}`);
}
return await response.json();
} catch (error) {
lastError = error;
console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`);
if (attempt < retries - 1) {
await new Promise(resolve => setTimeout(resolve, delay));
delay *= 2; // Exponential backoff
}
}
}
throw lastError;
};
return {
// Config endpoints
async getConfig() {
return request('/distributed/config');
},
async updateWorker(workerId, data) {
return request('/distributed/config/update_worker', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId, ...data })
});
},
async deleteWorker(workerId) {
return request('/distributed/config/delete_worker', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId })
});
},
async updateSetting(key, value) {
return request('/distributed/config/update_setting', {
method: 'POST',
body: JSON.stringify({ key, value })
});
},
async updateMaster(data) {
return request('/distributed/config/update_master', {
method: 'POST',
body: JSON.stringify(data)
});
},
// Worker management endpoints
async launchWorker(workerId) {
return request('/distributed/launch_worker', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId })
});
},
async stopWorker(workerId) {
return request('/distributed/stop_worker', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId })
});
},
async getManagedWorkers() {
return request('/distributed/managed_workers');
},
async getWorkerLog(workerId, lines = 1000) {
return request(`/distributed/worker_log/${workerId}?lines=${lines}`);
},
async clearLaunchingFlag(workerId) {
return request('/distributed/worker/clear_launching', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId })
});
},
// Job preparation
async prepareJob(multiJobId) {
return request('/distributed/prepare_job', {
method: 'POST',
body: JSON.stringify({ multi_job_id: multiJobId })
});
},
// Image loading
async loadImage(imagePath) {
return request('/distributed/load_image', {
method: 'POST',
body: JSON.stringify({ image_path: imagePath })
});
},
// Network info
async getNetworkInfo() {
return request('/distributed/network_info');
},
// Status checking (with timeout)
async checkStatus(url, timeout = TIMEOUTS.DEFAULT_FETCH) {
const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), timeout);
try {
const response = await fetch(url, {
method: 'GET',
mode: 'cors',
signal: controller.signal
});
clearTimeout(timeoutId);
if (!response.ok) throw new Error(`HTTP ${response.status}`);
return await response.json();
} catch (error) {
clearTimeout(timeoutId);
throw error;
}
},
// Batch status checking
async checkMultipleStatuses(urls) {
return Promise.allSettled(
urls.map(url => this.checkStatus(url))
);
}
};
}
+153
View File
@@ -0,0 +1,153 @@
export const BUTTON_STYLES = {
// Base styles with unified padding
base: "width: 100%; padding: 4px 14px; color: white; border: none; border-radius: 4px; cursor: pointer; transition: all 0.2s; font-size: 12px; font-weight: 500;",
// Context-specific combined styles
workerControl: "flex: 1; font-size: 11px;",
// Layout modifiers
hidden: "display: none;",
marginLeftAuto: "margin-left: auto;",
// Color variants
cancel: "background-color: #555;",
info: "background-color: #333;",
success: "background-color: #4a7c4a;",
error: "background-color: #7c4a4a;",
launch: "background-color: #4a7c4a;",
stop: "background-color: #7c4a4a;",
log: "background-color: #685434;",
clearMemory: "background-color: #555; padding: 6px 14px;",
interrupt: "background-color: #555; padding: 6px 14px;",
};
export const STATUS_COLORS = {
DISABLED_GRAY: "#666",
OFFLINE_RED: "#c04c4c",
ONLINE_GREEN: "#3ca03c",
PROCESSING_YELLOW: "#f0ad4e"
};
export const UI_COLORS = {
MUTED_TEXT: "#888",
SECONDARY_TEXT: "#ccc",
BORDER_LIGHT: "#555",
BORDER_DARK: "#444",
BORDER_DARKER: "#3a3a3a",
BACKGROUND_DARK: "#2a2a2a",
BACKGROUND_DARKER: "#1e1e1e",
ICON_COLOR: "#666",
ACCENT_COLOR: "#777"
};
export const PULSE_ANIMATION_CSS = `
@keyframes pulse {
0% {
opacity: 1;
transform: scale(0.8);
box-shadow: 0 0 0 0 rgba(240, 173, 78, 0.7);
}
50% {
opacity: 0.3;
transform: scale(1.1);
box-shadow: 0 0 0 6px rgba(240, 173, 78, 0);
}
100% {
opacity: 1;
transform: scale(0.8);
box-shadow: 0 0 0 0 rgba(240, 173, 78, 0);
}
}
.status-pulsing {
animation: pulse 1.2s ease-in-out infinite;
transform-origin: center;
}
/* Button hover effects */
.distributed-button:hover:not(:disabled) {
filter: brightness(1.2);
transition: filter 0.2s ease;
}
.distributed-button:disabled {
opacity: 0.6;
cursor: not-allowed;
}
/* Settings button animation */
.settings-btn {
transition: transform 0.2s ease;
}
/* Expanded settings panel */
.worker-settings {
max-height: 0;
overflow: hidden;
opacity: 0;
transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease;
}
.worker-settings.expanded {
max-height: 500px;
opacity: 1;
padding: 12px 0;
}
`;
export const UI_STYLES = {
statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;",
formGroup: "display: flex; flex-direction: column; gap: 5px;",
formLabel: "font-size: 12px; color: #ccc; font-weight: 500;",
formInput: "padding: 6px 10px; background: #2a2a2a; border: 1px solid #444; color: white; font-size: 12px; border-radius: 4px; transition: border-color 0.2s;",
// Card styles
cardBase: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;",
workerCard: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex; background: #2a2a2a;",
cardBlueprint: "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);",
cardAdd: "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;",
// Column styles
columnBase: "display: flex; align-items: center; justify-content: center;",
checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);",
contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;",
// Row and content styles
infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;",
workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;",
// Form and controls styles
buttonGroup: "display: flex; gap: 4px; margin-top: 10px;",
settingsForm: "display: flex; flex-direction: column; gap: 10px;",
checkboxGroup: "display: flex; align-items: center; gap: 8px; margin: 5px 0;",
formLabelClickable: "font-size: 12px; color: #ccc; cursor: pointer;",
settingsToggle: "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;",
controlsWrapper: "display: flex; gap: 6px; align-items: stretch; width: 100%;",
// Existing styles
settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto; padding: 4px;",
infoBox: "background-color: #333; color: #999; padding: 5px 14px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;",
workerSettings: "margin: 0 12px; padding: 0 12px; background: #1e1e1e; border-radius: 4px; border: 1px solid #2a2a2a;"
};
export const TIMEOUTS = {
DEFAULT_FETCH: 5000, // ms for general API calls
STATUS_CHECK: 1200, // ms for status checks
LAUNCH: 90000, // ms for worker launch (longer for model loading)
RETRY_DELAY: 1000, // initial delay for exponential backoff
MAX_RETRIES: 3, // max retry attempts
// UI feedback delays
BUTTON_RESET: 3000, // button text/state reset after actions
FLASH_SHORT: 1000, // brief success feedback
FLASH_MEDIUM: 1500, // medium error feedback
FLASH_LONG: 2000, // longer error feedback
// Operational delays
POST_ACTION_DELAY: 500, // delay after operations before status checks
STATUS_CHECK_DELAY: 100, // brief delay before status checks
// Background tasks
LOG_REFRESH: 2000, // log auto-refresh interval
IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache
};
+590
View File
@@ -0,0 +1,590 @@
import { api } from "../../scripts/api.js";
import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo } from './workerUtils.js';
import { TIMEOUTS } from './constants.js';
/**
* Convert paths in the API prompt to match the target platform's separator
* @param {Object} apiPrompt - The workflow API prompt
* @param {string} targetSeparator - The target path separator ('\\' or '/')
* @returns {Object} The converted API prompt
*/
function convertPathsForPlatform(apiPrompt, targetSeparator) {
// Validate target separator
if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) {
console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion');
return apiPrompt;
}
// Regex to identify likely file paths with extensions
const isLikelyFilename = (value) => {
return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i);
};
function convert(obj) {
if (typeof obj === 'string') {
// Only convert strings that look like file paths
if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) {
// Replace any path separator with the target one
return obj.replace(/[\\\/]/g, targetSeparator);
}
return obj;
} else if (Array.isArray(obj)) {
return obj.map(convert);
} else if (typeof obj === 'object' && obj !== null) {
const newObj = {};
for (const [key, value] of Object.entries(obj)) {
newObj[key] = convert(value);
}
return newObj;
}
return obj;
}
return convert(apiPrompt);
}
export function setupInterceptor(extension) {
api.queuePrompt = async (number, prompt) => {
if (extension.isEnabled) {
const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
if (hasCollector || hasDistUpscale) {
const result = await executeParallelDistributed(extension, prompt);
// Immediate status check for instant feedback
extension.checkAllWorkerStatuses();
// Another check after a short delay to catch state changes
setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
return result;
}
}
return extension.originalQueuePrompt(number, prompt);
};
}
export async function executeParallelDistributed(extension, promptWrapper) {
try {
const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution
const enabledWorkers = extension.enabledWorkers;
// Pre-flight health check on all enabled workers
const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
// Case: Enabled workers but all offline
if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
extension.log("No active workers found. All enabled workers are offline.");
if (extension.ui?.showToast) {
extension.ui.showToast(extension.app, "error", "All Workers Offline",
`${enabledWorkers.length} worker(s) enabled but all are offline or unreachable. Check worker connections and try again.`, 5000);
}
// Fall back to master-only execution
return extension.originalQueuePrompt(0, promptWrapper);
}
extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug");
// Check if master host might be unreachable by workers (cloudflare tunnel down)
const masterHost = extension.config?.master?.host || '';
const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost);
if (isCloudflareHost && activeWorkers.length > 0) {
// Try to verify if the cloudflare tunnel is actually up
try {
const testUrl = `${window.location.protocol}//${masterHost}/prompt`;
const response = await fetch(testUrl, {
method: 'GET',
mode: 'cors',
cache: 'no-cache',
signal: AbortSignal.timeout(3000) // 3 second timeout
});
if (!response.ok) {
throw new Error('Master not reachable');
}
} catch (error) {
// Cloudflare tunnel appears to be down
extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error");
if (extension.ui?.showCloudflareWarning) {
extension.ui.showCloudflareWarning(extension, masterHost);
}
// Stop execution - workers won't be able to send results back
extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error");
return null; // This will prevent the workflow from running
}
}
// Find all distributed nodes in the workflow
const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector");
const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed");
const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
// Map original node IDs to truly unique job IDs for this specific run
const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`]));
// Prepare a separate job queue on the backend for each unique job ID
const preparePromises = Array.from(job_id_map.values()).map(uniqueId => prepareDistributedJob(extension, uniqueId));
await Promise.all(preparePromises);
const jobs = [];
// Use only active workers
const participants = ['master', ...activeWorkers.map(w => w.id)];
for (const participantId of participants) {
const options = {
enabled_worker_ids: activeWorkers.map(w => w.id),
workflow: promptWrapper.workflow,
job_id_map: job_id_map // Pass the map of unique IDs
};
const jobApiPrompt = await prepareApiPromptForParticipant(
extension, promptWrapper.output, participantId, options
);
if (participantId === 'master') {
jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } });
} else {
const worker = activeWorkers.find(w => w.id === participantId);
if (worker) {
const job = {
type: 'worker',
worker,
prompt: jobApiPrompt,
workflow: promptWrapper.workflow
};
// Add image references if found for remote workers
if (options._imageReferences) {
job.imageReferences = options._imageReferences;
}
jobs.push(job);
}
}
}
const result = await executeJobs(extension, jobs);
return result;
} catch (error) {
extension.log("Parallel execution failed: " + error.message, "error");
throw error;
}
}
export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
const isMaster = participantId === 'master';
// Find all distributed nodes once (before pruning)
const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed");
const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
// For workers, handle platform-specific path conversion
if (!isMaster) {
const workerInfo = extension.config.workers.find(w => w.id === participantId);
if (workerInfo && workerInfo.host) {
// Remote or cloud worker - needs path translation
try {
const workerUrl = extension.getWorkerUrl(workerInfo);
const systemInfo = await getCachedWorkerSystemInfo(workerUrl);
const targetSeparator = systemInfo?.platform?.path_separator;
if (targetSeparator) {
// Convert paths to match worker's platform
jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator);
extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug");
} else {
extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug");
}
} catch (e) {
extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn");
// Continue without path conversion
}
}
// Prune the workflow to only include distributed node dependencies
if (allDistributedNodes.length > 0) {
jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes);
}
}
// Handle image references for remote workers
if (!isMaster && options.enabled_worker_ids) {
// Check if this is a remote worker
const workerId = participantId;
const workerInfo = extension.config.workers.find(w => w.id === workerId);
const isRemote = workerInfo && workerInfo.host;
if (isRemote) {
// Find all image/video references in the pruned workflow
const imageReferences = findImageReferences(extension, jobApiPrompt);
if (imageReferences.size > 0) {
extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug");
// Store image references for later processing
options._imageReferences = imageReferences;
}
}
}
// Handle Distributed seed nodes
const distributorNodes = findNodesByClass(jobApiPrompt, "DistributedSeed");
if (distributorNodes.length > 0) {
extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug");
}
for (const seedNode of distributorNodes) {
const { inputs } = jobApiPrompt[seedNode.id];
inputs.is_worker = !isMaster;
if (!isMaster) {
const workerIndex = options.enabled_worker_ids.indexOf(participantId);
inputs.worker_id = `worker_${workerIndex}`;
extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug");
}
}
// Handle Distributed collector nodes (already found above)
for (const collector of collectorNodes) {
const { inputs } = jobApiPrompt[collector.id];
// Check if this collector is downstream from a distributed upscaler
const hasUpstreamDistributedUpscaler = hasUpstreamNode(
jobApiPrompt,
collector.id,
'UltimateSDUpscaleDistributed'
);
if (hasUpstreamDistributedUpscaler) {
// Set pass_through mode for this collector
inputs.pass_through = true;
extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug");
} else {
// Normal collector behavior
// Get the unique job ID from the map created for this execution
const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id;
// Use the truly unique ID for this execution
inputs.multi_job_id = uniqueJobId;
inputs.is_worker = !isMaster;
if (isMaster) {
inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
} else {
inputs.master_url = extension.getMasterUrl();
// Also make the worker_job_id unique to prevent potential caching issues
inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`;
inputs.worker_id = participantId;
}
}
}
// Handle Ultimate SD Upscale Distributed nodes
for (const upscaleNode of upscaleNodes) {
const { inputs } = jobApiPrompt[upscaleNode.id];
// Get the unique job ID from the map
const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id;
inputs.multi_job_id = uniqueJobId;
inputs.is_worker = !isMaster;
if (isMaster) {
inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
} else {
inputs.master_url = extension.getMasterUrl();
inputs.worker_id = participantId;
// Workers also need the enabled_worker_ids to calculate tile distribution
inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
}
}
return jobApiPrompt;
}
export async function prepareDistributedJob(extension, multi_job_id) {
try {
await extension.api.prepareJob(multi_job_id);
} catch (error) {
extension.log("Error preparing job: " + error.message, "error");
throw error;
}
}
export async function executeJobs(extension, jobs) {
let masterPromptId = null;
// Pre-load all unique images before dispatching to workers
const allImageReferences = new Map();
for (const job of jobs) {
if (job.type === 'worker' && job.imageReferences) {
for (const [filename, info] of job.imageReferences) {
allImageReferences.set(filename, info);
}
}
}
if (allImageReferences.size > 0) {
extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug");
await loadImagesForWorker(extension, allImageReferences);
}
// Now dispatch jobs in parallel
const promises = jobs.map(job => {
if (job.type === 'master') {
return extension.originalQueuePrompt(0, job.promptWrapper).then(result => {
masterPromptId = result;
return result;
});
} else {
return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences);
}
});
await Promise.all(promises);
// Trigger immediate status check for instant feedback
extension.checkAllWorkerStatuses();
return masterPromptId || { "prompt_id": "distributed-job-dispatched" };
}
async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) {
const workerUrl = extension.getWorkerUrl(worker);
// Debug logging - always log to console for debugging
extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug");
extension.log('[Distributed] Worker URL: ' + workerUrl, "debug");
// Handle image uploads for remote workers
if (imageReferences && imageReferences.size > 0) {
// Check if this is a local worker (same host as master)
const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost');
if (isLocalWorker) {
extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug");
} else {
extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug");
try {
// Load images from master
const images = await loadImagesForWorker(extension, imageReferences);
// Upload images to worker
if (images.length > 0) {
await uploadImagesToWorker(extension, workerUrl, images);
extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug");
}
} catch (error) {
extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error");
// Continue with workflow execution even if image upload fails
}
}
}
const promptToSend = {
prompt,
extra_data: { extra_pnginfo: { workflow } },
client_id: api.clientId
};
extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");
try {
await fetch(`${workerUrl}/prompt`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
mode: 'cors',
body: JSON.stringify(promptToSend)
});
} catch (e) {
extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
}
}
export async function loadImagesForWorker(extension, imageReferences) {
const images = [];
// Use a cache to avoid loading the same image multiple times
if (!extension._imageCache) {
extension._imageCache = new Map();
}
for (const [filename, info] of imageReferences) {
try {
// Check cache first
if (extension._imageCache.has(filename)) {
images.push(extension._imageCache.get(filename));
extension.log(`Using cached image: ${filename}`, "debug");
continue;
}
// Limit cache size
if (extension._imageCache.size >= 10) {
const oldestKey = extension._imageCache.keys().next().value;
extension._imageCache.delete(oldestKey);
extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug");
}
// Load image from master's filesystem via API
try {
const data = await extension.api.loadImage(filename);
const imageData = {
name: filename,
image: data.image_data,
hash: data.hash // Include hash from the response
};
images.push(imageData);
// Cache the image for future use
extension._imageCache.set(filename, imageData);
extension.log(`Loaded and cached image: ${filename}`, "debug");
} catch (loadError) {
extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error");
throw loadError;
}
} catch (error) {
extension.log(`Error loading image ${filename}: ${error.message}`, "error");
}
}
// Clear cache after a reasonable time to avoid memory issues
setTimeout(() => {
if (extension._imageCache && extension._imageCache.size > 0) {
extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug");
extension._imageCache.clear();
}
}, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds
return images;
}
export async function uploadImagesToWorker(extension, workerUrl, images) {
// Upload images to worker's ComfyUI instance
for (const imageData of images) {
// Check if file already exists with matching hash
if (imageData.hash) {
try {
const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
mode: 'cors',
body: JSON.stringify({
filename: imageData.name,
hash: imageData.hash
})
});
if (checkResponse.ok) {
const result = await checkResponse.json();
if (result.exists && result.hash_matches) {
extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug");
continue;
}
}
} catch (error) {
// If check fails, proceed with upload
extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug");
}
}
const formData = new FormData();
// Detect MIME type from base64 header (supports both image and video)
const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\w+);base64,/);
const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected
// Convert base64 to blob
const base64Data = imageData.image.replace(/^data:(?:image|video)\/\w+;base64,/, '');
const byteCharacters = atob(base64Data);
const byteNumbers = new Array(byteCharacters.length);
for (let i = 0; i < byteCharacters.length; i++) {
byteNumbers[i] = byteCharacters.charCodeAt(i);
}
const byteArray = new Uint8Array(byteNumbers);
const blob = new Blob([byteArray], { type: mimeType });
// Use original filename without heavy cleaning
let cleanName = imageData.name;
let subfolder = '';
// Extract subfolder if present
if (cleanName.includes('/')) {
const parts = cleanName.split('/');
subfolder = parts.slice(0, -1).join('/');
cleanName = parts[parts.length - 1];
}
formData.append('image', blob, cleanName);
formData.append('type', 'input');
formData.append('subfolder', subfolder);
formData.append('overwrite', 'true');
try {
const response = await fetch(`${workerUrl}/upload/image`, {
method: 'POST',
mode: 'cors',
body: formData
});
if (!response.ok) {
throw new Error(`Upload failed: ${response.statusText}`);
}
extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug");
} catch (error) {
extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error");
// Continue with other images
}
}
}
export async function performPreflightCheck(extension, workers) {
if (workers.length === 0) return [];
extension.log(`Performing pre-flight health check on ${workers.length} workers...`, "debug");
const startTime = Date.now();
const checkPromises = workers.map(async (worker) => {
const url = extension.getWorkerUrl(worker, '/prompt');
extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug");
try {
const response = await fetch(url, {
method: 'GET',
mode: 'cors',
signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
});
if (response.ok) {
extension.log(`Worker ${worker.name} is active`, "debug");
return { worker, active: true };
} else {
extension.log(`Worker ${worker.name} returned ${response.status}`, "debug");
return { worker, active: false };
}
} catch (error) {
extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug");
return { worker, active: false };
}
});
const results = await Promise.all(checkPromises);
const activeWorkers = results.filter(r => r.active).map(r => r.worker);
const elapsed = Date.now() - startTime;
extension.log(`Pre-flight check completed in ${elapsed}ms. Active workers: ${activeWorkers.length}/${workers.length}`, "debug");
// Update UI status indicators for inactive workers
results.filter(r => !r.active).forEach(r => {
const statusDot = document.getElementById(`status-${r.worker.id}`);
if (statusDot) {
// Remove pulsing animation once status is determined
statusDot.classList.remove('status-pulsing');
statusDot.style.backgroundColor = "#c04c4c"; // Red for offline
statusDot.title = "Offline - Cannot connect";
}
});
return activeWorkers;
}
-2728
View File
File diff suppressed because it is too large Load Diff
+86
View File
@@ -0,0 +1,86 @@
import { app } from "/scripts/app.js";
app.registerExtension({
name: "Distributed.ImageBatchDivider",
async nodeCreated(node) {
if (node.comfyClass === "ImageBatchDivider") {
try {
const updateOutputs = () => {
if (!node.widgets) return;
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
if (!divideByWidget) return;
const divideBy = parseInt(divideByWidget.value, 10) || 1;
const totalOutputs = divideBy; // Direct divide by value
// Ensure outputs array exists
if (!node.outputs) node.outputs = [];
// Remove excess outputs
while (node.outputs.length > totalOutputs) {
node.removeOutput(node.outputs.length - 1);
}
// Add missing outputs
while (node.outputs.length < totalOutputs) {
const outputIndex = node.outputs.length + 1;
node.addOutput(`batch_${outputIndex}`, "IMAGE");
}
if (node.setDirty) node.setDirty(true); // Refresh canvas
};
// Initial update with delay to allow workflow loading
setTimeout(updateOutputs, 200);
// Find the widget and set up responsive handlers
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
if (divideByWidget) {
// Override callback for immediate trigger on value set
const originalCallback = divideByWidget.callback;
divideByWidget.callback = (value) => {
updateOutputs();
if (originalCallback) originalCallback.call(divideByWidget, value); // Preserve 'this' context
};
// Add event listener for real-time input changes (e.g., typing/dragging)
if (divideByWidget.inputEl) {
divideByWidget.inputEl.addEventListener('input', updateOutputs);
}
// Lightweight MutationObserver as fallback (observe attributes on widget element if available)
const observer = new MutationObserver(updateOutputs);
if (divideByWidget.element) {
observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true });
}
// Store cleanup function
node._batchDividerCleanup = () => {
observer.disconnect();
if (divideByWidget.inputEl) {
divideByWidget.inputEl.removeEventListener('input', updateOutputs);
}
divideByWidget.callback = originalCallback; // Restore original
};
}
// Add post-configure hook for reliable workflow loading
const originalConfigure = node.configure;
node.configure = function(data) {
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
updateOutputs(); // Re-run after config load
return result;
};
} catch (error) {
console.error("Error in ImageBatchDivider extension:", error);
}
}
},
nodeBeforeRemove(node) {
if (node.comfyClass === "ImageBatchDivider" && node._batchDividerCleanup) {
node._batchDividerCleanup();
}
}
});
+1336
View File
File diff suppressed because it is too large Load Diff
+277
View File
@@ -0,0 +1,277 @@
import { BUTTON_STYLES } from './constants.js';
export async function renderSidebarContent(extension, el) {
// Panel is being opened/rendered
extension.log("Panel opened", "debug");
if (!el) {
extension.log("No element provided to renderSidebarContent", "debug");
return;
}
// Prevent infinite recursion
if (extension._isRendering) {
extension.log("Already rendering, skipping", "debug");
return;
}
extension._isRendering = true;
try {
// Store reference to the panel element
extension.panelElement = el;
// Show loading indicator
el.innerHTML = '';
const loadingDiv = document.createElement("div");
loadingDiv.style.cssText = "display: flex; align-items: center; justify-content: center; height: calc(100vh - 100px); color: #888;";
loadingDiv.innerHTML = `<svg width="24" height="24" viewBox="0 0 24 24" style="color: #888;">
<circle cx="12" cy="12" r="10" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-dasharray="40 40"/>
</svg>`;
el.appendChild(loadingDiv);
// Add rotation animation
const style = document.createElement('style');
style.textContent = `
@keyframes rotate {
from { transform: rotate(0deg); }
to { transform: rotate(360deg); }
}
`;
document.head.appendChild(style);
loadingDiv.querySelector('svg').style.animation = 'rotate 1s linear infinite';
// Preload data outside render
await extension.loadConfig();
await extension.loadManagedWorkers();
el.innerHTML = '';
// Create toolbar header to match ComfyUI style
const toolbar = document.createElement("div");
toolbar.className = "p-toolbar p-component border-x-0 border-t-0 rounded-none px-2 py-1 min-h-8";
toolbar.style.cssText = "border-bottom: 1px solid #444; background: transparent; display: flex; align-items: center;";
const toolbarStart = document.createElement("div");
toolbarStart.className = "p-toolbar-start";
toolbarStart.style.cssText = "display: flex; align-items: center;";
const titleSpan = document.createElement("span");
titleSpan.className = "text-xs 2xl:text-sm truncate";
titleSpan.textContent = "COMFYUI DISTRIBUTED";
titleSpan.title = "ComfyUI Distributed";
toolbarStart.appendChild(titleSpan);
toolbar.appendChild(toolbarStart);
const toolbarCenter = document.createElement("div");
toolbarCenter.className = "p-toolbar-center";
toolbar.appendChild(toolbarCenter);
const toolbarEnd = document.createElement("div");
toolbarEnd.className = "p-toolbar-end";
toolbar.appendChild(toolbarEnd);
el.appendChild(toolbar);
// Main container with adjusted padding
const container = document.createElement("div");
container.style.cssText = "padding: 15px; display: flex; flex-direction: column; height: calc(100% - 32px);";
// Detect master info on panel open (in case CUDA info wasn't available at startup)
extension.log(`Panel opened. CUDA device count: ${extension.cudaDeviceCount}, Workers: ${extension.config?.workers?.length || 0}`, "debug");
if (!extension.cudaDeviceCount) {
await extension.detectMasterIP();
}
// Now render with guaranteed up-to-date config
// Master Node Section
const masterDiv = extension.ui.renderEntityCard('master', extension.config?.master, extension);
container.appendChild(masterDiv);
// Workers Section (no heading)
const gpuSection = document.createElement("div");
gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;";
const gpuList = document.createElement("div");
const workers = extension.config?.workers || [];
// If no workers exist, show a full blueprint placeholder first
if (workers.length === 0) {
const blueprintDiv = extension.ui.renderEntityCard('blueprint', { onClick: () => extension.addNewWorker() }, extension);
gpuList.appendChild(blueprintDiv);
}
// Show existing workers
workers.forEach(worker => {
const gpuDiv = extension.ui.renderEntityCard('worker', worker, extension);
gpuList.appendChild(gpuDiv);
});
gpuSection.appendChild(gpuList);
// Only show the minimal "Add Worker" box if there are existing workers
if (workers.length > 0) {
const addWorkerDiv = extension.ui.renderEntityCard('add', { onClick: () => extension.addNewWorker() }, extension);
gpuSection.appendChild(addWorkerDiv);
}
container.appendChild(gpuSection);
const actionsSection = document.createElement("div");
actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;";
// Create a row for both buttons
const buttonRow = document.createElement("div");
buttonRow.style.cssText = "display: flex; gap: 8px;";
const clearMemButton = extension.ui.createButtonHelper(
"Clear Worker VRAM",
(e) => extension._handleClearMemory(e.target),
BUTTON_STYLES.clearMemory
);
clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)";
clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory;
const interruptButton = extension.ui.createButtonHelper(
"Interrupt Workers",
(e) => extension._handleInterruptWorkers(e.target),
BUTTON_STYLES.interrupt
);
interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs";
interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt;
buttonRow.appendChild(clearMemButton);
buttonRow.appendChild(interruptButton);
actionsSection.appendChild(buttonRow);
container.appendChild(actionsSection);
// Settings section
const settingsSection = document.createElement("div");
settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;";
// Settings header with toggle
const settingsHeader = document.createElement("div");
settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;";
const workerSettingsTitle = document.createElement("h4");
workerSettingsTitle.textContent = "Settings";
workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;";
const workerSettingsToggle = document.createElement("span");
workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed
workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;";
settingsHeader.appendChild(workerSettingsTitle);
settingsHeader.appendChild(workerSettingsToggle);
// Hover effect for header
settingsHeader.onmouseover = () => {
workerSettingsToggle.style.color = "#fff";
};
settingsHeader.onmouseout = () => {
workerSettingsToggle.style.color = "#888";
};
// Collapsible settings content
const settingsContent = document.createElement("div");
settingsContent.style.cssText = "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;";
const settingsDiv = document.createElement("div");
settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;";
// Toggle functionality
let settingsExpanded = false;
settingsHeader.onclick = () => {
settingsExpanded = !settingsExpanded;
if (settingsExpanded) {
settingsContent.style.maxHeight = "200px";
settingsContent.style.opacity = "1";
workerSettingsToggle.style.transform = "rotate(90deg)";
} else {
settingsContent.style.maxHeight = "0";
settingsContent.style.opacity = "0";
workerSettingsToggle.style.transform = "rotate(0deg)";
}
};
// Debug mode setting
const debugGroup = document.createElement("div");
debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
const debugCheckbox = document.createElement("input");
debugCheckbox.type = "checkbox";
debugCheckbox.id = "setting-debug";
debugCheckbox.checked = extension.config?.settings?.debug || false;
debugCheckbox.onchange = (e) => extension._updateSetting('debug', e.target.checked);
const debugLabel = document.createElement("label");
debugLabel.htmlFor = "setting-debug";
debugLabel.textContent = "Debug Mode";
debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
debugGroup.appendChild(debugCheckbox);
debugGroup.appendChild(debugLabel);
// Auto-launch workers setting
const autoLaunchGroup = document.createElement("div");
autoLaunchGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
const autoLaunchCheckbox = document.createElement("input");
autoLaunchCheckbox.type = "checkbox";
autoLaunchCheckbox.id = "setting-auto-launch";
autoLaunchCheckbox.checked = extension.config?.settings?.auto_launch_workers || false;
autoLaunchCheckbox.onchange = (e) => extension._updateSetting('auto_launch_workers', e.target.checked);
const autoLaunchLabel = document.createElement("label");
autoLaunchLabel.htmlFor = "setting-auto-launch";
autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup";
autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
autoLaunchGroup.appendChild(autoLaunchCheckbox);
autoLaunchGroup.appendChild(autoLaunchLabel);
// Stop workers on exit setting
const stopOnExitGroup = document.createElement("div");
stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
const stopOnExitCheckbox = document.createElement("input");
stopOnExitCheckbox.type = "checkbox";
stopOnExitCheckbox.id = "setting-stop-on-exit";
stopOnExitCheckbox.checked = extension.config?.settings?.stop_workers_on_master_exit !== false; // Default true
stopOnExitCheckbox.onchange = (e) => extension._updateSetting('stop_workers_on_master_exit', e.target.checked);
const stopOnExitLabel = document.createElement("label");
stopOnExitLabel.htmlFor = "setting-stop-on-exit";
stopOnExitLabel.textContent = "Stop Local Workers on Master Exit";
stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
stopOnExitGroup.appendChild(stopOnExitCheckbox);
stopOnExitGroup.appendChild(stopOnExitLabel);
settingsDiv.appendChild(debugGroup);
settingsDiv.appendChild(autoLaunchGroup);
settingsDiv.appendChild(stopOnExitGroup);
settingsContent.appendChild(settingsDiv);
settingsSection.appendChild(settingsHeader);
settingsSection.appendChild(settingsContent);
container.appendChild(settingsSection);
const summarySection = document.createElement("div");
summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;";
const summary = document.createElement("div");
summary.id = "distributed-summary";
summary.style.cssText = "font-size: 11px; color: #888;";
summarySection.appendChild(summary);
container.appendChild(summarySection);
el.appendChild(container);
extension.updateSummary();
// Start checking worker statuses immediately in parallel
setTimeout(() => extension.checkAllWorkerStatuses(), 0);
} finally {
// Always reset the rendering flag
extension._isRendering = false;
}
}
+61
View File
@@ -0,0 +1,61 @@
export function createStateManager() {
const state = {
workers: new Map(), // Unified worker state: { status, managed, launching, expanded, ... }
masterStatus: 'online',
};
return {
// Worker state management
getWorker(workerId) {
return state.workers.get(String(workerId)) || {};
},
updateWorker(workerId, updates) {
const id = String(workerId);
const current = state.workers.get(id) || {};
state.workers.set(id, { ...current, ...updates });
return state.workers.get(id);
},
setWorkerStatus(workerId, status) {
return this.updateWorker(workerId, { status });
},
setWorkerManaged(workerId, info) {
return this.updateWorker(workerId, { managed: info });
},
setWorkerLaunching(workerId, launching) {
return this.updateWorker(workerId, { launching });
},
setWorkerExpanded(workerId, expanded) {
return this.updateWorker(workerId, { expanded });
},
isWorkerLaunching(workerId) {
return this.getWorker(workerId).launching || false;
},
isWorkerExpanded(workerId) {
return this.getWorker(workerId).expanded || false;
},
isWorkerManaged(workerId) {
return !!this.getWorker(workerId).managed;
},
getWorkerStatus(workerId) {
return this.getWorker(workerId).status || {};
},
// Master state
setMasterStatus(status) {
state.masterStatus = status;
},
getMasterStatus() {
return state.masterStatus;
}
};
}
+1236
View File
File diff suppressed because it is too large Load Diff
+325
View File
@@ -0,0 +1,325 @@
import { BUTTON_STYLES, TIMEOUTS } from './constants.js';
export async function handleWorkerOperation(extension, button, operation, successText, errorText) {
const originalText = button.textContent;
const originalStyle = button.style.cssText;
button.textContent = operation.loadingText;
button.disabled = true;
try {
const urlsToProcess = extension.enabledWorkers.map(w => ({
name: w.name,
url: extension.getWorkerUrl(w)
}));
if (urlsToProcess.length === 0) {
button.textContent = "No Workers";
button.style.backgroundColor = "#c04c4c";
setTimeout(() => {
button.textContent = originalText;
button.style.cssText = originalStyle;
button.disabled = false;
}, TIMEOUTS.BUTTON_RESET);
return;
}
const promises = urlsToProcess.map(target =>
fetch(`${target.url}${operation.endpoint}`, {
method: 'POST',
mode: 'cors'
})
.then(response => ({ ok: response.ok, name: target.name }))
.catch(() => ({ ok: false, name: target.name }))
);
const results = await Promise.all(promises);
const failures = results.filter(r => !r.ok);
if (failures.length === 0) {
button.textContent = successText;
button.style.backgroundColor = BUTTON_STYLES.success.split(':')[1].trim().replace(';', '');
if (operation.onSuccess) operation.onSuccess();
} else {
button.textContent = errorText;
button.style.backgroundColor = BUTTON_STYLES.error.split(':')[1].trim().replace(';', '');
extension.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`, "error");
}
setTimeout(() => {
button.textContent = originalText;
button.style.cssText = originalStyle;
}, TIMEOUTS.BUTTON_RESET);
} finally {
button.disabled = false;
}
}
export async function handleInterruptWorkers(extension, button) {
return handleWorkerOperation(extension, button, {
name: "Interrupt",
endpoint: "/interrupt",
loadingText: "Interrupting...",
onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY)
}, "Interrupted!", "Error! See Console");
}
export async function handleClearMemory(extension, button) {
return handleWorkerOperation(extension, button, {
name: "Clear memory",
endpoint: "/distributed/clear_memory",
loadingText: "Clearing..."
}, "Success!", "Error! See Console");
}
export function findNodesByClass(apiPrompt, className) {
return Object.entries(apiPrompt)
.filter(([, nodeData]) => nodeData.class_type === className)
.map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData }));
}
/**
* Find all image references in the workflow
* Looks for inputs named "image" that contain filename strings
*/
export function findImageReferences(extension, apiPrompt) {
const images = new Map();
// Updated regex to handle:
// - Standard files: "image.png"
// - Subfolder files: "subfolder/image.png"
// - ComfyUI special format: "clipspace/file.png [input]"
// - Video files: "video.mp4", "animation.avi", etc.
const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i;
for (const [nodeId, node] of Object.entries(apiPrompt)) {
// Check for both 'image' and 'video' inputs
const mediaInputs = [];
if (node.inputs && node.inputs.image) {
mediaInputs.push(node.inputs.image);
}
if (node.inputs && node.inputs.video) {
mediaInputs.push(node.inputs.video);
}
for (const mediaValue of mediaInputs) {
if (typeof mediaValue === 'string') {
// Clean special suffixes like [input] or [output]
const cleanValue = mediaValue.replace(/\s*\[\w+\]$/, '').trim();
if (imageExtensions.test(cleanValue)) {
images.set(cleanValue, {
nodeId,
nodeType: node.class_type,
inputName: 'image' // Keep as 'image' for compatibility
});
extension.log(`Found media reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug");
}
}
}
}
return images;
}
/**
* Find only upstream nodes (inputs) for distributed collector nodes
* This is used for workers to avoid executing downstream nodes like SaveImage
* @param {Object} apiPrompt - The API prompt containing the workflow
* @param {Array<string>} collectorIds - Array of collector node IDs
* @returns {Set<string>} Set of node IDs that feed into collectors
*/
export function findCollectorUpstreamNodes(apiPrompt, collectorIds) {
const connected = new Set(collectorIds); // Include all collectors
const toProcess = [...collectorIds];
// Only traverse upstream (inputs)
while (toProcess.length > 0) {
const nodeId = toProcess.pop();
const node = apiPrompt[nodeId];
// Traverse upstream (inputs) only
if (node && node.inputs) {
for (const [inputName, inputValue] of Object.entries(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2) {
const sourceNodeId = String(inputValue[0]);
if (!connected.has(sourceNodeId)) {
connected.add(sourceNodeId);
toProcess.push(sourceNodeId);
}
}
}
}
}
return connected;
}
/**
* Prune workflow to only include nodes connected to distributed nodes
* @param {Object} apiPrompt - The full workflow API prompt
* @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided)
* @returns {Object} Pruned API prompt with only required nodes
*/
export function pruneWorkflowForWorker(extension, apiPrompt, distributedNodes = null) {
// Find all distributed nodes if not provided
if (!distributedNodes) {
const collectorNodes = findNodesByClass(apiPrompt, "DistributedCollector");
const upscaleNodes = findNodesByClass(apiPrompt, "UltimateSDUpscaleDistributed");
distributedNodes = [...collectorNodes, ...upscaleNodes];
}
if (distributedNodes.length === 0) {
// No distributed nodes, return full workflow
return apiPrompt;
}
// Get all nodes connected to distributed nodes
const distributedIds = distributedNodes.map(node => node.id);
// For workers, only include upstream nodes (this removes ALL downstream nodes after collectors)
const connectedNodes = findCollectorUpstreamNodes(apiPrompt, distributedIds);
extension.log(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`, "debug");
// Create pruned prompt with only required nodes
const prunedPrompt = {};
for (const nodeId of connectedNodes) {
prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId]));
}
// Check if any distributed node has downstream SaveImage nodes that were removed
// If so, add a PreviewImage node after the collector
for (const distNode of distributedNodes) {
const distNodeId = distNode.id;
// Check if this distributed node had any downstream nodes in the original workflow
const originalOutputMap = new Map();
for (const [nodeId, node] of Object.entries(apiPrompt)) {
if (node.inputs) {
for (const [inputName, inputValue] of Object.entries(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2 && String(inputValue[0]) === distNodeId) {
if (!originalOutputMap.has(distNodeId)) {
originalOutputMap.set(distNodeId, []);
}
originalOutputMap.get(distNodeId).push({nodeId, inputName});
}
}
}
}
// If this distributed node had downstream nodes that were removed, add a PreviewImage
if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) {
// Generate unique numeric ID: max existing numeric key +1
const existingIds = Object.keys(prunedPrompt)
.filter(k => !isNaN(parseInt(k)))
.map(k => parseInt(k));
const maxId = existingIds.length > 0 ? Math.max(...existingIds) : 0;
const previewNodeId = String(maxId + 1);
// Add PreviewImage node connected to the distributed node
prunedPrompt[previewNodeId] = {
inputs: {
images: [distNodeId, 0] // Connect to first output of distributed node
},
class_type: "PreviewImage",
_meta: {
title: "Preview Image (auto-added)"
}
};
extension.log(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`, "debug");
}
}
return prunedPrompt;
}
/**
* Check if a node has an upstream node of a specific type
* @param {Object} apiPrompt - The workflow API prompt
* @param {string} nodeId - The node to check
* @param {string} upstreamType - The class_type to look for upstream
* @returns {boolean} True if an upstream node of the specified type exists
*/
export function hasUpstreamNode(apiPrompt, nodeId, upstreamType) {
const visited = new Set();
const toProcess = [nodeId];
while (toProcess.length > 0) {
const currentId = toProcess.pop();
if (visited.has(currentId)) continue;
visited.add(currentId);
const node = apiPrompt[currentId];
if (!node) continue;
// Check inputs for upstream connections
if (node.inputs) {
for (const [inputName, inputValue] of Object.entries(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2) {
const sourceNodeId = String(inputValue[0]);
const sourceNode = apiPrompt[sourceNodeId];
if (sourceNode && sourceNode.class_type === upstreamType) {
return true;
}
if (!visited.has(sourceNodeId)) {
toProcess.push(sourceNodeId);
}
}
}
}
}
return false;
}
/**
* Get system information from a worker
* @param {string} workerUrl - The worker URL
* @returns {Promise<Object>} System information including platform details
*/
export async function getWorkerSystemInfo(workerUrl) {
try {
const response = await fetch(`${workerUrl}/distributed/system_info`);
if (!response.ok) {
throw new Error(`HTTP ${response.status}`);
}
return await response.json();
} catch (error) {
console.warn(`Failed to get system info from ${workerUrl}:`, error);
// Return sensible defaults
return {
platform: {
os_name: 'posix', // Assume Linux
path_separator: '/',
system: 'Linux'
}
};
}
}
// Cache system info to avoid repeated calls
const systemInfoCache = new Map();
/**
* Get cached system information from a worker
* @param {string} workerUrl - The worker URL
* @returns {Promise<Object>} Cached or fresh system information
*/
export async function getCachedWorkerSystemInfo(workerUrl) {
if (systemInfoCache.has(workerUrl)) {
return systemInfoCache.get(workerUrl);
}
const info = await getWorkerSystemInfo(workerUrl);
systemInfoCache.set(workerUrl, info);
return info;
}
/**
* Clear the system info cache
*/
export function clearSystemInfoCache() {
systemInfoCache.clear();
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff