Compare commits

..
Author SHA1 Message Date
Robert Wojciechowski abb7f2dd31 test: remove redundant coverage and consolidate happy paths 2026-10-10 22:20:34 +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
Robert Wojciechowski db01824d6e Update pyproject.toml 2026-06-06 09:10:43 +10:00
Robert Wojciechowski 6e7818bee6 fix: preserve delegate-only master config inputs (#86)
Fixes #85
2026-06-05 17:12:25 +10:00
Robert Wojciechowski c246d965e7 Update pyproject.toml 2026-05-28 09:00:17 +10:00
Robert Wojciechowski d6ce26f143 Merge pull request #84 from robertvoy/fix/issue-83
Fix collector handling of ComfyUI list inputs
2026-05-28 08:59:02 +10:00
20 changed files with 1140 additions and 175 deletions
+2 -2
View File
@@ -147,7 +147,7 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
Control your distributed cluster programmatically without opening the browser. Control your distributed cluster programmatically without opening the browser.
* **Endpoint:** `POST /distributed/queue` * **Endpoint:** `POST /distributed/queue`
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID. * **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md) * **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN). > **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
@@ -175,7 +175,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
| **Distributed Value** | Outputs per-worker override values with fallback to default | | **Distributed Value** | Outputs per-worker override values with fallback to default |
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers | | **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
| **Image Batch Divider** | Splits image batches for multi-GPU output | | **Image Batch Divider** | Splits image batches for multi-GPU output |
| **Audio Batch Divider** | Splits audio batches for multi-GPU output | | **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode | | **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work | | **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
+11 -4
View File
@@ -295,17 +295,24 @@ async def job_complete_endpoint(request):
errors.append("worker_id: expected non-empty string") errors.append("worker_id: expected non-empty string")
if not isinstance(batch_idx, int) or batch_idx < 0: if not isinstance(batch_idx, int) or batch_idx < 0:
errors.append("batch_idx: expected non-negative integer") errors.append("batch_idx: expected non-negative integer")
if not isinstance(image_payload, str) or not image_payload.strip(): image_provided = image_payload is not None
audio_provided = audio_payload is not None
if not image_provided and not audio_provided:
errors.append("expected at least one of image or audio")
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
errors.append("image: expected non-empty base64 PNG string") errors.append("image: expected non-empty base64 PNG string")
if audio_payload is not None and not isinstance(audio_payload, dict): if audio_provided and not isinstance(audio_payload, dict):
errors.append("audio: expected object when provided") errors.append("audio: expected object when provided")
if not isinstance(is_last, bool): if not isinstance(is_last, bool):
errors.append("is_last: expected boolean") errors.append("is_last: expected boolean")
if errors: if errors:
return await handle_api_error(request, errors, 400) return await handle_api_error(request, errors, 400)
tensor = _decode_canonical_png_tensor(image_payload) try:
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
except ValueError as exc:
return await handle_api_error(request, exc, 400)
multi_job_id = job_id.strip() multi_job_id = job_id.strip()
worker_id = worker_id.strip() worker_id = worker_id.strip()
+237 -10
View File
@@ -125,6 +125,209 @@ def _find_upstream_nodes(prompt_obj, start_ids):
return connected return connected
_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
"PrimitiveBoolean",
"PrimitiveFloat",
"PrimitiveInt",
"PrimitiveNode",
"PrimitiveString",
}
_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
"LoadImage",
}
_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
# only after its connected inputs recursively prove safe.
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
"CreateList": ("inputs.",),
}
# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
def _get_delegate_master_node_class_mappings():
"""Return ComfyUI node-class mappings when available."""
if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
try:
import nodes as comfy_nodes # type: ignore
except Exception: # pragma: no cover - depends on ComfyUI runtime imports
return {}
return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
def _get_delegate_master_node_class(class_type):
mappings = _get_delegate_master_node_class_mappings()
return mappings.get(class_type) if isinstance(mappings, dict) else None
def _normalize_delegate_master_return_type(return_type):
if return_type is None:
return ""
return str(return_type).strip().upper()
def _delegate_master_type_is_safe_scalar(type_name):
return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
def _delegate_master_type_is_safe_config(type_name):
return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
def _delegate_master_output_is_safe_scalar(class_type, output_index):
"""Return True when a registered node output is lightweight config data."""
if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
return True
node_class = _get_delegate_master_node_class(class_type)
return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
try:
output_type = return_types[int(output_index)]
except (IndexError, TypeError, ValueError):
return False
return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
def _get_delegate_master_input_types(class_type):
node_class = _get_delegate_master_node_class(class_type)
input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
if callable(input_types):
try:
input_types = input_types()
except TypeError:
return {}
return input_types if isinstance(input_types, dict) else {}
def _normalize_delegate_master_input_type(input_spec):
if isinstance(input_spec, (list, tuple)) and input_spec:
return _normalize_delegate_master_return_type(input_spec[0])
return _normalize_delegate_master_return_type(input_spec)
def _delegate_master_input_is_safe_scalar(class_type, input_name):
"""Return True when a registered downstream input expects config data."""
for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
if input_name.startswith(prefix):
return True
input_types = _get_delegate_master_input_types(class_type)
for section_name in ("required", "optional"):
section = input_types.get(section_name, {})
if isinstance(section, dict) and input_name in section:
input_type = _normalize_delegate_master_input_type(section[input_name])
return _delegate_master_type_is_safe_config(input_type)
return False
def _is_delegate_master_always_retained_upstream_node(node):
if not isinstance(node, dict):
return False
class_type = node.get("class_type")
return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
def _is_delegate_master_retained_upstream_node(node, output_index=0):
"""Return True for lightweight upstream nodes safe to keep on the master."""
if not isinstance(node, dict):
return False
class_type = node.get("class_type")
if not isinstance(class_type, str):
return False
return (
class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
or class_type.startswith("Primitive")
or _delegate_master_output_is_safe_scalar(class_type, output_index)
)
def _collect_delegate_master_retained_upstream_branch(
prompt_obj,
node_id,
output_index,
memo,
visiting,
):
"""Return safe retained branch nodes, or None when the branch is not safe."""
node_id = str(node_id)
cache_key = (node_id, output_index)
if cache_key in memo:
cached = memo[cache_key]
return None if cached is None else set(cached)
if cache_key in visiting:
memo[cache_key] = None
return None
node = prompt_obj.get(node_id)
if not _is_delegate_master_retained_upstream_node(node, output_index):
memo[cache_key] = None
return None
visiting.add(cache_key)
retained = {node_id}
inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
class_type = node.get("class_type") if isinstance(node, dict) else None
for input_name, value in inputs.items():
if not (isinstance(value, list) and len(value) == 2):
continue
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
visiting.remove(cache_key)
memo[cache_key] = None
return None
source_id = str(value[0])
branch = _collect_delegate_master_retained_upstream_branch(
prompt_obj,
source_id,
value[1],
memo,
visiting,
)
if branch is None:
visiting.remove(cache_key)
memo[cache_key] = None
return None
retained.update(branch)
visiting.remove(cache_key)
memo[cache_key] = frozenset(retained)
return retained
def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
"""Return lightweight upstream nodes needed by kept delegate-master nodes."""
connected = set()
memo = {}
for node_id in start_ids:
node = prompt_obj.get(str(node_id)) or {}
inputs = node.get("inputs", {})
class_type = node.get("class_type") if isinstance(node, dict) else None
for input_name, value in inputs.items():
if not (isinstance(value, list) and len(value) == 2):
continue
source_node = prompt_obj.get(str(value[0]))
if _is_delegate_master_always_retained_upstream_node(source_node):
connected.add(str(value[0]))
continue
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
continue
branch = _collect_delegate_master_retained_upstream_branch(
prompt_obj,
value[0],
value[1],
memo,
set(),
)
if branch is not None:
connected.update(branch)
return connected
def prune_prompt_for_worker(prompt_obj): def prune_prompt_for_worker(prompt_obj):
"""Prune worker prompt to distributed nodes and their upstream dependencies.""" """Prune worker prompt to distributed nodes and their upstream dependencies."""
collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector") collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
@@ -148,16 +351,33 @@ def prune_prompt_for_worker(prompt_obj):
downstream = _find_downstream_nodes(prompt_obj, [dist_id]) downstream = _find_downstream_nodes(prompt_obj, [dist_id])
has_removed_downstream = any(node_id != dist_id for node_id in downstream) has_removed_downstream = any(node_id != dist_id for node_id in downstream)
if has_removed_downstream: if has_removed_downstream:
preview_id = next_id() original_node = prompt_obj.get(str(dist_id), {})
pruned_prompt[preview_id] = { class_type = original_node.get("class_type")
"inputs": { inputs = original_node.get("inputs", {})
"images": [dist_id, 0], image_connected = class_type != "DistributedCollector" or (
}, isinstance(inputs.get("images"), list)
"class_type": "PreviewImage", and len(inputs["images"]) == 2
"_meta": { )
"title": "Preview Image (auto-added)", audio_connected = (
}, class_type == "DistributedCollector"
} and isinstance(inputs.get("audio"), list)
and len(inputs["audio"]) == 2
)
if image_connected:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {"images": [dist_id, 0]},
"class_type": "PreviewImage",
"_meta": {"title": "Preview Image (auto-added)"},
}
elif audio_connected:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {"audio": [dist_id, 1]},
"class_type": "PreviewAudio",
"_meta": {"title": "Preview Audio (auto-added)"},
}
return pruned_prompt return pruned_prompt
@@ -167,6 +387,9 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
downstream = _find_downstream_nodes(prompt_obj, collector_ids) downstream = _find_downstream_nodes(prompt_obj, collector_ids)
nodes_to_keep = set(collector_ids) nodes_to_keep = set(collector_ids)
nodes_to_keep.update(downstream) nodes_to_keep.update(downstream)
nodes_to_keep.update(
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
)
pruned_prompt = {} pruned_prompt = {}
for node_id in nodes_to_keep: for node_id in nodes_to_keep:
@@ -194,6 +417,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
collector_entry = pruned_prompt.get(collector_id) collector_entry = pruned_prompt.get(collector_id)
if not collector_entry: if not collector_entry:
continue continue
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
original_images = original_inputs.get("images")
if not (isinstance(original_images, list) and len(original_images) == 2):
continue
placeholder_id = next_id() placeholder_id = next_id()
pruned_prompt[placeholder_id] = { pruned_prompt[placeholder_id] = {
"class_type": "DistributedEmptyImage", "class_type": "DistributedEmptyImage",
+18 -8
View File
@@ -19,8 +19,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI: - `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`). - Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
- Resolves enabled/selected workers. - Resolves enabled/selected workers.
- Pings workers (`GET /prompt`) to include only reachable ones. - Probes and dispatches workers through `/distributed/worker_ws` by default.
- Dispatches the workflow to workers (`POST /prompt`). - If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
- Queues the master workflow in ComfyUI’s prompt queue. - Queues the master workflow in ComfyUI’s prompt queue.
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run. - If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
@@ -61,7 +61,8 @@ Queue a workflow for distributed execution.
#### Fields #### Fields
- `prompt` (required unless `workflow.prompt` is present, object) - `prompt` (required unless `workflow.prompt` is present, object)
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`. - A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
- This is not the normal visual workflow export from the ComfyUI editor.
- `workflow` (optional, object) - `workflow` (optional, object)
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`. - Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
- If you don’t care about UI metadata, you can omit it. - If you don’t care about UI metadata, you can omit it.
@@ -123,10 +124,14 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
## Worker requirements (important) ## Worker requirements (important)
For a worker to participate, it must be reachable from the master: For a worker to participate, it must be reachable from the master. By default:
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
If `settings.websocket_orchestration=false`:
- Health check: `GET <worker-base>/prompt` must return HTTP 200. - Health check: `GET <worker-base>/prompt` must return HTTP 200.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow. - Dispatch: `POST <worker-base>/prompt` must accept the prompt.
Also, for collector-based flows: Also, for collector-based flows:
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
## Examples ## Examples
### 1) Minimal `curl` ### 1) Minimal `curl` request envelope
```bash ```bash
curl -X POST "http://127.0.0.1:8188/distributed/queue" \ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -209,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
-d @payload.json -d @payload.json
``` ```
Where `payload.json` contains at least: `payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
```json ```json
{ {
"prompt": { "prompt": {
"1": {"class_type": "KSampler", "inputs": {} } "<node_id>": {
"class_type": "<node class>",
"inputs": {"<required_input>": "<value or connection>"}
}
}, },
"enabled_worker_ids": [], "enabled_worker_ids": [],
"client_id": "external-client" "client_id": "external-client"
} }
``` ```
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
### 2) Python (`requests`) ### 2) Python (`requests`)
```python ```python
+2 -2
View File
@@ -79,7 +79,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk) 📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
**On Runpod:** **On Runpod:**
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!** > If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account. 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes). 2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
- SAGE_ATTENTION: optional optimisation (set to true/false) - SAGE_ATTENTION: optional optimisation (set to true/false)
5. Deploy your pod. 5. Deploy your pod.
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system. 6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below: 7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
``` ```
# Download from CivitAI # Download from CivitAI
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
+1 -1
View File
@@ -26,6 +26,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedModelName": "Distributed Model Name", "DistributedModelName": "Distributed Model Name",
"DistributedValue": "Distributed Value", "DistributedValue": "Distributed Value",
"ImageBatchDivider": "Image Batch Divider", "ImageBatchDivider": "Image Batch Divider",
"AudioBatchDivider": "Audio Batch Divider", "AudioBatchDivider": "Audio Segment Divider",
"DistributedEmptyImage": "Distributed Empty Image", "DistributedEmptyImage": "Distributed Empty Image",
} }
+73 -45
View File
@@ -29,7 +29,6 @@ class DistributedCollectorNode:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"images": ("IMAGE",),
"load_balance": ( "load_balance": (
"BOOLEAN", "BOOLEAN",
{ {
@@ -38,7 +37,10 @@ class DistributedCollectorNode:
}, },
), ),
}, },
"optional": { "audio": ("AUDIO",) }, "optional": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
},
"hidden": { "hidden": {
"multi_job_id": ("STRING", {"default": ""}), "multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}), "is_worker": ("BOOLEAN", {"default": False}),
@@ -102,8 +104,9 @@ class DistributedCollectorNode:
return None return None
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate} return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False): def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
images = self._normalize_images_input(images) if images is not None:
images = self._normalize_images_input(images)
audio = self._normalize_audio_input(audio) audio = self._normalize_audio_input(audio)
load_balance = self._unwrap_list_input(load_balance) load_balance = self._unwrap_list_input(load_balance)
multi_job_id = self._unwrap_list_input(multi_job_id) multi_job_id = self._unwrap_list_input(multi_job_id)
@@ -115,6 +118,14 @@ class DistributedCollectorNode:
pass_through = self._unwrap_list_input(pass_through) pass_through = self._unwrap_list_input(pass_through)
delegate_only = self._unwrap_list_input(delegate_only) delegate_only = self._unwrap_list_input(delegate_only)
remote_only_master = (
bool(multi_job_id)
and not is_worker
and (delegate_only or is_master_delegate_only())
)
if images is None and audio is None and not remote_only_master:
raise ValueError("DistributedCollector requires at least one image or audio input")
# Create empty audio if not provided # Create empty audio if not provided
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100} empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -141,41 +152,56 @@ class DistributedCollectorNode:
return result return result
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id): async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Send image batch to master via canonical JSON envelopes.""" """Send an image batch, optionally with audio, or an audio-only completion."""
batch_size = image_batch.shape[0]
if batch_size == 0:
return
encoded_audio = encode_audio_payload(audio) encoded_audio = encode_audio_payload(audio)
session = await get_client_session() session = await get_client_session()
url = f"{master_url}/distributed/job_complete" url = f"{master_url}/distributed/job_complete"
for batch_idx in range(batch_size):
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
byte_io = io.BytesIO()
img.save(byte_io, format='PNG', compress_level=0)
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
payload = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": int(batch_idx),
"image": f"data:image/png;base64,{encoded_image}",
"is_last": bool(batch_idx == batch_size - 1),
}
if payload["is_last"] and encoded_audio is not None:
payload["audio"] = encoded_audio
payloads = []
batch_size = 0 if image_batch is None else image_batch.shape[0]
if batch_size == 0:
if encoded_audio is None:
raise ValueError("Worker completion requires image or audio data")
payloads.append(
{
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": 0,
"audio": encoded_audio,
"is_last": True,
}
)
else:
for batch_idx in range(batch_size):
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
byte_io = io.BytesIO()
img.save(byte_io, format='PNG', compress_level=0)
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
payload = {
"job_id": str(multi_job_id),
"worker_id": str(worker_id),
"batch_idx": int(batch_idx),
"image": f"data:image/png;base64,{encoded_image}",
"is_last": bool(batch_idx == batch_size - 1),
}
if payload["is_last"] and encoded_audio is not None:
payload["audio"] = encoded_audio
payloads.append(payload)
for payload in payloads:
timeout_seconds = 60 if "image" in payload else 600
try: try:
async with session.post( async with session.post(
url, url,
json=payload, json=payload,
timeout=aiohttp.ClientTimeout(total=60), timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response: ) as response:
response.raise_for_status() response.raise_for_status()
except Exception as e: except Exception as e:
log(f"Worker - Failed to send canonical image envelope to master: {e}") media_type = "image/audio" if "image" in payload else "audio-only"
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
debug_log(f"Worker - Full error details: URL={url}") debug_log(f"Worker - Full error details: URL={url}")
raise # Re-raise to handle at caller level raise
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None): def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
"""Combine audio from master and workers into a single audio output. """Combine audio from master and workers into a single audio output.
@@ -257,8 +283,8 @@ class DistributedCollectorNode:
images_on_cpu, images_on_cpu,
delegate_mode: bool, delegate_mode: bool,
fallback_images, fallback_images,
) -> torch.Tensor: ):
"""Assemble final tensor: master first, then workers in enabled order.""" """Assemble final tensor, or return None when the job contains only audio."""
ordered_tensors = [] ordered_tensors = []
if not delegate_mode and images_on_cpu is not None: if not delegate_mode and images_on_cpu is not None:
for i in range(master_batch_size): for i in range(master_batch_size):
@@ -289,15 +315,15 @@ class DistributedCollectorNode:
if cpu_tensors: if cpu_tensors:
return ensure_contiguous(torch.cat(cpu_tensors, dim=0)) return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
elif fallback_images is not None: if fallback_images is not None:
return ensure_contiguous(fallback_images) return ensure_contiguous(fallback_images)
else: return None
raise ValueError("No image data collected from master or workers")
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False): async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
if is_worker: if is_worker:
# Worker mode: send images and audio to master in a single batch # Worker mode: send images and audio to master in a single batch
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") image_count = 0 if images is None else images.shape[0]
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id) await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
return (images, audio if audio is not None else self.EMPTY_AUDIO) return (images, audio if audio is not None else self.EMPTY_AUDIO)
else: else:
@@ -332,14 +358,15 @@ class DistributedCollectorNode:
master_audio = None master_audio = None
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers") debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
else: else:
images_on_cpu = images.cpu() if images is None:
master_batch_size = images.shape[0] images_on_cpu = None
master_batch_size = 0
else:
images_on_cpu = ensure_contiguous(images.cpu())
master_batch_size = images.shape[0]
master_audio = audio # Keep master's audio for later master_audio = audio # Keep master's audio for later
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...") debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
# Ensure master images are contiguous
images_on_cpu = ensure_contiguous(images_on_cpu)
# Initialize storage for collected images and audio # Initialize storage for collected images and audio
worker_images = {} # Dict to store images by worker_id and index worker_images = {} # Dict to store images by worker_id and index
@@ -511,18 +538,19 @@ class DistributedCollectorNode:
if multi_job_id in prompt_server.distributed_pending_jobs: if multi_job_id in prompt_server.distributed_pending_jobs:
del prompt_server.distributed_pending_jobs[multi_job_id] del prompt_server.distributed_pending_jobs[multi_job_id]
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
try: try:
combined = self._reorder_and_combine_tensors( combined = self._reorder_and_combine_tensors(
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
) )
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total " if combined is None:
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") debug_log(f"Master - Job {multi_job_id} complete with audio only")
else:
# Combine audio from master and workers debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers) f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
return (combined, combined_audio) return (combined, combined_audio)
except Exception as e: except Exception as e:
log(f"Master - Error combining images: {e}") log(f"Master - Error combining images: {e}")
# Return just the master images as fallback # Preserve collected audio even when image assembly fails.
return (images, audio if audio is not None else self.EMPTY_AUDIO) return (images, combined_audio)
+6 -6
View File
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
""" """
Distributed version of Ultimate SD Upscale (No Upscale). Distributed version of Ultimate SD Upscale (No Upscale).
Supports three processing modes: Supports two currently selected processing modes:
1. Single GPU: No workers available, process everything locally 1. Single GPU: No workers available, process everything locally
2. Static Mode: Small batches, distributes tiles across workers (flattened) 2. Distributed tile queue: Workers pull tiles from a shared queue
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
Features: Features:
- Multi-mode batch handling for efficient video/image upscaling - Tile-based batch handling for video/image upscaling
- Tiled VAE support for memory efficiency - Tiled VAE support for memory efficiency
- Dynamic load balancing for large batches - Shared work queue so faster workers can process more tiles
- Backward compatible with single-image workflows - Backward compatible with single-image workflows
Environment Variables: Environment Variables:
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20) - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB) - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
Threshold: dynamic_threshold input controls mode switch (default 8) The hidden dynamic_threshold input is retained for workflow compatibility but
does not affect the current mode-selection policy.
""" """
def __init__(self): def __init__(self):
+2 -2
View File
@@ -269,7 +269,7 @@ class ImageBatchDivider:
class AudioBatchDivider: class AudioBatchDivider:
"""Divides an audio waveform into multiple parts along the time/samples dimension.""" """Divides an audio waveform into sequential segments along the time dimension."""
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
@@ -282,7 +282,7 @@ class AudioBatchDivider:
"max": 10, "max": 10,
"step": 1, "step": 1,
"display": "number", "display": "number",
"tooltip": "Number of parts to divide the audio into" "tooltip": "Number of sequential time segments to create"
}), }),
} }
} }
+1 -1
View File
@@ -1,7 +1,7 @@
[project] [project]
name = "ComfyUI-Distributed" name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud" description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.4.4" version = "1.4.7"
license = {file = "LICENSE"} license = {file = "LICENSE"}
dependencies = [] dependencies = []
+71
View File
@@ -302,6 +302,77 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(queued["audio"]["sample_rate"], 44100) self.assertEqual(queued["audio"]["sample_rate"], 44100)
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4)) self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
async def test_job_complete_accepts_audio_without_image(self):
queue = asyncio.Queue()
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
request = _FakeRequest(
{
"job_id": "audio-only-job",
"worker_id": "worker-1",
"batch_idx": 0,
"audio": self._encoded_audio_payload(),
"is_last": True,
}
)
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 200)
decode_image.assert_not_called()
queued = await queue.get()
self.assertIsNone(queued["tensor"])
self.assertEqual(queued["audio"]["sample_rate"], 44100)
async def test_job_complete_rejects_payload_without_image_or_audio(self):
request = _FakeRequest(
{
"job_id": "job-1",
"worker_id": "worker-1",
"batch_idx": 0,
"is_last": True,
}
)
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 400)
self.assertIn("image or audio", response.payload.get("message", "").lower())
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
request = _FakeRequest(
{
"job_id": "job-1",
"worker_id": "worker-1",
"batch_idx": 0,
"image": 123,
"audio": self._encoded_audio_payload(),
"is_last": True,
}
)
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 400)
self.assertIn("image", response.payload.get("message", "").lower())
async def test_job_complete_returns_400_for_malformed_audio(self):
request = _FakeRequest(
{
"job_id": "job-1",
"worker_id": "worker-1",
"batch_idx": 0,
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
"is_last": True,
}
)
response = await job_routes.job_complete_endpoint(request)
self.assertEqual(response.status, 400)
self.assertIn("audio.shape", response.payload.get("message", "").lower())
def test_decode_audio_payload_rejects_bad_shape(self): def test_decode_audio_payload_rejects_bad_shape(self):
bad = { bad = {
"sample_rate": 44100, "sample_rate": 44100,
-15
View File
@@ -164,26 +164,11 @@ class ConvertPathsForPlatformTests(unittest.TestCase):
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class FindMediaReferencesTests(unittest.TestCase): class FindMediaReferencesTests(unittest.TestCase):
def test_finds_image_input(self):
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
refs = ms._find_media_references(prompt)
self.assertIn("photo.png", refs)
def test_finds_video_input(self):
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
refs = ms._find_media_references(prompt)
self.assertIn("clip.mp4", refs)
def test_finds_file_input_for_load_video(self): def test_finds_file_input_for_load_video(self):
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}} prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
refs = ms._find_media_references(prompt) refs = ms._find_media_references(prompt)
self.assertIn("1 - Copy.mp4", refs) self.assertIn("1 - Copy.mp4", refs)
def test_finds_audio_input(self):
prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
refs = ms._find_media_references(prompt)
self.assertIn("track.wav", refs)
def test_strips_annotation_suffix(self): def test_strips_annotation_suffix(self):
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}} prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
refs = ms._find_media_references(prompt) refs = ms._find_media_references(prompt)
+159
View File
@@ -124,6 +124,49 @@ def test_collector_opts_into_comfyui_list_inputs():
assert collector.INPUT_IS_LIST is True assert collector.INPUT_IS_LIST is True
def test_collector_exposes_images_as_optional_input():
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
assert "images" not in input_types["required"]
assert input_types["optional"]["images"] == ("IMAGE",)
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
collector = _load_collector_module().DistributedCollectorNode()
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
images, returned_audio = collector.run(images=None, audio=[audio])
assert images is None
assert returned_audio is audio
def test_collector_rejects_missing_images_and_audio():
collector = _load_collector_module().DistributedCollectorNode()
try:
collector.run(images=None, audio=None)
except ValueError as exc:
assert "image or audio" in str(exc).lower()
else:
raise AssertionError("Expected collector to reject a run with no media input")
def test_delegate_only_master_allows_no_local_media_input():
collector = _load_collector_module().DistributedCollectorNode()
images, audio = collector.run(
images=None,
audio=None,
multi_job_id=["delegate-audio-job"],
delegate_only=[True],
enabled_worker_ids=["[]"],
)
assert images is None
assert tuple(audio["waveform"].shape) == (1, 2, 1)
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs(): def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
collector = _load_collector_module().DistributedCollectorNode() collector = _load_collector_module().DistributedCollectorNode()
first = torch.zeros(1, 2, 2, 3) first = torch.zeros(1, 2, 2, 3)
@@ -203,3 +246,119 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
assert [payload["is_last"] for payload in posted_payloads] == [False, True] assert [payload["is_last"] for payload in posted_payloads] == [False, True]
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"} assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"} assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
def test_audio_only_worker_sends_one_completion_without_image():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
posted = []
class _FakeResponse:
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
def raise_for_status(self):
return None
class _FakeSession:
def post(self, url, json, timeout):
posted.append((url, json, timeout.total))
return _FakeResponse()
async def _fake_get_client_session():
return _FakeSession()
module.get_client_session = _fake_get_client_session
module.encode_audio_payload = lambda value: {"encoded": value is audio}
images, returned_audio = collector.run(
images=None,
audio=[audio],
multi_job_id=["audio-job"],
is_worker=[True],
master_url=["http://master"],
worker_id=["worker-a"],
)
assert images is None
assert returned_audio is audio
assert len(posted) == 1
assert posted[0][0] == "http://master/distributed/job_complete"
assert posted[0][2] == 600
assert posted[0][1] == {
"job_id": "audio-job",
"worker_id": "worker-a",
"batch_idx": 0,
"audio": {"encoded": True},
"is_last": True,
}
def test_audio_only_master_combines_local_and_worker_audio():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait(
{
"worker_id": "worker-a",
"image_index": 0,
"tensor": None,
"audio": worker_audio,
"is_last": True,
}
)
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
images, combined_audio = asyncio.run(
collector.execute(
images=None,
audio=master_audio,
multi_job_id="audio-job",
enabled_worker_ids='["worker-a"]',
)
)
assert images is None
assert combined_audio["sample_rate"] == 48000
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
module = _load_collector_module()
collector = module.DistributedCollectorNode()
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
queue = asyncio.Queue()
queue.put_nowait(
{
"worker_id": "worker-a",
"image_index": 0,
"tensor": None,
"audio": worker_audio,
"is_last": True,
}
)
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
images, combined_audio = asyncio.run(
collector.execute(
images=None,
audio=None,
multi_job_id="delegate-audio-job",
enabled_worker_ids='["worker-a"]',
delegate_only=True,
)
)
assert images is None
assert combined_audio["sample_rate"] == 48000
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
+10 -22
View File
@@ -86,27 +86,25 @@ class ParseTilesFromFormTests(unittest.TestCase):
# --- happy paths --- # --- happy paths ---
def test_single_tile_returns_one_entry(self): def test_single_tile_returns_image_and_metadata(self):
tiles = pp._parse_tiles_from_form(_make_form(1)) tiles = pp._parse_tiles_from_form(_make_form(1))
self.assertEqual(len(tiles), 1) self.assertEqual(len(tiles), 1)
def test_multiple_tiles_all_returned(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(len(tiles), 3)
def test_tile_image_is_pil_image(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
self.assertIsInstance(tiles[0]["image"], PILImage.Image)
def test_tile_metadata_fields_are_parsed(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
tile = tiles[0] tile = tiles[0]
self.assertIsInstance(tile["image"], PILImage.Image)
self.assertEqual(tile["tile_idx"], 0) self.assertEqual(tile["tile_idx"], 0)
self.assertEqual(tile["x"], 0) self.assertEqual(tile["x"], 0)
self.assertEqual(tile["y"], 0) self.assertEqual(tile["y"], 0)
self.assertEqual(tile["extracted_width"], 64) self.assertEqual(tile["extracted_width"], 64)
self.assertEqual(tile["extracted_height"], 64) self.assertEqual(tile["extracted_height"], 64)
def test_multiple_tiles_preserve_count_order_and_coordinates(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(len(tiles), 3)
for i, tile in enumerate(tiles):
self.assertEqual(tile["tile_idx"], i)
self.assertEqual(tiles[1]["x"], 64)
self.assertEqual(tiles[2]["x"], 128)
def test_padding_is_parsed_from_form(self): def test_padding_is_parsed_from_form(self):
tiles = pp._parse_tiles_from_form(_make_form(1, padding=16)) tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
self.assertEqual(tiles[0]["padding"], 16) self.assertEqual(tiles[0]["padding"], 16)
@@ -138,16 +136,6 @@ class ParseTilesFromFormTests(unittest.TestCase):
self.assertNotIn("batch_idx", tiles[0]) self.assertNotIn("batch_idx", tiles[0])
self.assertNotIn("global_idx", tiles[0]) self.assertNotIn("global_idx", tiles[0])
def test_tile_indices_match_metadata_order(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
for i, tile in enumerate(tiles):
self.assertEqual(tile["tile_idx"], i)
def test_x_coordinates_reflect_metadata(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(tiles[1]["x"], 64)
self.assertEqual(tiles[2]["x"], 128)
# --- error cases --- # --- error cases ---
def test_missing_tiles_metadata_raises_value_error(self): def test_missing_tiles_metadata_raises_value_error(self):
+541 -9
View File
@@ -71,6 +71,15 @@ def _collector_only_prompt():
} }
def _audio_only_prompt():
"""1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
return {
"1": {"class_type": "LoadAudio", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
"3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
}
def _delegate_prompt(): def _delegate_prompt():
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)""" """1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
return { return {
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
self.assertEqual(len(preview_nodes), 1) self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0]) self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
def test_injects_preview_audio_for_audio_only_collector(self):
result = pt.prune_prompt_for_worker(_audio_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
def test_prefers_image_preview_when_collector_has_images_and_audio(self):
prompt = _linear_prompt()
prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
prompt["4"]["inputs"]["audio"] = ["6", 0]
result = pt.prune_prompt_for_worker(prompt)
self.assertEqual(
len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
1,
)
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
def test_preserves_image_preview_for_distributed_upscale(self):
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
def test_no_preview_image_when_no_downstream(self): def test_no_preview_image_when_no_downstream(self):
result = pt.prune_prompt_for_worker(_collector_only_prompt()) result = pt.prune_prompt_for_worker(_collector_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"] preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
@@ -242,7 +280,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
def test_upscale_node_is_treated_as_distributed(self): def test_upscale_node_is_treated_as_distributed(self):
prompt = { prompt = {
"1": {"class_type": "KSampler", "inputs": {}}, "1": {"class_type": "KSampler", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}}, "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}}, "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
} }
result = pt.prune_prompt_for_worker(prompt) result = pt.prune_prompt_for_worker(prompt)
@@ -264,8 +302,8 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
self.assertNotIn("1", result) self.assertNotIn("1", result)
self.assertNotIn("2", result) self.assertNotIn("2", result)
def test_removes_dangling_upstream_refs(self): def test_replaces_dangling_upstream_ref_with_one_empty_image_placeholder(self):
"""Collector must not retain dangling refs to pruned upstream nodes.""" """Collector must point to exactly one valid placeholder after pruning."""
prompt = _delegate_prompt() prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"]) result = pt.prepare_delegate_master_prompt(prompt, ["3"])
collector_inputs = result["3"].get("inputs", {}) collector_inputs = result["3"].get("inputs", {})
@@ -276,26 +314,520 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
self.assertNotEqual(source_id, "2") self.assertNotEqual(source_id, "2")
self.assertIn(source_id, result) self.assertIn(source_id, result)
self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage") self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
def test_injects_empty_image_placeholder(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"] empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(len(empty_nodes), 1) self.assertEqual(len(empty_nodes), 1)
placeholder_id = empty_nodes[0][0] placeholder_id = empty_nodes[0][0]
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0]) self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
def test_audio_only_collector_does_not_get_image_placeholder(self):
prompt = _audio_only_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(empty_nodes, [])
self.assertNotIn("images", result["2"].get("inputs", {}))
self.assertNotIn("audio", result["2"].get("inputs", {}))
def test_one_placeholder_per_collector(self): def test_one_placeholder_per_collector(self):
"""Two collectors → two placeholders.""" """Two collectors → two placeholders."""
prompt = { prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}}, "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
"2": {"class_type": "DistributedCollector", "inputs": {}}, "2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}, "3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
"10": {"class_type": "LoadImage", "inputs": {}},
"11": {"class_type": "LoadImage", "inputs": {}},
} }
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"]) result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"] empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(len(empty_nodes), 2) self.assertEqual(len(empty_nodes), 2)
def test_preserves_primitive_string_for_downstream_required_input(self):
"""Delegate-only master keeps primitive inputs needed by SaveImage."""
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {"SaveImage": save_image}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test/input_bug"}},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["15", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertEqual(result["15"], prompt["15"])
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
def test_preserves_load_image_for_switch_alternate_required_input(self):
"""Delegate-only master keeps LoadImage inputs needed by switches."""
prompt = {
"848": {"class_type": "LoadImage", "inputs": {"image": "input_bug2_00003_.png"}},
"862": {"class_type": "VAEDecode", "inputs": {}},
"854": {"class_type": "DistributedCollector", "inputs": {"images": ["862", 0]}},
"850": {
"class_type": "ComfySwitchNode",
"inputs": {
"on_false": ["848", 0],
"on_true": ["854", 0],
},
},
"851": {"class_type": "PreviewImage", "inputs": {"images": ["850", 0]}},
}
result = pt.prepare_delegate_master_prompt(prompt, ["854"])
self.assertIn("848", result)
self.assertEqual(result["850"]["inputs"]["on_false"], ["848", 0])
self.assertEqual(result["850"]["inputs"]["on_true"], ["854", 0])
self.assertNotIn("862", result)
def test_preserves_registered_string_utility_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps scalar utility chains used by SaveImage."""
string_concat = type(
"StringConcatenate",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"string_a": ("STRING",),
"string_b": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"SaveImage": save_image,
"StringConcatenate": string_concat,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"16": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug3"}},
"17": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["15", 0],
"string_b": ["16", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("16", result)
self.assertIn("17", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["17", 0])
self.assertEqual(result["17"]["inputs"]["string_a"], ["15", 0])
self.assertEqual(result["17"]["inputs"]["string_b"], ["16", 0])
def test_preserves_registered_multi_string_join_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps multi-input scalar utility chains."""
join_string_multi = type(
"JoinStringMulti",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {"string_1": ("STRING",)},
"optional": {
"string_2": ("STRING",),
"string_3": ("STRING",),
},
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"JoinStringMulti": join_string_multi,
"SaveImage": save_image,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"16": {"class_type": "PrimitiveString", "inputs": {"value": "input"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "bug4"}},
"18": {
"class_type": "JoinStringMulti",
"inputs": {
"string_1": ["15", 0],
"string_2": ["16", 0],
"string_3": ["17", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["18", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("16", result)
self.assertIn("17", result)
self.assertIn("18", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["18", 0])
def test_does_not_preserve_scalar_utility_with_heavy_upstream_dependency(self):
"""Scalar utility nodes are retained only when their full input branch is safe."""
string_concat = type(
"StringConcatenate",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"string_a": ("STRING",),
"string_b": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"SaveImage": save_image,
"StringConcatenate": string_concat,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"17": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["15", 0],
"string_b": ["8", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertNotIn("8", result)
self.assertNotIn("17", result)
self.assertNotIn("filename_prefix", result["9"]["inputs"])
def test_does_not_preserve_scalar_output_for_non_scalar_downstream_input(self):
"""Scalar outputs are retained only for scalar/config downstream inputs."""
string_provider = type("StringProvider", (), {"RETURN_TYPES": ("STRING",)})
image_consumer = type(
"ImageConsumer",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"mask": ("IMAGE",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"ImageConsumer": image_consumer,
"StringProvider": string_provider,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"17": {"class_type": "StringProvider", "inputs": {}},
"9": {
"class_type": "ImageConsumer",
"inputs": {
"images": ["11", 0],
"mask": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertNotIn("17", result)
self.assertNotIn("mask", result["9"]["inputs"])
def test_preserves_scalar_list_join_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps list-shaped scalar config chains."""
create_list = type(
"CreateList",
(),
{
"RETURN_TYPES": ("LIST",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {"inputs.input0": ("STRING",)},
"optional": {
"inputs.input1": ("STRING",),
"inputs.input2": ("STRING",),
},
}
),
},
)
string_data_list_join = type(
"StringDataListJoin",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"strings": ("LIST",),
"delimiter": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"CreateList": create_list,
"SaveImage": save_image,
"StringDataListJoin": string_data_list_join,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
"29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
"28": {
"class_type": "CreateList",
"inputs": {
"inputs.input0": ["15", 0],
"inputs.input1": ["17", 0],
"inputs.input2": ["29", 0],
},
},
"32": {
"class_type": "StringDataListJoin",
"inputs": {
"strings": ["28", 0],
"delimiter": "/",
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["32", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("17", result)
self.assertIn("28", result)
self.assertIn("29", result)
self.assertIn("32", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["15", 0])
self.assertNotIn("8", result)
def test_preserves_builtin_create_list_string_data_list_join_subgraph(self):
"""Delegate-only master handles ComfyUI 0.23 CreateList data-list output."""
string_data_list_join = type(
"StringDataListJoin",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_IS_LIST": True,
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"strings": ("STRING", {"forceInput": True}),
"sep": ("STRING", {"default": " "}),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"Basic data handling: StringDataListJoin": string_data_list_join,
"SaveImage": save_image,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
"28": {
"class_type": "CreateList",
"inputs": {
"inputs.input0": ["17", 0],
"inputs.input1": ["15", 0],
"inputs.input2": ["29", 0],
},
},
"32": {
"class_type": "Basic data handling: StringDataListJoin",
"inputs": {
"strings": ["28", 0],
"sep": "/",
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["32", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
for node_id in ("15", "17", "28", "29", "32"):
self.assertIn(node_id, result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["17", 0])
self.assertNotIn("8", result)
def test_result_is_independent_copy(self): def test_result_is_independent_copy(self):
prompt = _delegate_prompt() prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"]) result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+2 -2
View File
@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
createCheckboxSetting( createCheckboxSetting(
"setting-debug", "setting-debug",
"Debug Mode", "Debug Mode",
"Enable verbose logging in the browser console.", "Enable verbose logging in the browser console and ComfyUI server output.",
extension.config?.settings?.debug || false, extension.config?.settings?.debug || false,
(event) => extension._updateSetting("debug", event.target.checked) (event) => extension._updateSetting("debug", event.target.checked)
) )
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
createNumberSetting( createNumberSetting(
"setting-worker-timeout", "setting-worker-timeout",
"Worker Timeout", "Worker Timeout",
"Seconds without a heartbeat before a worker is considered timed out. Default 60.", "Maximum result-wait and heartbeat inactivity period before recovery begins. Busy workers may receive additional grace. Default: 60 seconds.",
extension.config?.settings?.worker_timeout_seconds ?? 60, extension.config?.settings?.worker_timeout_seconds ?? 60,
10, 10,
1, 1,
-15
View File
@@ -1,15 +0,0 @@
import { describe, expect, it } from "vitest";
import { buildWorkerWebSocketUrl } from "../urlUtils.js";
describe("execution decision helpers", () => {
it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
"ws://worker.local:8188/distributed/worker_ws"
);
expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
"wss://worker.example.com/distributed/worker_ws"
);
});
});
-5
View File
@@ -106,11 +106,6 @@ describe("buildWorkerWebSocketUrl", () => {
"wss://worker.example.com/distributed/worker_ws" "wss://worker.example.com/distributed/worker_ws"
); );
}); });
it("always appends /distributed/worker_ws", () => {
const url = buildWorkerWebSocketUrl("http://worker.local:8188");
expect(url.endsWith("/distributed/worker_ws")).toBe(true);
});
}); });
+3 -25
View File
@@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { getWorkerUrl } from "../workerLifecycle.js"; import { getWorkerUrl } from "../workerLifecycle.js";
describe("workerLifecycle URL construction", () => { describe("workerLifecycle URL wiring", () => {
let originalWindow; let originalWindow;
beforeEach(() => { beforeEach(() => {
@@ -22,31 +22,9 @@ describe("workerLifecycle URL construction", () => {
globalThis.window = originalWindow; globalThis.window = originalWindow;
}); });
it("builds local worker URL with explicit local port", () => { // URL variants are covered in urlUtils.test.js; keep the wrapper wiring here.
it("forwards the worker and endpoint using window.location", () => {
const worker = { id: "w1", port: 8189, type: "local" }; const worker = { id: "w1", port: 8189, type: "local" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt"); expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
}); });
it("builds remote worker URL with host:port", () => {
const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
});
it("builds cloud worker URL as https", () => {
const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt");
});
it("rewrites runpod proxy hostname for local worker ports", () => {
globalThis.window = {
location: {
hostname: "podabc.proxy.runpod.net",
protocol: "https:",
port: "",
origin: "https://podabc.proxy.runpod.net",
},
};
const worker = { id: "w4", port: 8189, type: "local" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt");
});
}); });
+1 -1
View File
@@ -814,7 +814,7 @@
"outputs": [], "outputs": [],
"properties": {}, "properties": {},
"widgets_values": [ "widgets_values": [
"If all your GPUs are the same/similar, set static_distribution to true\n" "Workers pull tiles from a shared queue, so faster workers can process more tiles automatically.\n"
], ],
"color": "#432", "color": "#432",
"bgcolor": "#653" "bgcolor": "#653"