Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0e9ee7fbac | ||
|
|
5386de10e3 | ||
|
|
41f9e44945 | ||
|
|
6e698512b7 | ||
|
|
79cc9f5ad5 | ||
|
|
4eba87cba3 | ||
|
|
27e08de94c | ||
|
|
c8453a139b | ||
|
|
aae831e1e5 | ||
|
|
e7ab67733b | ||
|
|
dd55ff740e | ||
|
|
a6d0b82d35 |
@@ -19,6 +19,6 @@ jobs:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
|
||||
---
|
||||
|
||||
## Key Features
|
||||
## Key Features
|
||||
|
||||
#### Parallel Workflow Processing
|
||||
- Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
|
||||
@@ -27,22 +27,12 @@
|
||||
- Intelligent distribution
|
||||
- Handles single images and videos
|
||||
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
- JSON configuration with UI controls
|
||||
|
||||
---
|
||||
|
||||
## Current Architecture
|
||||
|
||||
- Workflow-level load balancing is controlled by **Distributed Collector** via the `load_balance` toggle.
|
||||
- There is **no Distributed Queue node** anymore.
|
||||
- With `load_balance=true`, orchestration selects one least-busy execution participant:
|
||||
- If master participation is enabled, master is included as a candidate.
|
||||
- If master is in orchestrator-only mode, only workers are considered.
|
||||
|
||||
---
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
- JSON configuration with UI controls
|
||||
|
||||
---
|
||||
|
||||
## Worker Types
|
||||
|
||||
@@ -61,7 +51,6 @@ ComfyUI Distributed supports three types of workers:
|
||||
## Requirements
|
||||
|
||||
- ComfyUI
|
||||
> Note: Desktop app not currently supported
|
||||
- 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
|
||||
@@ -92,19 +81,19 @@ Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a sp
|
||||
|
||||
## Workflow Examples
|
||||
|
||||
### Basic Parallel Generation
|
||||
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
|
||||
### Basic Parallel Generation
|
||||
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
|
||||
|
||||

|
||||
|
||||
> [Download workflow](/workflows/distributed-txt2img.json)
|
||||
|
||||
1. Open your ComfyUI workflow
|
||||
2. Add **Distributed Seed** → connect to sampler's seed
|
||||
3. Add **Distributed Collector** → after VAE Decode
|
||||
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
|
||||
5. Enable workers in the UI
|
||||
6. Run the workflow!
|
||||
2. Add **Distributed Seed** → connect to sampler's seed
|
||||
3. Add **Distributed Collector** → after VAE Decode
|
||||
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
|
||||
5. Enable workers in the UI
|
||||
6. Run the workflow!
|
||||
|
||||
### Parallel WAN Generation
|
||||
Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
|
||||
@@ -122,12 +111,12 @@ 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 Image Upscaling
|
||||
Accelerate Ultimate SD Upscaler 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.
|
||||
|
||||

|
||||
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
|
||||
1. Load your image
|
||||
2. Upscale with ESRGAN or similar
|
||||
@@ -153,7 +142,7 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
|
||||
|
||||
---
|
||||
|
||||
## Developer API
|
||||
## Developer API
|
||||
|
||||
Control your distributed cluster programmatically without opening the browser.
|
||||
|
||||
@@ -161,32 +150,32 @@ Control your distributed cluster programmatically without opening the browser.
|
||||
* **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).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
|
||||
> **⚠️ 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).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
|
||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||
|
||||
@@ -206,7 +195,7 @@ No, it does not speed up the generation of a single image or video. Instead, it
|
||||
|
||||
<details>
|
||||
<summary>Does it work with the ComfyUI desktop app?</summary>
|
||||
Currently, it is not compatible with the ComfyUI desktop app.
|
||||
Yes, it does now.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
@@ -249,3 +238,4 @@ Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
+3
-1
@@ -217,7 +217,7 @@ async def distributed_queue_endpoint(request):
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
try:
|
||||
prompt_id, worker_count = await orchestrate_distributed_execution(
|
||||
prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
|
||||
payload.prompt,
|
||||
payload.workflow_meta,
|
||||
payload.client_id,
|
||||
@@ -227,6 +227,8 @@ async def distributed_queue_endpoint(request):
|
||||
)
|
||||
return web.json_response({
|
||||
"prompt_id": prompt_id,
|
||||
"number": prompt_number,
|
||||
"node_errors": node_errors,
|
||||
"worker_count": worker_count,
|
||||
"auto_prepare_supported": True,
|
||||
})
|
||||
|
||||
@@ -13,7 +13,7 @@ from ..utils.constants import (
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_master_url
|
||||
from ..utils.network import build_master_url, build_master_callback_url
|
||||
from ..utils.trace_logger import trace_debug
|
||||
from .schemas import parse_positive_float, parse_positive_int
|
||||
from .orchestration.dispatch import (
|
||||
@@ -144,6 +144,7 @@ async def _prepare_worker_payload(
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
trace_execution_id,
|
||||
worker_prep_semaphore,
|
||||
@@ -153,6 +154,11 @@ async def _prepare_worker_payload(
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
worker_master_url = build_master_callback_url(
|
||||
worker,
|
||||
config=config,
|
||||
prompt_server_instance=prompt_server,
|
||||
)
|
||||
|
||||
worker_type = str(worker.get("type") or "local").strip().lower()
|
||||
is_remote_like = bool(worker.get("host")) and worker_type != "local"
|
||||
@@ -167,7 +173,7 @@ async def _prepare_worker_payload(
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
worker_master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
@@ -202,7 +208,7 @@ async def orchestrate_distributed_execution(
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int]: (prompt_id, worker_count)
|
||||
tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
@@ -311,8 +317,18 @@ async def orchestrate_distributed_execution(
|
||||
|
||||
if not job_id_map:
|
||||
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
|
||||
prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
|
||||
return prompt_id, 0
|
||||
queue_result = await queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
queue_result["number"],
|
||||
0,
|
||||
queue_result.get("node_errors", {}),
|
||||
)
|
||||
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
@@ -360,6 +376,7 @@ async def orchestrate_distributed_execution(
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_semaphore,
|
||||
@@ -385,9 +402,17 @@ async def orchestrate_distributed_execution(
|
||||
]
|
||||
)
|
||||
|
||||
prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
|
||||
queue_result = await queue_prompt_payload(
|
||||
master_prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
prompt_id = queue_result["prompt_id"]
|
||||
prompt_number = queue_result["number"]
|
||||
node_errors = queue_result.get("node_errors", {})
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
|
||||
)
|
||||
return prompt_id, len(worker_payloads)
|
||||
return prompt_id, prompt_number, len(worker_payloads), node_errors
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.4.1"
|
||||
version = "1.4.3"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
|
||||
@@ -162,7 +162,7 @@ def _load_job_routes_module():
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
|
||||
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -221,7 +221,7 @@ job_routes = _load_job_routes_module()
|
||||
|
||||
|
||||
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_id(self):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
@@ -233,12 +233,14 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
with patch.object(
|
||||
job_routes,
|
||||
"orchestrate_distributed_execution",
|
||||
new=AsyncMock(return_value=("prompt_123", 2)),
|
||||
new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
|
||||
):
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
||||
self.assertEqual(response.payload.get("number"), 42)
|
||||
self.assertEqual(response.payload.get("node_errors"), {})
|
||||
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
||||
|
||||
async def test_distributed_queue_missing_prompt_returns_400(self):
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
def _load_async_helpers_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
|
||||
package_name = "dist_async_helpers_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
execution_module = types.ModuleType("execution")
|
||||
|
||||
async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
|
||||
return (True, None, ["9"], {})
|
||||
|
||||
execution_module.validate_prompt = _validate_prompt
|
||||
execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
|
||||
sys.modules["execution"] = execution_module
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
trigger_on_prompt=lambda payload: payload,
|
||||
number=12,
|
||||
prompt_queue=_PromptQueue(),
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_loop = lambda: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module, prompt_server
|
||||
|
||||
|
||||
async_helpers, prompt_server = _load_async_helpers_module()
|
||||
|
||||
|
||||
class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
|
||||
result = await async_helpers.queue_prompt_payload(
|
||||
{"1": {"class_type": "Node"}},
|
||||
workflow_meta={"id": "workflow-1"},
|
||||
client_id="client-1",
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
|
||||
self.assertIsInstance(result["prompt_id"], str)
|
||||
self.assertTrue(result["prompt_id"])
|
||||
self.assertEqual(result["number"], 12)
|
||||
self.assertEqual(result["node_errors"], {})
|
||||
|
||||
self.assertEqual(prompt_server.number, 13)
|
||||
self.assertEqual(len(prompt_server.prompt_queue.items), 1)
|
||||
queued_item = prompt_server.prompt_queue.items[0]
|
||||
self.assertEqual(queued_item[0], 12)
|
||||
extra_data = queued_item[3]
|
||||
self.assertEqual(extra_data["client_id"], "client-1")
|
||||
self.assertIn("create_time", extra_data)
|
||||
self.assertIsInstance(extra_data["create_time"], int)
|
||||
self.assertGreater(extra_data["create_time"], 0)
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -90,14 +90,48 @@ class NetworkHelpersTests(unittest.TestCase):
|
||||
"https://master.example.com",
|
||||
)
|
||||
|
||||
def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
|
||||
cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8188",
|
||||
)
|
||||
|
||||
def test_build_master_url_keeps_explicit_port_in_host(self):
|
||||
cfg = {"master": {"host": "192.168.68.56:8001"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
def test_build_master_url_falls_back_to_server_address(self):
|
||||
cfg = {"master": {"host": ""}}
|
||||
cfg = {"master": {"host": "", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8190",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_uses_loopback_for_local_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8001",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_process_module(module_filename: str):
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
|
||||
package_name = "dist_proc_testpkg"
|
||||
module_name = module_filename[:-3]
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
process_pkg = types.ModuleType(f"{package_name}.workers.process")
|
||||
process_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers.process"] = process_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
process_module = types.ModuleType(f"{package_name}.utils.process")
|
||||
process_module.get_python_executable = lambda: "/usr/bin/test-python"
|
||||
sys.modules[f"{package_name}.utils.process"] = process_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.process.{module_name}",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
root_discovery_module = _load_process_module("root_discovery.py")
|
||||
launch_builder_module = _load_process_module("launch_builder.py")
|
||||
|
||||
|
||||
class ComfyRootDiscoveryTests(unittest.TestCase):
|
||||
def test_prefers_loaded_comfyui_module_path(self):
|
||||
discovery = root_discovery_module.ComfyRootDiscovery()
|
||||
server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py")
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/opt/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(sys.modules, {"server": server_module}, clear=False), \
|
||||
patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \
|
||||
patch.dict(root_discovery_module.os.environ, {}, clear=True):
|
||||
self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI")
|
||||
|
||||
|
||||
class LaunchCommandBuilderTests(unittest.TestCase):
|
||||
def test_inherits_runtime_layout_args_for_desktop(self):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
runtime_args = Namespace(
|
||||
listen="127.0.0.1",
|
||||
base_directory="C:/Users/test/ComfyUI",
|
||||
temp_directory=None,
|
||||
input_directory="C:/Users/test/ComfyUI/input",
|
||||
output_directory="C:/Users/test/ComfyUI/output",
|
||||
user_directory="C:/Users/test/ComfyUI/user",
|
||||
front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app",
|
||||
extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]],
|
||||
enable_manager=True,
|
||||
disable_manager_ui=False,
|
||||
enable_manager_legacy_ui=False,
|
||||
windows_standalone_build=True,
|
||||
log_stdout=True,
|
||||
verbose="INFO",
|
||||
enable_cors_header="*",
|
||||
)
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
comfy_cli_args = types.ModuleType("comfy.cli_args")
|
||||
comfy_cli_args.args = runtime_args
|
||||
|
||||
worker_config = {
|
||||
"port": 9001,
|
||||
"extra_args": "--preview-method auto",
|
||||
}
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/desktop/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"comfy": comfy_module, "comfy.cli_args": comfy_cli_args},
|
||||
clear=False,
|
||||
), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists):
|
||||
cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI")
|
||||
|
||||
self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"])
|
||||
self.assertIn("--listen", cmd)
|
||||
self.assertIn("127.0.0.1", cmd)
|
||||
self.assertIn("--base-directory", cmd)
|
||||
self.assertIn("C:/Users/test/ComfyUI", cmd)
|
||||
self.assertIn("--input-directory", cmd)
|
||||
self.assertIn("--output-directory", cmd)
|
||||
self.assertIn("--user-directory", cmd)
|
||||
self.assertIn("--front-end-root", cmd)
|
||||
self.assertIn("--extra-model-paths-config", cmd)
|
||||
self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd)
|
||||
self.assertIn("--enable-manager", cmd)
|
||||
self.assertIn("--windows-standalone-build", cmd)
|
||||
self.assertIn("--log-stdout", cmd)
|
||||
self.assertIn("--disable-auto-launch", cmd)
|
||||
self.assertIn("--enable-cors-header", cmd)
|
||||
self.assertIn("*", cmd)
|
||||
self.assertIn("--port", cmd)
|
||||
self.assertIn("9001", cmd)
|
||||
self.assertNotIn("--auto-launch", cmd)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+16
-2
@@ -3,6 +3,7 @@ Async helper utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import execution
|
||||
import server
|
||||
@@ -104,7 +105,12 @@ class PromptValidationError(RuntimeError):
|
||||
super().__init__(f"Invalid prompt: {merged}")
|
||||
|
||||
|
||||
async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
||||
async def queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta=None,
|
||||
client_id=None,
|
||||
include_queue_metadata=False,
|
||||
):
|
||||
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
||||
payload = {"prompt": prompt_obj}
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
@@ -117,7 +123,7 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
raise PromptValidationError(error_payload, node_errors)
|
||||
|
||||
extra_data = {}
|
||||
extra_data = {"create_time": int(time.time() * 1000)}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if client_id:
|
||||
@@ -132,4 +138,12 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
||||
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)
|
||||
|
||||
if include_queue_metadata:
|
||||
return {
|
||||
"prompt_id": prompt_id,
|
||||
"number": number,
|
||||
"node_errors": {},
|
||||
}
|
||||
|
||||
return prompt_id
|
||||
|
||||
+43
-8
@@ -66,6 +66,25 @@ def normalize_host(value):
|
||||
return host.split("/")[0]
|
||||
|
||||
|
||||
def _split_host_and_port(host):
|
||||
if not host:
|
||||
return host, None
|
||||
|
||||
if host.startswith("["):
|
||||
match = re.match(r"^(\[[^\]]+\])(?::(\d+))?$", host)
|
||||
if match:
|
||||
parsed_port = int(match.group(2)) if match.group(2) else None
|
||||
return match.group(1), parsed_port
|
||||
return host, None
|
||||
|
||||
if host.count(":") == 1:
|
||||
candidate_host, candidate_port = host.rsplit(":", 1)
|
||||
if candidate_port.isdigit():
|
||||
return candidate_host, int(candidate_port)
|
||||
|
||||
return host, None
|
||||
|
||||
|
||||
def build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
host = (worker.get("host") or "").strip()
|
||||
@@ -126,12 +145,7 @@ def build_master_url(config=None, prompt_server_instance=None):
|
||||
prompt_server_instance = prompt_server_instance or server.PromptServer.instance
|
||||
master_cfg = (config or {}).get("master", {}) or {}
|
||||
configured_host = (master_cfg.get("host") or "").strip()
|
||||
configured_port = master_cfg.get("port")
|
||||
default_port = getattr(prompt_server_instance, "port", 8188) or 8188
|
||||
try:
|
||||
port = int(configured_port or default_port)
|
||||
except (TypeError, ValueError):
|
||||
port = int(default_port)
|
||||
runtime_port = getattr(prompt_server_instance, "port", 8188) or 8188
|
||||
|
||||
def _needs_https(hostname):
|
||||
hostname = hostname.lower()
|
||||
@@ -149,10 +163,11 @@ def build_master_url(config=None, prompt_server_instance=None):
|
||||
if configured_host.startswith(("http://", "https://")):
|
||||
return configured_host.rstrip("/")
|
||||
|
||||
host = configured_host
|
||||
host, explicit_port = _split_host_and_port(configured_host)
|
||||
port = explicit_port if explicit_port is not None else int(runtime_port)
|
||||
scheme = "https" if _needs_https(host) or port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
if configured_port is None and scheme == "https" and _needs_https(host):
|
||||
if explicit_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}"
|
||||
@@ -160,7 +175,27 @@ def build_master_url(config=None, prompt_server_instance=None):
|
||||
address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1"
|
||||
if address in ("0.0.0.0", "::"):
|
||||
address = "127.0.0.1"
|
||||
port = int(runtime_port)
|
||||
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_master_callback_url(worker, config=None, prompt_server_instance=None):
|
||||
"""Build the callback URL a specific worker should use to reach the master."""
|
||||
prompt_server_instance = prompt_server_instance or server.PromptServer.instance
|
||||
|
||||
worker_type = str((worker or {}).get("type") or "").strip().lower()
|
||||
worker_host = normalize_host((worker or {}).get("host"))
|
||||
local_hosts = {"", "localhost", "127.0.0.1", "::1", "[::1]", "0.0.0.0"}
|
||||
|
||||
is_local_worker = worker_type == "local" or worker_host in local_hosts
|
||||
if is_local_worker:
|
||||
port = int(getattr(prompt_server_instance, "port", 8188) or 8188)
|
||||
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}://127.0.0.1{port_part}"
|
||||
|
||||
return build_master_url(config=config, prompt_server_instance=prompt_server_instance)
|
||||
|
||||
@@ -4,7 +4,7 @@ import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js';
|
||||
import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
|
||||
|
||||
export function setupInterceptor(extension) {
|
||||
api.queuePrompt = async (number, prompt) => {
|
||||
api.queuePrompt = async (number, prompt, ...rest) => {
|
||||
if (extension.isEnabled) {
|
||||
const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
|
||||
const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
|
||||
@@ -18,7 +18,7 @@ export function setupInterceptor(extension) {
|
||||
return result;
|
||||
}
|
||||
}
|
||||
return extension.originalQueuePrompt(number, prompt);
|
||||
return extension.originalQueuePrompt(number, prompt, ...rest);
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -10,6 +10,62 @@ from ...utils.process import get_python_executable
|
||||
class LaunchCommandBuilder:
|
||||
"""Build command-lines for launching worker ComfyUI processes."""
|
||||
|
||||
def _extend_arg(self, cmd, flag, value):
|
||||
if value in (None, "", [], ()):
|
||||
return
|
||||
cmd.extend([flag, str(value)])
|
||||
|
||||
def _extend_grouped_args(self, cmd, flag, values):
|
||||
for group in values or []:
|
||||
flattened = [str(item) for item in group if item]
|
||||
if flattened:
|
||||
cmd.append(flag)
|
||||
cmd.extend(flattened)
|
||||
|
||||
def _get_runtime_args(self):
|
||||
try:
|
||||
from comfy.cli_args import args
|
||||
return args
|
||||
except Exception as exc:
|
||||
debug_log(f"Could not read current ComfyUI CLI args for worker launch: {exc}")
|
||||
return None
|
||||
|
||||
def _build_runtime_launch_args(self):
|
||||
args = self._get_runtime_args()
|
||||
if args is None:
|
||||
return []
|
||||
|
||||
inherited = []
|
||||
self._extend_arg(inherited, "--listen", getattr(args, "listen", None))
|
||||
self._extend_arg(inherited, "--base-directory", getattr(args, "base_directory", None))
|
||||
self._extend_arg(inherited, "--temp-directory", getattr(args, "temp_directory", None))
|
||||
self._extend_arg(inherited, "--input-directory", getattr(args, "input_directory", None))
|
||||
self._extend_arg(inherited, "--output-directory", getattr(args, "output_directory", None))
|
||||
self._extend_arg(inherited, "--user-directory", getattr(args, "user_directory", None))
|
||||
self._extend_arg(inherited, "--front-end-root", getattr(args, "front_end_root", None))
|
||||
self._extend_grouped_args(
|
||||
inherited,
|
||||
"--extra-model-paths-config",
|
||||
getattr(args, "extra_model_paths_config", None),
|
||||
)
|
||||
|
||||
if getattr(args, "enable_manager", False):
|
||||
inherited.append("--enable-manager")
|
||||
if getattr(args, "disable_manager_ui", False):
|
||||
inherited.append("--disable-manager-ui")
|
||||
if getattr(args, "enable_manager_legacy_ui", False):
|
||||
inherited.append("--enable-manager-legacy-ui")
|
||||
if getattr(args, "windows_standalone_build", False):
|
||||
inherited.append("--windows-standalone-build")
|
||||
if getattr(args, "log_stdout", False):
|
||||
inherited.append("--log-stdout")
|
||||
|
||||
verbose = getattr(args, "verbose", None)
|
||||
if verbose and verbose != "INFO":
|
||||
inherited.extend(["--verbose", str(verbose)])
|
||||
|
||||
return inherited
|
||||
|
||||
def _find_windows_terminal(self):
|
||||
"""Find Windows Terminal executable."""
|
||||
possible_paths = [
|
||||
@@ -39,10 +95,19 @@ class LaunchCommandBuilder:
|
||||
cmd = [
|
||||
get_python_executable(),
|
||||
main_py,
|
||||
"--port",
|
||||
str(worker_config["port"]),
|
||||
"--enable-cors-header",
|
||||
]
|
||||
cmd.extend(self._build_runtime_launch_args())
|
||||
cmd.extend(["--port", str(worker_config["port"])])
|
||||
|
||||
current_args = self._get_runtime_args()
|
||||
current_cors = getattr(current_args, "enable_cors_header", None) if current_args else None
|
||||
cmd.append("--enable-cors-header")
|
||||
if current_cors is not None:
|
||||
cmd.append(str(current_cors))
|
||||
|
||||
if "--disable-auto-launch" not in cmd:
|
||||
cmd.append("--disable-auto-launch")
|
||||
|
||||
debug_log(f"Using main.py: {main_py}")
|
||||
else:
|
||||
error_msg = f"Could not find main.py in {comfy_root}\n"
|
||||
|
||||
@@ -33,6 +33,7 @@ class ProcessLifecycle:
|
||||
env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0))
|
||||
env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
env["COMFYUI_MASTER_PID"] = str(os.getpid())
|
||||
env["COMFYUI_IS_WORKER"] = "1"
|
||||
|
||||
cmd = self._manager.build_launch_command(worker_config, comfy_root)
|
||||
cwd = comfy_root
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
|
||||
@@ -6,6 +7,21 @@ from ...utils.logging import debug_log, log
|
||||
class ComfyRootDiscovery:
|
||||
"""Resolve the ComfyUI root directory across local and container layouts."""
|
||||
|
||||
def _find_root_from_loaded_modules(self):
|
||||
"""Use already-imported ComfyUI modules to locate the runtime root."""
|
||||
for module_name in ("server", "folder_paths", "main"):
|
||||
module = sys.modules.get(module_name)
|
||||
module_file = getattr(module, "__file__", None)
|
||||
if not module_file:
|
||||
continue
|
||||
|
||||
candidate = os.path.dirname(os.path.abspath(module_file))
|
||||
if os.path.exists(os.path.join(candidate, "main.py")):
|
||||
debug_log(f"Found ComfyUI root via loaded module {module_name}: {candidate}")
|
||||
return candidate
|
||||
|
||||
return None
|
||||
|
||||
def find_comfy_root(self):
|
||||
# Start from current file location.
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
@@ -17,12 +33,17 @@ class ComfyRootDiscovery:
|
||||
debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}")
|
||||
return env_root
|
||||
|
||||
# Method 2: Try going up from custom_nodes directory.
|
||||
# Method 2: Inspect the already-loaded ComfyUI runtime modules.
|
||||
runtime_root = self._find_root_from_loaded_modules()
|
||||
if runtime_root:
|
||||
return runtime_root
|
||||
|
||||
# Method 3: Try going up from custom_nodes directory.
|
||||
if os.path.exists(os.path.join(potential_root, "main.py")):
|
||||
debug_log(f"Found ComfyUI root via directory traversal: {potential_root}")
|
||||
return potential_root
|
||||
|
||||
# Method 3: Look for common Docker paths.
|
||||
# Method 4: Look for common Docker paths.
|
||||
docker_paths = [
|
||||
"/basedir",
|
||||
"/ComfyUI",
|
||||
@@ -37,7 +58,7 @@ class ComfyRootDiscovery:
|
||||
debug_log(f"Found ComfyUI root in Docker path: {path}")
|
||||
return path
|
||||
|
||||
# Method 4: Search upwards for main.py.
|
||||
# Method 5: Search upwards for main.py.
|
||||
search_dir = current_dir
|
||||
for _ in range(5):
|
||||
if os.path.exists(os.path.join(search_dir, "main.py")):
|
||||
@@ -48,7 +69,7 @@ class ComfyRootDiscovery:
|
||||
break
|
||||
search_dir = parent
|
||||
|
||||
# Method 5: Try to import and use folder_paths.
|
||||
# Method 6: Try to import and use folder_paths.
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
|
||||
Reference in New Issue
Block a user