Compare commits

...
Author SHA1 Message Date
robertvoy 1d4dcb985e Feat: Audio Batch Divider 2026-01-25 17:32:24 +11:00
robertvoy b6d4756392 collector supports audio 2026-01-24 18:32:24 +11:00
robertvoy 044c54b780 merged fork 2026-01-24 14:19:22 +11:00
robertvoy b768e7fc8a https handling in host field 2025-12-30 12:11:12 +11:00
robertvoy 9805049268 DistributedModelName 2025-12-30 11:17:44 +11:00
robertvoy 399ed7c7d1 feat: initial setup for cloudflare integration 2025-12-29 09:38:03 +11:00
Robert Wojciechowski 64ba4b75db Update pyproject.toml 2025-12-28 10:55:48 +11:00
Robert Wojciechowski d353d118d7 Update README.md 2025-12-28 10:54:16 +11:00
robertvoy 5dd10cef39 organise docs 2025-12-28 10:48:54 +11:00
Robert Wojciechowski 8eba225f5e Merge pull request #61 from umanets/feature/distributed-queue-api
Add public POST /distributed/queue API endpoint + docs
2025-12-28 10:36:59 +11:00
Igor Umanets 9bc2cfc43e Update README-API.md with links to video walkthrough and examples repository 2025-12-26 16:42:18 +01:00
Igor Umanets 370b48438b Add instructions for retrieving enabled worker IDs in API documentation 2025-12-26 15:12:05 +01:00
Igor Umanets 71c7a99e1e Add API documentation for distributed queue 2025-12-26 14:05:15 +01:00
Igor Umanets def7c0df1d Add distributed queue API orchestration 2025-12-26 10:47:34 +01:00
Robert Wojciechowski 5f5320b34e Update pyproject.toml 2025-10-28 08:01:06 +11:00
Robert Wojciechowski 5a33da828c Merge pull request #50 from IT-BillDeng/feature/disable-master
Feature/disable master
2025-10-28 08:00:30 +11:00
IT-BillDeng ef8e0542c8 Document master participation, orchestrator-only, and fallback modes in the worker setup guide 2025-10-27 03:03:17 +08:00
IT-BillDeng 89b7812fb4 Improve master participation logic with fallback handling and UI messaging 2025-10-27 03:00:57 +08:00
IT-BillDeng f13b953a82 Enable UI toggle for master participation with live status feedback and delegate-mode orchestration. 2025-10-27 02:29:51 +08:00
Robert Wojciechowski 1e3e094cdb Update README.md 2025-10-25 09:40:53 +11:00
Robert Wojciechowski 7bcf9290ba Update video-upscaler-runpod-preset.md 2025-09-15 12:18:44 +10:00
Robert Wojciechowski 6aec567071 Update README.md 2025-09-15 12:01:27 +10:00
Robert Wojciechowski 2fec61d567 Update video-upscaler-runpod-preset.md 2025-09-15 11:53:49 +10:00
Robert Wojciechowski 528cd2e895 Update pyproject.toml 2025-09-15 11:51:51 +10:00
Robert Wojciechowski 0572e084e3 Update video-upscaler-runpod-preset.md 2025-09-15 11:51:28 +10:00
Robert Wojciechowski bc9bba0c76 Update video-upscaler-runpod-preset.md 2025-09-15 11:50:33 +10:00
Robert Wojciechowski 701c6de9eb Update README.md 2025-09-15 11:44:08 +10:00
Robert Wojciechowski c846aac0e6 Update video-upscaler-runpod-preset.md 2025-09-15 11:41:32 +10:00
Robert Wojciechowski 0af3df071d Create video-upscaler-runpod-preset.md 2025-09-15 11:32:29 +10:00
robertvoy f7c7b960c7 change upscaler in video upscaler workflow 2025-09-14 09:54:10 +10:00
robertvoy c45b50830e Merge branch 'main' of https://github.com/robertvoy/ComfyUI-Distributed 2025-09-13 14:42:41 +10:00
robertvoy 213ac38227 Add new /prompt liveness probe (USDU + collector). when worker timeout value in exceeded, the master will probe the workers. 2025-09-13 14:42:33 +10:00
robertvoy 9ae354a788 upscaler workflow change 2025-09-13 12:21:26 +10:00
Robert Wojciechowski c1928ebd09 Update README.md 2025-09-13 10:41:52 +10:00
Robert Wojciechowski 98a63938c8 Update README.md 2025-09-13 10:40:03 +10:00
Robert Wojciechowski 27b73fd280 Update README.md 2025-09-13 10:38:55 +10:00
robertvoy 48c8792af2 video upscaler workflow 2025-09-13 10:34:06 +10:00
robertvoy 4b8836cd03 Unify worker timeout via UI across collector/upscaler (90s); add interruptible sliced waits + cleanup to collector 2025-09-13 10:26:03 +10:00
robertvoy b69fc40fbd updated video upscaler workflow 2025-09-12 09:58:53 +10:00
Robert Wojciechowski e5eb2c8b86 Delete workflows/distributed-upscale-batch.json 2025-09-12 09:49:48 +10:00
Robert Wojciechowski b90c3e42ac Update README.md 2025-09-12 09:49:27 +10:00
robertvoy 187b38d4f5 worker timeout in settings, fixed requeue video tiles 2025-09-11 16:40:19 +10:00
robertvoy 39fb4db486 warning for non-WAN compatible frame count and WanVideo Closest Frame Count 2025-09-11 14:26:53 +10:00
robertvoy 763e238120 fixed: if batch =! (4*n)+1 2025-09-11 12:51:13 +10:00
robertvoy 857aa80174 optimisations 2025-09-10 10:34:00 +10:00
robertvoy 287e4c05fa speed optimisation 2025-09-10 10:15:35 +10:00
robertvoy 58f26a106f clean up 2025-09-10 09:38:40 +10:00
robertvoy 138cac3eb7 remove static distribution 2025-09-09 19:52:34 +10:00
robertvoy 0e07868b98 initial 2025-09-09 19:18:57 +10:00
Robert Wojciechowski a8be7c7ed9 Update pyproject.toml 2025-08-25 13:38:01 +10:00
Robert Wojciechowski c9ca7f5390 Merge pull request #32 from doubletwisted/fix/pasted-paths
Update executionUtils.js
2025-08-25 13:37:34 +10:00
Dominik Bargiel c860df0aa7 Update executionUtils.js
Previous fix caused problems with loading models that are in subfolders like:
"'FLUX.1\\Shakker-Labs-ControlNet-Union-Pro-2.0.safetensors"

Now images and wideos are treated separetlly.
2025-08-24 19:46:40 +02:00
Robert Wojciechowski 05662eb7eb Update pyproject.toml 2025-08-16 10:49:56 +10:00
Robert Wojciechowski 8441c92d95 Merge pull request #29 from doubletwisted/fix/pasted-paths
fix: normalize pasted image paths and subfolder extraction
2025-08-16 10:49:27 +10:00
doubletwisted 4b513d49c3 fix: normalize pasted image paths and subfolder extraction (Windows-safe, upstream-only) 2025-08-15 21:45:19 +02:00
Robert Wojciechowski f91d2d579d Update pyproject.toml 2025-08-13 10:00:36 +10:00
Robert Wojciechowski 1ff863fb5d Add files via upload 2025-08-13 10:00:03 +10:00
Robert Wojciechowski 74bba258da Update worker-setup-guides.md 2025-08-09 12:40:53 +10:00
Robert Wojciechowski bb559e8ce0 Merge pull request #27 from ComfyNodePRs/update-publish-yaml
Update Github Action for Publishing to Comfy Registry
2025-08-06 06:53:17 +10:00
Robert Wojciechowski c032de224b Update README.md 2025-08-05 07:58:40 +10:00
Robert Wojciechowski 68f8627b0b Update README.md 2025-08-04 11:09:36 +10:00
Robert Wojciechowski 0e3c04fba7 Update README.md 2025-08-04 11:02:55 +10:00
robertvoy e076cf3455 Merge branch 'main' of https://github.com/robertvoy/ComfyUI-Distributed 2025-08-03 18:38:06 +10:00
robertvoy 5f3d358f61 wan 2.2 workflow 2025-08-03 18:38:00 +10:00
Robert Wojciechowski 1a6bff3ab1 Update pyproject.toml 2025-08-01 15:15:35 +10:00
Robert Wojciechowski d7aec0b0fe Update pyproject.toml 2025-08-01 15:13:11 +10:00
snomiao 26aaa1f1c8 chore(publish): update GitHub Actions workflow for node publishing
- Add permissions for issue writing
- Restrict job execution to specific repository owner
- Update action version for node publishing to v1
2025-07-10 00:44:29 +00:00
28 changed files with 5853 additions and 3805 deletions
+5 -1
View File
@@ -7,14 +7,18 @@ on:
paths:
- "pyproject.toml"
permissions:
issues: write
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
if: ${{ github.repository_owner == 'robertvoy' }}
steps:
- name: Check out code
uses: actions/checkout@v4
- name: Publish Custom Node
uses: Comfy-Org/publish-node-action@main
uses: Comfy-Org/publish-node-action@v1
with:
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+4
View File
@@ -0,0 +1,4 @@
bin/
logs/
gpu_config.json
__pycache__/
+47 -7
View File
@@ -1,11 +1,11 @@
<div align="center">
<img width="320" src="https://github.com/user-attachments/assets/533bb98d-0c4a-499f-9bca-5c937e361087" />
<img width="250" src="https://github.com/user-attachments/assets/533bb98d-0c4a-499f-9bca-5c937e361087" />
<br><br>
<a href="https://www.youtube.com/watch?v=p6eE3IlAbOs"><img src="https://img.shields.io/badge/Video_Tutorial-grey?style=flat&logo=youtube&logoColor=white" alt="Video Tutorial"></a>
<a href="/docs/worker-setup-guides.md"><img src="https://img.shields.io/badge/Setup_Guides-grey?style=flat&logo=gitbook&logoColor=white" alt="Setup Guides"></a>
<a href="/workflows"><img src="https://img.shields.io/badge/Workflows-grey?style=flat&logo=json&logoColor=white" alt="Workflows"></a>
<a href="https://buymeacoffee.com/robertvoy"><img src="https://img.shields.io/badge/Donation-grey?style=flat&logo=buymeacoffee&logoColor=white" alt="Donation"></a>
<a href="https://x.com/robertvoy3D"><img src="https://img.shields.io/twitter/follow/robertvoy3D" alt="Twitter"></a>
<a href="https://x.com/rbw_ai"><img src="https://img.shields.io/twitter/follow/rbw_ai" alt="Twitter"></a>
<br><br>
</div>
@@ -25,7 +25,7 @@
#### Distributed Upscaling
- Accelerate Ultimate SD Upscale by distributing tiles across GPUs
- Intelligent distribution
- Handles single images and batches
- Handles single images and videos
#### Ease of Use
- Auto-setup local workers; easily add remote/cloud ones
@@ -72,6 +72,14 @@ ComfyUI Distributed supports three types of workers:
---
## Official Sponsor
[<img width="1500" height="339" src="https://github.com/user-attachments/assets/c5f75e1f-3e19-4c57-b05d-151311cd1cf0" />](https://get.runpod.io/0bw29uf3ug0p)
Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a special bonus.
---
## Workflow Examples
### Basic Parallel Generation
@@ -103,7 +111,7 @@ Generate multiple videos in the time it takes to generate one. Each worker uses
7. Enable workers in the UI
8. Run the workflow!
### Distributed Upscaling
### Distributed Image Upscaling
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a)
@@ -114,9 +122,36 @@ Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, w
2. Upscale with ESRGAN or similar
3. Connect to **Ultimate SD Upscale Distributed**
4. Configure tile settings
> If your GPUs are similar, set `static_distribution` to true; otherwise, false
5. Enable workers for faster processing
### Distributed Video Upscaling
Accelerate Ultimate SD Upscaler by distributing video tiles across multiple workers, with speed scaling as you add more GPUs.
![Video Upscaler workflow](https://github.com/user-attachments/assets/3c3d61b1-0b5f-422e-8c58-7c1555fed765)
> [Download workflow](/workflows/distributed-upscale-video.json)
1. Load your video
2. Optional: upscale with ESRGAN or similar
3. Connect to **Ultimate SD Upscale Distributed**
4. Configure tile settings
5. Use RES4LYF (bong/res2) to get better results
6. Enable workers for faster processing
> You can run this workflow entirely on Runpod with minimal setup. [Check out the guide here.](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/video-upscaler-runpod-preset.md)
---
## Developer API
Control your distributed cluster programmatically without opening the browser.
* **Endpoint:** `POST /distributed/queue`
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
---
## FAQ
@@ -161,6 +196,11 @@ This software is provided "as is" without any warranties, express or implied, in
## Support the Project
<img width="200" src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
<img width="200" align="right" src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
If my custom nodes have added value to your workflow, consider fueling future development with a coffee!
Your support helps keep this project thriving.
Buy me a coffee at: https://buymeacoffee.com/robertvoy
If my custom nodes have added value to your workflow, consider fueling future development with a coffee! Your support helps keep this project thriving. Buy me a coffee at: https://buymeacoffee.com/robertvoy
+1 -1
View File
@@ -62,4 +62,4 @@ __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
debug_log("Loaded Distributed nodes.")
debug_log(f"Config file: {CONFIG_FILE}")
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
+838 -157
View File
File diff suppressed because it is too large Load Diff
+605
View File
@@ -0,0 +1,605 @@
import asyncio
import json
import time
import uuid
from collections import deque
import aiohttp
import execution
import server
from .utils.config import load_config
from .utils.logging import debug_log, log
prompt_server = server.PromptServer.instance
def ensure_distributed_state():
"""Ensure prompt_server has the state used by distributed queue orchestration."""
if not hasattr(prompt_server, "distributed_pending_jobs"):
prompt_server.distributed_pending_jobs = {}
if not hasattr(prompt_server, "distributed_jobs_lock"):
prompt_server.distributed_jobs_lock = asyncio.Lock()
if not hasattr(prompt_server, "distributed_worker_ws"):
prompt_server.distributed_worker_ws = {}
async def _get_client_session():
"""Get or create aiohttp client session (shared with distributed.py)."""
if not hasattr(prompt_server, "_distributed_session"):
prompt_server._distributed_session = aiohttp.ClientSession()
return prompt_server._distributed_session
class PromptIndex:
"""Cache prompt metadata for faster worker/master prompt preparation."""
def __init__(self, prompt_obj):
self._prompt_json = json.dumps(prompt_obj)
self.nodes_by_class = {}
self.class_by_node = {}
self.inputs_by_node = {}
for node_id, node in _iter_prompt_nodes(prompt_obj):
class_type = node.get("class_type")
node_id_str = str(node_id)
if class_type:
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
self.class_by_node[node_id_str] = class_type
self.inputs_by_node[node_id_str] = node.get("inputs", {})
self._upstream_cache = {}
def copy_prompt(self):
return json.loads(self._prompt_json)
def nodes_for_class(self, class_name):
return self.nodes_by_class.get(class_name, [])
def has_upstream(self, start_node_id, target_class):
cache_key = (str(start_node_id), target_class)
if cache_key in self._upstream_cache:
return self._upstream_cache[cache_key]
visited = set()
stack = [str(start_node_id)]
while stack:
node_id = stack.pop()
if node_id in visited:
continue
visited.add(node_id)
inputs = self.inputs_by_node.get(node_id, {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
upstream_id = str(value[0])
if self.class_by_node.get(upstream_id) == target_class:
self._upstream_cache[cache_key] = True
return True
if upstream_id in self.inputs_by_node:
stack.append(upstream_id)
self._upstream_cache[cache_key] = False
return False
def _iter_prompt_nodes(prompt_obj):
for node_id, node in prompt_obj.items():
if isinstance(node, dict):
yield str(node_id), node
def _find_nodes_by_class(prompt_obj, class_name):
nodes = []
for node_id, node in _iter_prompt_nodes(prompt_obj):
if node.get("class_type") == class_name:
nodes.append(node_id)
return nodes
def _find_downstream_nodes(prompt_obj, start_ids):
"""Return all nodes reachable downstream from the provided IDs."""
adjacency = {}
for node_id, node in _iter_prompt_nodes(prompt_obj):
inputs = node.get("inputs", {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
source_id = str(value[0])
adjacency.setdefault(source_id, set()).add(str(node_id))
connected = set(start_ids)
queue = deque(start_ids)
while queue:
current = queue.popleft()
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
if dependent not in connected:
connected.add(dependent)
queue.append(dependent)
return connected
def _create_numeric_id_generator(prompt_obj):
"""Return a closure that yields new numeric string IDs."""
max_id = 0
for node_id in prompt_obj.keys():
try:
numeric = int(node_id)
except (TypeError, ValueError):
continue
max_id = max(max_id, numeric)
counter = max_id
def _next_id():
nonlocal counter
counter += 1
return str(counter)
return _next_id
def _prepare_delegate_master_prompt(prompt_obj, collector_ids):
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
nodes_to_keep = set(collector_ids)
nodes_to_keep.update(downstream)
pruned_prompt = {}
for node_id in nodes_to_keep:
node = prompt_obj.get(node_id)
if node is not None:
pruned_prompt[node_id] = json.loads(json.dumps(node))
pruned_ids = set(pruned_prompt.keys())
for node_id, node in pruned_prompt.items():
inputs = node.get("inputs")
if not inputs:
continue
for input_name, input_value in list(inputs.items()):
if isinstance(input_value, list) and len(input_value) == 2:
source_id = str(input_value[0])
if source_id not in pruned_ids:
inputs.pop(input_name, None)
debug_log(
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
)
next_id = _create_numeric_id_generator(pruned_prompt)
for collector_id in collector_ids:
collector_entry = pruned_prompt.get(collector_id)
if not collector_entry:
continue
placeholder_id = next_id()
pruned_prompt[placeholder_id] = {
"class_type": "DistributedEmptyImage",
"inputs": {
"height": 64,
"width": 64,
"channels": 3,
},
"_meta": {
"title": "Distributed Empty Image (auto-added)",
},
}
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
debug_log(
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
)
return pruned_prompt
async def _ensure_distributed_queue(job_id):
"""Ensure a queue exists for the given distributed job ID."""
ensure_distributed_state()
async with prompt_server.distributed_jobs_lock:
if job_id not in prompt_server.distributed_pending_jobs:
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
def _generate_job_id_map(prompt_index, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
"UltimateSDUpscaleDistributed"
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
return job_map
def _resolve_enabled_workers(config, requested_ids=None):
"""Return a list of worker configs that should participate."""
workers = []
for worker in config.get("workers", []):
worker_id = str(worker.get("id") or "").strip()
if not worker_id:
continue
if requested_ids is not None:
if worker_id not in requested_ids:
continue
elif not worker.get("enabled", False):
continue
raw_port = worker.get("port", worker.get("listen_port", 8188))
try:
port = int(raw_port or 8188)
except (TypeError, ValueError):
log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
port = 8188
workers.append(
{
"id": worker_id,
"name": worker.get("name", worker_id),
"host": worker.get("host"),
"port": port,
"type": worker.get("type", "local"),
}
)
return workers
def _resolve_master_url():
"""Best-effort reconstruction of the master's public URL."""
cfg = load_config()
master_cfg = cfg.get("master", {}) or {}
configured_host = (master_cfg.get("host") or "").strip()
configured_port = master_cfg.get("port")
default_port = getattr(prompt_server, "port", 8188) or 8188
port = int(configured_port or default_port)
def _needs_https(hostname):
hostname = hostname.lower()
https_domains = (
".proxy.runpod.net",
".ngrok-free.app",
".ngrok-free.dev",
".ngrok.io",
".trycloudflare.com",
".cloudflare.dev",
)
return any(hostname.endswith(suffix) for suffix in https_domains)
if configured_host:
if configured_host.startswith(("http://", "https://")):
return configured_host.rstrip("/")
host = configured_host
scheme = "https" if _needs_https(host) or port == 443 else "http"
default_port_for_scheme = 443 if scheme == "https" else 80
# For ngrok/cloud domains without explicit port, default to the scheme's default
if configured_port is None and scheme == "https" and _needs_https(host):
port = default_port_for_scheme
port_part = "" if port == default_port_for_scheme else f":{port}"
return f"{scheme}://{host}{port_part}"
address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
if address in ("0.0.0.0", "::"):
address = "127.0.0.1"
scheme = "https" if port == 443 else "http"
default_port_for_scheme = 443 if scheme == "https" else 80
port_part = "" if port == default_port_for_scheme else f":{port}"
return f"{scheme}://{address}{port_part}"
def _build_worker_url(worker, endpoint=""):
"""Construct the worker base URL with optional endpoint."""
host = (worker.get("host") or "").strip()
port = int(worker.get("port", 8188) or 8188)
if not host:
host = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
if host.startswith(("http://", "https://")):
base = host.rstrip("/")
else:
is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
scheme = "https" if is_cloud else "http"
default_port = 443 if scheme == "https" else 80
port_part = "" if port == default_port else f":{port}"
base = f"{scheme}://{host}{port_part}"
return f"{base}{endpoint}"
async def _worker_is_active(worker):
"""Ping worker's /prompt endpoint to confirm it's reachable."""
url = _build_worker_url(worker, "/prompt")
session = await _get_client_session()
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
return resp.status == 200
except Exception:
return False
async def _worker_ws_is_active(worker):
"""Ping worker's websocket endpoint to confirm it's reachable."""
session = await _get_client_session()
url = _build_worker_url(worker, "/distributed/worker_ws")
try:
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
await ws.close()
return True
except Exception:
return False
async def _get_worker_ws_connection(worker):
"""Return a cached websocket connection to a worker (creating if needed)."""
ensure_distributed_state()
worker_id = worker.get("id")
existing = prompt_server.distributed_worker_ws.get(worker_id)
if existing and not existing.closed:
return existing
session = await _get_client_session()
url = _build_worker_url(worker, "/distributed/worker_ws")
ws = await session.ws_connect(url, heartbeat=20, timeout=10)
prompt_server.distributed_worker_ws[worker_id] = ws
return ws
async def _close_worker_ws(worker_id):
ws = prompt_server.distributed_worker_ws.pop(worker_id, None)
if ws and not ws.closed:
await ws.close()
async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False):
"""Send the prepared prompt to a worker ComfyUI instance."""
url = _build_worker_url(worker, "/prompt")
payload = {"prompt": prompt_obj}
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
if extra_data:
payload["extra_data"] = extra_data
if use_websocket:
request_id = uuid.uuid4().hex
ws_payload = {
"type": "dispatch_prompt",
"request_id": request_id,
"prompt": prompt_obj,
"workflow": workflow_meta,
"client_id": client_id,
}
try:
ws = await _get_worker_ws_connection(worker)
await ws.send_json(ws_payload)
response = await asyncio.wait_for(ws.receive(), timeout=60)
if response.type == aiohttp.WSMsgType.TEXT:
data = json.loads(response.data or "{}")
if (
data.get("type") == "dispatch_ack"
and data.get("request_id") == request_id
and data.get("ok")
):
return
error_msg = data.get("error") or "Worker rejected websocket dispatch."
raise RuntimeError(error_msg)
raise RuntimeError("Unexpected websocket response from worker.")
except Exception as exc:
worker_id = worker.get("id")
await _close_worker_ws(worker_id)
log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
session = await _get_client_session()
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
"""Queue the master prompt through ComfyUI's prompt queue."""
payload = {"prompt": prompt_obj}
payload = prompt_server.trigger_on_prompt(payload)
prompt = payload["prompt"]
prompt_id = str(uuid.uuid4())
valid = await execution.validate_prompt(prompt_id, prompt, None)
if not valid[0]:
raise RuntimeError(f"Invalid prompt: {valid[1]}")
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
if client_id:
extra_data["client_id"] = client_id
sensitive = {}
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
if key in extra_data:
sensitive[key] = extra_data.pop(key)
number = getattr(prompt_server, "number", 0)
prompt_server.number = number + 1
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
prompt_server.prompt_queue.put(prompt_queue_item)
return prompt_id
def _apply_participant_overrides(
prompt_copy,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
):
"""Return a prompt copy with hidden inputs configured for master/worker."""
is_master = participant_id == "master"
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
enabled_json = json.dumps(enabled_worker_ids)
# Distributed seed nodes
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
if not is_master:
idx = worker_index_map.get(participant_id, 0)
inputs["worker_id"] = f"worker_{idx}"
else:
inputs["worker_id"] = ""
# Distributed collectors
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
inputs["pass_through"] = True
continue
unique_id = job_id_map.get(node_id, node_id)
inputs["multi_job_id"] = unique_id
inputs["is_worker"] = not is_master
if is_master:
inputs["enabled_worker_ids"] = enabled_json
inputs["delegate_only"] = bool(delegate_master)
inputs.pop("master_url", None)
inputs.pop("worker_id", None)
else:
inputs["master_url"] = master_url
inputs["worker_id"] = participant_id
inputs["enabled_worker_ids"] = enabled_json
inputs["delegate_only"] = False
# Distributed upscaler nodes
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
unique_id = job_id_map.get(node_id, node_id)
inputs["multi_job_id"] = unique_id
inputs["is_worker"] = not is_master
if is_master:
inputs["enabled_worker_ids"] = enabled_json
inputs.pop("master_url", None)
inputs.pop("worker_id", None)
else:
inputs["master_url"] = master_url
inputs["worker_id"] = participant_id
inputs["enabled_worker_ids"] = enabled_json
return prompt_copy
async def orchestrate_distributed_execution(
prompt_obj,
workflow_meta,
client_id,
enabled_worker_ids=None,
delegate_master=None,
):
"""Core orchestration logic for the /distributed/queue endpoint.
Returns:
tuple[str, int]: (prompt_id, worker_count)
"""
ensure_distributed_state()
config = load_config()
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
workers = _resolve_enabled_workers(config, requested_ids)
prompt_index = PromptIndex(prompt_obj)
# Respect master delegate-only configuration
if delegate_master is None:
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
if not workers and delegate_master:
debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.")
delegate_master = False
# Filter to active workers
active_workers = []
for worker in workers:
if use_websocket:
is_active = await _worker_ws_is_active(worker)
else:
is_active = await _worker_is_active(worker)
if is_active:
active_workers.append(worker)
else:
log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
if not active_workers and delegate_master:
debug_log("All workers offline while delegate-only requested; enabling master participation.")
delegate_master = False
enabled_ids = [worker["id"] for worker in active_workers]
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
if not job_id_map:
prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
return prompt_id, 0
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
master_url = _resolve_master_url()
master_prompt = prompt_index.copy_prompt()
master_prompt = _apply_participant_overrides(
master_prompt,
"master",
enabled_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
)
if delegate_master:
collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector")
upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
if upscale_nodes:
debug_log(
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
)
elif not collector_ids:
debug_log(
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
)
else:
master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
if active_workers:
debug_log(
"Active distributed workers: "
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
)
worker_payloads = []
for worker in active_workers:
worker_prompt = prompt_index.copy_prompt()
worker_prompt = _apply_participant_overrides(
worker_prompt,
worker["id"],
enabled_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
)
worker_payloads.append((worker, worker_prompt))
if worker_payloads:
await asyncio.gather(
*[
_dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket)
for worker, wprompt in worker_payloads
]
)
prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
return prompt_id, len(worker_payloads)
+643 -879
View File
File diff suppressed because it is too large Load Diff
+201
View File
@@ -0,0 +1,201 @@
# ComfyUI-Distributed API (Experimental)
This document describes the **public HTTP API** added to ComfyUI-Distributed to allow queueing *distributed* workflows from external tools (scripts, services, CI jobs, render farms, etc.) without using the ComfyUI web UI.
## Demo
- Video walkthrough: https://youtu.be/yiQlPd0MzLk
## Examples Repository
- Examples repo: https://github.com/umanets/ComfyUI-Distributed-API-examples.git
---
## Overview
### What this adds
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
- Resolves enabled/selected workers.
- Pings workers (`GET /prompt`) to include only reachable ones.
- Dispatches the workflow to workers (`POST /prompt`).
- Queues the master workflow in ComfyUI’s prompt queue.
### What it does *not* add
- Authentication/authorization.
- A separate “job status” API for distributed results (you still use ComfyUI’s normal prompt history / websocket flow, and the existing `/distributed/queue_status/{job_id}` behavior for collector queues).
---
## Endpoint: `POST /distributed/queue`
Queue a workflow for distributed execution.
### URL
- `http://<master-host>:<master-port>/distributed/queue`
### Headers
- `Content-Type: application/json`
### Request Body
```json
{
"prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
"workflow": { },
"client_id": "optional",
"delegate_master": false,
"enabled_worker_ids": ["1", "2"]
}
```
#### Fields
- `prompt` (required, object)
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
- `workflow` (optional, object)
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
- If you don’t care about UI metadata, you can omit it.
- `client_id` (optional, string)
- Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events).
- `delegate_master` (optional, boolean)
- If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`.
- Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master.
- `enabled_worker_ids` (optional, array of strings)
- If provided, only these worker IDs will be considered.
- If omitted, the plugin uses workers marked as enabled in the UI config.
##### How to get `enabled_worker_ids`
Worker IDs come from the plugin config (`GET /distributed/config`) under `workers[].id`.
Example (bash + `jq`):
```bash
curl -s "http://127.0.0.1:8188/distributed/config" \
| jq -r '.workers[] | "id=\(.id)\tname=\(.name)\tenabled=\(.enabled)\thost=\(.host)\tport=\(.port)\ttype=\(.type)"'
```
Example (PowerShell):
```powershell
$cfg = Invoke-RestMethod "http://127.0.0.1:8188/distributed/config"
$cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -AutoSize
```
### Response Body
```json
{
"prompt_id": "<uuid>",
"worker_count": 2
}
```
- `prompt_id` — the master prompt id queued into ComfyUI.
- `worker_count` — number of workers that received a dispatched prompt (only those that passed the health check).
### Status Codes
- `200` — queued.
- `400` — invalid JSON or invalid body.
- `500` — orchestration/dispatch failure (see server logs for details).
---
## Worker requirements (important)
For a worker to participate, it must be reachable from the master:
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
Also, for collector-based flows:
- Workers will send results back to the master via `POST /distributed/job_complete` (that route must be reachable from workers).
### CORS note
If you call the API from a browser (not from a backend), ensure the master ComfyUI is started with `--enable-cors-header`.
---
## Examples
### 1) Minimal `curl`
```bash
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
-H "Content-Type: application/json" \
-d @payload.json
```
Where `payload.json` contains at least:
```json
{
"prompt": {
"1": {"class_type": "KSampler", "inputs": {} }
}
}
```
### 2) Python (`requests`)
```python
import requests
url = "http://127.0.0.1:8188/distributed/queue"
payload = {
"prompt": {...},
"workflow": {...},
"delegate_master": False,
"enabled_worker_ids": ["1", "2"],
}
r = requests.post(url, json=payload, timeout=60)
r.raise_for_status()
print(r.json())
```
### 3) JavaScript (`fetch`)
```js
const url = "http://127.0.0.1:8188/distributed/queue";
const payload = {
prompt: {/* ... */},
workflow: {/* ... */},
delegate_master: false,
enabled_worker_ids: ["1", "2"],
};
const resp = await fetch(url, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
if (!resp.ok) throw new Error(await resp.text());
console.log(await resp.json());
```
---
## Operational notes / gotchas
- If the workflow contains **no distributed nodes**, the endpoint falls back to normal master queueing and returns `worker_count: 0`.
- Worker selection is “best-effort”: offline workers are skipped.
- For public URLs/tunnels: prefer configuring `master.host` with an explicit scheme (`https://...`) to avoid ambiguity.
---
## Changelog (this feature)
- Added `POST /distributed/queue` endpoint.
- Added orchestration module used by the endpoint.
+20
View File
@@ -0,0 +1,20 @@
![Clipboard Image](https://github.com/user-attachments/assets/5dc5224f-3f47-442c-b94a-116afeb28132)
**Accelerated Creative Video Upscaler On Runpod:**
1. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=0bw29uf3ug0p) template.
2. Filter instances by CUDA 12.8 (add filter in Additional Filters at the top of the page).
3. Choose 4x 5090s
4. Press Edit Template to configure the pod's Environment Variables:
- CIVITAI_API_TOKEN: Not necessary for this workflow.
- HF_API_TOKEN: [get your token here](https://huggingface.co/settings/tokens)
- SAGE_ATTENTION: optional optimisation (set to true/false). Recommended for this workflow.
- PRESET_VIDEO_UPSCALER: set to true. This will download everything you need.
5. Deploy your pod.
6. Once pod setup is complete, connect to ComfyUI running on your pod.
7. In ComfyUI, open the GPU panel on the left.
> If you set SAGE_ATTENTION to true, add "--use-sage-attention" to Extra Args on the workers.
8. Launch the workers.
9. [Load the workflow.](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/workflows/distributed-upscale-video.json)
10. Upload video, add prompt and run workflow.
11. Right-click the Video Combine node and click Save Preview to save the video.
+9 -1
View File
@@ -6,6 +6,14 @@
<img width="600" src="https://github.com/user-attachments/assets/609c42aa-8a1c-4a3f-939e-f3552fa1d54f" />
### Master participation modes
The master can either contribute GPU work or stay in **orchestrator-only** mode:
- **Participating**: Master renders alongside workers, useful when you want every available GPU.
- **Orchestrator-only**: Master sends jobs to selected workers but skips local rendering. Enable this by opening the Distributed panel and unchecking the master toggle. The master card will display *“Master disabled: running as orchestrator only.”*
- **Fallback**: If orchestrator-only is enabled but no workers remain selected, the master automatically re-enables execution to guarantee the workflow still runs. The UI shows a green *“Master fallback execution active”* badge so you know work is executing locally again.
### Types of Workers
- **Local workers**: Additional GPUs on the same machine as the master
@@ -75,7 +83,7 @@
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. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
3. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template.
3. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=0bw29uf3ug0p) template.
4. Make sure your Network Volume is mounted and choose a suitable GPU.
> ⚠️ To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8 (add filter in Additional Filters).
6. Press Edit Template to configure the pod's Environment Variables:
+3 -3
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally and remotely "
version = "1.0.5"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.2.1"
license = {file = "LICENSE"}
dependencies = []
@@ -12,4 +12,4 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
[tool.comfy]
PublisherId = "robertvoy"
DisplayName = "ComfyUI-Distributed"
Icon = ""
Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
+397
View File
@@ -0,0 +1,397 @@
"""
Lightweight Cloudflare Tunnel manager.
Responsibilities:
- Locate or download the cloudflared binary for the current platform.
- Start/stop a quick tunnel pointing to the local ComfyUI server.
- Track status and persist minimal state in the shared config.
"""
import asyncio
import os
import platform
import re
import shutil
import signal
import stat
import subprocess
import threading
import time
from urllib import request, error as urlerror
from .config import load_config, save_config
from .logging import debug_log, log
from .network import get_server_port
from .process import is_process_alive, terminate_process
# Regex to capture the generated public URL from cloudflared output
PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE)
# Default timeout (seconds) to wait for a tunnel URL before giving up
TUNNEL_START_TIMEOUT = 25
def _normalize_host(value):
if not value or not isinstance(value, str):
return ""
value = value.strip()
value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE)
return value.split("/")[0]
class CloudflareTunnelManager:
def __init__(self):
self.process = None
self.pid = None
self.public_url = None
self.last_error = None
self.log_file = None
self.status = "stopped"
self.previous_master_host = None
self._lock = asyncio.Lock()
self._url_event = None
self._reader_thread = None
self._loop = None
self._recent_logs = []
self.binary_path = None
self._restore_state()
# --- Binary discovery & download ---
@property
def base_dir(self):
return os.path.dirname(os.path.dirname(__file__))
@property
def bin_dir(self):
return os.path.join(self.base_dir, "bin")
def _get_asset_name(self):
system = platform.system().lower()
machine = platform.machine().lower()
if system == "windows":
if "arm" in machine:
return "cloudflared-windows-arm64.exe"
return "cloudflared-windows-amd64.exe"
if system == "darwin":
if machine in ("arm64", "aarch64"):
return "cloudflared-darwin-arm64"
return "cloudflared-darwin-amd64"
if system == "linux":
if machine in ("arm64", "aarch64"):
return "cloudflared-linux-arm64"
return "cloudflared-linux-amd64"
raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
def _download_cloudflared(self):
asset = self._get_asset_name()
url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
os.makedirs(self.bin_dir, exist_ok=True)
target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared")
debug_log(f"Downloading cloudflared from {url}")
try:
with request.urlopen(url, timeout=30) as resp:
with open(target_path, "wb") as f:
shutil.copyfileobj(resp, f)
except urlerror.URLError as exc:
raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
# Make executable
st = os.stat(target_path)
os.chmod(target_path, st.st_mode | stat.S_IEXEC)
debug_log(f"Downloaded cloudflared to {target_path}")
return target_path
def ensure_binary(self):
# Allow explicit override
env_path = os.environ.get("CLOUDFLARED_PATH")
if env_path and os.path.exists(env_path):
self.binary_path = env_path
return env_path
# Check cached path
cached = self.binary_path
if cached and os.path.exists(cached):
return cached
# Local bin directory
candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared")
if os.path.exists(candidate):
self.binary_path = candidate
return candidate
# System PATH
path_binary = shutil.which("cloudflared")
if path_binary:
self.binary_path = path_binary
return path_binary
# Download on demand
self.binary_path = self._download_cloudflared()
return self.binary_path
# --- State helpers ---
def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None):
cfg = load_config()
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
if status is not None:
tunnel_cfg["status"] = status
if public_url is not None:
tunnel_cfg["public_url"] = public_url
if pid is not None:
tunnel_cfg["pid"] = pid
if log_file is not None:
tunnel_cfg["log_file"] = log_file
if previous_host is not None:
tunnel_cfg["previous_master_host"] = previous_host
if master_host is not None:
cfg.setdefault("master", {})["host"] = master_host
cfg["tunnel"] = tunnel_cfg
save_config(cfg)
def _restore_state(self):
cfg = load_config()
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
self.public_url = tunnel_cfg.get("public_url") or None
self.previous_master_host = tunnel_cfg.get("previous_master_host")
self.log_file = tunnel_cfg.get("log_file")
pid = tunnel_cfg.get("pid")
if pid and is_process_alive(pid):
self.pid = pid
self.status = tunnel_cfg.get("status", "running")
debug_log(f"Detected existing cloudflared process (pid={pid})")
else:
# Clear stale info
self._persist_state(status="stopped", public_url="", pid=None, log_file=None)
self.status = "stopped"
self.pid = None
# --- Process output handling ---
def _append_log(self, line):
if self.log_file:
try:
with open(self.log_file, "a", encoding="utf-8", errors="replace") as f:
f.write(line + "\n")
except Exception as exc: # pragma: no cover - best effort logging
debug_log(f"Failed to write tunnel log: {exc}")
self._recent_logs.append(line)
# Keep last ~200 lines to avoid unbounded growth
if len(self._recent_logs) > 200:
self._recent_logs = self._recent_logs[-200:]
def _reader(self):
assert self.process is not None
loop = self._loop
for raw_line in iter(self.process.stdout.readline, ""):
line = raw_line.strip()
if not line:
continue
self._append_log(line)
match = PUBLIC_URL_PATTERN.search(line)
if match and not self.public_url:
self.public_url = match.group(1).rstrip("/")
self.status = "running"
if self._url_event and loop:
loop.call_soon_threadsafe(self._url_event.set)
# Capture obvious error strings
if "error" in line.lower() and not self.last_error:
self.last_error = line
# Process exited
if self.status == "starting" and not self.public_url:
self.status = "error"
if not self.last_error:
self.last_error = "Cloudflare tunnel exited before becoming ready"
if self._url_event and loop:
loop.call_soon_threadsafe(self._url_event.set)
elif self.status == "running":
# Treat unexpected exit as stopped so UI can retry
self.status = "stopped"
elif self.status == "stopped":
# Normal stop path, nothing to do
pass
# --- Public API ---
async def start_tunnel(self):
async with self._lock:
if self.process and self.process.poll() is None:
return {
"status": self.status,
"public_url": self.public_url,
"pid": self.process.pid,
"log_file": self.log_file,
}
# If we had a stale pid, ensure it's not running
if self.pid and is_process_alive(self.pid):
debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
await self.stop_tunnel()
binary = await asyncio.to_thread(self.ensure_binary)
port = get_server_port()
self.status = "starting"
self.last_error = None
self.public_url = None
self._recent_logs = []
# Remember current master host so we can restore it if the tunnel stops
config = load_config()
master_host = (config.get("master") or {}).get("host") or ""
tunnel_cfg = config.get("tunnel") or {}
if tunnel_cfg.get("previous_master_host"):
self.previous_master_host = tunnel_cfg.get("previous_master_host")
else:
self.previous_master_host = master_host
os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
timestamp = time.strftime("%Y%m%d-%H%M%S")
self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
cmd = [
binary,
"tunnel",
"--no-autoupdate",
"--url",
f"http://127.0.0.1:{port}",
]
debug_log(f"Starting cloudflared: {' '.join(cmd)}")
try:
self.process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
except FileNotFoundError:
self.status = "error"
raise RuntimeError("cloudflared binary not found")
except Exception as exc:
self.status = "error"
raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
self.pid = self.process.pid
self._persist_state(
status="starting",
pid=self.pid,
log_file=self.log_file,
previous_host=self.previous_master_host
)
# Kick off background reader
self._loop = asyncio.get_running_loop()
self._url_event = asyncio.Event()
self._reader_thread = threading.Thread(target=self._reader, daemon=True)
self._reader_thread.start()
try:
await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT)
except asyncio.TimeoutError:
self.last_error = "Timed out waiting for Cloudflare to assign a URL"
await self.stop_tunnel()
raise RuntimeError(self.last_error)
if self.status != "running" or not self.public_url:
error_msg = self.last_error or "Cloudflare tunnel failed to start"
await self.stop_tunnel()
raise RuntimeError(error_msg)
debug_log(f"Cloudflare tunnel ready at {self.public_url}")
master_host = _normalize_host(self.public_url)
self._persist_state(
status="running",
public_url=self.public_url,
pid=self.pid,
log_file=self.log_file,
previous_host=self.previous_master_host or "",
master_host=master_host
)
return {
"status": self.status,
"public_url": self.public_url,
"pid": self.pid,
"log_file": self.log_file,
}
async def stop_tunnel(self):
async with self._lock:
pid = self.process.pid if self.process else self.pid
if not pid:
self._persist_state(status="stopped", public_url=None, pid=None, log_file=None)
self.status = "stopped"
return {"status": "stopped"}
debug_log(f"Stopping cloudflared (pid={pid})")
if self.process:
terminate_process(self.process, timeout=5)
else:
try:
os.kill(pid, signal.SIGTERM)
time.sleep(0.5)
except Exception as exc: # pragma: no cover - best effort
debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
config = load_config()
tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {}
self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host)
active_url = tunnel_cfg.get("public_url")
current_master_host = (config.get("master") or {}).get("host")
restore_host = None
if active_url:
active_host = _normalize_host(active_url)
current_host = _normalize_host(current_master_host)
if current_host == active_host:
restore_host = self.previous_master_host or ""
self.status = "stopped"
self.public_url = None
self.pid = None
self.process = None
self.last_error = None
self._url_event = None
if self._reader_thread and self._reader_thread.is_alive():
self._reader_thread.join(timeout=1)
self._reader_thread = None
self._persist_state(
status="stopped",
public_url="",
pid=None,
log_file=self.log_file,
previous_host=self.previous_master_host,
master_host=restore_host
)
return {"status": "stopped"}
def get_status(self):
alive = False
pid = self.process.pid if self.process else self.pid
if pid:
alive = is_process_alive(pid)
if not alive and self.status == "running":
self.status = "stopped"
return {
"status": self.status,
"public_url": self.public_url,
"pid": pid,
"log_file": self.log_file,
"last_error": self.last_error,
"binary_path": self.binary_path or shutil.which("cloudflared"),
"recent_logs": self._recent_logs[-20:],
"previous_master_host": self.previous_master_host,
}
# Singleton tunnel manager
cloudflare_tunnel_manager = CloudflareTunnelManager()
+47 -3
View File
@@ -5,6 +5,9 @@ import os
import json
from .logging import log
# Import defaults for timeout fallbacks
from .constants import HEARTBEAT_TIMEOUT
CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
def get_default_config():
@@ -15,7 +18,16 @@ def get_default_config():
"settings": {
"debug": False,
"auto_launch_workers": False,
"stop_workers_on_master_exit": True
"stop_workers_on_master_exit": True,
"master_delegate_only": False,
"websocket_orchestration": True
},
"tunnel": {
"status": "stopped",
"public_url": "",
"pid": None,
"log_file": "",
"previous_master_host": ""
}
}
@@ -24,7 +36,12 @@ def load_config():
if os.path.exists(CONFIG_FILE):
try:
with open(CONFIG_FILE, 'r') as f:
return json.load(f)
data = json.load(f)
defaults = get_default_config()
for key, value in defaults.items():
if key not in data:
data[key] = value
return data
except Exception as e:
log(f"Error loading config, using defaults: {e}")
return get_default_config()
@@ -47,4 +64,31 @@ def ensure_config_exists():
from .logging import debug_log
debug_log("Created default config file")
else:
log("Could not create default config file")
log("Could not create default config file")
def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
"""Return the unified worker timeout (seconds).
Priority:
1) UI-configured setting `settings.worker_timeout_seconds`
2) Fallback to provided `default` (defaults to HEARTBEAT_TIMEOUT which itself
can be overridden via the COMFYUI_HEARTBEAT_TIMEOUT env var)
This value should be used anywhere we consider a worker "timed out" from the
master's perspective (e.g., collector waits, upscaler result collection).
"""
try:
cfg = load_config()
val = int(cfg.get('settings', {}).get('worker_timeout_seconds', default))
return max(1, val)
except Exception:
return max(1, int(default))
def is_master_delegate_only() -> bool:
"""Returns True when master should skip local workload and act as orchestrator only."""
try:
cfg = load_config()
return bool(cfg.get('settings', {}).get('master_delegate_only', False))
except Exception:
return False
+1 -1
View File
@@ -38,4 +38,4 @@ MEMORY_CLEAR_DELAY = 0.5
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
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds (default 60s)
+277 -411
View File
@@ -4,22 +4,99 @@ import json
import copy
import os
import io
from aiohttp import web
from aiohttp import web, ClientTimeout
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
# We avoid converting to tensors on the master for tiles; blending uses 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
from .config import load_config
def _parse_tiles_from_form(data):
"""Parse tiles submitted via multipart/form-data into a list of tile dicts.
Expects the following fields in the aiohttp form data:
- 'tiles_metadata': JSON list with per-tile metadata items containing at least
'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional
'batch_idx' and 'global_idx' are included when available.
- 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG).
- 'padding': integer padding used during extraction (optional; defaults 0).
Returns: list of dicts with keys: 'image', 'tile_idx', 'x', 'y',
'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx',
plus 'padding'.
"""
try:
# Parse padding if present
padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
except Exception:
padding = 0
# Parse tiles metadata (JSON list)
meta_raw = data.get('tiles_metadata')
if meta_raw is None:
raise ValueError("Missing tiles_metadata")
try:
metadata = json.loads(meta_raw)
except Exception as e:
raise ValueError(f"Invalid tiles_metadata JSON: {e}")
if not isinstance(metadata, list):
raise ValueError("tiles_metadata must be a list")
tiles = []
# Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ...
for i, meta in enumerate(metadata):
file_field = data.get(f'tile_{i}')
if file_field is None or not hasattr(file_field, 'file'):
raise ValueError(f"Missing tile data for index {i}")
# Read image bytes and decode to PIL
raw = file_field.file.read()
try:
img = Image.open(io.BytesIO(raw)).convert("RGB")
except Exception as e:
raise ValueError(f"Invalid image data for tile {i}: {e}")
# Build tile dictionary (store PIL only; master blends via PIL)
try:
tile_info = {
'image': img,
'tile_idx': int(meta.get('tile_idx', i)),
'x': int(meta.get('x', 0)),
'y': int(meta.get('y', 0)),
'extracted_width': int(meta.get('extracted_width', img.width)),
'extracted_height': int(meta.get('extracted_height', img.height)),
'padding': int(padding),
}
except Exception as e:
raise ValueError(f"Invalid metadata values for tile {i}: {e}")
# Optional fields
if 'batch_idx' in meta:
try:
tile_info['batch_idx'] = int(meta['batch_idx'])
except Exception:
pass
if 'global_idx' in meta:
try:
tile_info['global_idx'] = int(meta['global_idx'])
except Exception:
pass
tiles.append(tile_info)
return tiles
# Unified Job Data Structure Keys
@@ -36,7 +113,46 @@ JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static
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):
from typing import List, Optional
async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None):
"""Initialize queue for dynamic mode (per-image), with collector fields.
- Creates JOB_PENDING_TASKS with image indices
- Adds 'completed_images' dict and 'pending_images' alias used by collectors
"""
await _init_job_queue(
multi_job_id,
'dynamic',
batch_size=batch_size,
all_indices=all_indices or list(range(batch_size)),
enabled_workers=enabled_workers,
)
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
job_data['completed_images'] = {}
job_data['pending_images'] = job_data[JOB_PENDING_TASKS]
debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]):
"""Initialize queue for static mode (batched-per-tile).
- Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1]
"""
await _init_job_queue(
multi_job_id,
'static',
batch_size=batch_size,
num_tiles_per_image=num_tiles_per_image,
enabled_workers=enabled_workers,
batched_static=True,
)
# Initialization handled by master; avoid duplicate init logs here
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, batched_static: bool = False):
"""Unified initialization for job queues in static and dynamic modes."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
@@ -56,17 +172,22 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
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):
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")
debug_log(f"Initialized image queue with {batch_size} pending items")
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
job_data[JOB_BATCH_SIZE] = batch_size
job_data['batched_static'] = bool(batched_static)
# For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
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")
if batched_static and num_tiles_per_image is not None:
for i in range(num_tiles_per_image):
await pending_queue.put(i)
else:
total_tiles = batch_size * num_tiles_per_image
for i in range(total_tiles):
await pending_queue.put(i)
# Keep backward compatibility - if task assignments provided, still track them
if task_assignments and enabled_workers:
@@ -80,37 +201,13 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
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
# Note: legacy task distribution and queue pull helpers removed
async def _drain_results_queue(multi_job_id):
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
"""Drain pending results from queue and update completed_tasks. Returns count drained.
Uses non-blocking get_nowait to avoid await timeouts and reduce latency.
"""
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)
@@ -120,46 +217,47 @@ async def _drain_results_queue(multi_job_id):
completed_tasks = job_data[JOB_COMPLETED_TASKS]
collected = 0
while not q.empty():
while True:
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:
result = q.get_nowait()
except asyncio.QueueEmpty:
break
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]
return collected
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
@@ -174,13 +272,92 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
requeued_count = 0
completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {})
# Allow override via config setting 'worker_timeout_seconds'
cfg = load_config()
hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT))
for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()):
if current_time - last_heartbeat > HEARTBEAT_TIMEOUT:
age = current_time - last_heartbeat
debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
if age > hb_timeout:
# Busy-only grace policy: require positive signal from worker (/prompt)
# We also log assignment state for diagnostics but do not grace on it alone.
assigned = job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, [])
incomplete_assigned = 0
try:
if assigned:
batched_static = bool(job_data.get('batched_static', False))
if batched_static:
num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
batch_size = job_data.get(JOB_BATCH_SIZE, 1)
for task_id in assigned:
for b in range(batch_size):
gidx = b * num_tiles_per_image + task_id
if gidx not in completed_tasks:
incomplete_assigned += 1
break
else:
for task_id in assigned:
if task_id not in completed_tasks:
incomplete_assigned += 1
debug_log(f"Assigned diagnostics: total_assigned={len(assigned)} incomplete_assigned={incomplete_assigned}")
except Exception as e:
debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
busy = False
probe_status = None
probe_queue = None
try:
cfg_workers = load_config().get('workers', [])
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None)
if wrec:
host = wrec.get('host') or 'localhost'
port = int(wrec.get('port', 8188))
url = f"http://{host}:{port}/prompt"
debug_log(f"Probing worker {worker} at {url}")
session = await get_client_session()
async with session.get(url, timeout=ClientTimeout(total=2.0)) as resp:
probe_status = resp.status
if resp.status == 200:
try:
payload = await resp.json()
probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0))
except Exception:
probe_queue = 0
busy = probe_queue is not None and probe_queue > 0
except Exception as e:
debug_log(f"Probe failed for worker {worker}: {e}")
finally:
debug_log(f"Probe diagnostics: http_status={probe_status} queue_remaining={probe_queue}")
if busy:
job_data[JOB_WORKER_STATUS][worker] = current_time
debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
continue
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 batched_static, task_id is a tile_idx; consider it complete only if
# all corresponding global_idx entries are present in completed_tasks.
batched_static = bool(job_data.get('batched_static', False))
if batched_static:
num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
batch_size = job_data.get(JOB_BATCH_SIZE, 1)
# Check all global indices for this tile across the batch
all_done = True
for b in range(batch_size):
gidx = b * num_tiles_per_image + task_id
if gidx not in completed_tasks:
all_done = False
break
if not all_done:
await job_data[JOB_PENDING_TASKS].put(task_id)
requeued_count += 1
else:
# Legacy/global-idx mode: task_id is a global index key
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:
@@ -226,42 +403,6 @@ async def _cleanup_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:
@@ -314,7 +455,7 @@ async def submit_tiles_endpoint(request):
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)
return await handle_api_error(request, "Job not configured for tile submissions", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
@@ -324,115 +465,17 @@ async def submit_tiles_endpoint(request):
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
}]
try:
tiles = _parse_tiles_from_form(data)
except ValueError as e:
return await handle_api_error(request, str(e), 400)
# 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)
return await handle_api_error(request, "Job not configured for tile submissions", 400)
q = job_data[JOB_QUEUE]
if batch_size > 0 or len(tiles) > 0:
@@ -441,6 +484,7 @@ async def submit_tiles_endpoint(request):
'tiles': tiles,
'is_last': is_last
})
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
else:
await q.put({
'worker_id': worker_id,
@@ -484,7 +528,7 @@ async def submit_image_endpoint(request):
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)
return await handle_api_error(request, "Job not configured for image submissions", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
@@ -500,7 +544,7 @@ async def submit_image_endpoint(request):
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)
return await handle_api_error(request, "Job not configured for image submissions", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
@@ -516,188 +560,7 @@ async def submit_image_endpoint(request):
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)
# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles.
@@ -708,7 +571,8 @@ def clone_control_chain(control, clone_hint=True):
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()
hint = getattr(control, 'cond_hint_original', None)
new_control.cond_hint_original = hint.clone() if hint is not None else None
if hasattr(control, 'previous_controlnet'):
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
return new_control
@@ -722,10 +586,12 @@ def clone_conditioning(cond_list, clone_hints=True):
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()
if new_dict['mask'] is not None:
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 new_dict['pooled_output'] is not None:
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])
@@ -776,7 +642,7 @@ async def request_image_endpoint(request):
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)
return await handle_api_error(request, "Invalid job configuration", 400)
try:
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
@@ -795,7 +661,7 @@ async def request_image_endpoint(request):
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})
return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.get('batched_static', False)})
except asyncio.TimeoutError:
if mode == 'dynamic':
return web.json_response({"image_idx": None})
+56 -1
View File
@@ -4,7 +4,6 @@ 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
@@ -442,6 +441,61 @@ def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad
cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W)
# Added Flux-Kontext Support crop_reference_latents by TBG ETUR
def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
"""
1. Resize each latent to `canvas_size` in latent units.
2. Crop the rectangle `region` (pixel coordinates).
3. Down-sample the crop to latent-space `tile_size`.
Expects a list of BCHW tensors under "reference_latents".
"""
latents = cond_dict.get("reference_latents")
if not isinstance(latents, list):
return # nothing to do
k = 8 # down-sample factor from pixel space → latent space (SD-type models)
W_can_px, H_can_px = canvas_size
# canvas size expressed in latent units
W_can_lat, H_can_lat = W_can_px // k, H_can_px // k
W_tile_px, H_tile_px = tile_size
W_tile_lat, H_tile_lat = max(1, W_tile_px // k), max(1, H_tile_px // k)
x1_px, y1_px, x2_px, y2_px = region
new_latents = []
for t in latents: # (B,C,H_lat_in,W_lat_in)
if t.ndim != 4:
raise ValueError(f"expected BCHW, got {t.shape}")
# 1. Resize to canvas resolution in latent units only if needed
if t.shape[-2:] != (H_can_lat, W_can_lat):
t = F.interpolate(t,
size=(H_can_lat, W_can_lat),
mode="bilinear",
align_corners=False)
# 2. Convert pixel crop → latent slice
w0_lat = int(round(x1_px / k))
w1_lat = int(round(x2_px / k))
h0_lat = int(round(y1_px / k))
h1_lat = int(round(y2_px / k))
cropped = t[:, :, h0_lat:h1_lat, w0_lat:w1_lat] # view
# 3. Down-sample to latent-tile size
cropped = F.interpolate(cropped,
size=(H_tile_lat, W_tile_lat),
mode="bilinear",
align_corners=False)
new_latents.append(cropped)
cond_dict["reference_latents"] = new_latents
def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0):
cropped = []
@@ -452,5 +506,6 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
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)
crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
cropped.append(n)
return cropped
+20 -1
View File
@@ -141,6 +141,25 @@ export function createApiClient(baseUrl) {
return Promise.allSettled(
urls.map(url => this.checkStatus(url))
);
},
// Cloudflare tunnel management
async startTunnel() {
return request('/distributed/tunnel/start', {
method: 'POST',
body: JSON.stringify({})
});
},
async stopTunnel() {
return request('/distributed/tunnel/stop', {
method: 'POST',
body: JSON.stringify({})
});
},
async getTunnelStatus() {
return request('/distributed/tunnel/status');
}
};
}
}
+18 -1
View File
@@ -91,6 +91,23 @@ export const PULSE_ANIMATION_CSS = `
opacity: 1;
padding: 12px 0;
}
/* Cloudflare tunnel spinner */
@keyframes tunnel-spin {
from { transform: rotate(0deg); }
to { transform: rotate(360deg); }
}
.tunnel-spinner {
width: 14px;
height: 14px;
border: 2px solid rgba(255, 255, 255, 0.35);
border-top-color: #fff;
border-radius: 50%;
display: inline-block;
animation: tunnel-spin 0.9s linear infinite;
margin-right: 8px;
vertical-align: middle;
}
`;
export const UI_STYLES = {
@@ -150,4 +167,4 @@ export const TIMEOUTS = {
// Background tasks
LOG_REFRESH: 2000, // log auto-refresh interval
IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache
};
};
Binary file not shown.

After

Width:  |  Height:  |  Size: 4.7 KiB

+350 -25
View File
@@ -1,5 +1,5 @@
import { api } from "../../scripts/api.js";
import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo } from './workerUtils.js';
import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo, findCollectorDownstreamNodes } from './workerUtils.js';
import { TIMEOUTS } from './constants.js';
/**
@@ -19,13 +19,25 @@ function convertPathsForPlatform(apiPrompt, targetSeparator) {
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);
};
const isImageOrVideo = (value) => {
return value.match(/\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\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);
const trimmed = obj.trim();
const hasDrive = /^[A-Za-z]:\\\\|^[A-Za-z]:\//.test(trimmed);
const isAbsolute = trimmed.startsWith('/') || trimmed.startsWith('\\\\');
const hasProtocol = /^\w+:\/\//.test(trimmed);
// For annotated relative image/video paths, keep forward slashes
if (!hasDrive && !isAbsolute && !hasProtocol && isImageOrVideo(trimmed)) {
return trimmed.replace(/[\\\\]/g, '/');
}
// Otherwise replace any path separator with the worker's target separator
return trimmed.replace(/[\\\\\/]/g, targetSeparator);
}
return obj;
} else if (Array.isArray(obj)) {
@@ -48,7 +60,8 @@ export function setupInterceptor(extension) {
if (extension.isEnabled) {
const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0;
if (hasCollector || hasDistUpscale) {
const result = await executeParallelDistributed(extension, prompt);
// Immediate status check for instant feedback
@@ -57,6 +70,14 @@ export function setupInterceptor(extension) {
setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
return result;
}
// DistributedQueue: route entire workflow to least-busy worker
if (hasDistQueue) {
const result = await executeQueueDistributed(extension, prompt);
extension.checkAllWorkerStatuses();
setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
return result;
}
}
return extension.originalQueuePrompt(number, prompt);
};
@@ -82,11 +103,11 @@ export async function executeParallelDistributed(extension, 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 {
@@ -97,28 +118,43 @@ export async function executeParallelDistributed(extension, promptWrapper) {
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];
const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only);
let masterDelegateActive = false;
if (wantsDelegate) {
if (activeWorkers.length === 0) {
extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug");
} else if (upscaleNodes.length > 0) {
extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn");
} else if (collectorNodes.length === 0) {
extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug");
} else {
masterDelegateActive = true;
extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug");
}
}
// 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}`]));
@@ -135,7 +171,8 @@ export async function executeParallelDistributed(extension, promptWrapper) {
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
job_id_map: job_id_map, // Pass the map of unique IDs
delegate_master: masterDelegateActive && participantId === 'master'
};
const jobApiPrompt = await prepareApiPromptForParticipant(
@@ -172,14 +209,147 @@ export async function executeParallelDistributed(extension, promptWrapper) {
}
}
/**
* Execute a workflow via DistributedQueue - routes to the least-busy worker.
* The workflow runs ONLY on the selected worker, not on the master.
*/
export async function executeQueueDistributed(extension, promptWrapper) {
try {
const enabledWorkers = extension.enabledWorkers;
if (enabledWorkers.length === 0) {
extension.log("DistributedQueue: No enabled workers. Falling back to local execution.", "warn");
return extension.originalQueuePrompt(0, promptWrapper);
}
// Fetch queue status from all enabled workers
const workerStatuses = await fetchWorkerQueueStatuses(extension, enabledWorkers);
if (workerStatuses.length === 0) {
extension.log("DistributedQueue: No reachable workers. Falling back to local execution.", "warn");
return extension.originalQueuePrompt(0, promptWrapper);
}
// Select the least-busy worker
const selectedWorker = selectLeastBusyWorker(extension, workerStatuses);
const worker = selectedWorker.worker;
const queueRemaining = selectedWorker.queueRemaining;
extension.log(`DistributedQueue: Routing to ${worker.name} (queue_remaining=${queueRemaining})`, "info");
// Prepare the prompt with skip_dispatch=True for DistributedQueue nodes
const promptToSend = JSON.parse(JSON.stringify(promptWrapper.output));
markSkipDispatch(promptToSend);
// Dispatch to the selected worker
const workerUrl = extension.getWorkerUrl(worker);
try {
const response = await fetch(`${workerUrl}/prompt`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
mode: 'cors',
body: JSON.stringify({
prompt: promptToSend,
extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
client_id: api.clientId
})
});
if (!response.ok) {
throw new Error(`Worker returned ${response.status}`);
}
const result = await response.json();
extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");
return { prompt_id: result.prompt_id, worker_id: worker.id };
} catch (error) {
extension.log(`DistributedQueue: Failed to dispatch to ${worker.name}: ${error.message}`, "error");
// Fall back to local execution
return extension.originalQueuePrompt(0, promptWrapper);
}
} catch (error) {
extension.log("DistributedQueue execution failed: " + error.message, "error");
throw error;
}
}
async function fetchWorkerQueueStatuses(extension, workers) {
const statuses = [];
const checkPromises = workers.map(async (worker) => {
const url = extension.getWorkerUrl(worker, '/prompt');
try {
const response = await fetch(url, {
method: 'GET',
mode: 'cors',
cache: 'no-store',
signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
});
if (response.ok) {
const data = await response.json();
const queueRemaining = data.exec_info?.queue_remaining || 0;
return { worker, queueRemaining, online: true };
}
} catch (error) {
extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
}
return null;
});
const results = await Promise.all(checkPromises);
return results.filter(r => r !== null);
}
function selectLeastBusyWorker(extension, statuses) {
// Find idle workers (queue_remaining == 0)
const idleWorkers = statuses.filter(s => s.queueRemaining === 0);
if (idleWorkers.length > 0) {
// Round-robin among idle workers
if (!extension._distributedQueueRRIndex) {
extension._distributedQueueRRIndex = 0;
}
const index = extension._distributedQueueRRIndex % idleWorkers.length;
extension._distributedQueueRRIndex++;
return idleWorkers[index];
}
// No idle workers - pick the one with the shortest queue
return statuses.reduce((min, s) => s.queueRemaining < min.queueRemaining ? s : min);
}
function markSkipDispatch(promptObj) {
for (const node of Object.values(promptObj)) {
if (node && typeof node === 'object' && node.class_type === 'DistributedQueue') {
node.inputs = node.inputs || {};
node.inputs.skip_dispatch = true;
}
}
}
export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
const isMaster = participantId === 'master';
const delegateMaster = Boolean(options.delegate_master);
// Find all distributed nodes once (before pruning)
const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
let collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed");
const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
let allDistributedNodes = [...collectorNodes, ...upscaleNodes];
if (isMaster && delegateMaster) {
if (upscaleNodes.length > 0) {
extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn");
} else if (collectorNodes.length === 0) {
extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug");
} else {
jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes);
collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
allDistributedNodes = [...collectorNodes, ...upscaleNodes];
}
}
// For workers, handle platform-specific path conversion
if (!isMaster) {
@@ -270,6 +440,9 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p
inputs.is_worker = !isMaster;
if (isMaster) {
inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
if (delegateMaster) {
inputs.delegate_only = true;
}
} else {
inputs.master_url = extension.getMasterUrl();
// Also make the worker_job_id unique to prevent potential caching issues
@@ -311,6 +484,75 @@ export async function prepareDistributedJob(extension, multi_job_id) {
}
}
function createNumericIdGenerator(promptObj) {
let maxId = 0;
for (const key of Object.keys(promptObj)) {
const numeric = parseInt(key, 10);
if (!Number.isNaN(numeric)) {
maxId = Math.max(maxId, numeric);
}
}
return () => {
maxId += 1;
return String(maxId);
};
}
function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) {
const collectorIds = collectorNodes.map((node) => node.id);
const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds);
// Ensure collectors themselves are present
collectorIds.forEach((id) => nodesToKeep.add(id));
const prunedPrompt = {};
nodesToKeep.forEach((nodeId) => {
const nodeData = apiPrompt[nodeId];
if (nodeData) {
prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData));
}
});
const prunedIds = new Set(Object.keys(prunedPrompt));
// Remove dangling references to trimmed nodes
for (const [nodeId, node] of Object.entries(prunedPrompt)) {
if (!node.inputs) continue;
for (const [inputName, inputValue] of Object.entries(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2) {
const sourceId = String(inputValue[0]);
if (!prunedIds.has(sourceId)) {
delete node.inputs[inputName];
extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug");
}
}
}
}
const nextId = createNumericIdGenerator(prunedPrompt);
collectorIds.forEach((collectorId) => {
const collectorEntry = prunedPrompt[collectorId];
if (!collectorEntry) return;
const placeholderId = nextId();
prunedPrompt[placeholderId] = {
class_type: "DistributedEmptyImage",
inputs: {
height: 64,
width: 64,
channels: 3
}
};
prunedIds.add(placeholderId);
collectorEntry.inputs = collectorEntry.inputs || {};
collectorEntry.inputs.images = [placeholderId, 0];
});
extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug");
return prunedPrompt;
}
export async function executeJobs(extension, jobs) {
let masterPromptId = null;
@@ -383,24 +625,102 @@ async function dispatchToWorker(extension, worker, prompt, workflow, imageRefere
const promptToSend = {
prompt,
extra_data: { extra_pnginfo: { workflow } },
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' },
if (extension.config?.settings?.websocket_orchestration) {
const wsUrl = buildWorkerWebSocketUrl(workerUrl);
await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend);
return;
}
await fetch(`${workerUrl}/prompt`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
mode: 'cors',
body: JSON.stringify(promptToSend)
body: JSON.stringify({
prompt,
extra_data: { extra_pnginfo: { workflow } },
client_id: api.clientId
})
});
} catch (e) {
extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
}
}
function buildWorkerWebSocketUrl(workerUrl) {
const url = new URL(workerUrl);
url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
url.pathname = '/distributed/worker_ws';
url.search = '';
return url.toString();
}
function generateRequestId() {
if (crypto.randomUUID) {
return crypto.randomUUID();
}
return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
const r = Math.random() * 16 | 0;
const v = c === 'x' ? r : (r & 0x3 | 0x8);
return v.toString(16);
});
}
async function dispatchPromptViaWebSocket(extension, wsUrl, payload) {
const requestId = generateRequestId();
const message = {
type: 'dispatch_prompt',
request_id: requestId,
prompt: payload.prompt,
workflow: payload.workflow,
client_id: payload.client_id
};
await new Promise((resolve, reject) => {
const ws = new WebSocket(wsUrl);
const timeoutId = setTimeout(() => {
ws.close();
reject(new Error('Websocket dispatch timed out.'));
}, 60000);
ws.addEventListener('open', () => {
ws.send(JSON.stringify(message));
});
ws.addEventListener('message', (event) => {
try {
const data = JSON.parse(event.data || '{}');
if (data.type !== 'dispatch_ack' || data.request_id !== requestId) {
return;
}
clearTimeout(timeoutId);
ws.close();
if (data.ok) {
resolve();
} else {
reject(new Error(data.error || 'Worker rejected websocket dispatch.'));
}
} catch (error) {
clearTimeout(timeoutId);
ws.close();
reject(error);
}
});
ws.addEventListener('error', () => {
clearTimeout(timeoutId);
ws.close();
reject(new Error('Websocket connection failed.'));
});
});
extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug");
}
export async function loadImagesForWorker(extension, imageReferences) {
const images = [];
@@ -507,9 +827,9 @@ export async function uploadImagesToWorker(extension, workerUrl, images) {
let cleanName = imageData.name;
let subfolder = '';
// Extract subfolder if present
if (cleanName.includes('/')) {
const parts = cleanName.split('/');
// Extract subfolder if present (handle both slash styles)
if (cleanName.includes('/') || cleanName.includes('\\')) {
const parts = cleanName.replace(/\\/g, '/').split('/');
subfolder = parts.slice(0, -1).join('/');
cleanName = parts[parts.length - 1];
}
@@ -553,9 +873,10 @@ export async function performPreflightCheck(extension, workers) {
const response = await fetch(url, {
method: 'GET',
mode: 'cors',
cache: 'no-store',
signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
});
if (response.ok) {
extension.log(`Worker ${worker.name} is active`, "debug");
return { worker, active: true };
@@ -564,6 +885,10 @@ export async function performPreflightCheck(extension, workers) {
return { worker, active: false };
}
} catch (error) {
if (error?.name === 'AbortError') {
extension.log(`Worker ${worker.name} pre-flight check timed out; assuming active`, "debug");
return { worker, active: true, uncertain: true };
}
extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug");
return { worker, active: false };
}
@@ -587,4 +912,4 @@ export async function performPreflightCheck(extension, workers) {
});
return activeWorkers;
}
}
+74 -72
View File
@@ -1,86 +1,88 @@
import { app } from "/scripts/app.js";
// Configuration for each batch divider node type
const BATCH_DIVIDER_NODES = {
"ImageBatchDivider": { outputPrefix: "batch_", outputType: "IMAGE" },
"AudioBatchDivider": { outputPrefix: "audio_", outputType: "AUDIO" }
};
app.registerExtension({
name: "Distributed.ImageBatchDivider",
name: "Distributed.BatchDividers",
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 config = BATCH_DIVIDER_NODES[node.comfyClass];
if (!config) return;
try {
const updateOutputs = () => {
if (!node.widgets) return;
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
};
if (!divideByWidget) return;
const divideBy = parseInt(divideByWidget.value, 10) || 1;
const totalOutputs = divideBy;
// Ensure outputs array exists
if (!node.outputs) node.outputs = [];
// Remove excess outputs
while (node.outputs.length > totalOutputs) {
node.removeOutput(node.outputs.length - 1);
}
// 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;
// Add missing outputs
while (node.outputs.length < totalOutputs) {
const outputIndex = node.outputs.length + 1;
node.addOutput(`${config.outputPrefix}${outputIndex}`, config.outputType);
}
if (node.setDirty) node.setDirty(true);
};
// 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) {
const originalCallback = divideByWidget.callback;
divideByWidget.callback = (value) => {
updateOutputs();
if (originalCallback) originalCallback.call(divideByWidget, value);
};
if (divideByWidget.inputEl) {
divideByWidget.inputEl.addEventListener('input', updateOutputs);
}
const observer = new MutationObserver(updateOutputs);
if (divideByWidget.element) {
observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true });
}
node._batchDividerCleanup = () => {
observer.disconnect();
if (divideByWidget.inputEl) {
divideByWidget.inputEl.removeEventListener('input', updateOutputs);
}
divideByWidget.callback = originalCallback;
};
} catch (error) {
console.error("Error in ImageBatchDivider extension:", error);
}
const originalConfigure = node.configure;
node.configure = function(data) {
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
updateOutputs();
return result;
};
} catch (error) {
console.error(`Error in ${node.comfyClass} extension:`, error);
}
},
nodeBeforeRemove(node) {
if (node.comfyClass === "ImageBatchDivider" && node._batchDividerCleanup) {
if (BATCH_DIVIDER_NODES[node.comfyClass] && node._batchDividerCleanup) {
node._batchDividerCleanup();
}
}
});
});
+238 -35
View File
@@ -17,6 +17,8 @@ class DistributedExtension {
this.logAutoRefreshInterval = null;
this.masterSettingsExpanded = false;
this.app = app; // Store app reference for toast notifications
this.tunnelStatus = { status: "unknown" };
this.tunnelElements = {};
// Initialize centralized state
this.state = createStateManager();
@@ -90,6 +92,34 @@ class DistributedExtension {
return this.enabledWorkers.length > 0;
}
isMasterParticipationEnabled() {
return !Boolean(this.config?.settings?.master_delegate_only);
}
isMasterFallbackActive() {
return Boolean(this.config?.settings?.master_delegate_only) && this.enabledWorkers.length === 0;
}
isMasterParticipating() {
return this.isMasterParticipationEnabled() || this.isMasterFallbackActive();
}
async updateMasterParticipation(enabled) {
if (!this.config?.settings) {
this.config.settings = {};
}
const delegateOnly = !enabled;
if (this.config.settings.master_delegate_only === delegateOnly) {
return;
}
await this._updateSetting('master_delegate_only', delegateOnly);
if (this.panelElement) {
renderSidebarContent(this, this.panelElement);
}
}
async loadConfig() {
try {
this.config = await this.api.getConfig();
@@ -118,6 +148,162 @@ class DistributedExtension {
}
}
_applyMasterHost(host) {
if (!host || !this.config) return;
if (!this.config.master) this.config.master = {};
this.config.master.host = host;
const hostInput = document.getElementById('master-host');
if (hostInput) {
hostInput.value = host;
}
}
_parseHostInput(value) {
if (!value) {
return { host: "", port: null };
}
let cleaned = value.trim().replace(/^https?:\/\//i, "");
cleaned = cleaned.split("/")[0];
try {
const url = new URL(`http://${cleaned}`);
const port = url.port ? parseInt(url.port, 10) : null;
return {
host: url.hostname || cleaned,
port: Number.isFinite(port) ? port : null,
};
} catch (error) {
return { host: cleaned, port: null };
}
}
updateTunnelUIElements() {
const elements = this.tunnelElements || {};
const status = (this.tunnelStatus?.status || "stopped").toLowerCase();
const enableColor = "#665533"; // requested yellow/brown tone
const disableColor = "#7c4a4a"; // match worker delete button
const colors = {
running: disableColor,
starting: enableColor,
stopped: enableColor,
error: disableColor,
unknown: enableColor,
stopping: disableColor
};
if (elements.button) {
elements.button.disabled = status === "starting" || status === "stopping";
if (status === "starting") {
elements.button.innerHTML = `<span class="tunnel-spinner"></span> Starting...`;
elements.button.style.backgroundColor = enableColor;
} else if (status === "running") {
elements.button.textContent = "Disable Cloudflare Tunnel";
elements.button.style.backgroundColor = disableColor;
} else if (status === "error") {
elements.button.textContent = "Retry Cloudflare Tunnel";
elements.button.style.backgroundColor = disableColor;
} else {
elements.button.textContent = "Enable Cloudflare Tunnel";
elements.button.style.backgroundColor = enableColor;
}
}
if (elements.status) {
elements.status.textContent = status.toUpperCase();
elements.status.style.backgroundColor = colors[status] || colors.stopped;
}
if (elements.url) {
const url = this.tunnelStatus?.public_url;
if (url) {
elements.url.innerHTML = `<a href="${url}" target="_blank" style="color: #eee; text-decoration: none;">${url}</a>`;
} else {
elements.url.textContent = status === "starting" ? "Requesting public URL..." : "No tunnel active";
}
}
if (elements.copyBtn) {
const hasUrl = Boolean(this.tunnelStatus?.public_url);
elements.copyBtn.disabled = !hasUrl;
elements.copyBtn.style.opacity = hasUrl ? "1" : "0.5";
}
}
async refreshTunnelStatus() {
try {
const data = await this.api.getTunnelStatus();
this.tunnelStatus = data.tunnel || { status: "stopped" };
if (data.master_host !== undefined) {
this._applyMasterHost(data.master_host);
}
return this.tunnelStatus;
} catch (error) {
this.tunnelStatus = { status: "error", last_error: error.message };
this.log("Failed to fetch tunnel status: " + error.message, "error");
return this.tunnelStatus;
} finally {
this.updateTunnelUIElements();
}
}
async handleTunnelToggle(button) {
const currentStatus = (this.tunnelStatus?.status || "stopped").toLowerCase();
if (currentStatus === "starting" || currentStatus === "stopping") {
return;
}
const setStatus = (status) => {
this.tunnelStatus = { ...(this.tunnelStatus || {}), status };
this.updateTunnelUIElements();
};
if (currentStatus === "running") {
setStatus("stopping");
try {
if (button) {
button.innerHTML = `<span class="tunnel-spinner"></span> Stopping...`;
button.disabled = true;
}
const data = await this.api.stopTunnel();
this.tunnelStatus = data.tunnel || { status: "stopped" };
if (data.master_host !== undefined) {
this._applyMasterHost(data.master_host);
}
this.updateTunnelUIElements();
this.ui.showToast(this.app, "info", "Cloudflare Tunnel Disabled", "Master address restored", 4000);
} catch (error) {
this.tunnelStatus = { status: "error", last_error: error.message };
this.updateTunnelUIElements();
this.ui.showToast(this.app, "error", "Failed to stop tunnel", error.message, 5000);
} finally {
if (button) button.disabled = false;
}
return;
}
// Start tunnel
setStatus("starting");
if (button) {
button.innerHTML = `<span class="tunnel-spinner"></span> Starting...`;
button.disabled = true;
}
try {
const data = await this.api.startTunnel();
this.tunnelStatus = data.tunnel || { status: "running" };
if (data.master_host !== undefined) {
this._applyMasterHost(data.master_host);
}
this.updateTunnelUIElements();
const url = data.tunnel?.public_url || data.master_host;
this.ui.showToast(this.app, "success", "Cloudflare Tunnel Ready", url || "Public URL created", 4500);
} catch (error) {
this.tunnelStatus = { status: "error", last_error: error.message };
this.updateTunnelUIElements();
this.ui.showToast(this.app, "error", "Failed to start tunnel", error.message, 5000);
} finally {
if (button) button.disabled = false;
}
}
async updateWorkerEnabled(workerId, enabled) {
const worker = this.config.workers.find(w => w.id === workerId);
if (worker) {
@@ -143,6 +329,10 @@ class DistributedExtension {
} catch (error) {
this.log("Error updating worker: " + error.message, "error");
}
if (this.panelElement) {
await renderSidebarContent(this, this.panelElement);
}
}
async _updateSetting(key, value) {
@@ -154,11 +344,22 @@ class DistributedExtension {
try {
await this.api.updateSetting(key, value);
const prettyKey = key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase());
let detail;
if (key === 'worker_timeout_seconds') {
const secs = parseInt(value, 10);
detail = `Worker Timeout set to ${Number.isFinite(secs) ? secs : value}s`;
} else if (typeof value === 'boolean') {
detail = `${prettyKey} ${value ? 'enabled' : 'disabled'}`;
} else {
detail = `${prettyKey} set to ${value}`;
}
app.extensionManager.toast.add({
severity: "success",
summary: "Setting Updated",
detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`,
detail,
life: 2000
});
} catch (error) {
@@ -217,22 +418,7 @@ class DistributedExtension {
this.panelElement = null;
}
updateSummary() {
const summaryEl = document.getElementById('distributed-summary');
if (summaryEl) {
// Count active workers from state
const activeWorkers = this.config.workers.filter(w => {
const status = this.state.getWorkerStatus(w.id);
return status.online;
}).length;
const totalGPUs = activeWorkers + 1;
if (this.isEnabled) {
summaryEl.textContent = `If Collector node is present, total generation = (${totalGPUs} GPUs × Batch Size)`;
} else {
summaryEl.textContent = "Only the master GPU will be used.";
}
}
}
// updateSummary removed
// --- Core Logic & Execution ---
@@ -304,11 +490,19 @@ class DistributedExtension {
// Update master status dot
const statusDot = document.getElementById('master-status');
if (statusDot) {
if (isProcessing) {
statusDot.style.backgroundColor = "#f0ad4e";
if (!this.isMasterParticipating()) {
if (isProcessing) {
statusDot.style.backgroundColor = STATUS_COLORS.PROCESSING_YELLOW;
statusDot.title = `Orchestrating (${queueRemaining} in queue)`;
} else {
statusDot.style.backgroundColor = STATUS_COLORS.DISABLED_GRAY;
statusDot.title = "Master orchestrator only";
}
} else if (isProcessing) {
statusDot.style.backgroundColor = STATUS_COLORS.PROCESSING_YELLOW;
statusDot.title = `Processing (${queueRemaining} in queue)`;
} else {
statusDot.style.backgroundColor = "#4CAF50";
statusDot.style.backgroundColor = STATUS_COLORS.ONLINE_GREEN;
statusDot.title = "Online";
}
}
@@ -317,15 +511,17 @@ class DistributedExtension {
// Master is always online (we're running on it), so keep it green
const statusDot = document.getElementById('master-status');
if (statusDot) {
statusDot.style.backgroundColor = "#4CAF50";
statusDot.title = "Online";
statusDot.style.backgroundColor = this.isMasterParticipating() ? STATUS_COLORS.ONLINE_GREEN : STATUS_COLORS.DISABLED_GRAY;
statusDot.title = this.isMasterParticipating() ? "Online" : "Master orchestrator only";
}
}
}
// Helper to build worker URL
getWorkerUrl(worker, endpoint = '') {
const host = worker.host || window.location.hostname;
const parsed = this._parseHostInput(worker.host || window.location.hostname);
const host = parsed.host || window.location.hostname;
const resolvedPort = parsed.port || worker.port;
// Cloud workers always use HTTPS
const isCloud = worker.type === 'cloud';
@@ -348,13 +544,13 @@ class DistributedExtension {
}
// Determine protocol: HTTPS for cloud, Runpod proxies, or port 443
const useHttps = isCloud || isRunpodProxy || worker.port === 443;
const useHttps = isCloud || isRunpodProxy || resolvedPort === 443;
const protocol = useHttps ? 'https' : 'http';
// Only add port if non-standard
const defaultPort = useHttps ? 443 : 80;
const needsPort = !isRunpodProxy && worker.port !== defaultPort;
const portStr = needsPort ? `:${worker.port}` : '';
const needsPort = !isRunpodProxy && resolvedPort !== defaultPort;
const portStr = needsPort ? `:${resolvedPort}` : '';
return `${protocol}://${finalHost}${portStr}${endpoint}`;
}
@@ -440,21 +636,21 @@ class DistributedExtension {
async launchWorker(workerId) {
const worker = this.config.workers.find(w => w.id === workerId);
const launchBtn = document.querySelector(`#controls-${workerId} button`);
// If worker is disabled, enable it first
if (!worker.enabled) {
await this.updateWorkerEnabled(workerId, true);
// Update the checkbox UI
const checkbox = document.getElementById(`gpu-${workerId}`);
if (checkbox) {
checkbox.checked = true;
}
this.updateSummary();
}
// Re-query button AFTER updateWorkerEnabled (which may re-render sidebar)
const launchBtn = document.querySelector(`#controls-${workerId} button`);
this.ui.updateStatusDot(workerId, "#f0ad4e", "Launching...", true);
this.state.setWorkerLaunching(workerId, true);
@@ -798,7 +994,8 @@ class DistributedExtension {
return true;
}
// Otherwise check by host (backward compatibility)
const host = worker.host || window.location.hostname;
const parsed = this._parseHostInput(worker.host || window.location.hostname);
const host = parsed.host || window.location.hostname;
return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname;
}
@@ -1018,10 +1215,16 @@ class DistributedExtension {
const workerType = document.getElementById(`worker-type-${workerId}`).value;
const isRemote = workerType === 'remote' || workerType === 'cloud';
const isCloud = workerType === 'cloud';
const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
const port = parseInt(document.getElementById(`port-${workerId}`).value);
const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
const parsedHost = isRemote ? this._parseHostInput(rawHost) : { host: window.location.hostname, port: null };
const host = isRemote ? parsedHost.host : window.location.hostname;
let port = parseInt(document.getElementById(`port-${workerId}`).value);
const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value);
const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value;
if (isRemote && Number.isFinite(parsedHost.port)) {
port = parsedHost.port;
}
// Validate
if (!name.trim()) {
@@ -1333,4 +1536,4 @@ app.registerExtension({
async setup() {
new DistributedExtension();
}
});
});
+73 -31
View File
@@ -43,6 +43,8 @@ export async function renderSidebarContent(extension, el) {
// Preload data outside render
await extension.loadConfig();
await extension.loadManagedWorkers();
await extension.refreshTunnelStatus();
extension.tunnelElements = {};
el.innerHTML = '';
@@ -148,57 +150,67 @@ export async function renderSidebarContent(extension, el) {
// Settings section
const settingsSection = document.createElement("div");
settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;";
// Top separator only; spacing handled by the clickable toggle area for equal top/bottom spacing
settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;";
// Settings header with toggle
// Settings header with toggle (full-area clickable between separators)
const settingsToggleArea = document.createElement("div");
// Equal spacing above header (to top separator) and below header (to bottom separator)
settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;";
const settingsHeader = document.createElement("div");
settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;";
settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;";
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);
settingsToggleArea.appendChild(settingsHeader);
// Hover effect for toggle area
settingsToggleArea.onmouseover = () => { workerSettingsToggle.style.color = "#fff"; };
settingsToggleArea.onmouseout = () => { workerSettingsToggle.style.color = "#888"; };
// Hover effect for header
settingsHeader.onmouseover = () => {
workerSettingsToggle.style.color = "#fff";
};
settingsHeader.onmouseout = () => {
workerSettingsToggle.style.color = "#888";
};
// A small separator shown only when collapsed (to make the section boundary obvious)
const settingsSeparator = document.createElement("div");
// No margin so the bottom spacing is controlled by settingsToggleArea padding-bottom
settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;";
// 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;";
settingsDiv.style.cssText = "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;";
// Toggle functionality
let settingsExpanded = false;
settingsHeader.onclick = () => {
settingsToggleArea.onclick = () => {
settingsExpanded = !settingsExpanded;
if (settingsExpanded) {
settingsContent.style.maxHeight = "200px";
settingsContent.style.opacity = "1";
workerSettingsToggle.style.transform = "rotate(90deg)";
settingsSeparator.style.display = "none";
} else {
settingsContent.style.maxHeight = "0";
settingsContent.style.opacity = "0";
workerSettingsToggle.style.transform = "rotate(0deg)";
settingsSeparator.style.display = "block";
}
};
// Debug mode setting
// Section: General
const generalLabel = document.createElement("div");
generalLabel.textContent = "GENERAL";
generalLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;";
const debugGroup = document.createElement("div");
debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
debugGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;";
const debugCheckbox = document.createElement("input");
debugCheckbox.type = "checkbox";
debugCheckbox.id = "setting-debug";
@@ -209,13 +221,14 @@ export async function renderSidebarContent(extension, el) {
debugLabel.htmlFor = "setting-debug";
debugLabel.textContent = "Debug Mode";
debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
debugLabel.title = "Enable verbose logging in the browser console.";
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;";
autoLaunchGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;";
const autoLaunchCheckbox = document.createElement("input");
autoLaunchCheckbox.type = "checkbox";
@@ -227,13 +240,14 @@ export async function renderSidebarContent(extension, el) {
autoLaunchLabel.htmlFor = "setting-auto-launch";
autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup";
autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
autoLaunchLabel.title = "Start local worker processes automatically when the master starts.";
autoLaunchGroup.appendChild(autoLaunchCheckbox);
autoLaunchGroup.appendChild(autoLaunchLabel);
// Stop workers on exit setting
// Stop workers on exit setting (under General)
const stopOnExitGroup = document.createElement("div");
stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
stopOnExitGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;";
const stopOnExitCheckbox = document.createElement("input");
stopOnExitCheckbox.type = "checkbox";
@@ -245,28 +259,56 @@ export async function renderSidebarContent(extension, el) {
stopOnExitLabel.htmlFor = "setting-stop-on-exit";
stopOnExitLabel.textContent = "Stop Local Workers on Master Exit";
stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
stopOnExitLabel.title = "Stop local worker processes automatically when the master exits.";
stopOnExitGroup.appendChild(stopOnExitCheckbox);
stopOnExitGroup.appendChild(stopOnExitLabel);
settingsDiv.appendChild(generalLabel);
settingsDiv.appendChild(debugGroup);
settingsDiv.appendChild(autoLaunchGroup);
settingsDiv.appendChild(stopOnExitGroup);
// Worker Timeout setting (seconds)
// Section: Timeouts
const timeoutsLabel = document.createElement("div");
timeoutsLabel.textContent = "TIMEOUTS";
timeoutsLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;";
const timeoutGroup = document.createElement("div");
timeoutGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 6px;";
const timeoutLabel = document.createElement("label");
timeoutLabel.htmlFor = "setting-worker-timeout";
timeoutLabel.textContent = "Worker Timeout";
timeoutLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: default;";
timeoutLabel.title = "Seconds without a heartbeat before a worker is considered timed out and its tasks are requeued. Default 60. For WAN, consider 300–600.";
const timeoutInput = document.createElement("input");
timeoutInput.type = "number";
timeoutInput.id = "setting-worker-timeout";
timeoutInput.min = "10";
timeoutInput.step = "1";
timeoutInput.style.cssText = "width: 80px; padding: 2px 6px; background: #222; color: #ddd; border: 1px solid #333; border-radius: 3px;";
timeoutInput.value = (extension.config?.settings?.worker_timeout_seconds ?? 60);
timeoutInput.onchange = (e) => {
const v = parseInt(e.target.value, 10);
if (!Number.isFinite(v) || v <= 0) return;
extension._updateSetting('worker_timeout_seconds', v);
};
timeoutGroup.appendChild(timeoutLabel);
timeoutGroup.appendChild(timeoutInput);
settingsDiv.appendChild(timeoutsLabel);
settingsDiv.appendChild(timeoutGroup);
settingsContent.appendChild(settingsDiv);
settingsSection.appendChild(settingsHeader);
settingsSection.appendChild(settingsToggleArea);
settingsSection.appendChild(settingsSeparator);
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);
@@ -274,4 +316,4 @@ export async function renderSidebarContent(extension, el) {
// Always reset the rendering flag
extension._isRendering = false;
}
}
}
+78 -15
View File
@@ -3,28 +3,32 @@ import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS, UI_COLORS, TIMEOUTS } from './
const cardConfigs = {
master: {
checkbox: {
enabled: false,
checked: true,
disabled: true,
opacity: 0.6,
title: "Master node is always enabled"
enabled: true,
masterToggle: true,
title: "Toggle master participation in workloads"
},
statusDot: {
color: STATUS_COLORS.ONLINE_GREEN,
title: 'Online',
id: 'master-status',
initialColor: (_, extension) => extension.isMasterParticipating() ? STATUS_COLORS.ONLINE_GREEN : STATUS_COLORS.DISABLED_GRAY,
initialTitle: (_, extension) => extension.isMasterParticipating() ? 'Master participating' : 'Master orchestrator only',
dynamic: true
},
infoText: (data, extension) => {
const cudaDevice = extension.config?.master?.cuda_device ?? extension.masterCudaDevice;
const cudaInfo = cudaDevice !== undefined ? `CUDA ${cudaDevice} • ` : '';
const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80');
return `<strong id="master-name-display">${data?.name || extension.config?.master?.name || "Master"}</strong><br><small style="color: ${UI_COLORS.MUTED_TEXT};"><span id="master-cuda-info">${cudaInfo}Port ${port}</span></small>`;
const participationEnabled = extension.isMasterParticipationEnabled();
const fallbackActive = extension.isMasterFallbackActive();
let delegateBadge = '';
if (!participationEnabled) {
delegateBadge = fallbackActive
? `<br><small style="color: #6bd06b;">Fallback active • Master executing</small>`
: `<br><small style="color: ${UI_COLORS.SECONDARY_TEXT};">Orchestrator-only mode</small>`;
}
return `<strong id="master-name-display">${data?.name || extension.config?.master?.name || "Master"}</strong><br><small style="color: ${UI_COLORS.MUTED_TEXT};"><span id="master-cuda-info">${cudaInfo}Port ${port}</span></small>${delegateBadge}`;
},
controls: {
type: 'info',
text: 'Master',
style: "background-color: #333; color: #999;"
type: 'master'
},
settings: {
formType: 'master',
@@ -868,7 +872,36 @@ export class DistributedUI {
if (config?.opacity) checkbox.style.opacity = config.opacity;
if (config?.title) column.title = config.title;
if (config?.enabled && !config?.disabled && data?.id) {
const isMasterToggle = config?.masterToggle && typeof extension.isMasterParticipating === 'function';
if (isMasterToggle) {
const participationEnabled = extension.isMasterParticipationEnabled();
const fallbackActive = extension.isMasterFallbackActive();
const buildTitle = (enabled, fallback) => {
if (enabled) {
return "Master participating • Click to switch to orchestrator-only";
}
if (fallback) {
return "No workers selected • Master fallback execution active";
}
return "Master orchestrator-only • Click to re-enable participation";
};
checkbox.checked = participationEnabled;
checkbox.style.pointerEvents = "none";
column.style.cursor = "pointer";
column.title = buildTitle(participationEnabled, fallbackActive);
column.onclick = async (event) => {
if (event) {
event.stopPropagation();
event.preventDefault();
}
const nextState = !extension.isMasterParticipationEnabled();
const nextFallback = !nextState && extension.enabledWorkers.length === 0;
checkbox.checked = nextState;
column.title = buildTitle(nextState, nextFallback);
await extension.updateMasterParticipation(nextState);
};
} else if (config?.enabled && !config?.disabled && data?.id) {
checkbox.style.pointerEvents = "none";
column.style.cursor = "pointer";
column.onclick = async () => {
@@ -890,10 +923,10 @@ export class DistributedUI {
let id = config.id;
if (typeof config.initialColor === 'function') {
color = config.initialColor(data);
color = config.initialColor(data, extension);
}
if (typeof config.initialTitle === 'function') {
title = config.initialTitle(data);
title = config.initialTitle(data, extension);
}
if (typeof config.id === 'function') {
id = config.id(data);
@@ -940,7 +973,29 @@ export class DistributedUI {
const controlsWrapper = document.createElement("div");
controlsWrapper.style.cssText = this.styles.controlsWrapper;
if (config.dynamic && data) {
if (config.type === 'master') {
const participationEnabled = extension.isMasterParticipationEnabled();
const fallbackActive = extension.isMasterFallbackActive();
let message;
const badge = document.createElement("div");
badge.style.cssText = this.styles.infoBox;
if (fallbackActive) {
message = "No workers selected. Master fallback execution active.";
badge.textContent = message;
badge.style.backgroundColor = "#243024";
badge.style.color = "#6bd06b";
badge.style.border = "1px solid #335533";
} else if (!participationEnabled) {
message = "Master disabled: running as orchestrator only.";
badge.textContent = message;
badge.style.backgroundColor = "#3a3a3a";
badge.style.color = "#ffcc66";
} else {
message = "Master participating in workflows.";
badge.textContent = message;
}
controlsWrapper.appendChild(badge);
} else if (config.dynamic && data) {
if (isRemote) {
const isCloud = data.type === 'cloud';
const workerTypeText = isCloud ? "Cloud worker" : "Remote worker";
@@ -1029,6 +1084,14 @@ export class DistributedUI {
const hostResult = this.createFormGroup("Host:", extension.config?.master?.host || "", "master-host", "text", "Auto-detect if empty");
settingsForm.appendChild(hostResult.group);
// Cloudflare tunnel toggle (simple button inside master settings)
const tunnelBtn = this.createButton("Enable Cloudflare Tunnel", (e) => extension.handleTunnelToggle(e.target), "background-color: #665533;");
tunnelBtn.id = "cloudflare-tunnel-button";
tunnelBtn.style.cssText = BUTTON_STYLES.base + " background-color: #665533; margin: 4px 0 -5px 0;";
settingsForm.appendChild(tunnelBtn);
extension.tunnelElements = { button: tunnelBtn };
extension.updateTunnelUIElements();
const saveBtn = this.createButton("Save", async () => {
const nameInput = document.getElementById('master-name');
+46 -3
View File
@@ -105,13 +105,15 @@ export function findImageReferences(extension, apiPrompt) {
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, {
// Normalize to forward slashes so subfolder/filename derivation is consistent on Windows
const normalizedValue = cleanValue.replace(/\\/g, '/');
if (imageExtensions.test(normalizedValue)) {
images.set(normalizedValue, {
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");
extension.log(`Found media reference: ${normalizedValue} in node ${nodeId} (${node.class_type})`, "debug");
}
}
}
@@ -153,6 +155,47 @@ export function findCollectorUpstreamNodes(apiPrompt, collectorIds) {
return connected;
}
/**
* Find all downstream nodes from the provided collector IDs.
* Used to keep only post-collection nodes when master runs in orchestrator mode.
* @param {Object} apiPrompt
* @param {Array<string>} collectorIds
* @returns {Set<string>}
*/
export function findCollectorDownstreamNodes(apiPrompt, collectorIds) {
const adjacency = new Map();
for (const [nodeId, node] of Object.entries(apiPrompt)) {
if (!node.inputs) continue;
for (const inputValue of Object.values(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2) {
const sourceId = String(inputValue[0]);
if (!adjacency.has(sourceId)) {
adjacency.set(sourceId, new Set());
}
adjacency.get(sourceId).add(String(nodeId));
}
}
}
const connected = new Set(collectorIds);
const queue = [...collectorIds];
while (queue.length > 0) {
const current = queue.shift();
const dependents = adjacency.get(current);
if (!dependents) continue;
dependents.forEach((dependentId) => {
if (!connected.has(dependentId)) {
connected.add(dependentId);
queue.push(dependentId);
}
});
}
return connected;
}
/**
* Prune workflow to only include nodes connected to distributed nodes
* @param {Object} apiPrompt - The full workflow API prompt
+1
View File
@@ -39,6 +39,7 @@ except ImportError:
def monitor_and_run(master_pid, command):
"""Run command and monitor master process."""
# Start the actual worker process
print(f"[Distributed] Launching worker command: {' '.join(command)}")
worker_process = subprocess.Popen(command)
print(f"[Distributed] Started worker PID: {worker_process.pid}")
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff