Compare commits
25
Commits
@@ -1,4 +1,3 @@
|
||||
# These are supported funding model platforms
|
||||
|
||||
github: robertvoy
|
||||
buy_me_a_coffee: robertvoy
|
||||
|
||||
@@ -19,6 +19,6 @@ jobs:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
|
||||
---
|
||||
|
||||
## Key Features
|
||||
## Key Features
|
||||
|
||||
#### Parallel Workflow Processing
|
||||
- Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
|
||||
@@ -27,22 +27,12 @@
|
||||
- Intelligent distribution
|
||||
- Handles single images and videos
|
||||
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
- JSON configuration with UI controls
|
||||
|
||||
---
|
||||
|
||||
## Current Architecture
|
||||
|
||||
- Workflow-level load balancing is controlled by **Distributed Collector** via the `load_balance` toggle.
|
||||
- There is **no Distributed Queue node** anymore.
|
||||
- With `load_balance=true`, orchestration selects one least-busy execution participant:
|
||||
- If master participation is enabled, master is included as a candidate.
|
||||
- If master is in orchestrator-only mode, only workers are considered.
|
||||
|
||||
---
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
- JSON configuration with UI controls
|
||||
|
||||
---
|
||||
|
||||
## Worker Types
|
||||
|
||||
@@ -61,7 +51,6 @@ ComfyUI Distributed supports three types of workers:
|
||||
## Requirements
|
||||
|
||||
- ComfyUI
|
||||
> Note: Desktop app not currently supported
|
||||
- Multiple NVIDIA GPUs
|
||||
> No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
|
||||
- That's it
|
||||
@@ -92,19 +81,19 @@ Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a sp
|
||||
|
||||
## Workflow Examples
|
||||
|
||||
### Basic Parallel Generation
|
||||
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
|
||||
### Basic Parallel Generation
|
||||
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
|
||||
|
||||

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

|
||||
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
|
||||
1. Load your image
|
||||
2. Upscale with ESRGAN or similar
|
||||
@@ -153,40 +142,40 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
|
||||
|
||||
---
|
||||
|
||||
## Developer API
|
||||
## Developer API
|
||||
|
||||
Control your distributed cluster programmatically without opening the browser.
|
||||
|
||||
* **Endpoint:** `POST /distributed/queue`
|
||||
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
|
||||
* **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).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
|
||||
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio 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 |
|
||||
|
||||
@@ -206,7 +195,7 @@ No, it does not speed up the generation of a single image or video. Instead, it
|
||||
|
||||
<details>
|
||||
<summary>Does it work with the ComfyUI desktop app?</summary>
|
||||
Currently, it is not compatible with the ComfyUI desktop app.
|
||||
Yes, it does now.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
@@ -249,3 +238,4 @@ Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
+15
-22
@@ -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']
|
||||
|
||||
+14
-5
@@ -217,7 +217,7 @@ async def distributed_queue_endpoint(request):
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
try:
|
||||
prompt_id, worker_count = await orchestrate_distributed_execution(
|
||||
prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
|
||||
payload.prompt,
|
||||
payload.workflow_meta,
|
||||
payload.client_id,
|
||||
@@ -227,6 +227,8 @@ async def distributed_queue_endpoint(request):
|
||||
)
|
||||
return web.json_response({
|
||||
"prompt_id": prompt_id,
|
||||
"number": prompt_number,
|
||||
"node_errors": node_errors,
|
||||
"worker_count": worker_count,
|
||||
"auto_prepare_supported": True,
|
||||
})
|
||||
@@ -293,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()
|
||||
|
||||
|
||||
@@ -125,6 +125,209 @@ def _find_upstream_nodes(prompt_obj, start_ids):
|
||||
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):
|
||||
"""Prune worker prompt to distributed nodes and their upstream dependencies."""
|
||||
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])
|
||||
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
|
||||
|
||||
@@ -167,6 +387,9 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
nodes_to_keep.update(
|
||||
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
|
||||
)
|
||||
|
||||
pruned_prompt = {}
|
||||
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)
|
||||
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",
|
||||
|
||||
@@ -13,7 +13,7 @@ from ..utils.constants import (
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_master_url
|
||||
from ..utils.network import build_master_url, build_master_callback_url
|
||||
from ..utils.trace_logger import trace_debug
|
||||
from .schemas import parse_positive_float, parse_positive_int
|
||||
from .orchestration.dispatch import (
|
||||
@@ -144,6 +144,7 @@ async def _prepare_worker_payload(
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
trace_execution_id,
|
||||
worker_prep_semaphore,
|
||||
@@ -153,6 +154,11 @@ async def _prepare_worker_payload(
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
worker_master_url = build_master_callback_url(
|
||||
worker,
|
||||
config=config,
|
||||
prompt_server_instance=prompt_server,
|
||||
)
|
||||
|
||||
worker_type = str(worker.get("type") or "local").strip().lower()
|
||||
is_remote_like = bool(worker.get("host")) and worker_type != "local"
|
||||
@@ -167,7 +173,7 @@ async def _prepare_worker_payload(
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
worker_master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
@@ -202,7 +208,7 @@ async def orchestrate_distributed_execution(
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int]: (prompt_id, worker_count)
|
||||
tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
@@ -311,8 +317,18 @@ async def orchestrate_distributed_execution(
|
||||
|
||||
if not job_id_map:
|
||||
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
|
||||
prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
|
||||
return prompt_id, 0
|
||||
queue_result = await queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
queue_result["number"],
|
||||
0,
|
||||
queue_result.get("node_errors", {}),
|
||||
)
|
||||
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
@@ -360,6 +376,7 @@ async def orchestrate_distributed_execution(
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_semaphore,
|
||||
@@ -385,9 +402,17 @@ async def orchestrate_distributed_execution(
|
||||
]
|
||||
)
|
||||
|
||||
prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
|
||||
queue_result = await queue_prompt_payload(
|
||||
master_prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
prompt_id = queue_result["prompt_id"]
|
||||
prompt_number = queue_result["number"]
|
||||
node_errors = queue_result.get("node_errors", {})
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
|
||||
)
|
||||
return prompt_id, len(worker_payloads)
|
||||
return prompt_id, prompt_number, len(worker_payloads), node_errors
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
# 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.
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
@@ -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
-31
@@ -1,31 +1 @@
|
||||
from .utilities import (
|
||||
DistributedSeed,
|
||||
DistributedModelName,
|
||||
DistributedValue,
|
||||
ImageBatchDivider,
|
||||
AudioBatchDivider,
|
||||
DistributedEmptyImage,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
)
|
||||
from .collector import DistributedCollectorNode
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DistributedCollector": DistributedCollectorNode,
|
||||
"DistributedSeed": DistributedSeed,
|
||||
"DistributedModelName": DistributedModelName,
|
||||
"DistributedValue": DistributedValue,
|
||||
"ImageBatchDivider": ImageBatchDivider,
|
||||
"AudioBatchDivider": AudioBatchDivider,
|
||||
"DistributedEmptyImage": DistributedEmptyImage,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DistributedCollector": "Distributed Collector",
|
||||
"DistributedSeed": "Distributed Seed",
|
||||
"DistributedModelName": "Distributed Model Name",
|
||||
"DistributedValue": "Distributed Value",
|
||||
"ImageBatchDivider": "Image Batch Divider",
|
||||
"AudioBatchDivider": "Audio Batch Divider",
|
||||
"DistributedEmptyImage": "Distributed Empty Image",
|
||||
}
|
||||
"""Private execution helpers; public registration lives in nodes.v3."""
|
||||
|
||||
+131
-44
@@ -22,13 +22,13 @@ prompt_server = _server.PromptServer.instance
|
||||
|
||||
|
||||
class DistributedCollectorNode:
|
||||
INPUT_IS_LIST = True
|
||||
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"load_balance": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
@@ -37,7 +37,10 @@ class DistributedCollectorNode:
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": { "audio": ("AUDIO",) },
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
"audio": ("AUDIO",),
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
@@ -55,7 +58,74 @@ class DistributedCollectorNode:
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "image"
|
||||
|
||||
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):
|
||||
@staticmethod
|
||||
def _unwrap_list_input(value):
|
||||
"""Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
|
||||
if isinstance(value, (list, tuple)) and len(value) == 1:
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
def _normalize_images_input(self, images):
|
||||
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
|
||||
if isinstance(images, (list, tuple)):
|
||||
if not images:
|
||||
raise ValueError("Collector received an empty image list")
|
||||
if not all(isinstance(image, torch.Tensor) for image in images):
|
||||
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
|
||||
if len(images) == 1:
|
||||
return ensure_contiguous(images[0])
|
||||
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
|
||||
return ensure_contiguous(images)
|
||||
|
||||
def _normalize_audio_input(self, audio):
|
||||
"""Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
|
||||
if not isinstance(audio, (list, tuple)):
|
||||
return audio
|
||||
|
||||
audio_items = [item for item in audio if item is not None]
|
||||
if not audio_items:
|
||||
return None
|
||||
if len(audio_items) == 1:
|
||||
return audio_items[0]
|
||||
|
||||
waveforms = []
|
||||
sample_rate = 44100
|
||||
for item in audio_items:
|
||||
if not isinstance(item, dict):
|
||||
raise TypeError("Collector expected AUDIO list items to be dictionaries")
|
||||
waveform = item.get("waveform")
|
||||
if waveform is None or waveform.numel() == 0:
|
||||
continue
|
||||
waveforms.append(waveform)
|
||||
if sample_rate == 44100:
|
||||
sample_rate = item.get("sample_rate", 44100)
|
||||
|
||||
if not waveforms:
|
||||
return None
|
||||
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
|
||||
|
||||
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)
|
||||
is_worker = self._unwrap_list_input(is_worker)
|
||||
master_url = self._unwrap_list_input(master_url)
|
||||
enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
|
||||
worker_batch_size = self._unwrap_list_input(worker_batch_size)
|
||||
worker_id = self._unwrap_list_input(worker_id)
|
||||
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}
|
||||
|
||||
@@ -82,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.
|
||||
@@ -198,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):
|
||||
@@ -230,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:
|
||||
@@ -273,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
|
||||
@@ -452,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)
|
||||
|
||||
@@ -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)",
|
||||
}
|
||||
|
||||
+2
-2
@@ -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
@@ -0,0 +1,181 @@
|
||||
"""Native V3 schemas with the existing execution algorithms kept intact.
|
||||
|
||||
Runtime objects are private, per-execution helpers, not registered V1 nodes.
|
||||
This avoids sharing mutable instance state through V3's sanitized class clones.
|
||||
"""
|
||||
import comfy.samplers
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
|
||||
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
|
||||
max=1125899906842624, force_input=False)],
|
||||
outputs=[io.Int.Output(display_name='seed')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seed, is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedValue(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedValue', display_name='Distributed Value', category='utils',
|
||||
inputs=[io.String.Input('default_value', default=''),
|
||||
io.String.Input('worker_values', default='{}')],
|
||||
outputs=[io.AnyType.Output(display_name='value')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedValue().distribute(
|
||||
default_value, worker_values, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedModelName(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
|
||||
inputs=[io.String.Input('text', default='')],
|
||||
outputs=[io.AnyType.Output(display_name='output')],
|
||||
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, text):
|
||||
result = _utilities.DistributedModelName().log_input(
|
||||
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
|
||||
return io.NodeOutput(*result['result'], ui=result['ui'])
|
||||
|
||||
|
||||
class ImageBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
|
||||
inputs=[io.Image.Input('images'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of parts to divide the batch into')],
|
||||
# The existing frontend still displays only divide_by sockets.
|
||||
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, divide_by):
|
||||
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
|
||||
|
||||
|
||||
class AudioBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
|
||||
inputs=[io.Audio.Input('audio'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of sequential time segments to create')],
|
||||
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, divide_by):
|
||||
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
|
||||
|
||||
|
||||
class DistributedEmptyImage(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
|
||||
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('width', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
|
||||
outputs=[io.Image.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, height, width, channels):
|
||||
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
|
||||
|
||||
|
||||
class DistributedCollector(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
|
||||
inputs=[io.Boolean.Input('load_balance', default=False,
|
||||
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
|
||||
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)],
|
||||
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
|
||||
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
|
||||
worker_id='', pass_through=False, delegate_only=False):
|
||||
return io.NodeOutput(*_CollectorRuntime().run(
|
||||
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
|
||||
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
|
||||
worker_batch_size=worker_batch_size, worker_id=worker_id,
|
||||
pass_through=pass_through, delegate_only=delegate_only))
|
||||
|
||||
|
||||
class UltimateSDUpscaleDistributed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='UltimateSDUpscaleDistributed',
|
||||
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
|
||||
inputs=[
|
||||
io.Image.Input('upscaled_image'), io.Model.Input('model'),
|
||||
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
|
||||
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
|
||||
io.Int.Input('steps', default=20, min=1, max=10000),
|
||||
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
|
||||
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
|
||||
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
|
||||
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
|
||||
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('padding', default=32, min=0, max=256, step=8),
|
||||
io.Int.Input('mask_blur', default=8, min=0, max=256),
|
||||
io.Boolean.Input('force_uniform_tiles', default=True),
|
||||
io.Boolean.Input('tiled_decode', default=False),
|
||||
], outputs=[io.Image.Output()], accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, **kwargs):
|
||||
return _UpscaleRuntime.IS_CHANGED(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
|
||||
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
|
||||
return io.NodeOutput(*_UpscaleRuntime().run(
|
||||
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
|
||||
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
|
||||
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.4.1"
|
||||
version = "1.4.7"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Internal extension lifecycle support (not a node provider)."""
|
||||
@@ -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}')
|
||||
@@ -162,7 +162,7 @@ def _load_job_routes_module():
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
|
||||
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -221,7 +221,7 @@ job_routes = _load_job_routes_module()
|
||||
|
||||
|
||||
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_id(self):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
@@ -233,12 +233,14 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
with patch.object(
|
||||
job_routes,
|
||||
"orchestrate_distributed_execution",
|
||||
new=AsyncMock(return_value=("prompt_123", 2)),
|
||||
new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
|
||||
):
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
||||
self.assertEqual(response.payload.get("number"), 42)
|
||||
self.assertEqual(response.payload.get("node_errors"), {})
|
||||
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
||||
|
||||
async def test_distributed_queue_missing_prompt_returns_400(self):
|
||||
@@ -300,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
@@ -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)"
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
@@ -0,0 +1,91 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
def _load_async_helpers_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
|
||||
package_name = "dist_async_helpers_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
execution_module = types.ModuleType("execution")
|
||||
|
||||
async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
|
||||
return (True, None, ["9"], {})
|
||||
|
||||
execution_module.validate_prompt = _validate_prompt
|
||||
execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
|
||||
sys.modules["execution"] = execution_module
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
trigger_on_prompt=lambda payload: payload,
|
||||
number=12,
|
||||
prompt_queue=_PromptQueue(),
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_loop = lambda: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module, prompt_server
|
||||
|
||||
|
||||
async_helpers, prompt_server = _load_async_helpers_module()
|
||||
|
||||
|
||||
class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
|
||||
result = await async_helpers.queue_prompt_payload(
|
||||
{"1": {"class_type": "Node"}},
|
||||
workflow_meta={"id": "workflow-1"},
|
||||
client_id="client-1",
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
|
||||
self.assertIsInstance(result["prompt_id"], str)
|
||||
self.assertTrue(result["prompt_id"])
|
||||
self.assertEqual(result["number"], 12)
|
||||
self.assertEqual(result["node_errors"], {})
|
||||
|
||||
self.assertEqual(prompt_server.number, 13)
|
||||
self.assertEqual(len(prompt_server.prompt_queue.items), 1)
|
||||
queued_item = prompt_server.prompt_queue.items[0]
|
||||
self.assertEqual(queued_item[0], 12)
|
||||
extra_data = queued_item[3]
|
||||
self.assertEqual(extra_data["client_id"], "client-1")
|
||||
self.assertIn("create_time", extra_data)
|
||||
self.assertIsInstance(extra_data["create_time"], int)
|
||||
self.assertGreater(extra_data["create_time"], 0)
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,364 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_collector_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
||||
package_name = "dist_collector_list_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
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
def get(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
routes=_Routes(),
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_management = types.ModuleType("comfy.model_management")
|
||||
|
||||
class InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_management.InterruptProcessingException = InterruptProcessingException
|
||||
model_management.throw_exception_if_processing_interrupted = lambda: None
|
||||
comfy_module.model_management = model_management
|
||||
|
||||
comfy_utils = types.ModuleType("comfy.utils")
|
||||
|
||||
class ProgressBar:
|
||||
def __init__(self, _total):
|
||||
self.total = _total
|
||||
self.updates = []
|
||||
|
||||
def update(self, value):
|
||||
self.updates.append(value)
|
||||
|
||||
comfy_utils.ProgressBar = ProgressBar
|
||||
comfy_module.utils = comfy_utils
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_management
|
||||
sys.modules["comfy.utils"] = comfy_utils
|
||||
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 0.1
|
||||
config_module.load_config = lambda: {"workers": []}
|
||||
config_module.is_master_delegate_only = lambda: False
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 1.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
def _ensure_contiguous(tensor):
|
||||
return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
|
||||
|
||||
image_module.ensure_contiguous = _ensure_contiguous
|
||||
image_module.tensor_to_pil = lambda *_args, **_kwargs: None
|
||||
image_module.pil_to_tensor = lambda value: value
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker: "http://worker"
|
||||
network_module.get_client_session = lambda: None
|
||||
network_module.probe_worker = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
||||
audio_payload_module.encode_audio_payload = lambda audio: audio
|
||||
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", 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
|
||||
|
||||
|
||||
def test_collector_opts_into_comfyui_list_inputs():
|
||||
collector = _load_collector_module().DistributedCollectorNode
|
||||
|
||||
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)
|
||||
second = torch.ones(1, 2, 2, 3)
|
||||
|
||||
images, audio = collector.run(
|
||||
images=[first, second],
|
||||
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],
|
||||
)
|
||||
|
||||
assert tuple(images.shape) == (2, 2, 2, 3)
|
||||
assert torch.equal(images[0:1], first)
|
||||
assert torch.equal(images[1:2], second)
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
|
||||
|
||||
def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
first = torch.zeros(1, 2, 2, 3)
|
||||
second = torch.ones(1, 2, 2, 3)
|
||||
posted_payloads = []
|
||||
|
||||
class _FakeImage:
|
||||
def save(self, fp, format=None, compress_level=None):
|
||||
fp.write(b"png-bytes")
|
||||
|
||||
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_payloads.append(json)
|
||||
return _FakeResponse()
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return _FakeSession()
|
||||
|
||||
module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
|
||||
module.get_client_session = _fake_get_client_session
|
||||
module.encode_audio_payload = lambda _audio: None
|
||||
|
||||
images, audio = collector.run(
|
||||
images=[first, second],
|
||||
load_balance=[False],
|
||||
audio=[None],
|
||||
multi_job_id=["job-list-1"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
enabled_worker_ids=["[]"],
|
||||
worker_batch_size=[1],
|
||||
worker_id=["worker-a"],
|
||||
pass_through=[False],
|
||||
delegate_only=[False],
|
||||
)
|
||||
|
||||
assert tuple(images.shape) == (2, 2, 2, 3)
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
assert len(posted_payloads) == 2
|
||||
assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
|
||||
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"])
|
||||
@@ -90,14 +90,48 @@ class NetworkHelpersTests(unittest.TestCase):
|
||||
"https://master.example.com",
|
||||
)
|
||||
|
||||
def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
|
||||
cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8188",
|
||||
)
|
||||
|
||||
def test_build_master_url_keeps_explicit_port_in_host(self):
|
||||
cfg = {"master": {"host": "192.168.68.56:8001"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
def test_build_master_url_falls_back_to_server_address(self):
|
||||
cfg = {"master": {"host": ""}}
|
||||
cfg = {"master": {"host": "", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8190",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_uses_loopback_for_local_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8001",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -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,17 +323,522 @@ 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"]
|
||||
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_does_not_preserve_non_primitive_upstream_for_collector(self):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
|
||||
self.assertNotIn("2", result)
|
||||
self.assertNotEqual(result["3"]["inputs"]["images"], ["2", 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):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,134 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_process_module(module_filename: str):
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
|
||||
package_name = "dist_proc_testpkg"
|
||||
module_name = module_filename[:-3]
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
process_pkg = types.ModuleType(f"{package_name}.workers.process")
|
||||
process_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers.process"] = process_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
process_module = types.ModuleType(f"{package_name}.utils.process")
|
||||
process_module.get_python_executable = lambda: "/usr/bin/test-python"
|
||||
sys.modules[f"{package_name}.utils.process"] = process_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.process.{module_name}",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
root_discovery_module = _load_process_module("root_discovery.py")
|
||||
launch_builder_module = _load_process_module("launch_builder.py")
|
||||
|
||||
|
||||
class ComfyRootDiscoveryTests(unittest.TestCase):
|
||||
def test_prefers_loaded_comfyui_module_path(self):
|
||||
discovery = root_discovery_module.ComfyRootDiscovery()
|
||||
server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py")
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/opt/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(sys.modules, {"server": server_module}, clear=False), \
|
||||
patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \
|
||||
patch.dict(root_discovery_module.os.environ, {}, clear=True):
|
||||
self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI")
|
||||
|
||||
|
||||
class LaunchCommandBuilderTests(unittest.TestCase):
|
||||
def test_inherits_runtime_layout_args_for_desktop(self):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
runtime_args = Namespace(
|
||||
listen="127.0.0.1",
|
||||
base_directory="C:/Users/test/ComfyUI",
|
||||
temp_directory=None,
|
||||
input_directory="C:/Users/test/ComfyUI/input",
|
||||
output_directory="C:/Users/test/ComfyUI/output",
|
||||
user_directory="C:/Users/test/ComfyUI/user",
|
||||
front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app",
|
||||
extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]],
|
||||
enable_manager=True,
|
||||
disable_manager_ui=False,
|
||||
enable_manager_legacy_ui=False,
|
||||
windows_standalone_build=True,
|
||||
log_stdout=True,
|
||||
verbose="INFO",
|
||||
enable_cors_header="*",
|
||||
)
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
comfy_cli_args = types.ModuleType("comfy.cli_args")
|
||||
comfy_cli_args.args = runtime_args
|
||||
|
||||
worker_config = {
|
||||
"port": 9001,
|
||||
"extra_args": "--preview-method auto",
|
||||
}
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/desktop/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"comfy": comfy_module, "comfy.cli_args": comfy_cli_args},
|
||||
clear=False,
|
||||
), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists):
|
||||
cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI")
|
||||
|
||||
self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"])
|
||||
self.assertIn("--listen", cmd)
|
||||
self.assertIn("127.0.0.1", cmd)
|
||||
self.assertIn("--base-directory", cmd)
|
||||
self.assertIn("C:/Users/test/ComfyUI", cmd)
|
||||
self.assertIn("--input-directory", cmd)
|
||||
self.assertIn("--output-directory", cmd)
|
||||
self.assertIn("--user-directory", cmd)
|
||||
self.assertIn("--front-end-root", cmd)
|
||||
self.assertIn("--extra-model-paths-config", cmd)
|
||||
self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd)
|
||||
self.assertIn("--enable-manager", cmd)
|
||||
self.assertIn("--windows-standalone-build", cmd)
|
||||
self.assertIn("--log-stdout", cmd)
|
||||
self.assertIn("--disable-auto-launch", cmd)
|
||||
self.assertIn("--enable-cors-header", cmd)
|
||||
self.assertIn("*", cmd)
|
||||
self.assertIn("--port", cmd)
|
||||
self.assertIn("9001", cmd)
|
||||
self.assertNotIn("--auto-launch", cmd)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+16
-2
@@ -3,6 +3,7 @@ Async helper utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import execution
|
||||
import server
|
||||
@@ -104,7 +105,12 @@ class PromptValidationError(RuntimeError):
|
||||
super().__init__(f"Invalid prompt: {merged}")
|
||||
|
||||
|
||||
async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
||||
async def queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta=None,
|
||||
client_id=None,
|
||||
include_queue_metadata=False,
|
||||
):
|
||||
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
||||
payload = {"prompt": prompt_obj}
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
@@ -117,7 +123,7 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
raise PromptValidationError(error_payload, node_errors)
|
||||
|
||||
extra_data = {}
|
||||
extra_data = {"create_time": int(time.time() * 1000)}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if client_id:
|
||||
@@ -132,4 +138,12 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_server.prompt_queue.put(prompt_queue_item)
|
||||
|
||||
if include_queue_metadata:
|
||||
return {
|
||||
"prompt_id": prompt_id,
|
||||
"number": number,
|
||||
"node_errors": {},
|
||||
}
|
||||
|
||||
return prompt_id
|
||||
|
||||
+43
-8
@@ -66,6 +66,25 @@ def normalize_host(value):
|
||||
return host.split("/")[0]
|
||||
|
||||
|
||||
def _split_host_and_port(host):
|
||||
if not host:
|
||||
return host, None
|
||||
|
||||
if host.startswith("["):
|
||||
match = re.match(r"^(\[[^\]]+\])(?::(\d+))?$", host)
|
||||
if match:
|
||||
parsed_port = int(match.group(2)) if match.group(2) else None
|
||||
return match.group(1), parsed_port
|
||||
return host, None
|
||||
|
||||
if host.count(":") == 1:
|
||||
candidate_host, candidate_port = host.rsplit(":", 1)
|
||||
if candidate_port.isdigit():
|
||||
return candidate_host, int(candidate_port)
|
||||
|
||||
return host, None
|
||||
|
||||
|
||||
def build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
host = (worker.get("host") or "").strip()
|
||||
@@ -126,12 +145,7 @@ def build_master_url(config=None, prompt_server_instance=None):
|
||||
prompt_server_instance = prompt_server_instance or server.PromptServer.instance
|
||||
master_cfg = (config or {}).get("master", {}) or {}
|
||||
configured_host = (master_cfg.get("host") or "").strip()
|
||||
configured_port = master_cfg.get("port")
|
||||
default_port = getattr(prompt_server_instance, "port", 8188) or 8188
|
||||
try:
|
||||
port = int(configured_port or default_port)
|
||||
except (TypeError, ValueError):
|
||||
port = int(default_port)
|
||||
runtime_port = getattr(prompt_server_instance, "port", 8188) or 8188
|
||||
|
||||
def _needs_https(hostname):
|
||||
hostname = hostname.lower()
|
||||
@@ -149,10 +163,11 @@ def build_master_url(config=None, prompt_server_instance=None):
|
||||
if configured_host.startswith(("http://", "https://")):
|
||||
return configured_host.rstrip("/")
|
||||
|
||||
host = configured_host
|
||||
host, explicit_port = _split_host_and_port(configured_host)
|
||||
port = explicit_port if explicit_port is not None else int(runtime_port)
|
||||
scheme = "https" if _needs_https(host) or port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
if configured_port is None and scheme == "https" and _needs_https(host):
|
||||
if explicit_port is None and scheme == "https" and _needs_https(host):
|
||||
port = default_port_for_scheme
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{host}{port_part}"
|
||||
@@ -160,7 +175,27 @@ def build_master_url(config=None, prompt_server_instance=None):
|
||||
address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1"
|
||||
if address in ("0.0.0.0", "::"):
|
||||
address = "127.0.0.1"
|
||||
port = int(runtime_port)
|
||||
scheme = "https" if port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{address}{port_part}"
|
||||
|
||||
|
||||
def build_master_callback_url(worker, config=None, prompt_server_instance=None):
|
||||
"""Build the callback URL a specific worker should use to reach the master."""
|
||||
prompt_server_instance = prompt_server_instance or server.PromptServer.instance
|
||||
|
||||
worker_type = str((worker or {}).get("type") or "").strip().lower()
|
||||
worker_host = normalize_host((worker or {}).get("host"))
|
||||
local_hosts = {"", "localhost", "127.0.0.1", "::1", "[::1]", "0.0.0.0"}
|
||||
|
||||
is_local_worker = worker_type == "local" or worker_host in local_hosts
|
||||
if is_local_worker:
|
||||
port = int(getattr(prompt_server_instance, "port", 8188) or 8188)
|
||||
scheme = "https" if port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://127.0.0.1{port_part}"
|
||||
|
||||
return build_master_url(config=config, prompt_server_instance=prompt_server_instance)
|
||||
|
||||
+9
-12
@@ -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) {
|
||||
|
||||
@@ -4,7 +4,7 @@ import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js';
|
||||
import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
|
||||
|
||||
export function setupInterceptor(extension) {
|
||||
api.queuePrompt = async (number, prompt) => {
|
||||
api.queuePrompt = async (number, prompt, ...rest) => {
|
||||
if (extension.isEnabled) {
|
||||
const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
|
||||
const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
|
||||
@@ -18,7 +18,7 @@ export function setupInterceptor(extension) {
|
||||
return result;
|
||||
}
|
||||
}
|
||||
return extension.originalQueuePrompt(number, prompt);
|
||||
return extension.originalQueuePrompt(number, prompt, ...rest);
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -10,6 +10,62 @@ from ...utils.process import get_python_executable
|
||||
class LaunchCommandBuilder:
|
||||
"""Build command-lines for launching worker ComfyUI processes."""
|
||||
|
||||
def _extend_arg(self, cmd, flag, value):
|
||||
if value in (None, "", [], ()):
|
||||
return
|
||||
cmd.extend([flag, str(value)])
|
||||
|
||||
def _extend_grouped_args(self, cmd, flag, values):
|
||||
for group in values or []:
|
||||
flattened = [str(item) for item in group if item]
|
||||
if flattened:
|
||||
cmd.append(flag)
|
||||
cmd.extend(flattened)
|
||||
|
||||
def _get_runtime_args(self):
|
||||
try:
|
||||
from comfy.cli_args import args
|
||||
return args
|
||||
except Exception as exc:
|
||||
debug_log(f"Could not read current ComfyUI CLI args for worker launch: {exc}")
|
||||
return None
|
||||
|
||||
def _build_runtime_launch_args(self):
|
||||
args = self._get_runtime_args()
|
||||
if args is None:
|
||||
return []
|
||||
|
||||
inherited = []
|
||||
self._extend_arg(inherited, "--listen", getattr(args, "listen", None))
|
||||
self._extend_arg(inherited, "--base-directory", getattr(args, "base_directory", None))
|
||||
self._extend_arg(inherited, "--temp-directory", getattr(args, "temp_directory", None))
|
||||
self._extend_arg(inherited, "--input-directory", getattr(args, "input_directory", None))
|
||||
self._extend_arg(inherited, "--output-directory", getattr(args, "output_directory", None))
|
||||
self._extend_arg(inherited, "--user-directory", getattr(args, "user_directory", None))
|
||||
self._extend_arg(inherited, "--front-end-root", getattr(args, "front_end_root", None))
|
||||
self._extend_grouped_args(
|
||||
inherited,
|
||||
"--extra-model-paths-config",
|
||||
getattr(args, "extra_model_paths_config", None),
|
||||
)
|
||||
|
||||
if getattr(args, "enable_manager", False):
|
||||
inherited.append("--enable-manager")
|
||||
if getattr(args, "disable_manager_ui", False):
|
||||
inherited.append("--disable-manager-ui")
|
||||
if getattr(args, "enable_manager_legacy_ui", False):
|
||||
inherited.append("--enable-manager-legacy-ui")
|
||||
if getattr(args, "windows_standalone_build", False):
|
||||
inherited.append("--windows-standalone-build")
|
||||
if getattr(args, "log_stdout", False):
|
||||
inherited.append("--log-stdout")
|
||||
|
||||
verbose = getattr(args, "verbose", None)
|
||||
if verbose and verbose != "INFO":
|
||||
inherited.extend(["--verbose", str(verbose)])
|
||||
|
||||
return inherited
|
||||
|
||||
def _find_windows_terminal(self):
|
||||
"""Find Windows Terminal executable."""
|
||||
possible_paths = [
|
||||
@@ -39,10 +95,19 @@ class LaunchCommandBuilder:
|
||||
cmd = [
|
||||
get_python_executable(),
|
||||
main_py,
|
||||
"--port",
|
||||
str(worker_config["port"]),
|
||||
"--enable-cors-header",
|
||||
]
|
||||
cmd.extend(self._build_runtime_launch_args())
|
||||
cmd.extend(["--port", str(worker_config["port"])])
|
||||
|
||||
current_args = self._get_runtime_args()
|
||||
current_cors = getattr(current_args, "enable_cors_header", None) if current_args else None
|
||||
cmd.append("--enable-cors-header")
|
||||
if current_cors is not None:
|
||||
cmd.append(str(current_cors))
|
||||
|
||||
if "--disable-auto-launch" not in cmd:
|
||||
cmd.append("--disable-auto-launch")
|
||||
|
||||
debug_log(f"Using main.py: {main_py}")
|
||||
else:
|
||||
error_msg = f"Could not find main.py in {comfy_root}\n"
|
||||
|
||||
@@ -33,6 +33,7 @@ class ProcessLifecycle:
|
||||
env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0))
|
||||
env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
env["COMFYUI_MASTER_PID"] = str(os.getpid())
|
||||
env["COMFYUI_IS_WORKER"] = "1"
|
||||
|
||||
cmd = self._manager.build_launch_command(worker_config, comfy_root)
|
||||
cwd = comfy_root
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
|
||||
@@ -6,6 +7,21 @@ from ...utils.logging import debug_log, log
|
||||
class ComfyRootDiscovery:
|
||||
"""Resolve the ComfyUI root directory across local and container layouts."""
|
||||
|
||||
def _find_root_from_loaded_modules(self):
|
||||
"""Use already-imported ComfyUI modules to locate the runtime root."""
|
||||
for module_name in ("server", "folder_paths", "main"):
|
||||
module = sys.modules.get(module_name)
|
||||
module_file = getattr(module, "__file__", None)
|
||||
if not module_file:
|
||||
continue
|
||||
|
||||
candidate = os.path.dirname(os.path.abspath(module_file))
|
||||
if os.path.exists(os.path.join(candidate, "main.py")):
|
||||
debug_log(f"Found ComfyUI root via loaded module {module_name}: {candidate}")
|
||||
return candidate
|
||||
|
||||
return None
|
||||
|
||||
def find_comfy_root(self):
|
||||
# Start from current file location.
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
@@ -17,12 +33,17 @@ class ComfyRootDiscovery:
|
||||
debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}")
|
||||
return env_root
|
||||
|
||||
# Method 2: Try going up from custom_nodes directory.
|
||||
# Method 2: Inspect the already-loaded ComfyUI runtime modules.
|
||||
runtime_root = self._find_root_from_loaded_modules()
|
||||
if runtime_root:
|
||||
return runtime_root
|
||||
|
||||
# Method 3: Try going up from custom_nodes directory.
|
||||
if os.path.exists(os.path.join(potential_root, "main.py")):
|
||||
debug_log(f"Found ComfyUI root via directory traversal: {potential_root}")
|
||||
return potential_root
|
||||
|
||||
# Method 3: Look for common Docker paths.
|
||||
# Method 4: Look for common Docker paths.
|
||||
docker_paths = [
|
||||
"/basedir",
|
||||
"/ComfyUI",
|
||||
@@ -37,7 +58,7 @@ class ComfyRootDiscovery:
|
||||
debug_log(f"Found ComfyUI root in Docker path: {path}")
|
||||
return path
|
||||
|
||||
# Method 4: Search upwards for main.py.
|
||||
# Method 5: Search upwards for main.py.
|
||||
search_dir = current_dir
|
||||
for _ in range(5):
|
||||
if os.path.exists(os.path.join(search_dir, "main.py")):
|
||||
@@ -48,7 +69,7 @@ class ComfyRootDiscovery:
|
||||
break
|
||||
search_dir = parent
|
||||
|
||||
# Method 5: Try to import and use folder_paths.
|
||||
# Method 6: Try to import and use folder_paths.
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user