Compare commits

...
Author SHA1 Message Date
Robert Wojciechowski a2b341fb99 fix: hide Distributed Value backing fields in Nodes 2.0
Use native widget visibility without converting backing fields into sockets. Preserve widget types, serializers and options identity; cover edits, reload, worker refresh and linked widgets.
2026-10-10 22:43:49 +00:00
Robert Wojciechowski 9de7e158e1 feat: migrate distributed nodes to native ComfyUI V3 API 2026-10-10 22:06:46 +00:00
Robert Wojciechowski 32ac027e0e docs: align selected copy with current behavior 2026-07-12 11:13:22 +00:00
Robert Wojciechowski fec959e1ac chore: bump version to 1.4.7 2026-07-12 10:40:27 +00:00
Robert Wojciechowski cedc8d4591 feat: support audio-only distributed workflows (#87)
Support image, audio, or combined distributed collection with guarded audio-only worker, master, endpoint, preview, and delegate-only paths.
2026-07-12 20:36:43 +10:00
27 changed files with 1810 additions and 212 deletions
+2 -2
View File
@@ -147,7 +147,7 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
Control your distributed cluster programmatically without opening the browser. Control your distributed cluster programmatically without opening the browser.
* **Endpoint:** `POST /distributed/queue` * **Endpoint:** `POST /distributed/queue`
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID. * **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md) * **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). > **⚠️ 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).
@@ -175,7 +175,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
| **Distributed Value** | Outputs per-worker override values with fallback to default | | **Distributed Value** | Outputs per-worker override values with fallback to default |
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers | | **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
| **Image Batch Divider** | Splits image batches for multi-GPU output | | **Image Batch Divider** | Splits image batches for multi-GPU output |
| **Audio Batch Divider** | Splits audio batches for multi-GPU output | | **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode | | **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 | | **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
+15 -22
View File
@@ -1,29 +1,22 @@
# Import everything needed from the main module """ComfyUI-Distributed's native V3 extension entrypoint."""
from .distributed import ( from comfy_api.v0_0_2 import ComfyExtension, io
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
)
# Import utilities from .nodes.v3 import NODES
from .utils.config import ensure_config_exists, CONFIG_FILE from .runtime.bootstrap import initialize
from .utils.logging import debug_log
# Import distributed upscale nodes WEB_DIRECTORY = './web'
from .nodes.distributed_upscale import (
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
)
WEB_DIRECTORY = "./web"
ensure_config_exists() class DistributedExtension(ComfyExtension):
async def on_load(self) -> None:
initialize()
# Merge node mappings async def get_node_list(self) -> list[type[io.ComfyNode]]:
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS} return list(NODES)
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
debug_log("Loaded Distributed nodes.") async def comfy_entrypoint() -> DistributedExtension:
debug_log(f"Config file: {CONFIG_FILE}") return DistributedExtension()
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
+11 -4
View File
@@ -295,17 +295,24 @@ async def job_complete_endpoint(request):
errors.append("worker_id: expected non-empty string") errors.append("worker_id: expected non-empty string")
if not isinstance(batch_idx, int) or batch_idx < 0: if not isinstance(batch_idx, int) or batch_idx < 0:
errors.append("batch_idx: expected non-negative integer") errors.append("batch_idx: expected non-negative integer")
if not isinstance(image_payload, str) or not image_payload.strip(): image_provided = image_payload is not None
audio_provided = audio_payload is not None
if not image_provided and not audio_provided:
errors.append("expected at least one of image or audio")
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
errors.append("image: expected non-empty base64 PNG string") errors.append("image: expected non-empty base64 PNG string")
if audio_payload is not None and not isinstance(audio_payload, dict): if audio_provided and not isinstance(audio_payload, dict):
errors.append("audio: expected object when provided") errors.append("audio: expected object when provided")
if not isinstance(is_last, bool): if not isinstance(is_last, bool):
errors.append("is_last: expected boolean") errors.append("is_last: expected boolean")
if errors: if errors:
return await handle_api_error(request, errors, 400) return await handle_api_error(request, errors, 400)
tensor = _decode_canonical_png_tensor(image_payload) try:
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
except ValueError as exc:
return await handle_api_error(request, exc, 400)
multi_job_id = job_id.strip() multi_job_id = job_id.strip()
worker_id = worker_id.strip() worker_id = worker_id.strip()
+31 -10
View File
@@ -351,16 +351,33 @@ def prune_prompt_for_worker(prompt_obj):
downstream = _find_downstream_nodes(prompt_obj, [dist_id]) downstream = _find_downstream_nodes(prompt_obj, [dist_id])
has_removed_downstream = any(node_id != dist_id for node_id in downstream) has_removed_downstream = any(node_id != dist_id for node_id in downstream)
if has_removed_downstream: if has_removed_downstream:
preview_id = next_id() original_node = prompt_obj.get(str(dist_id), {})
pruned_prompt[preview_id] = { class_type = original_node.get("class_type")
"inputs": { inputs = original_node.get("inputs", {})
"images": [dist_id, 0], image_connected = class_type != "DistributedCollector" or (
}, isinstance(inputs.get("images"), list)
"class_type": "PreviewImage", and len(inputs["images"]) == 2
"_meta": { )
"title": "Preview Image (auto-added)", audio_connected = (
}, class_type == "DistributedCollector"
} and isinstance(inputs.get("audio"), list)
and len(inputs["audio"]) == 2
)
if image_connected:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {"images": [dist_id, 0]},
"class_type": "PreviewImage",
"_meta": {"title": "Preview Image (auto-added)"},
}
elif audio_connected:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {"audio": [dist_id, 1]},
"class_type": "PreviewAudio",
"_meta": {"title": "Preview Audio (auto-added)"},
}
return pruned_prompt return pruned_prompt
@@ -400,6 +417,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
collector_entry = pruned_prompt.get(collector_id) collector_entry = pruned_prompt.get(collector_id)
if not collector_entry: if not collector_entry:
continue continue
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
original_images = original_inputs.get("images")
if not (isinstance(original_images, list) and len(original_images) == 2):
continue
placeholder_id = next_id() placeholder_id = next_id()
pruned_prompt[placeholder_id] = { pruned_prompt[placeholder_id] = {
"class_type": "DistributedEmptyImage", "class_type": "DistributedEmptyImage",
+1 -1
View File
@@ -1,7 +1,7 @@
# conftest.py — project-level pytest configuration. # conftest.py — project-level pytest configuration.
# #
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports # Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
# (from .distributed import ...) that fail when pytest tries to import it as a # (from .nodes.v3 import ...) that fail when pytest tries to import it as a
# standalone module during Package.setup() for the root package node. # standalone module during Package.setup() for the root package node.
# #
# Fix: patch Package.setup() to skip the root-package's __init__.py import. # Fix: patch Package.setup() to skip the root-package's __init__.py import.
-51
View File
@@ -1,51 +0,0 @@
"""
ComfyUI-Distributed: thin entry point.
All implementation lives in workers/, nodes/, api/.
"""
import atexit
import os
import server
from .utils.config import ensure_config_exists
from .utils.logging import debug_log
from .utils.network import cleanup_client_session
from .workers import get_worker_manager
from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
from .upscale.job_store import ensure_tile_jobs_initialized
from .nodes import (
NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS,
ImageBatchDivider,
DistributedCollectorNode,
DistributedSeed,
DistributedModelName,
DistributedValue,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from . import api # noqa: F401 - triggers all @routes.* registrations
from .api.queue_orchestration import ensure_distributed_state
ensure_config_exists()
# Aiohttp session cleanup
async def _cleanup_session():
await cleanup_client_session()
atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup
# Initialize distributed job state on prompt_server
prompt_server = server.PromptServer.instance
ensure_distributed_state(prompt_server)
ensure_tile_jobs_initialized()
# Worker startup
if not os.environ.get('COMFYUI_IS_WORKER'):
atexit.register(sync_cleanup)
delayed_auto_launch()
register_async_signals()
+18 -8
View File
@@ -19,8 +19,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI: - `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`). - Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
- Resolves enabled/selected workers. - Resolves enabled/selected workers.
- Pings workers (`GET /prompt`) to include only reachable ones. - Probes and dispatches workers through `/distributed/worker_ws` by default.
- Dispatches the workflow to workers (`POST /prompt`). - If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
- Queues the master workflow in ComfyUI’s prompt queue. - Queues the master workflow in ComfyUI’s prompt queue.
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run. - If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
@@ -61,7 +61,8 @@ Queue a workflow for distributed execution.
#### Fields #### Fields
- `prompt` (required unless `workflow.prompt` is present, object) - `prompt` (required unless `workflow.prompt` is present, object)
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`. - A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
- This is not the normal visual workflow export from the ComfyUI editor.
- `workflow` (optional, object) - `workflow` (optional, object)
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`. - Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
- If you don’t care about UI metadata, you can omit it. - If you don’t care about UI metadata, you can omit it.
@@ -123,10 +124,14 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
## Worker requirements (important) ## Worker requirements (important)
For a worker to participate, it must be reachable from the master: For a worker to participate, it must be reachable from the master. By default:
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
If `settings.websocket_orchestration=false`:
- Health check: `GET <worker-base>/prompt` must return HTTP 200. - Health check: `GET <worker-base>/prompt` must return HTTP 200.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow. - Dispatch: `POST <worker-base>/prompt` must accept the prompt.
Also, for collector-based flows: Also, for collector-based flows:
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
## Examples ## Examples
### 1) Minimal `curl` ### 1) Minimal `curl` request envelope
```bash ```bash
curl -X POST "http://127.0.0.1:8188/distributed/queue" \ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -209,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
-d @payload.json -d @payload.json
``` ```
Where `payload.json` contains at least: `payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
```json ```json
{ {
"prompt": { "prompt": {
"1": {"class_type": "KSampler", "inputs": {} } "<node_id>": {
"class_type": "<node class>",
"inputs": {"<required_input>": "<value or connection>"}
}
}, },
"enabled_worker_ids": [], "enabled_worker_ids": [],
"client_id": "external-client" "client_id": "external-client"
} }
``` ```
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
### 2) Python (`requests`) ### 2) Python (`requests`)
```python ```python
+53
View File
@@ -0,0 +1,53 @@
# Native V3 node API
The package registers one `ComfyExtension` through `comfy_entrypoint()` and
imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
that API; there is no V1 registration fallback.
## Compatibility
- All eight node IDs, display names, categories, visible input order, defaults
and output order are retained. Existing execution algorithms remain in the
private collector, utilities and upscale modules.
- The image/audio dividers explicitly declare ten typed outputs. Their existing
frontend extensions still show only the selected number of outputs. This
replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
IMAGE/AUDIO socket indices or the ten returned values.
- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
its existing prompt-input names and Python defaults via `accept_all_inputs`;
it does not become visible widgets.
- The collector retains list-input handling. Upscale retains its always-changing
fingerprint and creates a private runtime helper for each V3 execution, so
mutable helper state is not attached to sanitized V3 class clones.
- Routes, distributed state and the existing worker startup/shutdown hooks are
initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
Worker mode still suppresses automatic worker launch. `distributed.py` is no
longer an entrypoint.
## Verification
Run ordinary unit tests from this repository:
```bash
python -m pytest tests -q -o addopts=
```
Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
```bash
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
/path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
```
The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
all node schemas with the V1 fixture, checks all five bundled workflows' node
IDs and socket/link contracts, exercises injected worker metadata, image/audio
lists, collector aggregation and divider outputs, rejects invalid upscale enums,
and decodes an actual preview PNG referenced by executor history.
The upscale GPU/model boundary is mocked to check argument forwarding and
per-execution helper isolation. Full checkpoint inference, browser canvas
acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
workers or model downloads are started; preview files use a temporary scratch
directory. The test does not install this branch into a live custom-node folder.
+2 -2
View File
@@ -79,7 +79,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk) 📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
**On Runpod:** **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!** > If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account. 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes). 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).
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
- SAGE_ATTENTION: optional optimisation (set to true/false) - SAGE_ATTENTION: optional optimisation (set to true/false)
5. Deploy your pod. 5. Deploy your pod.
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system. 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: 7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
``` ```
# Download from CivitAI # 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 comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
+1 -31
View File
@@ -1,31 +1 @@
from .utilities import ( """Private execution helpers; public registration lives in nodes.v3."""
DistributedSeed,
DistributedModelName,
DistributedValue,
ImageBatchDivider,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from .collector import DistributedCollectorNode
NODE_CLASS_MAPPINGS = {
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed,
"DistributedModelName": DistributedModelName,
"DistributedValue": DistributedValue,
"ImageBatchDivider": ImageBatchDivider,
"AudioBatchDivider": AudioBatchDivider,
"DistributedEmptyImage": DistributedEmptyImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedCollector": "Distributed Collector",
"DistributedSeed": "Distributed Seed",
"DistributedModelName": "Distributed Model Name",
"DistributedValue": "Distributed Value",
"ImageBatchDivider": "Image Batch Divider",
"AudioBatchDivider": "Audio Batch Divider",
"DistributedEmptyImage": "Distributed Empty Image",
}
+73 -45
View File
@@ -29,7 +29,6 @@ class DistributedCollectorNode:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"images": ("IMAGE",),
"load_balance": ( "load_balance": (
"BOOLEAN", "BOOLEAN",
{ {
@@ -38,7 +37,10 @@ class DistributedCollectorNode:
}, },
), ),
}, },
"optional": { "audio": ("AUDIO",) }, "optional": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
},
"hidden": { "hidden": {
"multi_job_id": ("STRING", {"default": ""}), "multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}), "is_worker": ("BOOLEAN", {"default": False}),
@@ -102,8 +104,9 @@ class DistributedCollectorNode:
return None return None
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate} return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False): def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
images = self._normalize_images_input(images) if images is not None:
images = self._normalize_images_input(images)
audio = self._normalize_audio_input(audio) audio = self._normalize_audio_input(audio)
load_balance = self._unwrap_list_input(load_balance) load_balance = self._unwrap_list_input(load_balance)
multi_job_id = self._unwrap_list_input(multi_job_id) multi_job_id = self._unwrap_list_input(multi_job_id)
@@ -115,6 +118,14 @@ class DistributedCollectorNode:
pass_through = self._unwrap_list_input(pass_through) pass_through = self._unwrap_list_input(pass_through)
delegate_only = self._unwrap_list_input(delegate_only) delegate_only = self._unwrap_list_input(delegate_only)
remote_only_master = (
bool(multi_job_id)
and not is_worker
and (delegate_only or is_master_delegate_only())
)
if images is None and audio is None and not remote_only_master:
raise ValueError("DistributedCollector requires at least one image or audio input")
# Create empty audio if not provided # Create empty audio if not provided
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100} empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -141,41 +152,56 @@ class DistributedCollectorNode:
return result return result
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id): async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Send image batch to master via canonical JSON envelopes.""" """Send an image batch, optionally with audio, or an audio-only completion."""
batch_size = image_batch.shape[0]
if batch_size == 0:
return
encoded_audio = encode_audio_payload(audio) encoded_audio = encode_audio_payload(audio)
session = await get_client_session() session = await get_client_session()
url = f"{master_url}/distributed/job_complete" url = f"{master_url}/distributed/job_complete"
for batch_idx in range(batch_size):
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
byte_io = io.BytesIO()
img.save(byte_io, format='PNG', compress_level=0)
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
payload = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": int(batch_idx),
"image": f"data:image/png;base64,{encoded_image}",
"is_last": bool(batch_idx == batch_size - 1),
}
if payload["is_last"] and encoded_audio is not None:
payload["audio"] = encoded_audio
payloads = []
batch_size = 0 if image_batch is None else image_batch.shape[0]
if batch_size == 0:
if encoded_audio is None:
raise ValueError("Worker completion requires image or audio data")
payloads.append(
{
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": 0,
"audio": encoded_audio,
"is_last": True,
}
)
else:
for batch_idx in range(batch_size):
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
byte_io = io.BytesIO()
img.save(byte_io, format='PNG', compress_level=0)
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
payload = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": int(batch_idx),
"image": f"data:image/png;base64,{encoded_image}",
"is_last": bool(batch_idx == batch_size - 1),
}
if payload["is_last"] and encoded_audio is not None:
payload["audio"] = encoded_audio
payloads.append(payload)
for payload in payloads:
timeout_seconds = 60 if "image" in payload else 600
try: try:
async with session.post( async with session.post(
url, url,
json=payload, json=payload,
timeout=aiohttp.ClientTimeout(total=60), timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response: ) as response:
response.raise_for_status() response.raise_for_status()
except Exception as e: except Exception as e:
log(f"Worker - Failed to send canonical image envelope to master: {e}") media_type = "image/audio" if "image" in payload else "audio-only"
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
debug_log(f"Worker - Full error details: URL={url}") debug_log(f"Worker - Full error details: URL={url}")
raise # Re-raise to handle at caller level raise
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None): def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
"""Combine audio from master and workers into a single audio output. """Combine audio from master and workers into a single audio output.
@@ -257,8 +283,8 @@ class DistributedCollectorNode:
images_on_cpu, images_on_cpu,
delegate_mode: bool, delegate_mode: bool,
fallback_images, fallback_images,
) -> torch.Tensor: ):
"""Assemble final tensor: master first, then workers in enabled order.""" """Assemble final tensor, or return None when the job contains only audio."""
ordered_tensors = [] ordered_tensors = []
if not delegate_mode and images_on_cpu is not None: if not delegate_mode and images_on_cpu is not None:
for i in range(master_batch_size): for i in range(master_batch_size):
@@ -289,15 +315,15 @@ class DistributedCollectorNode:
if cpu_tensors: if cpu_tensors:
return ensure_contiguous(torch.cat(cpu_tensors, dim=0)) return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
elif fallback_images is not None: if fallback_images is not None:
return ensure_contiguous(fallback_images) return ensure_contiguous(fallback_images)
else: return None
raise ValueError("No image data collected from master or workers")
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False): async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
if is_worker: if is_worker:
# Worker mode: send images and audio to master in a single batch # Worker mode: send images and audio to master in a single batch
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") image_count = 0 if images is None else images.shape[0]
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id) await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
return (images, audio if audio is not None else self.EMPTY_AUDIO) return (images, audio if audio is not None else self.EMPTY_AUDIO)
else: else:
@@ -332,14 +358,15 @@ class DistributedCollectorNode:
master_audio = None master_audio = None
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers") debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
else: else:
images_on_cpu = images.cpu() if images is None:
master_batch_size = images.shape[0] images_on_cpu = None
master_batch_size = 0
else:
images_on_cpu = ensure_contiguous(images.cpu())
master_batch_size = images.shape[0]
master_audio = audio # Keep master's audio for later master_audio = audio # Keep master's audio for later
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...") debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
# Ensure master images are contiguous
images_on_cpu = ensure_contiguous(images_on_cpu)
# Initialize storage for collected images and audio # Initialize storage for collected images and audio
worker_images = {} # Dict to store images by worker_id and index worker_images = {} # Dict to store images by worker_id and index
@@ -511,18 +538,19 @@ class DistributedCollectorNode:
if multi_job_id in prompt_server.distributed_pending_jobs: if multi_job_id in prompt_server.distributed_pending_jobs:
del prompt_server.distributed_pending_jobs[multi_job_id] del prompt_server.distributed_pending_jobs[multi_job_id]
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
try: try:
combined = self._reorder_and_combine_tensors( combined = self._reorder_and_combine_tensors(
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
) )
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total " if combined is None:
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") debug_log(f"Master - Job {multi_job_id} complete with audio only")
else:
# Combine audio from master and workers debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers) f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
return (combined, combined_audio) return (combined, combined_audio)
except Exception as e: except Exception as e:
log(f"Master - Error combining images: {e}") log(f"Master - Error combining images: {e}")
# Return just the master images as fallback # Preserve collected audio even when image assembly fails.
return (images, audio if audio is not None else self.EMPTY_AUDIO) return (images, combined_audio)
+6 -15
View File
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
""" """
Distributed version of Ultimate SD Upscale (No Upscale). Distributed version of Ultimate SD Upscale (No Upscale).
Supports three processing modes: Supports two currently selected processing modes:
1. Single GPU: No workers available, process everything locally 1. Single GPU: No workers available, process everything locally
2. Static Mode: Small batches, distributes tiles across workers (flattened) 2. Distributed tile queue: Workers pull tiles from a shared queue
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
Features: Features:
- Multi-mode batch handling for efficient video/image upscaling - Tile-based batch handling for video/image upscaling
- Tiled VAE support for memory efficiency - Tiled VAE support for memory efficiency
- Dynamic load balancing for large batches - Shared work queue so faster workers can process more tiles
- Backward compatible with single-image workflows - Backward compatible with single-image workflows
Environment Variables: Environment Variables:
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20) - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB) - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
Threshold: dynamic_threshold input controls mode switch (default 8) The hidden dynamic_threshold input is retained for workflow compatibility but
does not affect the current mode-selection policy.
""" """
def __init__(self): def __init__(self):
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
# Ensure initialization before registering routes # Ensure initialization before registering routes
ensure_tile_jobs_initialized() ensure_tile_jobs_initialized()
# Node registration
NODE_CLASS_MAPPINGS = {
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
}
+2 -2
View File
@@ -269,7 +269,7 @@ class ImageBatchDivider:
class AudioBatchDivider: class AudioBatchDivider:
"""Divides an audio waveform into multiple parts along the time/samples dimension.""" """Divides an audio waveform into sequential segments along the time dimension."""
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
@@ -282,7 +282,7 @@ class AudioBatchDivider:
"max": 10, "max": 10,
"step": 1, "step": 1,
"display": "number", "display": "number",
"tooltip": "Number of parts to divide the audio into" "tooltip": "Number of sequential time segments to create"
}), }),
} }
} }
+181
View File
@@ -0,0 +1,181 @@
"""Native V3 schemas with the existing execution algorithms kept intact.
Runtime objects are private, per-execution helpers, not registered V1 nodes.
This avoids sharing mutable instance state through V3's sanitized class clones.
"""
import comfy.samplers
from comfy_api.v0_0_2 import io
from . import utilities as _utilities
from .collector import DistributedCollectorNode as _CollectorRuntime
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
class DistributedSeed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
max=1125899906842624, force_input=False)],
outputs=[io.Int.Output(display_name='seed')],
accept_all_inputs=True,
)
@classmethod
def execute(cls, seed, is_worker=False, worker_id=''):
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
class DistributedValue(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedValue', display_name='Distributed Value', category='utils',
inputs=[io.String.Input('default_value', default=''),
io.String.Input('worker_values', default='{}')],
outputs=[io.AnyType.Output(display_name='value')],
accept_all_inputs=True,
)
@classmethod
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
return io.NodeOutput(*_utilities.DistributedValue().distribute(
default_value, worker_values, is_worker, worker_id))
class DistributedModelName(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
inputs=[io.String.Input('text', default='')],
outputs=[io.AnyType.Output(display_name='output')],
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
)
@classmethod
def execute(cls, text):
result = _utilities.DistributedModelName().log_input(
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
return io.NodeOutput(*result['result'], ui=result['ui'])
class ImageBatchDivider(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
inputs=[io.Image.Input('images'),
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
display_mode=io.NumberDisplay.number,
tooltip='Number of parts to divide the batch into')],
# The existing frontend still displays only divide_by sockets.
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
is_output_node=True,
)
@classmethod
def execute(cls, images, divide_by):
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
class AudioBatchDivider(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
inputs=[io.Audio.Input('audio'),
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
display_mode=io.NumberDisplay.number,
tooltip='Number of sequential time segments to create')],
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
is_output_node=True,
)
@classmethod
def execute(cls, audio, divide_by):
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
class DistributedEmptyImage(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
io.Int.Input('width', default=64, min=1, max=4096, step=1),
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
outputs=[io.Image.Output()],
)
@classmethod
def execute(cls, height, width, channels):
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
class DistributedCollector(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
inputs=[io.Boolean.Input('load_balance', default=False,
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)],
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
is_input_list=True, accept_all_inputs=True,
)
@classmethod
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
worker_id='', pass_through=False, delegate_only=False):
return io.NodeOutput(*_CollectorRuntime().run(
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
worker_batch_size=worker_batch_size, worker_id=worker_id,
pass_through=pass_through, delegate_only=delegate_only))
class UltimateSDUpscaleDistributed(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id='UltimateSDUpscaleDistributed',
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
inputs=[
io.Image.Input('upscaled_image'), io.Model.Input('model'),
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
io.Int.Input('steps', default=20, min=1, max=10000),
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
io.Int.Input('padding', default=32, min=0, max=256, step=8),
io.Int.Input('mask_blur', default=8, min=0, max=256),
io.Boolean.Input('force_uniform_tiles', default=True),
io.Boolean.Input('tiled_decode', default=False),
], outputs=[io.Image.Output()], accept_all_inputs=True,
)
@classmethod
def fingerprint_inputs(cls, **kwargs):
return _UpscaleRuntime.IS_CHANGED(**kwargs)
@classmethod
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
return io.NodeOutput(*_UpscaleRuntime().run(
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
+1 -1
View File
@@ -1,7 +1,7 @@
[project] [project]
name = "ComfyUI-Distributed" name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud" description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.4.6" version = "1.4.7"
license = {file = "LICENSE"} license = {file = "LICENSE"}
dependencies = [] dependencies = []
+1
View File
@@ -0,0 +1 @@
"""Internal extension lifecycle support (not a node provider)."""
+35
View File
@@ -0,0 +1,35 @@
"""Initialize routes, distributed state and the existing worker lifecycle."""
import atexit
import os
import server
from ..utils.config import CONFIG_FILE, ensure_config_exists
from ..utils.logging import debug_log
from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
from ..upscale.job_store import ensure_tile_jobs_initialized
_initialized = False
def initialize():
"""Called by ComfyExtension.on_load; initialize once per loaded package."""
global _initialized
if _initialized:
return
ensure_config_exists()
from .. import api # noqa: F401 - registers the existing @routes.* handlers
from ..api.queue_orchestration import ensure_distributed_state
ensure_distributed_state(server.PromptServer.instance)
ensure_tile_jobs_initialized()
if not os.environ.get('COMFYUI_IS_WORKER'):
atexit.register(sync_cleanup)
delayed_auto_launch()
register_async_signals()
_initialized = True
debug_log('Loaded Distributed nodes.')
debug_log(f'Config file: {CONFIG_FILE}')
+71
View File
@@ -302,6 +302,77 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(queued["audio"]["sample_rate"], 44100) self.assertEqual(queued["audio"]["sample_rate"], 44100)
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4)) self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
async def test_job_complete_accepts_audio_without_image(self):
queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
request = _FakeRequest(
{
"job_id": "audio-only-job",
"worker_id": "worker-1",
"batch_idx": 0,
"audio": self._encoded_audio_payload(),
"is_last": True,
}
)
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 200)
decode_image.assert_not_called()
queued = await queue.get()
self.assertIsNone(queued["tensor"])
self.assertEqual(queued["audio"]["sample_rate"], 44100)
async def test_job_complete_rejects_payload_without_image_or_audio(self):
request = _FakeRequest(
{
"job_id": "job-1",
"worker_id": "worker-1",
"batch_idx": 0,
"is_last": True,
}
)
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 400)
self.assertIn("image or audio", response.payload.get("message", "").lower())
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
request = _FakeRequest(
{
"job_id": "job-1",
"worker_id": "worker-1",
"batch_idx": 0,
"image": 123,
"audio": self._encoded_audio_payload(),
"is_last": True,
}
)
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 400)
self.assertIn("image", response.payload.get("message", "").lower())
async def test_job_complete_returns_400_for_malformed_audio(self):
request = _FakeRequest(
{
"job_id": "job-1",
"worker_id": "worker-1",
"batch_idx": 0,
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
"is_last": True,
}
)
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 400)
self.assertIn("audio.shape", response.payload.get("message", "").lower())
def test_decode_audio_payload_rejects_bad_shape(self): def test_decode_audio_payload_rejects_bad_shape(self):
bad = { bad = {
"sample_rate": 44100, "sample_rate": 44100,
+691
View File
@@ -0,0 +1,691 @@
{
"DistributedCollector": {
"input": {
"required": {
"load_balance": [
"BOOLEAN",
{
"default": false,
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
}
]
},
"optional": {
"images": [
"IMAGE"
],
"audio": [
"AUDIO"
]
},
"hidden": {
"multi_job_id": [
"STRING",
{
"default": ""
}
],
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"master_url": [
"STRING",
{
"default": ""
}
],
"enabled_worker_ids": [
"STRING",
{
"default": "[]"
}
],
"worker_batch_size": [
"INT",
{
"default": 1,
"min": 1,
"max": 1024
}
],
"worker_id": [
"STRING",
{
"default": ""
}
],
"pass_through": [
"BOOLEAN",
{
"default": false
}
],
"delegate_only": [
"BOOLEAN",
{
"default": false
}
]
}
},
"input_order": {
"required": [
"load_balance"
],
"optional": [
"images",
"audio"
],
"hidden": [
"multi_job_id",
"is_worker",
"master_url",
"enabled_worker_ids",
"worker_batch_size",
"worker_id",
"pass_through",
"delegate_only"
]
},
"output": [
"IMAGE",
"AUDIO"
],
"output_name": [
"images",
"audio"
],
"output_is_list": [
false,
false
],
"is_input_list": true,
"output_node": false,
"category": "image",
"display_name": "Distributed Collector"
},
"DistributedSeed": {
"input": {
"required": {
"seed": [
"INT",
{
"default": 1125899906842,
"min": 0,
"max": 1125899906842624,
"forceInput": false
}
]
},
"hidden": {
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"worker_id": [
"STRING",
{
"default": ""
}
]
}
},
"input_order": {
"required": [
"seed"
],
"hidden": [
"is_worker",
"worker_id"
]
},
"output": [
"INT"
],
"output_name": [
"seed"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "utils",
"display_name": "Distributed Seed"
},
"DistributedModelName": {
"input": {
"required": {
"text": [
"STRING",
{
"default": ""
}
]
},
"hidden": {
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO"
}
},
"input_order": {
"required": [
"text"
],
"hidden": [
"unique_id",
"extra_pnginfo"
]
},
"output": [
"*"
],
"output_name": [
"output"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": true,
"category": "utils",
"display_name": "Distributed Model Name"
},
"DistributedValue": {
"input": {
"required": {
"default_value": [
"STRING",
{
"default": ""
}
],
"worker_values": [
"STRING",
{
"default": "{}"
}
]
},
"hidden": {
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"worker_id": [
"STRING",
{
"default": ""
}
]
}
},
"input_order": {
"required": [
"default_value",
"worker_values"
],
"hidden": [
"is_worker",
"worker_id"
]
},
"output": [
"*"
],
"output_name": [
"value"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "utils",
"display_name": "Distributed Value"
},
"ImageBatchDivider": {
"input": {
"required": {
"images": [
"IMAGE"
],
"divide_by": [
"INT",
{
"default": 2,
"min": 1,
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of parts to divide the batch into"
}
]
}
},
"input_order": {
"required": [
"images",
"divide_by"
]
},
"output": [
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*"
],
"output_name": [
"batch_1",
"batch_2",
"batch_3",
"batch_4",
"batch_5",
"batch_6",
"batch_7",
"batch_8",
"batch_9",
"batch_10"
],
"output_is_list": [
false,
false,
false,
false,
false,
false,
false,
false,
false,
false
],
"is_input_list": false,
"output_node": true,
"category": "image",
"display_name": "Image Batch Divider"
},
"AudioBatchDivider": {
"input": {
"required": {
"audio": [
"AUDIO"
],
"divide_by": [
"INT",
{
"default": 2,
"min": 1,
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of sequential time segments to create"
}
]
}
},
"input_order": {
"required": [
"audio",
"divide_by"
]
},
"output": [
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*",
"*"
],
"output_name": [
"audio_1",
"audio_2",
"audio_3",
"audio_4",
"audio_5",
"audio_6",
"audio_7",
"audio_8",
"audio_9",
"audio_10"
],
"output_is_list": [
false,
false,
false,
false,
false,
false,
false,
false,
false,
false
],
"is_input_list": false,
"output_node": true,
"category": "audio",
"display_name": "Audio Segment Divider"
},
"DistributedEmptyImage": {
"input": {
"required": {
"height": [
"INT",
{
"default": 64,
"min": 1,
"max": 4096,
"step": 1
}
],
"width": [
"INT",
{
"default": 64,
"min": 1,
"max": 4096,
"step": 1
}
],
"channels": [
"INT",
{
"default": 3,
"min": 1,
"max": 4,
"step": 1
}
]
}
},
"input_order": {
"required": [
"height",
"width",
"channels"
]
},
"output": [
"IMAGE"
],
"output_name": [
"IMAGE"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "image",
"display_name": "Distributed Empty Image"
},
"UltimateSDUpscaleDistributed": {
"input": {
"required": {
"upscaled_image": [
"IMAGE"
],
"model": [
"MODEL"
],
"positive": [
"CONDITIONING"
],
"negative": [
"CONDITIONING"
],
"vae": [
"VAE"
],
"seed": [
"INT",
{
"default": 0,
"min": 0,
"max": 18446744073709551615
}
],
"steps": [
"INT",
{
"default": 20,
"min": 1,
"max": 10000
}
],
"cfg": [
"FLOAT",
{
"default": 8.0,
"min": 0.0,
"max": 100.0
}
],
"sampler_name": [
[
"euler",
"euler_cfg_pp",
"euler_ancestral",
"euler_ancestral_cfg_pp",
"heun",
"heunpp2",
"exp_heun_2_x0",
"exp_heun_2_x0_sde",
"dpm_2",
"dpm_2_ancestral",
"lms",
"dpm_fast",
"dpm_adaptive",
"dpmpp_2s_ancestral",
"dpmpp_2s_ancestral_cfg_pp",
"dpmpp_sde",
"dpmpp_sde_gpu",
"dpmpp_2m",
"dpmpp_2m_cfg_pp",
"dpmpp_2m_sde",
"dpmpp_2m_sde_gpu",
"dpmpp_2m_sde_heun",
"dpmpp_2m_sde_heun_gpu",
"dpmpp_3m_sde",
"dpmpp_3m_sde_gpu",
"ddpm",
"lcm",
"ipndm",
"ipndm_v",
"deis",
"cfgpp_ud10_ab",
"res_multistep",
"res_multistep_cfg_pp",
"res_multistep_ancestral",
"res_multistep_ancestral_cfg_pp",
"gradient_estimation",
"gradient_estimation_cfg_pp",
"er_sde",
"seeds_2",
"seeds_3",
"sa_solver",
"sa_solver_pece",
"ddim",
"uni_pc",
"uni_pc_bh2"
]
],
"scheduler": [
[
"simple",
"sgm_uniform",
"karras",
"exponential",
"ddim_uniform",
"beta",
"normal",
"linear_quadratic",
"kl_optimal"
]
],
"denoise": [
"FLOAT",
{
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01
}
],
"tile_width": [
"INT",
{
"default": 512,
"min": 64,
"max": 2048,
"step": 8
}
],
"tile_height": [
"INT",
{
"default": 512,
"min": 64,
"max": 2048,
"step": 8
}
],
"padding": [
"INT",
{
"default": 32,
"min": 0,
"max": 256,
"step": 8
}
],
"mask_blur": [
"INT",
{
"default": 8,
"min": 0,
"max": 256
}
],
"force_uniform_tiles": [
"BOOLEAN",
{
"default": true
}
],
"tiled_decode": [
"BOOLEAN",
{
"default": false
}
]
},
"hidden": {
"multi_job_id": [
"STRING",
{
"default": ""
}
],
"is_worker": [
"BOOLEAN",
{
"default": false
}
],
"master_url": [
"STRING",
{
"default": ""
}
],
"enabled_worker_ids": [
"STRING",
{
"default": "[]"
}
],
"worker_id": [
"STRING",
{
"default": ""
}
],
"tile_indices": [
"STRING",
{
"default": ""
}
],
"dynamic_threshold": [
"INT",
{
"default": 8,
"min": 1,
"max": 64
}
]
}
},
"input_order": {
"required": [
"upscaled_image",
"model",
"positive",
"negative",
"vae",
"seed",
"steps",
"cfg",
"sampler_name",
"scheduler",
"denoise",
"tile_width",
"tile_height",
"padding",
"mask_blur",
"force_uniform_tiles",
"tiled_decode"
],
"hidden": [
"multi_job_id",
"is_worker",
"master_url",
"enabled_worker_ids",
"worker_id",
"tile_indices",
"dynamic_threshold"
]
},
"output": [
"IMAGE"
],
"output_name": [
"IMAGE"
],
"output_is_list": [
false
],
"is_input_list": false,
"output_node": false,
"category": "image/upscaling",
"display_name": "Ultimate SD Upscale Distributed (No Upscale)"
}
}
+252
View File
@@ -0,0 +1,252 @@
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
No HTTP listener, workers, model downloads or live installation changes.
The original contracts were captured from 32ac027 using the same core.
"""
import asyncio
import inspect
import json
import os
from pathlib import Path
import sys
import tempfile
from unittest.mock import patch
from PIL import Image
ROOT = Path(__file__).resolve().parents[2]
COMFY_ROOT = Path(sys.argv[1]).resolve()
sys.path.insert(0, str(COMFY_ROOT))
os.environ['COMFYUI_IS_WORKER'] = '1'
import comfy.cli_args
comfy.cli_args.args.cpu = True
comfy.cli_args.args.disable_assets = True
from app.assets.manager import default_asset_manager
from comfy_api.v0_0_2 import io
import execution
import nodes
import server
import torch
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
def normalized_input(value):
kind = value[0]
opts = dict(value[1]) if len(value) > 1 else {}
if kind == 'STRING':
# V3 serializes the same default single-line widget explicitly.
opts.setdefault('multiline', False)
if kind == 'COMBO':
kind = opts['options']
opts = {key: val for key, val in opts.items() if key != 'options'}
if isinstance(kind, list):
# Single-selection is the V1 dropdown default as well.
opts.setdefault('multiselect', False)
return [kind, opts]
async def map_node(cls, values, extra=None):
"""Use the actual input parser and sanitized V3 executor clones."""
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
assert not missing, missing
# Use a worker thread for synchronous nodes: collector/upscale bridge back
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
def run():
return asyncio.run(execution._async_map_node_over_list(
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
return await asyncio.to_thread(run)
def check_saved_workflows(mapping):
checked = 0
occurrences = 0
for path in sorted((ROOT / 'workflows').glob('*.json')):
workflow = json.loads(path.read_text())
local_nodes = {node['id']: node for node in workflow['nodes']}
for saved in workflow['nodes']:
if saved['type'] not in mapping:
continue
occurrences += 1
info = mapping[saved['type']].GET_NODE_INFO_V1()
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
for socket in saved.get('inputs', []):
assert socket['name'] in inputs, (path.name, saved['id'], socket)
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
for index, output in enumerate(saved.get('outputs', [])):
assert output['type'] == info['output'][index], (path.name, index)
for link in workflow['links']:
if link[1] == saved['id']:
assert link[2] < len(info['output']), (path.name, link)
target = local_nodes[link[3]]['inputs'][link[4]]
assert target['type'] == info['output'][link[2]], (path.name, link)
checked += 1
assert checked == 5 and occurrences == 10, (checked, occurrences)
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
async def check_execution(mapping, module, prompt_server, asset_manager):
result = await map_node(mapping['DistributedSeed'],
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
assert result[0].result == (126,)
values = json.dumps({'_type': 'INT', '3': '17'})
result = await map_node(mapping['DistributedValue'],
{'default_value': '4', 'worker_values': values,
'is_worker': True, 'worker_id': 'worker_2'})
assert result[0].result == (17,)
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
{'extra_pnginfo': metadata})
assert result[0].result == ('model.ckpt',)
assert result[0].ui == {'text': ['model.ckpt']}
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
result = await map_node(mapping['DistributedEmptyImage'],
{'height': 8, 'width': 8, 'channels': 3})
empty = result[0].result[0]
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
for count in (1, 3, 10):
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
assert len(result[0].result) == 10
assert torch.equal(torch.cat(result[0].result[:count]), images)
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
assert len(result[0].result) == 10
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
collector = mapping['DistributedCollector']
prepared, missing, hidden = execution.get_input_data(
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
assert not missing
prepared['images'] = [images[:2], images[2:]]
prepared['audio'] = [audio, audio]
result = await execution._async_map_node_over_list(
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
assert torch.equal(result[0].result[0], images)
assert result[0].result[1]['waveform'].shape[-1] == 66
# Exercise real async aggregation without an HTTP endpoint or worker.
queue = asyncio.Queue()
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
result = await map_node(collector,
{'images': images[:2], 'load_balance': False,
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
assert torch.equal(result[0].result[0], images[:3])
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
# Prove V3's per-execution class clones do not leak helper instance state;
# replace only the GPU/model boundary, not the input parser or executor.
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
seen = []
def fake_upscale(self, *args):
assert not hasattr(self, 'acceptance_marker')
self.acceptance_marker = True
seen.append(args)
return (args[0],)
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
'force_uniform_tiles': True, 'tiled_decode': False,
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
with patch.object(upscale_runtime, 'run', fake_upscale):
for _ in range(2):
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
assert result[0].result[0] is images
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
'tile_indices', 'dynamic_threshold')), seen[0]
import math
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
graph = {
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
'pass_through': True}},
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
}
valid = await execution.validate_prompt('v3-graph', graph, None)
assert valid[0], valid
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
assert not rejected[0]
errors = [error for entry in rejected[3].values() for error in entry['errors']]
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
import folder_paths
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
with patch.object(folder_paths, 'temp_directory', temp):
executor = execution.PromptExecutor(
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
assert executor.success, executor.status_messages
history = executor.history_result
record = history['outputs']['4']['images'][0]
preview = Path(temp) / record.get('subfolder', '') / record['filename']
with Image.open(preview) as image:
assert image.size == (8, 8) and image.mode == 'RGB'
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
async def main():
asset_manager = default_asset_manager()
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
module = sys.modules[str(ROOT).replace('.', '_x_')]
assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
extension = await module.comfy_entrypoint()
classes = await extension.get_node_list()
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
assert len(classes) == len(mapping) == len(BASELINE) == 8
assert set(mapping) == set(BASELINE)
for node_id, cls in mapping.items():
assert issubclass(cls, io.ComfyNode)
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
old = dict(BASELINE[node_id])
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
# V1's ByPassTypeTuple advertises '*' when indexed, while its
# underlying tuple and existing frontend declare IMAGE/AUDIO.
# Native V3 declares all ten existing typed sockets explicitly.
assert old['output'] == ['*'] * 10
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
for group in ('required', 'optional'):
old_inputs = old['input'].get(group, {})
new_inputs = new['input'].get(group, {})
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
for name, original in old_inputs.items():
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
assert old[key] == new[key], (node_id, key, old[key], new[key])
# Standard context lives in cls.hidden; orchestrator metadata remains
# accepted by its original kwarg name, without creating new widgets.
signature = inspect.signature(cls.execute)
for name, field in old['input'].get('hidden', {}).items():
if isinstance(field, list):
assert cls.GET_SCHEMA().accept_all_inputs, node_id
assert name in signature.parameters, (node_id, name)
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
else:
assert name in new['input']['hidden'], (node_id, name)
expected_hidden = []
if node_id == 'DistributedModelName':
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
if old['output_node']:
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
print('SCHEMA_PARITY_OK', len(mapping))
check_saved_workflows(mapping)
await check_execution(mapping, module, prompt_server, asset_manager)
print('V3_ACCEPTANCE_OK')
asyncio.run(main())
+159
View File
@@ -124,6 +124,49 @@ def test_collector_opts_into_comfyui_list_inputs():
assert collector.INPUT_IS_LIST is True assert collector.INPUT_IS_LIST is True
def test_collector_exposes_images_as_optional_input():
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
assert "images" not in input_types["required"]
assert input_types["optional"]["images"] == ("IMAGE",)
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
collector = _load_collector_module().DistributedCollectorNode()
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
images, returned_audio = collector.run(images=None, audio=[audio])
assert images is None
assert returned_audio is audio
def test_collector_rejects_missing_images_and_audio():
collector = _load_collector_module().DistributedCollectorNode()
try:
collector.run(images=None, audio=None)
except ValueError as exc:
assert "image or audio" in str(exc).lower()
else:
raise AssertionError("Expected collector to reject a run with no media input")
def test_delegate_only_master_allows_no_local_media_input():
collector = _load_collector_module().DistributedCollectorNode()
images, audio = collector.run(
images=None,
audio=None,
multi_job_id=["delegate-audio-job"],
delegate_only=[True],
enabled_worker_ids=["[]"],
)
assert images is None
assert tuple(audio["waveform"].shape) == (1, 2, 1)
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs(): def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
collector = _load_collector_module().DistributedCollectorNode() collector = _load_collector_module().DistributedCollectorNode()
first = torch.zeros(1, 2, 2, 3) first = torch.zeros(1, 2, 2, 3)
@@ -203,3 +246,119 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
assert [payload["is_last"] for payload in posted_payloads] == [False, True] assert [payload["is_last"] for payload in posted_payloads] == [False, True]
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"} assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"} assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
def test_audio_only_worker_sends_one_completion_without_image():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
posted = []
class _FakeResponse:
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
def raise_for_status(self):
return None
class _FakeSession:
def post(self, url, json, timeout):
posted.append((url, json, timeout.total))
return _FakeResponse()
async def _fake_get_client_session():
return _FakeSession()
module.get_client_session = _fake_get_client_session
module.encode_audio_payload = lambda value: {"encoded": value is audio}
images, returned_audio = collector.run(
images=None,
audio=[audio],
multi_job_id=["audio-job"],
is_worker=[True],
master_url=["http://master"],
worker_id=["worker-a"],
)
assert images is None
assert returned_audio is audio
assert len(posted) == 1
assert posted[0][0] == "http://master/distributed/job_complete"
assert posted[0][2] == 600
assert posted[0][1] == {
"job_id": "audio-job",
"worker_id": "worker-a",
"batch_idx": 0,
"audio": {"encoded": True},
"is_last": True,
}
def test_audio_only_master_combines_local_and_worker_audio():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait(
{
"worker_id": "worker-a",
"image_index": 0,
"tensor": None,
"audio": worker_audio,
"is_last": True,
}
)
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
images, combined_audio = asyncio.run(
collector.execute(
images=None,
audio=master_audio,
multi_job_id="audio-job",
enabled_worker_ids='["worker-a"]',
)
)
assert images is None
assert combined_audio["sample_rate"] == 48000
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait(
{
"worker_id": "worker-a",
"image_index": 0,
"tensor": None,
"audio": worker_audio,
"is_last": True,
}
)
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
images, combined_audio = asyncio.run(
collector.execute(
images=None,
audio=None,
multi_job_id="delegate-audio-job",
enabled_worker_ids='["worker-a"]',
delegate_only=True,
)
)
assert images is None
assert combined_audio["sample_rate"] == 48000
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
+51 -3
View File
@@ -71,6 +71,15 @@ def _collector_only_prompt():
} }
def _audio_only_prompt():
"""1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
return {
"1": {"class_type": "LoadAudio", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
"3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
}
def _delegate_prompt(): def _delegate_prompt():
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)""" """1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
return { return {
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
self.assertEqual(len(preview_nodes), 1) self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0]) self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
def test_injects_preview_audio_for_audio_only_collector(self):
result = pt.prune_prompt_for_worker(_audio_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
def test_prefers_image_preview_when_collector_has_images_and_audio(self):
prompt = _linear_prompt()
prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
prompt["4"]["inputs"]["audio"] = ["6", 0]
result = pt.prune_prompt_for_worker(prompt)
self.assertEqual(
len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
1,
)
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
def test_preserves_image_preview_for_distributed_upscale(self):
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
def test_no_preview_image_when_no_downstream(self): def test_no_preview_image_when_no_downstream(self):
result = pt.prune_prompt_for_worker(_collector_only_prompt()) result = pt.prune_prompt_for_worker(_collector_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"] preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
@@ -242,7 +280,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
def test_upscale_node_is_treated_as_distributed(self): def test_upscale_node_is_treated_as_distributed(self):
prompt = { prompt = {
"1": {"class_type": "KSampler", "inputs": {}}, "1": {"class_type": "KSampler", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}}, "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}}, "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
} }
result = pt.prune_prompt_for_worker(prompt) result = pt.prune_prompt_for_worker(prompt)
@@ -285,12 +323,22 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
placeholder_id = empty_nodes[0][0] placeholder_id = empty_nodes[0][0]
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0]) self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
def test_audio_only_collector_does_not_get_image_placeholder(self):
prompt = _audio_only_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(empty_nodes, [])
self.assertNotIn("images", result["2"].get("inputs", {}))
self.assertNotIn("audio", result["2"].get("inputs", {}))
def test_one_placeholder_per_collector(self): def test_one_placeholder_per_collector(self):
"""Two collectors → two placeholders.""" """Two collectors → two placeholders."""
prompt = { prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}}, "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
"2": {"class_type": "DistributedCollector", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}, "3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
"10": {"class_type": "LoadImage", "inputs": {}},
"11": {"class_type": "LoadImage", "inputs": {}},
} }
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"]) result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"] empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
+21
View File
@@ -0,0 +1,21 @@
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
import os
from pathlib import Path
import subprocess
import sys
import pytest
def test_native_v3_runtime():
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
if not comfy_root:
pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance')
helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py'
result = subprocess.run(
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
capture_output=True, text=True, timeout=120,
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
)
assert result.returncode == 0, result.stdout + '\n' + result.stderr
assert 'V3_ACCEPTANCE_OK' in result.stdout
+9 -12
View File
@@ -2,7 +2,6 @@ import { app } from "/scripts/app.js";
import { ENDPOINTS } from "./constants.js"; import { ENDPOINTS } from "./constants.js";
const NODE_CLASS = "DistributedValue"; const NODE_CLASS = "DistributedValue";
const CONVERTED_WIDGET = "converted-widget";
const DYNAMIC_DEFAULT_WIDGET = "_dv_default"; const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_"; const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
const WORKERS_CHANGED_EVENT = "distributed:workers-changed"; const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
@@ -42,15 +41,13 @@ function getDynamicWorkerWidgets(node) {
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX)); return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
} }
function hideWidgetForGood(node, widget, suffix = "") { function hideWidgetForGood(widget) {
if (!widget) return; if (!widget) return;
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return; // These are serialized backing fields, not widgets converted into sockets.
// Keep their native type/serializer and mutate the shared options in place
widget.origType = widget.type; // so both the classic canvas and Nodes 2.0 suppress the whole widget row.
widget.origComputeSize = widget.computeSize; widget.options ??= {};
widget.origSerializeValue = widget.serializeValue; widget.options.hidden = true;
widget.computeSize = () => [0, -4];
widget.type = `${CONVERTED_WIDGET}${suffix}`;
// Hide any attached DOM element (multiline widgets). // Hide any attached DOM element (multiline widgets).
if (widget.element) widget.element.style.display = "none"; if (widget.element) widget.element.style.display = "none";
@@ -58,14 +55,14 @@ function hideWidgetForGood(node, widget, suffix = "") {
if (widget.linkedWidgets) { if (widget.linkedWidgets) {
for (const linked of widget.linkedWidgets) { for (const linked of widget.linkedWidgets) {
hideWidgetForGood(node, linked, `:${widget.name}`); hideWidgetForGood(linked);
} }
} }
} }
function hideRawWidgets(node) { function hideRawWidgets(node) {
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value"); hideWidgetForGood(getRawDefaultWidget(node));
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values"); hideWidgetForGood(getRawWorkerValuesWidget(node));
} }
function removeDynamicDefaultWidget(node) { function removeDynamicDefaultWidget(node) {
+2 -2
View File
@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
createCheckboxSetting( createCheckboxSetting(
"setting-debug", "setting-debug",
"Debug Mode", "Debug Mode",
"Enable verbose logging in the browser console.", "Enable verbose logging in the browser console and ComfyUI server output.",
extension.config?.settings?.debug || false, extension.config?.settings?.debug || false,
(event) => extension._updateSetting("debug", event.target.checked) (event) => extension._updateSetting("debug", event.target.checked)
) )
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
createNumberSetting( createNumberSetting(
"setting-worker-timeout", "setting-worker-timeout",
"Worker Timeout", "Worker Timeout",
"Seconds without a heartbeat before a worker is considered timed out. Default 60.", "Maximum result-wait and heartbeat inactivity period before recovery begins. Busy workers may receive additional grace. Default: 60 seconds.",
extension.config?.settings?.worker_timeout_seconds ?? 60, extension.config?.settings?.worker_timeout_seconds ?? 60,
10, 10,
1, 1,
+120
View File
@@ -0,0 +1,120 @@
import { readFileSync } from "node:fs";
import vm from "node:vm";
import { describe, expect, it, vi } from "vitest";
// Evaluate this repository's browser entrypoint without needing a ComfyUI server.
// The real-renderer acceptance additionally checks the input rows in Chromium.
const source = readFileSync(new URL("../distributedValue.js", import.meta.url), "utf8")
.replace(/^import .*;\r?\n/gm, "");
function setup({ connected = false, targetType = "STRING" } = {}) {
let extension;
const listeners = new Map();
const tasks = [];
const workers = [
{ id: "gpu-a", name: "GPU A", enabled: true },
{ id: "gpu-b", name: "GPU B", enabled: true },
{ id: "off", name: "Disabled", enabled: false },
];
const target = {
inputs: [{ name: "value" }],
widgets: [{ name: "value", type: targetType === "INT" ? "number" : "string", options: { step: 1, precision: 0 } }],
};
const graph = { links: { 1: { target_id: 2, target_slot: 0 } }, getNodeById: () => target };
const raw = ["default_value", "worker_values"].map((name, index) => ({
name,
type: "text",
value: index === 0 ? "saved default" : '{"1":"saved worker"}',
options: { multiline: false },
computeSize: vi.fn(() => [200, 20]),
serializeValue: vi.fn(function () { return this.value; }),
}));
const node = {
comfyClass: "DistributedValue", graph, widgets: [...raw], size: [200, 100],
inputs: raw.map(w => ({ name: w.name, widget: { name: w.name }, link: null })),
outputs: [{ links: connected ? [1] : [] }],
computeSize: () => [200, 100], setSize: vi.fn(), setDirtyCanvas: vi.fn(),
addWidget(type, name, value, callback, options) {
const widget = { type, name, value, callback, options };
this.widgets.push(widget);
return widget;
},
configure(data) {
data.widgets_values.forEach((value, i) => { this.widgets[i].value = value; });
return "configured";
},
};
vm.runInNewContext(source, {
app: { graph, registerExtension: ext => { extension = ext; } },
ENDPOINTS: { CONFIG: "/distributed/config" },
fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
window: { addEventListener: (name, handler) => listeners.set(name, handler) },
setTimeout: callback => { tasks.push(callback); }, console,
}, { filename: "distributedValue.js" });
return { node, raw, extension, listeners, tasks, flush: () => { while (tasks.length) tasks.shift()(); } };
}
describe("Distributed Value internal widget visibility", () => {
it("hides raw fields natively without converting their types or replacing their options", async () => {
const h = setup();
const options = h.raw.map(w => w.options);
const sizes = h.raw.map(w => w.computeSize);
const serializers = h.raw.map(w => w.serializeValue);
const slots = h.node.inputs;
await h.extension.nodeCreated(h.node);
for (const [i, widget] of h.raw.entries()) {
expect(widget.type).toBe("text");
expect(widget.options).toBe(options[i]);
expect(widget.options.hidden).toBe(true);
expect(widget.computeSize).toBe(sizes[i]);
expect(widget.serializeValue).toBe(serializers[i]);
expect(h.node.widgets[i]).toBe(widget);
}
expect(h.node.inputs).toBe(slots);
expect(h.node.widgets.slice(2).map(w => w.label)).toEqual(["default_value", "GPU A", "GPU B"]);
expect(h.node.widgets.slice(2).every(w => !w.options.hidden)).toBe(true);
});
it("keeps typed edits in the original serialized fields across rebuild/configure", async () => {
const h = setup({ connected: true, targetType: "INT" });
await h.extension.nodeCreated(h.node);
h.node.widgets.find(w => w.name === "_dv_default").callback(42);
h.node.widgets.find(w => w.name === "_dv_worker_1").callback(7);
h.node.widgets.find(w => w.name === "_dv_worker_2").callback(9);
const values = h.raw.map(w => w.serializeValue());
expect(values[0]).toBe(42);
expect(JSON.parse(values[1])).toMatchObject({
_type: "INT", 1: "7", 2: "9", _by_worker_id: { "gpu-a": "7", "gpu-b": "9" },
});
expect(h.node.configure({ widgets_values: values })).toBe("configured");
h.flush();
expect(h.raw.map(w => w.serializeValue())).toEqual(values);
expect(h.raw.every(w => w.options.hidden && w.type === "text")).toBe(true);
expect(h.node.widgets.slice(2).map(w => w.value)).toEqual([42, 7, 9]);
});
it("reapplies visibility on worker refresh without duplicating raw or dynamic fields", async () => {
const h = setup();
await h.extension.nodeCreated(h.node);
h.raw[0].options.hidden = false;
await h.listeners.get("distributed:workers-changed")({ detail: { workers: [
{ id: "gpu-b", name: "Renamed B", enabled: true },
] } });
expect(h.raw.every(w => w.options.hidden)).toBe(true);
expect(h.node.widgets.map(w => w.name)).toEqual(["default_value", "worker_values", "_dv_default", "_dv_worker_1"]);
expect(h.node.widgets[3].label).toBe("Renamed B");
});
it("hides linked backing widgets and handles absent options without altering serializers", async () => {
const h = setup();
delete h.raw[1].options;
const linked = { type: "text", options: {}, serializeValue: vi.fn(() => "linked") };
h.raw[0].linkedWidgets = [linked];
await h.extension.nodeCreated(h.node);
expect(h.raw[1].options.hidden).toBe(true);
expect(linked.options.hidden).toBe(true);
expect(linked.type).toBe("text");
expect(linked.serializeValue()).toBe("linked");
});
});
+1 -1
View File
@@ -814,7 +814,7 @@
"outputs": [], "outputs": [],
"properties": {}, "properties": {},
"widgets_values": [ "widgets_values": [
"If all your GPUs are the same/similar, set static_distribution to true\n" "Workers pull tiles from a shared queue, so faster workers can process more tiles automatically.\n"
], ],
"color": "#432", "color": "#432",
"bgcolor": "#653" "bgcolor": "#653"