Compare commits
117
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1d4dcb985e | ||
|
|
b6d4756392 | ||
|
|
044c54b780 | ||
|
|
b768e7fc8a | ||
|
|
9805049268 | ||
|
|
399ed7c7d1 | ||
|
|
64ba4b75db | ||
|
|
d353d118d7 | ||
|
|
5dd10cef39 | ||
|
|
8eba225f5e | ||
|
|
9bc2cfc43e | ||
|
|
370b48438b | ||
|
|
71c7a99e1e | ||
|
|
def7c0df1d | ||
|
|
5f5320b34e | ||
|
|
5a33da828c | ||
|
|
ef8e0542c8 | ||
|
|
89b7812fb4 | ||
|
|
f13b953a82 | ||
|
|
1e3e094cdb | ||
|
|
7bcf9290ba | ||
|
|
6aec567071 | ||
|
|
2fec61d567 | ||
|
|
528cd2e895 | ||
|
|
0572e084e3 | ||
|
|
bc9bba0c76 | ||
|
|
701c6de9eb | ||
|
|
c846aac0e6 | ||
|
|
0af3df071d | ||
|
|
f7c7b960c7 | ||
|
|
c45b50830e | ||
|
|
213ac38227 | ||
|
|
9ae354a788 | ||
|
|
c1928ebd09 | ||
|
|
98a63938c8 | ||
|
|
27b73fd280 | ||
|
|
48c8792af2 | ||
|
|
4b8836cd03 | ||
|
|
b69fc40fbd | ||
|
|
e5eb2c8b86 | ||
|
|
b90c3e42ac | ||
|
|
187b38d4f5 | ||
|
|
39fb4db486 | ||
|
|
763e238120 | ||
|
|
857aa80174 | ||
|
|
287e4c05fa | ||
|
|
58f26a106f | ||
|
|
138cac3eb7 | ||
|
|
0e07868b98 | ||
|
|
a8be7c7ed9 | ||
|
|
c9ca7f5390 | ||
|
|
c860df0aa7 | ||
|
|
05662eb7eb | ||
|
|
8441c92d95 | ||
|
|
4b513d49c3 | ||
|
|
f91d2d579d | ||
|
|
1ff863fb5d | ||
|
|
74bba258da | ||
|
|
bb559e8ce0 | ||
|
|
c032de224b | ||
|
|
68f8627b0b | ||
|
|
0e3c04fba7 | ||
|
|
e076cf3455 | ||
|
|
5f3d358f61 | ||
|
|
1a6bff3ab1 | ||
|
|
d7aec0b0fe | ||
|
|
7cd014fc0e | ||
|
|
51e9a03dd5 | ||
|
|
421ca2e023 | ||
|
|
ae3201accb | ||
|
|
ae1d1a7f89 | ||
|
|
407e68816d | ||
|
|
e4b0fe11fd | ||
|
|
960c44b095 | ||
|
|
35925c1e1b | ||
|
|
feb56e234b | ||
|
|
f5d7307aef | ||
|
|
386f7ae2d4 | ||
|
|
a161213a8f | ||
|
|
a884527ad3 | ||
|
|
53be3e0c18 | ||
|
|
8f87d1ac29 | ||
|
|
b54f2be19c | ||
|
|
1e9ba2383d | ||
|
|
33fbe4cb62 | ||
|
|
5882a1064e | ||
|
|
b646f0b41c | ||
|
|
3c988daad1 | ||
|
|
500a5b38b2 | ||
|
|
a49c7a8ed5 | ||
|
|
8c41749ea6 | ||
|
|
394a2d8a9a | ||
|
|
e13ea47a93 | ||
|
|
81e644e661 | ||
|
|
b13c78ad71 | ||
|
|
520393b150 | ||
|
|
8c029e5702 | ||
|
|
0391605c80 | ||
|
|
784aa3a06b | ||
|
|
f589a723fe | ||
|
|
d247746e22 | ||
|
|
0fe8485c6c | ||
|
|
ccb00a7142 | ||
|
|
0acf557ba5 | ||
|
|
9926253c8b | ||
|
|
285710bb28 | ||
|
|
a58ca1fc1d | ||
|
|
f5a1e59d65 | ||
|
|
ff02c88224 | ||
|
|
b7c7c033d5 | ||
|
|
f7efde4eb5 | ||
|
|
f4f7e48d89 | ||
|
|
0a480c143a | ||
|
|
fb46d32e39 | ||
|
|
acf166fe6c | ||
|
|
e5a2ab8fc5 | ||
|
|
26aaa1f1c8 |
@@ -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 }}
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
bin/
|
||||
logs/
|
||||
gpu_config.json
|
||||
__pycache__/
|
||||
@@ -1,15 +1,15 @@
|
||||
<div align="center">
|
||||
<img width="320" src="https://github.com/user-attachments/assets/537d38cc-2e34-4527-bca7-0d79f4120cce" />
|
||||
<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>
|
||||
|
||||
> **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network.**
|
||||
> **A powerful extension for ComfyUI that enables distributed and parallel processing across multiple GPUs and machines. Generate more images and videos and accelerate your upscaling workflows by leveraging all available GPU resources in your network and cloud.**
|
||||
|
||||

|
||||
|
||||
@@ -18,14 +18,14 @@
|
||||
## Key Features
|
||||
|
||||
#### Parallel Workflow Processing
|
||||
- Run workflows on multiple GPUs simultaneously with varied seeds
|
||||
- Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
|
||||
- Scale output with more workers
|
||||
- Supports images and videos
|
||||
|
||||
#### 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
|
||||
@@ -34,11 +34,26 @@
|
||||
|
||||
---
|
||||
|
||||
## Worker Types
|
||||
|
||||
<img width="200" align="right" alt="ComfyUI_temp_khvcc_00034_@0 25x" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
|
||||
|
||||
ComfyUI Distributed supports three types of workers:
|
||||
|
||||
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
|
||||
- **Remote Workers** - GPUs on other computers in your network
|
||||
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
|
||||
|
||||
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
|
||||
|
||||
---
|
||||
|
||||
## Requirements
|
||||
|
||||
- ComfyUI
|
||||
> Note: Desktop app not currently supported
|
||||
- Multiple GPUs
|
||||
- Multiple NVIDIA GPUs
|
||||
> No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
|
||||
- That's it
|
||||
|
||||
---
|
||||
@@ -50,9 +65,18 @@
|
||||
git clone https://github.com/robertvoy/ComfyUI-Distributed.git
|
||||
```
|
||||
|
||||
2. **Restart ComfyUI** - If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master
|
||||
2. **Restart ComfyUI**
|
||||
- If you'll be using remote/cloud workers, add `--enable-cors-header` to your launch arguments on the master
|
||||
|
||||
3. Read the [setup guide](https://github.com/robertvoy/ComfyUI-Distributed/blob/feature/cloud-workers/docs/worker-setup-guides.md) for adding workers
|
||||
3. Read the [setup guide](/docs/worker-setup-guides.md) for adding 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.
|
||||
|
||||
---
|
||||
|
||||
@@ -87,8 +111,8 @@ 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
|
||||
Accelerate the process by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
|
||||
### Distributed Image Upscaling
|
||||
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
|
||||
|
||||

|
||||
|
||||
@@ -100,6 +124,34 @@ Accelerate the process by distributing tiles across multiple workers, with speed
|
||||
4. Configure tile settings
|
||||
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.
|
||||
|
||||

|
||||
|
||||
> [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
|
||||
@@ -143,4 +195,12 @@ This software is provided "as is" without any warranties, express or implied, in
|
||||
---
|
||||
|
||||
## Support the Project
|
||||
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
|
||||
|
||||
<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
|
||||
|
||||
|
||||
+1
-1
@@ -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())}")
|
||||
|
||||
+844
-152
File diff suppressed because it is too large
Load Diff
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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.
|
||||
@@ -1,15 +1,14 @@
|
||||
## Automating ComfyUI Model Downloads
|
||||
> This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM).
|
||||
|
||||
This guide will walk you through creating a shell script to automatically download the necessary models for your ComfyUI workflow, leveraging an advanced Large Language Model (LLM).
|
||||
1. In ComfyUI (on your local machine), export your workflow as an API workflow
|
||||
2. Copy the below prompt and upload the API workflow to an LLM **that has access to the internet**
|
||||
|
||||
<details>
|
||||
<summary><strong>📋 Click to expand the full prompt</strong></summary>
|
||||
|
||||
1. Launched the Pod with these Environment Variables:
|
||||
- CIVITAI_API_TOKEN ([get your token here](https://civitai.com/user/account))
|
||||
- HF_API_TOKEN ([get your token here](https://huggingface.co/settings/tokens))
|
||||
2. In ComfyUI, export your workflow as an API workflow
|
||||
3. Copy the below prompt and upload the API workflow to a LLM **that has access to the internet**
|
||||
```
|
||||
Create a sh script that will download the models from this workflow into the correct folders. For reference, these are the paths:
|
||||
|
||||
base_path: /workspace/ComfyUI
|
||||
checkpoints: models/checkpoints/
|
||||
clip: models/clip/
|
||||
@@ -25,32 +24,30 @@ text_encoders: models/text_encoders/
|
||||
unet: models/unet/
|
||||
upscale_models: models/upscale_models/
|
||||
vae: models/vae/
|
||||
|
||||
---
|
||||
|
||||
Important:
|
||||
Make sure you find the correct URLs for the models online.
|
||||
Use comfy cli to download the models: `comfy model download --url <URL> [--relative-path <PATH>] [--set-civitai-api-token <TOKEN>] [--set-hf-api-token <TOKEN>]`
|
||||
Make sure you add `--set-civitai-api-token $CIVITAI_API_TOKEN` for CivitAI download and `--set-hf-api-token $HF_API_TOKEN` for Hugging Face downloads.
|
||||
|
||||
---
|
||||
|
||||
Example:
|
||||
|
||||
#!/bin/bash
|
||||
|
||||
# Download from CivitAI
|
||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||
|
||||
# Download model from Hugging Face
|
||||
comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN
|
||||
|
||||
# If a model in the workflow was in a subfolder
|
||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints/SDXL --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||
|
||||
```
|
||||
4. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh`
|
||||
5. Upload it onto your Runpod instance, into `/workspace`
|
||||
|
||||
</details>
|
||||
|
||||
3. Review the LLMs output to make sure all download links are correct and save it as a .sh file, for example `download_models.sh`
|
||||
4. Launch the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) with these Environment Variables:
|
||||
- `CIVITAI_API_TOKEN`: [get your token here](https://civitai.com/user/account)
|
||||
- `HF_API_TOKEN`: [get your token here](https://huggingface.co/settings/tokens)
|
||||
5. Upload the .sh file to your Runpod instance, into `/workspace`
|
||||
6. Then run these commands:
|
||||
- `chmod 755 /workspace/download_models.sh`
|
||||
- `/workspace/download_models.sh`
|
||||
7. Confirm each model name (sometimes you might need to rename them to match the name on your local machine)
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||

|
||||
|
||||
**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.
|
||||
+62
-25
@@ -1,11 +1,19 @@
|
||||
## ComfyUI Distributed Worker Setup Guide
|
||||
|
||||
ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs, whether they're in the same machine, on your local network, or in the cloud.
|
||||
## Worker Setup Guide
|
||||
|
||||
**Master**: The main ComfyUI instance that coordinates and distributes work. This is where you load workflows, manage the queue, and view results.
|
||||
|
||||
**Worker**: A ComfyUI instance that receives and processes tasks from the master. Workers handle just the GPU computation and send results back to the master. You can have multiple workers connected to a single master, each utilizing their own GPU.
|
||||
|
||||
<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
|
||||
@@ -13,8 +21,14 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
|
||||
- **Cloud workers**: GPUs hosted on cloud services like Runpod
|
||||
|
||||
## Local workers
|
||||
|
||||
<img align="right" width="200" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
|
||||
|
||||
> These are added automatically on first launch, but you can add them manually if you need to.
|
||||
|
||||
|
||||
📺 [Watch Tutorial](https://youtu.be/p6eE3IlAbOs?si=K7Km0_flmPHwRQwz&t=43)
|
||||
|
||||
1. **Open** the Distributed GPU panel.
|
||||
2. **Click** "Add Worker" in the UI.
|
||||
3. **Configure** your local worker:
|
||||
@@ -25,8 +39,14 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
|
||||
4. **Save** and launch the local worker.
|
||||
|
||||
## Remote workers
|
||||
|
||||
<img align="right" width="200" src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
|
||||
|
||||
|
||||
> ComfyUI instances running on completely different computers on your network. These allow you to harness GPU power from other machines. Remote workers must be manually started on their respective computers and are connected via IP address.
|
||||
|
||||
📺 [Watch Tutorial](https://youtu.be/p6eE3IlAbOs?si=Oxj3EzPyf4jKDvfG&t=140)
|
||||
|
||||
**On the Remote Worker Machine:**
|
||||
1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!**
|
||||
- This ComfyUI instance will serve as a worker for your main master.
|
||||
@@ -49,38 +69,55 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
|
||||
6. **Save** the remote worker configuration.
|
||||
|
||||
## Cloud workers
|
||||
> ComfyUI instances running on a cloud service like [Runpod](https://get.runpod.io/0bw29uf3ug0p).
|
||||
|
||||
<img align="right" width="200" src="https://github.com/user-attachments/assets/a053f3ae-22f0-4e1c-8f2e-f26a1f660adf" />
|
||||
|
||||
> ComfyUI instances running on a cloud service like Runpod.
|
||||
|
||||
### Deploy Cloud Worker on Runpod
|
||||
|
||||
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
|
||||
|
||||
**On Runpod:**
|
||||
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
|
||||
|
||||
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
||||
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later.
|
||||
3. Now go to Pods and find a suitable GPU for your workflows.
|
||||
4. Choose the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template and make sure your network drive is mounted.
|
||||
> To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8.
|
||||
6. Launch your pod.
|
||||
7. Access your pod using JupyterLabs.
|
||||
8. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod).
|
||||
> You can use [this guide](model-download-script.md) to make this process easy for you. It will generate a shell script that will automatically download the models you need for a given workflow.
|
||||
8. If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you git clone ComfyUI-Distributed into custom_nodes. ⚠️ **Required!**
|
||||
9. Download any additional custom nodes using the ComfyUI Manager.
|
||||
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=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:
|
||||
- CIVITAI_API_TOKEN: [get your token here](https://civitai.com/user/account)
|
||||
- HF_API_TOKEN: [get your token here](https://huggingface.co/settings/tokens)
|
||||
- SAGE_ATTENTION: optional optimisation (set to true/false)
|
||||
5. Deploy your pod.
|
||||
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
|
||||
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
```
|
||||
# Download from CivitAI
|
||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||
# Download model from Hugging Face
|
||||
comfy model download --url https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors --relative-path /workspace/ComfyUI/models/unet --set-hf-api-token $HF_API_TOKEN
|
||||
```
|
||||
> ℹ️ Use [this guide](model-download-script.md) to make this process easy. It will generate a shell script that automatically downloads the models for a given workflow.
|
||||
9. Access ComfyUI through the Runpod URL.
|
||||
10. Download any additional custom nodes you need using the ComfyUI Manager.
|
||||
|
||||
**On the Main Machine:**
|
||||
1. Launch a Cloudflare tunnel
|
||||
1. **Launch** a Cloudflare tunnel.
|
||||
- Download from here: [https://github.com/cloudflare/cloudflared/releases](https://github.com/cloudflare/cloudflared/releases)
|
||||
- Then run, for example: `cloudflared-windows-amd64.exe tunnel --url http://localhost:8188`
|
||||
> Cloudflare tunnels create secure connections without exposing ports directly to the internet
|
||||
2. Copy the Cloudflare address
|
||||
> ℹ️ Cloudflare tunnels create secure connections without exposing ports directly to the internet and are required for Cloud Workers.
|
||||
2. **Copy** the Cloudflare address
|
||||
3. **Launch** ComfyUI with `--enable-cors-header` launch argument.
|
||||
4. **Open** the Distributed GPU panel (sidebar on the left).
|
||||
5. **Edit** the Master's host address and replace it with the Cloudflare address.
|
||||
- **Click** "Add Worker."
|
||||
- **Choose** "Cloud".
|
||||
- **Configure** your cloud worker:
|
||||
- **Host**: The Runpod address
|
||||
- **Port**: 443
|
||||
- **Save** the remote worker configuration.
|
||||
5. **Edit** the Master's settings to change the host address to the Cloudflare address.
|
||||
6. **Click** "Add Worker."
|
||||
7. **Choose** "Cloud".
|
||||
8. **Configure** your cloud worker:
|
||||
- **Host**: The ComfyUI Runpod address. For example: `wcegfo9tbbml9l-8188.proxy.runpod.net`
|
||||
- **Port**: 443
|
||||
9. **Save** the remote worker configuration.
|
||||
|
||||
---
|
||||
|
||||
@@ -91,7 +128,7 @@ ComfyUI-Distributed allows you to distribute your workflows across multiple GPUs
|
||||
- If your cloud platform doesn't provide a secure connection, use Cloudflare to create a tunnel for the worker. Each GPU needs their own tunnel for their respective port.
|
||||
- For example: `./cloudflared tunnel --url http://localhost:8188`
|
||||
1. **Launch** ComfyUI with the `--listen --enable-cors-header` arguments. ⚠️ **Required!**
|
||||
2. **Add** workers in the UI panel if the remote machine has more than one GPU.
|
||||
2. **Add** workers in the UI panel if the cloud machine has more than one GPU.
|
||||
- Make sure that they also have `--listen` set in `Extra Args`.
|
||||
- Then launch them.
|
||||
|
||||
|
||||
+3
-3
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally and remotely "
|
||||
version = "1.0.1"
|
||||
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"
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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',
|
||||
@@ -697,7 +701,7 @@ export class DistributedUI {
|
||||
// Create powered by Runpod text (initially hidden)
|
||||
const runpodText = document.createElement("a");
|
||||
runpodText.id = `runpod-text-${worker.id}`;
|
||||
runpodText.href = "https://get.runpod.io/0bw29uf3ug0p";
|
||||
runpodText.href = "https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers";
|
||||
runpodText.target = "_blank";
|
||||
runpodText.textContent = "Deploy Cloud Worker with Runpod";
|
||||
runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;";
|
||||
@@ -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');
|
||||
@@ -1233,4 +1296,4 @@ export class DistributedUI {
|
||||
|
||||
return card;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+46
-3
@@ -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
|
||||
|
||||
@@ -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
+1038
-1088
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user