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
31 changed files with 1811 additions and 1121 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.
* **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)
> **⚠️ 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 |
| **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 |
| **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 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
from .distributed import (
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
)
"""ComfyUI-Distributed's native V3 extension entrypoint."""
from comfy_api.v0_0_2 import ComfyExtension, io
# Import utilities
from .utils.config import ensure_config_exists, CONFIG_FILE
from .utils.logging import debug_log
from .nodes.v3 import NODES
from .runtime.bootstrap import initialize
# Import distributed upscale nodes
from .nodes.distributed_upscale import (
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
)
WEB_DIRECTORY = './web'
WEB_DIRECTORY = "./web"
ensure_config_exists()
class DistributedExtension(ComfyExtension):
async def on_load(self) -> None:
initialize()
# Merge node mappings
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return list(NODES)
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
debug_log("Loaded Distributed nodes.")
debug_log(f"Config file: {CONFIG_FILE}")
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
async def comfy_entrypoint() -> DistributedExtension:
return DistributedExtension()
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
-1
View File
@@ -3,4 +3,3 @@ from . import tunnel_routes # noqa: F401
from . import worker_routes # noqa: F401
from . import job_routes # noqa: F401
from . import usdu_routes # noqa: F401
from . import ltx_routes # noqa: F401
+11 -4
View File
@@ -295,17 +295,24 @@ async def job_complete_endpoint(request):
errors.append("worker_id: expected non-empty string")
if not isinstance(batch_idx, int) or batch_idx < 0:
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")
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")
if not isinstance(is_last, bool):
errors.append("is_last: expected boolean")
if errors:
return await handle_api_error(request, errors, 400)
tensor = _decode_canonical_png_tensor(image_payload)
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
try:
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()
worker_id = worker_id.strip()
-76
View File
@@ -1,76 +0,0 @@
"""HTTP routes for distributed LTX latent tile results."""
from __future__ import annotations
import time
from aiohttp import web
import server
from ..upscale.job_models import TileJobState
from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
from ..utils.logging import debug_log
from ..utils.network import handle_api_error
@server.PromptServer.instance.routes.post("/distributed/ltx/submit_tile")
async def submit_ltx_tile_endpoint(request):
"""Accept one processed LTX latent tile from a worker.
Body is JSON:
- multi_job_id: distributed job id
- worker_id: worker id
- tile_index: spatial tile id, or -1 for completion-only messages
- payload: base64 safetensors payload produced by nodes.ltx_tiled_sampler
- is_last: optional worker completion flag
"""
try:
content_length = request.headers.get("content-length")
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
data = await request.json()
multi_job_id = data.get("multi_job_id")
worker_id = str(data.get("worker_id") or "")
is_last = bool(data.get("is_last", False))
payload = data.get("payload")
tile_index = data.get("tile_index")
if not multi_job_id or not worker_id:
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not isinstance(job_data, TileJobState):
return await handle_api_error(request, "Job not configured for tile submissions", 404)
job_data.worker_status[worker_id] = time.time()
if is_last and (payload is None or int(tile_index or -1) < 0):
await job_data.queue.put({
"worker_id": worker_id,
"is_last": True,
"tiles": [],
})
debug_log(f"Received LTX completion signal from worker {worker_id}")
return web.json_response({"status": "success"})
if payload is None or tile_index is None:
return await handle_api_error(request, "Missing tile_index or payload", 400)
tile_index = int(tile_index)
await job_data.queue.put({
"worker_id": worker_id,
"is_last": is_last,
"tiles": [{
"tile_idx": tile_index,
"global_idx": tile_index,
"payload": payload,
"worker_id": worker_id,
}],
})
debug_log(f"Received LTX tile {tile_index} from worker {worker_id} (is_last={is_last})")
return web.json_response({"status": "success"})
except Exception as exc:
return await handle_api_error(request, exc, 500)
+31 -10
View File
@@ -351,16 +351,33 @@ def prune_prompt_for_worker(prompt_obj):
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
if has_removed_downstream:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {
"images": [dist_id, 0],
},
"class_type": "PreviewImage",
"_meta": {
"title": "Preview Image (auto-added)",
},
}
original_node = prompt_obj.get(str(dist_id), {})
class_type = original_node.get("class_type")
inputs = original_node.get("inputs", {})
image_connected = class_type != "DistributedCollector" or (
isinstance(inputs.get("images"), list)
and len(inputs["images"]) == 2
)
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
@@ -400,6 +417,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
collector_entry = pruned_prompt.get(collector_id)
if not collector_entry:
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()
pruned_prompt[placeholder_id] = {
"class_type": "DistributedEmptyImage",
+1 -1
View File
@@ -1,7 +1,7 @@
# conftest.py — project-level pytest configuration.
#
# 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.
#
# 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:
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
- Resolves enabled/selected workers.
- Pings workers (`GET /prompt`) to include only reachable ones.
- Dispatches the workflow to workers (`POST /prompt`).
- Probes and dispatches workers through `/distributed/worker_ws` by default.
- If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
- Queues the master workflow in ComfyUI’s prompt queue.
- 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
- `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 metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
- 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)
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.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
Also, for collector-based flows:
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
## Examples
### 1) Minimal `curl`
### 1) Minimal `curl` request envelope
```bash
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
```
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
{
"prompt": {
"1": {"class_type": "KSampler", "inputs": {} }
"<node_id>": {
"class_type": "<node class>",
"inputs": {"<required_input>": "<value or connection>"}
}
},
"enabled_worker_ids": [],
"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`)
```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)
**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.
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)
5. Deploy your pod.
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
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
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 -37
View File
@@ -1,37 +1 @@
from .utilities import (
DistributedSeed,
DistributedModelName,
DistributedValue,
ImageBatchDivider,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from .collector import DistributedCollectorNode
from .ltx_tiled_sampler import (
NODE_CLASS_MAPPINGS as LTX_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as LTX_DISPLAY_NAME_MAPPINGS,
)
NODE_CLASS_MAPPINGS = {
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed,
"DistributedModelName": DistributedModelName,
"DistributedValue": DistributedValue,
"ImageBatchDivider": ImageBatchDivider,
"AudioBatchDivider": AudioBatchDivider,
"DistributedEmptyImage": DistributedEmptyImage,
**LTX_CLASS_MAPPINGS,
}
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",
**LTX_DISPLAY_NAME_MAPPINGS,
}
"""Private execution helpers; public registration lives in nodes.v3."""
+73 -45
View File
@@ -29,7 +29,6 @@ class DistributedCollectorNode:
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"load_balance": (
"BOOLEAN",
{
@@ -38,7 +37,10 @@ class DistributedCollectorNode:
},
),
},
"optional": { "audio": ("AUDIO",) },
"optional": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
},
"hidden": {
"multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}),
@@ -102,8 +104,9 @@ class DistributedCollectorNode:
return None
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):
images = self._normalize_images_input(images)
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):
if images is not None:
images = self._normalize_images_input(images)
audio = self._normalize_audio_input(audio)
load_balance = self._unwrap_list_input(load_balance)
multi_job_id = self._unwrap_list_input(multi_job_id)
@@ -115,6 +118,14 @@ class DistributedCollectorNode:
pass_through = self._unwrap_list_input(pass_through)
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
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -141,41 +152,56 @@ class DistributedCollectorNode:
return result
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."""
batch_size = image_batch.shape[0]
if batch_size == 0:
return
"""Send an image batch, optionally with audio, or an audio-only completion."""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
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:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
response.raise_for_status()
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}")
raise # Re-raise to handle at caller level
raise
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
"""Combine audio from master and workers into a single audio output.
@@ -257,8 +283,8 @@ class DistributedCollectorNode:
images_on_cpu,
delegate_mode: bool,
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 = []
if not delegate_mode and images_on_cpu is not None:
for i in range(master_batch_size):
@@ -289,15 +315,15 @@ class DistributedCollectorNode:
if cpu_tensors:
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)
else:
raise ValueError("No image data collected from master or workers")
return None
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:
# 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)
return (images, audio if audio is not None else self.EMPTY_AUDIO)
else:
@@ -332,14 +358,15 @@ class DistributedCollectorNode:
master_audio = None
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
else:
images_on_cpu = images.cpu()
master_batch_size = images.shape[0]
if images is None:
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
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
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:
del prompt_server.distributed_pending_jobs[multi_job_id]
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
try:
combined = self._reorder_and_combine_tensors(
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 "
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
# Combine audio from master and workers
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
if combined is None:
debug_log(f"Master - Job {multi_job_id} complete with audio only")
else:
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
return (combined, combined_audio)
except Exception as e:
log(f"Master - Error combining images: {e}")
# Return just the master images as fallback
return (images, audio if audio is not None else self.EMPTY_AUDIO)
# Preserve collected audio even when image assembly fails.
return (images, combined_audio)
+6 -15
View File
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
"""
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
2. Static Mode: Small batches, distributes tiles across workers (flattened)
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
2. Distributed tile queue: Workers pull tiles from a shared queue
Features:
- Multi-mode batch handling for efficient video/image upscaling
- Tile-based batch handling for video/image upscaling
- 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
Environment Variables:
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
- 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):
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
# Ensure initialization before registering routes
ensure_tile_jobs_initialized()
# Node registration
NODE_CLASS_MAPPINGS = {
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
}
-668
View File
@@ -1,668 +0,0 @@
"""Distributed LTX tiled sampler node.
This node adapts the upstream LTX tiled sampler interface to the existing
ComfyUI-Distributed job queue. It distributes spatial latent tiles across the
same worker/master infrastructure used by the distributed upscaler instead of
creating a parallel queue implementation.
"""
from __future__ import annotations
import asyncio
import base64
import copy
import math
from typing import Callable, Dict, Iterable, List, Optional, Tuple
import aiohttp
import torch
from safetensors.torch import load as safetensors_load
from safetensors.torch import save as safetensors_save
from ..upscale.job_models import TileJobState
from ..upscale.job_store import (
_cleanup_job,
_drain_results_queue,
_get_completed_count,
_mark_task_completed,
ensure_tile_jobs_initialized,
init_static_job_batched,
)
from ..upscale.job_timeout import _check_and_requeue_timed_out_workers
from ..utils.async_helpers import run_async_in_server_loop
from ..utils.constants import HEARTBEAT_INTERVAL, TILE_WAIT_TIMEOUT
from ..utils.logging import debug_log, log
from ..utils.network import get_client_session
from ..utils.usdu_managment import _send_heartbeat_to_master
try: # Optional external LTX tiled sampler dependency.
from latent_tiled_sampler import LTXTiledSampler # type: ignore
except Exception: # pragma: no cover - depends on optional custom node install.
LTXTiledSampler = None # type: ignore
def _default_ltx_inputs() -> dict:
"""Fallback schema used when the source LTX sampler is not importable."""
return {
"required": {
"noise": ("NOISE",),
"guider": ("GUIDER",),
"sampler": ("SAMPLER",),
"sigmas": ("SIGMAS",),
"latent_image": ("LATENT",),
"h_tiles": ("INT", {"default": 1, "min": 1, "max": 64}),
"w_tiles": ("INT", {"default": 1, "min": 1, "max": 64}),
"overlap": ("INT", {"default": 0, "min": 0, "max": 4096}),
},
"optional": {
"denoise_mask": ("MASK",),
},
}
def _copy_base_inputs() -> dict:
if LTXTiledSampler is None:
return _default_ltx_inputs()
try:
return copy.deepcopy(LTXTiledSampler.INPUT_TYPES())
except Exception:
return _default_ltx_inputs()
def _compute_tile_starts(total_size: int, n_tiles: int, overlap: int) -> Tuple[List[int], int]:
"""Compute evenly spaced tile starts and tile size for one dimension."""
total_size = int(total_size)
n_tiles = max(1, int(n_tiles))
overlap = max(0, int(overlap))
if n_tiles <= 1 or total_size <= 1:
return [0], total_size
tile_size = math.ceil((total_size + (n_tiles - 1) * overlap) / n_tiles)
tile_size = min(tile_size, total_size)
if n_tiles == 2:
return [0, max(0, total_size - tile_size)], tile_size
stride = (total_size - tile_size) / (n_tiles - 1)
return [int(round(i * stride)) for i in range(n_tiles)], tile_size
def _make_window_1d(size: int, fade_left: int, fade_right: int, dtype, device):
"""Create a one-dimensional cosine edge window for tile blending."""
win = torch.ones(int(size), dtype=dtype, device=device)
if fade_left > 0:
fl = min(int(fade_left), int(size))
i = torch.arange(fl, dtype=dtype, device=device)
# Use fl rather than fl - 1 to preserve a nonzero center plateau for small fades.
win[:fl] = 0.5 * (1.0 - torch.cos(math.pi * i / max(1, fl)))
if fade_right > 0:
fr = min(int(fade_right), int(size))
i = torch.arange(fr, dtype=dtype, device=device)
win[int(size) - fr:] = 0.5 * (1.0 + torch.cos(math.pi * i / max(1, fr)))
return win
def _pack_tensors(tensors: Dict[str, torch.Tensor]) -> str:
"""Serialize tensors to a base64 safetensors payload for JSON transport."""
cpu_tensors = {
name: tensor.detach().contiguous().to("cpu")
for name, tensor in tensors.items()
if tensor is not None
}
payload = safetensors_save(cpu_tensors)
return base64.b64encode(payload).decode("ascii")
def _unpack_tensors(payload_b64: str, device) -> Dict[str, torch.Tensor]:
payload = base64.b64decode(payload_b64.encode("ascii"))
return {name: tensor.to(device) for name, tensor in safetensors_load(payload).items()}
def _extract_tensor_and_rebuilder(value) -> Tuple[torch.Tensor, Callable[[torch.Tensor], object]]:
"""Extract the tensor from common Comfy latent containers and preserve shape."""
if isinstance(value, dict) and "samples" in value:
original = value
def rebuild(tensor):
updated = dict(original)
updated["samples"] = tensor
return updated
return value["samples"], rebuild
if isinstance(value, tuple) and value and torch.is_tensor(value[0]):
original = value
def rebuild(tensor):
return (tensor, *original[1:])
return value[0], rebuild
if isinstance(value, list) and value and torch.is_tensor(value[0]):
original = value
def rebuild(tensor):
updated = list(original)
updated[0] = tensor
return updated
return value[0], rebuild
if torch.is_tensor(value):
return value, lambda tensor: tensor
raise TypeError("Expected latent_image to be a tensor or latent container with a tensor samples field")
def _extract_optional_mask(mask):
if mask is None:
return None
if isinstance(mask, dict) and "samples" in mask:
return mask["samples"]
return mask
def _tile_latent(value: torch.Tensor, tile: dict) -> torch.Tensor:
return value[:, :, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
def _tile_mask(mask, tile: dict):
if mask is None:
return None
if torch.is_tensor(mask) and mask.ndim >= 5:
return mask[:, :, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
if torch.is_tensor(mask) and mask.ndim >= 4:
return mask[:, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
return mask
def _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask):
result = guider.sample(
tile_noise,
tile_latent,
sampler,
sigmas,
denoise_mask=tile_mask,
disable_pbar=True,
)
if isinstance(result, tuple):
tile_samples = result[0]
tile_denoised = result[1] if len(result) > 1 else None
else:
tile_samples = result
tile_denoised = None
return tile_samples, tile_denoised
def _build_tile_queue(latent_shape: Iterable[int], n_h_tiles: int, n_w_tiles: int, overlap: int) -> List[dict]:
_batch, _channels, _frames, height, width = [int(v) for v in latent_shape]
starts_y, tile_h = _compute_tile_starts(height, n_h_tiles, overlap)
starts_x, tile_w = _compute_tile_starts(width, n_w_tiles, overlap)
queue = []
tile_idx = 0
for grid_y, y_start in enumerate(starts_y):
y_end = min(y_start + tile_h, height)
for grid_x, x_start in enumerate(starts_x):
x_end = min(x_start + tile_w, width)
queue.append(
{
"tile_index": tile_idx,
"h_start": y_start,
"h_end": y_end,
"w_start": x_start,
"w_end": x_end,
"grid_y": grid_y,
"grid_x": grid_x,
}
)
tile_idx += 1
return queue
def _blend_tile_into_sums(
tile: dict,
samples: torch.Tensor,
denoised: Optional[torch.Tensor],
output_sum: torch.Tensor,
weight_sum: torch.Tensor,
denoised_sum: Optional[torch.Tensor],
n_h_tiles: int,
n_w_tiles: int,
overlap: int,
):
y_start = tile["h_start"]
y_end = tile["h_end"]
x_start = tile["w_start"]
x_end = tile["w_end"]
grid_y = tile["grid_y"]
grid_x = tile["grid_x"]
fade_top = overlap if grid_y > 0 else 0
fade_bottom = overlap if grid_y < n_h_tiles - 1 else 0
fade_left = overlap if grid_x > 0 else 0
fade_right = overlap if grid_x < n_w_tiles - 1 else 0
dtype = output_sum.dtype
device = output_sum.device
wy = _make_window_1d(y_end - y_start, fade_top, fade_bottom, dtype, device)
wx = _make_window_1d(x_end - x_start, fade_left, fade_right, dtype, device)
weights = wy.view(1, 1, 1, y_end - y_start, 1) * wx.view(1, 1, 1, 1, x_end - x_start)
output_sum[:, :, :, y_start:y_end, x_start:x_end] += samples.to(device=device, dtype=dtype) * weights
weight_sum[:, :, :, y_start:y_end, x_start:x_end] += weights
if denoised_sum is not None and denoised is not None:
denoised_sum[:, :, :, y_start:y_end, x_start:x_end] += denoised.to(device=device, dtype=dtype) * weights
async def _get_next_ltx_tile_index(multi_job_id: str) -> Optional[int]:
prompt_server = ensure_tile_jobs_initialized()
pending_queue = None
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if isinstance(job_data, TileJobState):
pending_queue = job_data.pending_tasks
if pending_queue is None:
return None
try:
return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
except asyncio.TimeoutError:
return None
async def _get_ltx_completed_tasks(multi_job_id: str) -> dict:
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if isinstance(job_data, TileJobState):
return dict(job_data.completed_tasks)
return {}
async def _check_ltx_job_status(multi_job_id: str, master_url: str) -> bool:
session = await get_client_session()
async with session.get(f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}") as response:
if response.status != 200:
return False
data = await response.json()
return bool(data.get("ready", False))
async def _request_tile_from_master(multi_job_id: str, master_url: str, worker_id: str):
session = await get_client_session()
async with session.post(
f"{master_url}/distributed/request_image",
json={"multi_job_id": multi_job_id, "worker_id": str(worker_id)},
timeout=aiohttp.ClientTimeout(total=TILE_WAIT_TIMEOUT),
) as response:
if response.status == 200:
return await response.json()
if response.status == 404:
return None
response.raise_for_status()
async def _submit_ltx_tile_to_master(
multi_job_id: str,
master_url: str,
worker_id: str,
tile_index: int,
payload_b64: Optional[str] = None,
is_last: bool = False,
):
session = await get_client_session()
body = {
"multi_job_id": multi_job_id,
"worker_id": str(worker_id),
"tile_index": int(tile_index) if tile_index is not None else None,
"payload": payload_b64,
"is_last": bool(is_last),
}
async with session.post(f"{master_url}/distributed/ltx/submit_tile", json=body) as response:
response.raise_for_status()
return await response.json()
async def _ltx_worker_process(
node,
noise,
guider,
sampler,
sigmas,
latent_tensor,
denoise_mask,
n_h_tiles,
n_w_tiles,
tile_overlap,
):
job_id = node.multi_job_id
master_url = node.master_url
worker_id = node.worker_id or "ltx-worker"
tiles_by_index = {
tile["tile_index"]: tile
for tile in _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
}
for attempt in range(120):
if await _check_ltx_job_status(job_id, master_url):
break
await asyncio.sleep(1.0)
else:
log(f"LTX Dist Worker[{worker_id[:8]}]: job {job_id} was not ready; aborting")
return
full_noise = noise.generate_noise({"samples": latent_tensor})
last_heartbeat = 0.0
while True:
data = await _request_tile_from_master(job_id, master_url, worker_id)
if not data:
log(f"LTX Dist Worker[{worker_id[:8]}]: master job unavailable; stopping")
break
tile_index = data.get("tile_idx")
if tile_index is None:
await _submit_ltx_tile_to_master(job_id, master_url, worker_id, -1, is_last=True)
break
tile = tiles_by_index.get(int(tile_index))
if tile is None:
raise RuntimeError(f"Master assigned unknown LTX tile index {tile_index}")
tile_latent = _tile_latent(latent_tensor, tile)
tile_noise = _tile_latent(full_noise, tile)
tile_mask = _tile_mask(denoise_mask, tile)
tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
tensors = {"samples": tile_samples}
if tile_denoised is not None:
tensors["denoised"] = tile_denoised
await _submit_ltx_tile_to_master(
job_id,
master_url,
worker_id,
int(tile_index),
_pack_tensors(tensors),
)
now = asyncio.get_running_loop().time()
if now - last_heartbeat >= HEARTBEAT_INTERVAL:
await _send_heartbeat_to_master(job_id, master_url, worker_id)
last_heartbeat = now
async def _ltx_master_process(
node,
noise,
guider,
sampler,
sigmas,
latent_tensor,
denoise_mask,
n_h_tiles,
n_w_tiles,
tile_overlap,
):
job_id = node.multi_job_id
enabled_workers = node.enabled_worker_ids
tiles = _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
tiles_by_index = {tile["tile_index"]: tile for tile in tiles}
total_tiles = len(tiles)
await init_static_job_batched(job_id, batch_size=1, num_tiles_per_image=total_tiles, enabled_workers=enabled_workers)
full_noise = noise.generate_noise({"samples": latent_tensor})
output_sum = torch.zeros_like(latent_tensor)
weight_sum = torch.zeros_like(latent_tensor)
denoised_sum = torch.zeros_like(latent_tensor) if denoise_mask is not None else None
integrated = set()
last_timeout_check = 0.0
try:
while len(integrated) < total_tiles:
await _drain_results_queue(job_id)
completed = await _get_ltx_completed_tasks(job_id)
for tile_index, result in list(completed.items()):
if tile_index in integrated:
continue
if not isinstance(result, dict) or "payload" not in result:
# Locally completed tiles are blended immediately and then marked complete.
continue
tile = tiles_by_index.get(int(tile_index))
if tile is None:
continue
tensors = _unpack_tensors(result["payload"], latent_tensor.device)
_blend_tile_into_sums(
tile,
tensors["samples"],
tensors.get("denoised"),
output_sum,
weight_sum,
denoised_sum,
n_h_tiles,
n_w_tiles,
tile_overlap,
)
integrated.add(int(tile_index))
if len(integrated) >= total_tiles:
break
tile_index = await _get_next_ltx_tile_index(job_id)
if tile_index is not None and int(tile_index) not in integrated:
tile = tiles_by_index[int(tile_index)]
tile_latent = _tile_latent(latent_tensor, tile)
tile_noise = _tile_latent(full_noise, tile)
tile_mask = _tile_mask(denoise_mask, tile)
tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
_blend_tile_into_sums(
tile,
tile_samples,
tile_denoised,
output_sum,
weight_sum,
denoised_sum,
n_h_tiles,
n_w_tiles,
tile_overlap,
)
integrated.add(int(tile_index))
await _mark_task_completed(job_id, int(tile_index), {"tile_idx": int(tile_index), "source": "master"})
continue
now = asyncio.get_running_loop().time()
if now - last_timeout_check >= 10.0:
requeued = await _check_and_requeue_timed_out_workers(job_id, total_tiles)
if requeued:
log(f"LTX Dist: Requeued {requeued} timed-out tile(s)")
last_timeout_check = now
completed_count = await _get_completed_count(job_id)
debug_log(f"LTX Dist: Progress {completed_count}/{total_tiles} completed, {len(integrated)} integrated")
await asyncio.sleep(0.1)
finally:
await _cleanup_job(job_id)
weight_sum = torch.where(weight_sum == 0, torch.ones_like(weight_sum), weight_sum)
samples = output_sum / weight_sum
denoised = denoised_sum / weight_sum if denoised_sum is not None else None
return samples, denoised
def _ltx_local_tiled_process(
noise,
guider,
sampler,
sigmas,
latent_tensor,
denoise_mask,
n_h_tiles,
n_w_tiles,
tile_overlap,
):
"""Process all LTX latent tiles locally in the current ComfyUI execution context.
This path is used when no workers are enabled. Keeping it synchronous avoids
crossing ComfyUI's execution/inference context boundary while still using the
same tile extraction and blend logic as the distributed master/worker path.
"""
tiles = _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
full_noise = noise.generate_noise({"samples": latent_tensor})
output_sum = torch.zeros_like(latent_tensor)
weight_sum = torch.zeros_like(latent_tensor)
denoised_sum = None
for tile in tiles:
tile_latent = _tile_latent(latent_tensor, tile)
tile_noise = _tile_latent(full_noise, tile)
tile_mask = _tile_mask(denoise_mask, tile)
tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
if denoised_sum is None and tile_denoised is not None:
denoised_sum = torch.zeros_like(latent_tensor)
_blend_tile_into_sums(
tile,
tile_samples,
tile_denoised,
output_sum,
weight_sum,
denoised_sum,
n_h_tiles,
n_w_tiles,
tile_overlap,
)
weight_sum = torch.where(weight_sum == 0, torch.ones_like(weight_sum), weight_sum)
samples = output_sum / weight_sum
denoised = denoised_sum / weight_sum if denoised_sum is not None else None
return samples, denoised
class LTXTiledSamplerDistributed:
"""Distributed wrapper around the LTX tiled sampler."""
@classmethod
def INPUT_TYPES(cls):
base_inputs = _copy_base_inputs()
base_inputs.setdefault("hidden", {})
base_inputs["hidden"].update(
{
"multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}),
"master_url": ("STRING", {"default": ""}),
"worker_id": ("STRING", {"default": ""}),
"enabled_worker_ids": ("STRING", {"default": "[]"}),
}
)
return base_inputs
RETURN_TYPES = getattr(LTXTiledSampler, "RETURN_TYPES", ("LATENT", "LATENT", "LATENT"))
RETURN_NAMES = getattr(LTXTiledSampler, "RETURN_NAMES", ("samples", "denoised", "preview"))
FUNCTION = "sample_tiled_distributed"
CATEGORY = "latent/diffusion"
@classmethod
def IS_CHANGED(cls, **kwargs):
return float("nan")
def __init__(self):
self.base = LTXTiledSampler() if LTXTiledSampler is not None else None
def _run_base(self, noise, guider, sampler, sigmas, latent_image, denoise_mask, h_tiles, w_tiles, overlap):
if self.base is None:
raise RuntimeError(
"LTXTiledSamplerDistributed requires the base LTXTiledSampler node for local fallback. "
"Install the LTX tiled sampler dependency or run with distributed multi_job_id."
)
return self.base.sample_tiled(
noise,
guider,
sampler,
sigmas,
latent_image,
denoise_mask,
h_tiles,
w_tiles,
overlap,
)
def sample_tiled_distributed(
self,
noise,
guider,
sampler,
sigmas,
latent_image,
denoise_mask=None,
h_tiles: int = 1,
w_tiles: int = 1,
overlap: int = 0,
multi_job_id: str = "",
is_worker: bool = False,
master_url: str = "",
worker_id: str = "",
enabled_worker_ids: str = "[]",
):
if not multi_job_id:
return self._run_base(noise, guider, sampler, sigmas, latent_image, denoise_mask, h_tiles, w_tiles, overlap)
latent_tensor, rebuild_latent = _extract_tensor_and_rebuilder(latent_image)
if latent_tensor.ndim != 5:
raise ValueError(
f"LTXTiledSamplerDistributed expects a 5D latent tensor [B,C,F,H,W]; got shape {tuple(latent_tensor.shape)}"
)
mask_tensor = _extract_optional_mask(denoise_mask)
self.multi_job_id = str(multi_job_id)
self.master_url = str(master_url or "").rstrip("/")
self.worker_id = str(worker_id or "")
if isinstance(enabled_worker_ids, str):
try:
import json
parsed_workers = json.loads(enabled_worker_ids) if enabled_worker_ids else []
except Exception:
parsed_workers = []
else:
parsed_workers = enabled_worker_ids or []
self.enabled_worker_ids = [str(worker) for worker in parsed_workers]
if is_worker:
run_async_in_server_loop(
_ltx_worker_process(
self,
noise,
guider,
sampler,
sigmas,
latent_tensor,
mask_tensor,
int(h_tiles),
int(w_tiles),
int(overlap),
),
timeout=None,
)
placeholder = rebuild_latent(latent_tensor)
return (placeholder, placeholder, placeholder)
if not self.enabled_worker_ids:
samples, denoised = _ltx_local_tiled_process(
noise,
guider,
sampler,
sigmas,
latent_tensor,
mask_tensor,
int(h_tiles),
int(w_tiles),
int(overlap),
)
else:
samples, denoised = run_async_in_server_loop(
_ltx_master_process(
self,
noise,
guider,
sampler,
sigmas,
latent_tensor,
mask_tensor,
int(h_tiles),
int(w_tiles),
int(overlap),
),
timeout=None,
)
samples_latent = rebuild_latent(samples)
denoised_latent = rebuild_latent(denoised if denoised is not None else samples)
return (samples_latent, denoised_latent, samples_latent)
NODE_CLASS_MAPPINGS = {"LTXTiledSamplerDistributed": LTXTiledSamplerDistributed}
NODE_DISPLAY_NAME_MAPPINGS = {"LTXTiledSamplerDistributed": "LTX Tiled Sampler (Distributed)"}
+2 -2
View File
@@ -269,7 +269,7 @@ class ImageBatchDivider:
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
def INPUT_TYPES(s):
@@ -282,7 +282,7 @@ class AudioBatchDivider:
"max": 10,
"step": 1,
"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]
+2 -2
View File
@@ -1,9 +1,9 @@
[project]
name = "ComfyUI-Distributed"
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"}
dependencies = ["safetensors"]
dependencies = []
[project.urls]
Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
+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(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):
bad = {
"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
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():
collector = _load_collector_module().DistributedCollectorNode()
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["job_id"] for payload in posted_payloads} == {"job-list-1"}
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"])
-157
View File
@@ -1,157 +0,0 @@
import asyncio
import importlib.util
import sys
import types
import unittest
from pathlib import Path
import torch
def _load_ltx_module():
module_path = Path(__file__).resolve().parents[1] / "nodes" / "ltx_tiled_sampler.py"
package_name = "dist_ltx_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
for subpkg in ["nodes", "upscale", "utils"]:
pkg = types.ModuleType(f"{package_name}.{subpkg}")
pkg.__path__ = []
sys.modules[f"{package_name}.{subpkg}"] = pkg
job_models = types.ModuleType(f"{package_name}.upscale.job_models")
class TileJobState:
pass
job_models.TileJobState = TileJobState
sys.modules[f"{package_name}.upscale.job_models"] = job_models
async def _noop(*_args, **_kwargs):
return None
job_store = types.ModuleType(f"{package_name}.upscale.job_store")
job_store._cleanup_job = _noop
job_store._drain_results_queue = _noop
job_store._get_completed_count = _noop
job_store._mark_task_completed = _noop
job_store.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
distributed_tile_jobs_lock=asyncio.Lock(),
distributed_pending_tile_jobs={},
)
job_store.init_static_job_batched = _noop
sys.modules[f"{package_name}.upscale.job_store"] = job_store
job_timeout = types.ModuleType(f"{package_name}.upscale.job_timeout")
job_timeout._check_and_requeue_timed_out_workers = _noop
sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout
async_helpers = types.ModuleType(f"{package_name}.utils.async_helpers")
async_helpers.run_async_in_server_loop = lambda coro, timeout=None: asyncio.run(coro)
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers
constants = types.ModuleType(f"{package_name}.utils.constants")
constants.HEARTBEAT_INTERVAL = 10.0
constants.TILE_WAIT_TIMEOUT = 1.0
sys.modules[f"{package_name}.utils.constants"] = constants
logging = types.ModuleType(f"{package_name}.utils.logging")
logging.debug_log = lambda *_args, **_kwargs: None
logging.log = lambda *_args, **_kwargs: None
sys.modules[f"{package_name}.utils.logging"] = logging
network = types.ModuleType(f"{package_name}.utils.network")
async def _missing_session():
raise RuntimeError("network not available in unit test")
network.get_client_session = _missing_session
sys.modules[f"{package_name}.utils.network"] = network
usdu = types.ModuleType(f"{package_name}.utils.usdu_managment")
usdu._send_heartbeat_to_master = _noop
sys.modules[f"{package_name}.utils.usdu_managment"] = usdu
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.ltx_tiled_sampler", 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
ltx = _load_ltx_module()
class LtxTiledSamplerTests(unittest.TestCase):
def test_compute_tile_starts_with_overlap_covers_edges(self):
starts, tile_size = ltx._compute_tile_starts(total_size=10, n_tiles=3, overlap=2)
self.assertEqual(starts[0], 0)
self.assertEqual(starts[-1] + tile_size, 10)
self.assertGreaterEqual(tile_size, 4)
def test_pack_unpack_tensors_round_trips_with_safetensors(self):
tensor = torch.arange(12, dtype=torch.float32).reshape(1, 1, 3, 2, 2)
payload = ltx._pack_tensors({"samples": tensor})
restored = ltx._unpack_tensors(payload, device="cpu")
self.assertTrue(torch.equal(restored["samples"], tensor))
def test_latent_dict_rebuilder_preserves_metadata(self):
samples = torch.zeros((1, 2, 3, 4, 5))
latent = {"samples": samples, "noise_mask": "keep-me"}
extracted, rebuild = ltx._extract_tensor_and_rebuilder(latent)
rebuilt = rebuild(torch.ones_like(extracted))
self.assertIs(extracted, samples)
self.assertEqual(rebuilt["noise_mask"], "keep-me")
self.assertTrue(torch.equal(rebuilt["samples"], torch.ones_like(samples)))
def test_node_schema_includes_distributed_hidden_inputs(self):
schema = ltx.LTXTiledSamplerDistributed.INPUT_TYPES()
self.assertIn("multi_job_id", schema["hidden"])
self.assertIn("is_worker", schema["hidden"])
self.assertIn("master_url", schema["hidden"])
self.assertIn("enabled_worker_ids", schema["hidden"])
self.assertIn("LTXTiledSamplerDistributed", ltx.NODE_CLASS_MAPPINGS)
def test_local_tiled_process_runs_without_async_master_path(self):
class _Noise:
def generate_noise(self, latent):
return torch.ones_like(latent['samples']) * 0.25
class _Guider:
def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, disable_pbar=True):
return latent_image + noise
latent = torch.zeros((1, 1, 2, 4, 4), dtype=torch.float32)
samples, denoised = ltx._ltx_local_tiled_process(
_Noise(),
_Guider(),
sampler=object(),
sigmas=torch.tensor([1.0, 0.0]),
latent_tensor=latent,
denoise_mask=None,
n_h_tiles=2,
n_w_tiles=2,
tile_overlap=0,
)
self.assertIsNone(denoised)
self.assertEqual(tuple(samples.shape), tuple(latent.shape))
self.assertTrue(torch.allclose(samples, torch.full_like(latent, 0.25)))
if __name__ == "__main__":
unittest.main()
+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():
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
return {
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
self.assertEqual(len(preview_nodes), 1)
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):
result = pt.prune_prompt_for_worker(_collector_only_prompt())
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):
prompt = {
"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]}},
}
result = pt.prune_prompt_for_worker(prompt)
@@ -285,12 +323,22 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
placeholder_id = empty_nodes[0][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):
"""Two collectors → two placeholders."""
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {}},
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 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"])
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";
const NODE_CLASS = "DistributedValue";
const CONVERTED_WIDGET = "converted-widget";
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
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));
}
function hideWidgetForGood(node, widget, suffix = "") {
function hideWidgetForGood(widget) {
if (!widget) return;
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
widget.origType = widget.type;
widget.origComputeSize = widget.computeSize;
widget.origSerializeValue = widget.serializeValue;
widget.computeSize = () => [0, -4];
widget.type = `${CONVERTED_WIDGET}${suffix}`;
// These are serialized backing fields, not widgets converted into sockets.
// Keep their native type/serializer and mutate the shared options in place
// so both the classic canvas and Nodes 2.0 suppress the whole widget row.
widget.options ??= {};
widget.options.hidden = true;
// Hide any attached DOM element (multiline widgets).
if (widget.element) widget.element.style.display = "none";
@@ -58,14 +55,14 @@ function hideWidgetForGood(node, widget, suffix = "") {
if (widget.linkedWidgets) {
for (const linked of widget.linkedWidgets) {
hideWidgetForGood(node, linked, `:${widget.name}`);
hideWidgetForGood(linked);
}
}
}
function hideRawWidgets(node) {
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
hideWidgetForGood(getRawDefaultWidget(node));
hideWidgetForGood(getRawWorkerValuesWidget(node));
}
function removeDynamicDefaultWidget(node) {
+2 -2
View File
@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
createCheckboxSetting(
"setting-debug",
"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,
(event) => extension._updateSetting("debug", event.target.checked)
)
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
createNumberSetting(
"setting-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,
10,
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": [],
"properties": {},
"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",
"bgcolor": "#653"