Merge branch 'feature/cloud-workers'
This commit is contained in:
+30
@@ -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
File diff suppressed because it is too large
Load Diff
+1855
-418
File diff suppressed because it is too large
Load Diff
@@ -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`
|
||||
@@ -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
@@ -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
|
||||
@@ -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})
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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))
|
||||
);
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -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
|
||||
};
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -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
+951
-284
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user