Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a2b341fb99 | ||
|
|
9de7e158e1 | ||
|
|
32ac027e0e | ||
|
|
fec959e1ac | ||
|
|
cedc8d4591 |
@@ -147,7 +147,7 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
|
|||||||
Control your distributed cluster programmatically without opening the browser.
|
Control your distributed cluster programmatically without opening the browser.
|
||||||
|
|
||||||
* **Endpoint:** `POST /distributed/queue`
|
* **Endpoint:** `POST /distributed/queue`
|
||||||
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
|
* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
|
||||||
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
|
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
|
||||||
|
|
||||||
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
||||||
@@ -175,7 +175,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
|
|||||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||||
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
|
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
|
||||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||||
|
|
||||||
|
|||||||
+15
-22
@@ -1,29 +1,22 @@
|
|||||||
# Import everything needed from the main module
|
"""ComfyUI-Distributed's native V3 extension entrypoint."""
|
||||||
from .distributed import (
|
from comfy_api.v0_0_2 import ComfyExtension, io
|
||||||
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
|
|
||||||
)
|
|
||||||
|
|
||||||
# Import utilities
|
from .nodes.v3 import NODES
|
||||||
from .utils.config import ensure_config_exists, CONFIG_FILE
|
from .runtime.bootstrap import initialize
|
||||||
from .utils.logging import debug_log
|
|
||||||
|
|
||||||
# Import distributed upscale nodes
|
WEB_DIRECTORY = './web'
|
||||||
from .nodes.distributed_upscale import (
|
|
||||||
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
|
|
||||||
)
|
|
||||||
|
|
||||||
WEB_DIRECTORY = "./web"
|
|
||||||
|
|
||||||
ensure_config_exists()
|
class DistributedExtension(ComfyExtension):
|
||||||
|
async def on_load(self) -> None:
|
||||||
|
initialize()
|
||||||
|
|
||||||
# Merge node mappings
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||||
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
|
return list(NODES)
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
|
|
||||||
|
|
||||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
|
||||||
|
|
||||||
debug_log("Loaded Distributed nodes.")
|
async def comfy_entrypoint() -> DistributedExtension:
|
||||||
debug_log(f"Config file: {CONFIG_FILE}")
|
return DistributedExtension()
|
||||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
|
||||||
|
|
||||||
|
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
|
||||||
|
|||||||
+11
-4
@@ -295,17 +295,24 @@ async def job_complete_endpoint(request):
|
|||||||
errors.append("worker_id: expected non-empty string")
|
errors.append("worker_id: expected non-empty string")
|
||||||
if not isinstance(batch_idx, int) or batch_idx < 0:
|
if not isinstance(batch_idx, int) or batch_idx < 0:
|
||||||
errors.append("batch_idx: expected non-negative integer")
|
errors.append("batch_idx: expected non-negative integer")
|
||||||
if not isinstance(image_payload, str) or not image_payload.strip():
|
image_provided = image_payload is not None
|
||||||
|
audio_provided = audio_payload is not None
|
||||||
|
if not image_provided and not audio_provided:
|
||||||
|
errors.append("expected at least one of image or audio")
|
||||||
|
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
|
||||||
errors.append("image: expected non-empty base64 PNG string")
|
errors.append("image: expected non-empty base64 PNG string")
|
||||||
if audio_payload is not None and not isinstance(audio_payload, dict):
|
if audio_provided and not isinstance(audio_payload, dict):
|
||||||
errors.append("audio: expected object when provided")
|
errors.append("audio: expected object when provided")
|
||||||
if not isinstance(is_last, bool):
|
if not isinstance(is_last, bool):
|
||||||
errors.append("is_last: expected boolean")
|
errors.append("is_last: expected boolean")
|
||||||
if errors:
|
if errors:
|
||||||
return await handle_api_error(request, errors, 400)
|
return await handle_api_error(request, errors, 400)
|
||||||
|
|
||||||
tensor = _decode_canonical_png_tensor(image_payload)
|
try:
|
||||||
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
|
tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
|
||||||
|
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
|
||||||
|
except ValueError as exc:
|
||||||
|
return await handle_api_error(request, exc, 400)
|
||||||
multi_job_id = job_id.strip()
|
multi_job_id = job_id.strip()
|
||||||
worker_id = worker_id.strip()
|
worker_id = worker_id.strip()
|
||||||
|
|
||||||
|
|||||||
@@ -351,16 +351,33 @@ def prune_prompt_for_worker(prompt_obj):
|
|||||||
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
||||||
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
|
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
|
||||||
if has_removed_downstream:
|
if has_removed_downstream:
|
||||||
preview_id = next_id()
|
original_node = prompt_obj.get(str(dist_id), {})
|
||||||
pruned_prompt[preview_id] = {
|
class_type = original_node.get("class_type")
|
||||||
"inputs": {
|
inputs = original_node.get("inputs", {})
|
||||||
"images": [dist_id, 0],
|
image_connected = class_type != "DistributedCollector" or (
|
||||||
},
|
isinstance(inputs.get("images"), list)
|
||||||
"class_type": "PreviewImage",
|
and len(inputs["images"]) == 2
|
||||||
"_meta": {
|
)
|
||||||
"title": "Preview Image (auto-added)",
|
audio_connected = (
|
||||||
},
|
class_type == "DistributedCollector"
|
||||||
}
|
and isinstance(inputs.get("audio"), list)
|
||||||
|
and len(inputs["audio"]) == 2
|
||||||
|
)
|
||||||
|
|
||||||
|
if image_connected:
|
||||||
|
preview_id = next_id()
|
||||||
|
pruned_prompt[preview_id] = {
|
||||||
|
"inputs": {"images": [dist_id, 0]},
|
||||||
|
"class_type": "PreviewImage",
|
||||||
|
"_meta": {"title": "Preview Image (auto-added)"},
|
||||||
|
}
|
||||||
|
elif audio_connected:
|
||||||
|
preview_id = next_id()
|
||||||
|
pruned_prompt[preview_id] = {
|
||||||
|
"inputs": {"audio": [dist_id, 1]},
|
||||||
|
"class_type": "PreviewAudio",
|
||||||
|
"_meta": {"title": "Preview Audio (auto-added)"},
|
||||||
|
}
|
||||||
|
|
||||||
return pruned_prompt
|
return pruned_prompt
|
||||||
|
|
||||||
@@ -400,6 +417,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
|||||||
collector_entry = pruned_prompt.get(collector_id)
|
collector_entry = pruned_prompt.get(collector_id)
|
||||||
if not collector_entry:
|
if not collector_entry:
|
||||||
continue
|
continue
|
||||||
|
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
|
||||||
|
original_images = original_inputs.get("images")
|
||||||
|
if not (isinstance(original_images, list) and len(original_images) == 2):
|
||||||
|
continue
|
||||||
placeholder_id = next_id()
|
placeholder_id = next_id()
|
||||||
pruned_prompt[placeholder_id] = {
|
pruned_prompt[placeholder_id] = {
|
||||||
"class_type": "DistributedEmptyImage",
|
"class_type": "DistributedEmptyImage",
|
||||||
|
|||||||
+1
-1
@@ -1,7 +1,7 @@
|
|||||||
# conftest.py — project-level pytest configuration.
|
# conftest.py — project-level pytest configuration.
|
||||||
#
|
#
|
||||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||||
# (from .distributed import ...) that fail when pytest tries to import it as a
|
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
|
||||||
# standalone module during Package.setup() for the root package node.
|
# standalone module during Package.setup() for the root package node.
|
||||||
#
|
#
|
||||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||||
|
|||||||
@@ -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:
|
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
|
||||||
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
|
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
|
||||||
- Resolves enabled/selected workers.
|
- Resolves enabled/selected workers.
|
||||||
- Pings workers (`GET /prompt`) to include only reachable ones.
|
- Probes and dispatches workers through `/distributed/worker_ws` by default.
|
||||||
- Dispatches the workflow to workers (`POST /prompt`).
|
- If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
|
||||||
- Queues the master workflow in ComfyUI’s prompt queue.
|
- Queues the master workflow in ComfyUI’s prompt queue.
|
||||||
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
|
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
|
||||||
|
|
||||||
@@ -61,7 +61,8 @@ Queue a workflow for distributed execution.
|
|||||||
#### Fields
|
#### Fields
|
||||||
|
|
||||||
- `prompt` (required unless `workflow.prompt` is present, object)
|
- `prompt` (required unless `workflow.prompt` is present, object)
|
||||||
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
|
- A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
|
||||||
|
- This is not the normal visual workflow export from the ComfyUI editor.
|
||||||
- `workflow` (optional, object)
|
- `workflow` (optional, object)
|
||||||
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
|
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
|
||||||
- If you don’t care about UI metadata, you can omit it.
|
- If you don’t care about UI metadata, you can omit it.
|
||||||
@@ -123,10 +124,14 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
|||||||
|
|
||||||
## Worker requirements (important)
|
## Worker requirements (important)
|
||||||
|
|
||||||
For a worker to participate, it must be reachable from the master:
|
For a worker to participate, it must be reachable from the master. By default:
|
||||||
|
|
||||||
|
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
|
||||||
|
|
||||||
|
If `settings.websocket_orchestration=false`:
|
||||||
|
|
||||||
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
|
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
|
||||||
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
|
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
|
||||||
|
|
||||||
Also, for collector-based flows:
|
Also, for collector-based flows:
|
||||||
|
|
||||||
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
|
|||||||
|
|
||||||
## Examples
|
## Examples
|
||||||
|
|
||||||
### 1) Minimal `curl`
|
### 1) Minimal `curl` request envelope
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||||
@@ -209,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
|||||||
-d @payload.json
|
-d @payload.json
|
||||||
```
|
```
|
||||||
|
|
||||||
Where `payload.json` contains at least:
|
`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
|
||||||
|
|
||||||
```json
|
```json
|
||||||
{
|
{
|
||||||
"prompt": {
|
"prompt": {
|
||||||
"1": {"class_type": "KSampler", "inputs": {} }
|
"<node_id>": {
|
||||||
|
"class_type": "<node class>",
|
||||||
|
"inputs": {"<required_input>": "<value or connection>"}
|
||||||
|
}
|
||||||
},
|
},
|
||||||
"enabled_worker_ids": [],
|
"enabled_worker_ids": [],
|
||||||
"client_id": "external-client"
|
"client_id": "external-client"
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
|
||||||
|
|
||||||
### 2) Python (`requests`)
|
### 2) Python (`requests`)
|
||||||
|
|
||||||
```python
|
```python
|
||||||
|
|||||||
@@ -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)
|
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
|
||||||
|
|
||||||
**On Runpod:**
|
**On Runpod:**
|
||||||
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
|
> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
|
||||||
|
|
||||||
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
||||||
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
|
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
|
||||||
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
|||||||
- SAGE_ATTENTION: optional optimisation (set to true/false)
|
- SAGE_ATTENTION: optional optimisation (set to true/false)
|
||||||
5. Deploy your pod.
|
5. Deploy your pod.
|
||||||
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
|
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
|
||||||
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
|
7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||||
```
|
```
|
||||||
# Download from CivitAI
|
# Download from CivitAI
|
||||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||||
|
|||||||
+1
-31
@@ -1,31 +1 @@
|
|||||||
from .utilities import (
|
"""Private execution helpers; public registration lives in nodes.v3."""
|
||||||
DistributedSeed,
|
|
||||||
DistributedModelName,
|
|
||||||
DistributedValue,
|
|
||||||
ImageBatchDivider,
|
|
||||||
AudioBatchDivider,
|
|
||||||
DistributedEmptyImage,
|
|
||||||
AnyType,
|
|
||||||
ByPassTypeTuple,
|
|
||||||
any_type,
|
|
||||||
)
|
|
||||||
from .collector import DistributedCollectorNode
|
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
|
||||||
"DistributedCollector": DistributedCollectorNode,
|
|
||||||
"DistributedSeed": DistributedSeed,
|
|
||||||
"DistributedModelName": DistributedModelName,
|
|
||||||
"DistributedValue": DistributedValue,
|
|
||||||
"ImageBatchDivider": ImageBatchDivider,
|
|
||||||
"AudioBatchDivider": AudioBatchDivider,
|
|
||||||
"DistributedEmptyImage": DistributedEmptyImage,
|
|
||||||
}
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
||||||
"DistributedCollector": "Distributed Collector",
|
|
||||||
"DistributedSeed": "Distributed Seed",
|
|
||||||
"DistributedModelName": "Distributed Model Name",
|
|
||||||
"DistributedValue": "Distributed Value",
|
|
||||||
"ImageBatchDivider": "Image Batch Divider",
|
|
||||||
"AudioBatchDivider": "Audio Batch Divider",
|
|
||||||
"DistributedEmptyImage": "Distributed Empty Image",
|
|
||||||
}
|
|
||||||
|
|||||||
+73
-45
@@ -29,7 +29,6 @@ class DistributedCollectorNode:
|
|||||||
def INPUT_TYPES(s):
|
def INPUT_TYPES(s):
|
||||||
return {
|
return {
|
||||||
"required": {
|
"required": {
|
||||||
"images": ("IMAGE",),
|
|
||||||
"load_balance": (
|
"load_balance": (
|
||||||
"BOOLEAN",
|
"BOOLEAN",
|
||||||
{
|
{
|
||||||
@@ -38,7 +37,10 @@ class DistributedCollectorNode:
|
|||||||
},
|
},
|
||||||
),
|
),
|
||||||
},
|
},
|
||||||
"optional": { "audio": ("AUDIO",) },
|
"optional": {
|
||||||
|
"images": ("IMAGE",),
|
||||||
|
"audio": ("AUDIO",),
|
||||||
|
},
|
||||||
"hidden": {
|
"hidden": {
|
||||||
"multi_job_id": ("STRING", {"default": ""}),
|
"multi_job_id": ("STRING", {"default": ""}),
|
||||||
"is_worker": ("BOOLEAN", {"default": False}),
|
"is_worker": ("BOOLEAN", {"default": False}),
|
||||||
@@ -102,8 +104,9 @@ class DistributedCollectorNode:
|
|||||||
return None
|
return None
|
||||||
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
|
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
|
||||||
|
|
||||||
def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
|
def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
|
||||||
images = self._normalize_images_input(images)
|
if images is not None:
|
||||||
|
images = self._normalize_images_input(images)
|
||||||
audio = self._normalize_audio_input(audio)
|
audio = self._normalize_audio_input(audio)
|
||||||
load_balance = self._unwrap_list_input(load_balance)
|
load_balance = self._unwrap_list_input(load_balance)
|
||||||
multi_job_id = self._unwrap_list_input(multi_job_id)
|
multi_job_id = self._unwrap_list_input(multi_job_id)
|
||||||
@@ -115,6 +118,14 @@ class DistributedCollectorNode:
|
|||||||
pass_through = self._unwrap_list_input(pass_through)
|
pass_through = self._unwrap_list_input(pass_through)
|
||||||
delegate_only = self._unwrap_list_input(delegate_only)
|
delegate_only = self._unwrap_list_input(delegate_only)
|
||||||
|
|
||||||
|
remote_only_master = (
|
||||||
|
bool(multi_job_id)
|
||||||
|
and not is_worker
|
||||||
|
and (delegate_only or is_master_delegate_only())
|
||||||
|
)
|
||||||
|
if images is None and audio is None and not remote_only_master:
|
||||||
|
raise ValueError("DistributedCollector requires at least one image or audio input")
|
||||||
|
|
||||||
# Create empty audio if not provided
|
# Create empty audio if not provided
|
||||||
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||||
|
|
||||||
@@ -141,41 +152,56 @@ class DistributedCollectorNode:
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
||||||
"""Send image batch to master via canonical JSON envelopes."""
|
"""Send an image batch, optionally with audio, or an audio-only completion."""
|
||||||
batch_size = image_batch.shape[0]
|
|
||||||
if batch_size == 0:
|
|
||||||
return
|
|
||||||
|
|
||||||
encoded_audio = encode_audio_payload(audio)
|
encoded_audio = encode_audio_payload(audio)
|
||||||
|
|
||||||
session = await get_client_session()
|
session = await get_client_session()
|
||||||
url = f"{master_url}/distributed/job_complete"
|
url = f"{master_url}/distributed/job_complete"
|
||||||
for batch_idx in range(batch_size):
|
|
||||||
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
|
|
||||||
byte_io = io.BytesIO()
|
|
||||||
img.save(byte_io, format='PNG', compress_level=0)
|
|
||||||
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
|
|
||||||
payload = {
|
|
||||||
"job_id": str(multi_job_id),
|
|
||||||
"worker_id": str(worker_id),
|
|
||||||
"batch_idx": int(batch_idx),
|
|
||||||
"image": f"data:image/png;base64,{encoded_image}",
|
|
||||||
"is_last": bool(batch_idx == batch_size - 1),
|
|
||||||
}
|
|
||||||
if payload["is_last"] and encoded_audio is not None:
|
|
||||||
payload["audio"] = encoded_audio
|
|
||||||
|
|
||||||
|
payloads = []
|
||||||
|
batch_size = 0 if image_batch is None else image_batch.shape[0]
|
||||||
|
if batch_size == 0:
|
||||||
|
if encoded_audio is None:
|
||||||
|
raise ValueError("Worker completion requires image or audio data")
|
||||||
|
payloads.append(
|
||||||
|
{
|
||||||
|
"job_id": str(multi_job_id),
|
||||||
|
"worker_id": str(worker_id),
|
||||||
|
"batch_idx": 0,
|
||||||
|
"audio": encoded_audio,
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
for batch_idx in range(batch_size):
|
||||||
|
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
|
||||||
|
byte_io = io.BytesIO()
|
||||||
|
img.save(byte_io, format='PNG', compress_level=0)
|
||||||
|
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
|
||||||
|
payload = {
|
||||||
|
"job_id": str(multi_job_id),
|
||||||
|
"worker_id": str(worker_id),
|
||||||
|
"batch_idx": int(batch_idx),
|
||||||
|
"image": f"data:image/png;base64,{encoded_image}",
|
||||||
|
"is_last": bool(batch_idx == batch_size - 1),
|
||||||
|
}
|
||||||
|
if payload["is_last"] and encoded_audio is not None:
|
||||||
|
payload["audio"] = encoded_audio
|
||||||
|
payloads.append(payload)
|
||||||
|
|
||||||
|
for payload in payloads:
|
||||||
|
timeout_seconds = 60 if "image" in payload else 600
|
||||||
try:
|
try:
|
||||||
async with session.post(
|
async with session.post(
|
||||||
url,
|
url,
|
||||||
json=payload,
|
json=payload,
|
||||||
timeout=aiohttp.ClientTimeout(total=60),
|
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
|
||||||
) as response:
|
) as response:
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log(f"Worker - Failed to send canonical image envelope to master: {e}")
|
media_type = "image/audio" if "image" in payload else "audio-only"
|
||||||
|
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
|
||||||
debug_log(f"Worker - Full error details: URL={url}")
|
debug_log(f"Worker - Full error details: URL={url}")
|
||||||
raise # Re-raise to handle at caller level
|
raise
|
||||||
|
|
||||||
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
|
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
|
||||||
"""Combine audio from master and workers into a single audio output.
|
"""Combine audio from master and workers into a single audio output.
|
||||||
@@ -257,8 +283,8 @@ class DistributedCollectorNode:
|
|||||||
images_on_cpu,
|
images_on_cpu,
|
||||||
delegate_mode: bool,
|
delegate_mode: bool,
|
||||||
fallback_images,
|
fallback_images,
|
||||||
) -> torch.Tensor:
|
):
|
||||||
"""Assemble final tensor: master first, then workers in enabled order."""
|
"""Assemble final tensor, or return None when the job contains only audio."""
|
||||||
ordered_tensors = []
|
ordered_tensors = []
|
||||||
if not delegate_mode and images_on_cpu is not None:
|
if not delegate_mode and images_on_cpu is not None:
|
||||||
for i in range(master_batch_size):
|
for i in range(master_batch_size):
|
||||||
@@ -289,15 +315,15 @@ class DistributedCollectorNode:
|
|||||||
|
|
||||||
if cpu_tensors:
|
if cpu_tensors:
|
||||||
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
||||||
elif fallback_images is not None:
|
if fallback_images is not None:
|
||||||
return ensure_contiguous(fallback_images)
|
return ensure_contiguous(fallback_images)
|
||||||
else:
|
return None
|
||||||
raise ValueError("No image data collected from master or workers")
|
|
||||||
|
|
||||||
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
|
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
|
||||||
if is_worker:
|
if is_worker:
|
||||||
# Worker mode: send images and audio to master in a single batch
|
# Worker mode: send images and audio to master in a single batch
|
||||||
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
|
image_count = 0 if images is None else images.shape[0]
|
||||||
|
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
|
||||||
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
|
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
|
||||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||||
else:
|
else:
|
||||||
@@ -332,14 +358,15 @@ class DistributedCollectorNode:
|
|||||||
master_audio = None
|
master_audio = None
|
||||||
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
|
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
|
||||||
else:
|
else:
|
||||||
images_on_cpu = images.cpu()
|
if images is None:
|
||||||
master_batch_size = images.shape[0]
|
images_on_cpu = None
|
||||||
|
master_batch_size = 0
|
||||||
|
else:
|
||||||
|
images_on_cpu = ensure_contiguous(images.cpu())
|
||||||
|
master_batch_size = images.shape[0]
|
||||||
master_audio = audio # Keep master's audio for later
|
master_audio = audio # Keep master's audio for later
|
||||||
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
|
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
|
||||||
|
|
||||||
# Ensure master images are contiguous
|
|
||||||
images_on_cpu = ensure_contiguous(images_on_cpu)
|
|
||||||
|
|
||||||
|
|
||||||
# Initialize storage for collected images and audio
|
# Initialize storage for collected images and audio
|
||||||
worker_images = {} # Dict to store images by worker_id and index
|
worker_images = {} # Dict to store images by worker_id and index
|
||||||
@@ -511,18 +538,19 @@ class DistributedCollectorNode:
|
|||||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||||
del prompt_server.distributed_pending_jobs[multi_job_id]
|
del prompt_server.distributed_pending_jobs[multi_job_id]
|
||||||
|
|
||||||
|
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
||||||
try:
|
try:
|
||||||
combined = self._reorder_and_combine_tensors(
|
combined = self._reorder_and_combine_tensors(
|
||||||
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
|
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
|
||||||
)
|
)
|
||||||
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
if combined is None:
|
||||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
debug_log(f"Master - Job {multi_job_id} complete with audio only")
|
||||||
|
else:
|
||||||
# Combine audio from master and workers
|
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
||||||
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||||
|
|
||||||
return (combined, combined_audio)
|
return (combined, combined_audio)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log(f"Master - Error combining images: {e}")
|
log(f"Master - Error combining images: {e}")
|
||||||
# Return just the master images as fallback
|
# Preserve collected audio even when image assembly fails.
|
||||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
return (images, combined_audio)
|
||||||
|
|||||||
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
|
|||||||
"""
|
"""
|
||||||
Distributed version of Ultimate SD Upscale (No Upscale).
|
Distributed version of Ultimate SD Upscale (No Upscale).
|
||||||
|
|
||||||
Supports three processing modes:
|
Supports two currently selected processing modes:
|
||||||
1. Single GPU: No workers available, process everything locally
|
1. Single GPU: No workers available, process everything locally
|
||||||
2. Static Mode: Small batches, distributes tiles across workers (flattened)
|
2. Distributed tile queue: Workers pull tiles from a shared queue
|
||||||
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
|
|
||||||
|
|
||||||
Features:
|
Features:
|
||||||
- Multi-mode batch handling for efficient video/image upscaling
|
- Tile-based batch handling for video/image upscaling
|
||||||
- Tiled VAE support for memory efficiency
|
- Tiled VAE support for memory efficiency
|
||||||
- Dynamic load balancing for large batches
|
- Shared work queue so faster workers can process more tiles
|
||||||
- Backward compatible with single-image workflows
|
- Backward compatible with single-image workflows
|
||||||
|
|
||||||
Environment Variables:
|
Environment Variables:
|
||||||
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
|
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
|
||||||
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
|
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
|
||||||
|
|
||||||
Threshold: dynamic_threshold input controls mode switch (default 8)
|
The hidden dynamic_threshold input is retained for workflow compatibility but
|
||||||
|
does not affect the current mode-selection policy.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
|
|||||||
|
|
||||||
# Ensure initialization before registering routes
|
# Ensure initialization before registering routes
|
||||||
ensure_tile_jobs_initialized()
|
ensure_tile_jobs_initialized()
|
||||||
|
|
||||||
# Node registration
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
|
||||||
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
|
|
||||||
}
|
|
||||||
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
||||||
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
|
|
||||||
}
|
|
||||||
|
|||||||
+2
-2
@@ -269,7 +269,7 @@ class ImageBatchDivider:
|
|||||||
|
|
||||||
|
|
||||||
class AudioBatchDivider:
|
class AudioBatchDivider:
|
||||||
"""Divides an audio waveform into multiple parts along the time/samples dimension."""
|
"""Divides an audio waveform into sequential segments along the time dimension."""
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def INPUT_TYPES(s):
|
||||||
@@ -282,7 +282,7 @@ class AudioBatchDivider:
|
|||||||
"max": 10,
|
"max": 10,
|
||||||
"step": 1,
|
"step": 1,
|
||||||
"display": "number",
|
"display": "number",
|
||||||
"tooltip": "Number of parts to divide the audio into"
|
"tooltip": "Number of sequential time segments to create"
|
||||||
}),
|
}),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+181
@@ -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]
|
[project]
|
||||||
name = "ComfyUI-Distributed"
|
name = "ComfyUI-Distributed"
|
||||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||||
version = "1.4.6"
|
version = "1.4.7"
|
||||||
license = {file = "LICENSE"}
|
license = {file = "LICENSE"}
|
||||||
dependencies = []
|
dependencies = []
|
||||||
|
|
||||||
|
|||||||
@@ -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}')
|
||||||
@@ -302,6 +302,77 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
|||||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||||
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
||||||
|
|
||||||
|
async def test_job_complete_accepts_audio_without_image(self):
|
||||||
|
queue = asyncio.Queue()
|
||||||
|
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||||
|
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
|
||||||
|
request = _FakeRequest(
|
||||||
|
{
|
||||||
|
"job_id": "audio-only-job",
|
||||||
|
"worker_id": "worker-1",
|
||||||
|
"batch_idx": 0,
|
||||||
|
"audio": self._encoded_audio_payload(),
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
|
||||||
|
response = await job_routes.job_complete_endpoint(request)
|
||||||
|
|
||||||
|
self.assertEqual(response.status, 200)
|
||||||
|
decode_image.assert_not_called()
|
||||||
|
queued = await queue.get()
|
||||||
|
self.assertIsNone(queued["tensor"])
|
||||||
|
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||||
|
|
||||||
|
async def test_job_complete_rejects_payload_without_image_or_audio(self):
|
||||||
|
request = _FakeRequest(
|
||||||
|
{
|
||||||
|
"job_id": "job-1",
|
||||||
|
"worker_id": "worker-1",
|
||||||
|
"batch_idx": 0,
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await job_routes.job_complete_endpoint(request)
|
||||||
|
|
||||||
|
self.assertEqual(response.status, 400)
|
||||||
|
self.assertIn("image or audio", response.payload.get("message", "").lower())
|
||||||
|
|
||||||
|
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
|
||||||
|
request = _FakeRequest(
|
||||||
|
{
|
||||||
|
"job_id": "job-1",
|
||||||
|
"worker_id": "worker-1",
|
||||||
|
"batch_idx": 0,
|
||||||
|
"image": 123,
|
||||||
|
"audio": self._encoded_audio_payload(),
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await job_routes.job_complete_endpoint(request)
|
||||||
|
|
||||||
|
self.assertEqual(response.status, 400)
|
||||||
|
self.assertIn("image", response.payload.get("message", "").lower())
|
||||||
|
|
||||||
|
async def test_job_complete_returns_400_for_malformed_audio(self):
|
||||||
|
request = _FakeRequest(
|
||||||
|
{
|
||||||
|
"job_id": "job-1",
|
||||||
|
"worker_id": "worker-1",
|
||||||
|
"batch_idx": 0,
|
||||||
|
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await job_routes.job_complete_endpoint(request)
|
||||||
|
|
||||||
|
self.assertEqual(response.status, 400)
|
||||||
|
self.assertIn("audio.shape", response.payload.get("message", "").lower())
|
||||||
|
|
||||||
def test_decode_audio_payload_rejects_bad_shape(self):
|
def test_decode_audio_payload_rejects_bad_shape(self):
|
||||||
bad = {
|
bad = {
|
||||||
"sample_rate": 44100,
|
"sample_rate": 44100,
|
||||||
|
|||||||
+691
@@ -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())
|
||||||
@@ -124,6 +124,49 @@ def test_collector_opts_into_comfyui_list_inputs():
|
|||||||
assert collector.INPUT_IS_LIST is True
|
assert collector.INPUT_IS_LIST is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_collector_exposes_images_as_optional_input():
|
||||||
|
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
|
||||||
|
|
||||||
|
assert "images" not in input_types["required"]
|
||||||
|
assert input_types["optional"]["images"] == ("IMAGE",)
|
||||||
|
|
||||||
|
|
||||||
|
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
|
||||||
|
collector = _load_collector_module().DistributedCollectorNode()
|
||||||
|
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||||
|
|
||||||
|
images, returned_audio = collector.run(images=None, audio=[audio])
|
||||||
|
|
||||||
|
assert images is None
|
||||||
|
assert returned_audio is audio
|
||||||
|
|
||||||
|
|
||||||
|
def test_collector_rejects_missing_images_and_audio():
|
||||||
|
collector = _load_collector_module().DistributedCollectorNode()
|
||||||
|
|
||||||
|
try:
|
||||||
|
collector.run(images=None, audio=None)
|
||||||
|
except ValueError as exc:
|
||||||
|
assert "image or audio" in str(exc).lower()
|
||||||
|
else:
|
||||||
|
raise AssertionError("Expected collector to reject a run with no media input")
|
||||||
|
|
||||||
|
|
||||||
|
def test_delegate_only_master_allows_no_local_media_input():
|
||||||
|
collector = _load_collector_module().DistributedCollectorNode()
|
||||||
|
|
||||||
|
images, audio = collector.run(
|
||||||
|
images=None,
|
||||||
|
audio=None,
|
||||||
|
multi_job_id=["delegate-audio-job"],
|
||||||
|
delegate_only=[True],
|
||||||
|
enabled_worker_ids=["[]"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert images is None
|
||||||
|
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||||
|
|
||||||
|
|
||||||
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
|
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
|
||||||
collector = _load_collector_module().DistributedCollectorNode()
|
collector = _load_collector_module().DistributedCollectorNode()
|
||||||
first = torch.zeros(1, 2, 2, 3)
|
first = torch.zeros(1, 2, 2, 3)
|
||||||
@@ -203,3 +246,119 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
|
|||||||
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
|
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
|
||||||
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
|
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
|
||||||
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
|
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_audio_only_worker_sends_one_completion_without_image():
|
||||||
|
module = _load_collector_module()
|
||||||
|
collector = module.DistributedCollectorNode()
|
||||||
|
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||||
|
posted = []
|
||||||
|
|
||||||
|
class _FakeResponse:
|
||||||
|
async def __aenter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, exc_type, exc, tb):
|
||||||
|
return False
|
||||||
|
|
||||||
|
def raise_for_status(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class _FakeSession:
|
||||||
|
def post(self, url, json, timeout):
|
||||||
|
posted.append((url, json, timeout.total))
|
||||||
|
return _FakeResponse()
|
||||||
|
|
||||||
|
async def _fake_get_client_session():
|
||||||
|
return _FakeSession()
|
||||||
|
|
||||||
|
module.get_client_session = _fake_get_client_session
|
||||||
|
module.encode_audio_payload = lambda value: {"encoded": value is audio}
|
||||||
|
|
||||||
|
images, returned_audio = collector.run(
|
||||||
|
images=None,
|
||||||
|
audio=[audio],
|
||||||
|
multi_job_id=["audio-job"],
|
||||||
|
is_worker=[True],
|
||||||
|
master_url=["http://master"],
|
||||||
|
worker_id=["worker-a"],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert images is None
|
||||||
|
assert returned_audio is audio
|
||||||
|
assert len(posted) == 1
|
||||||
|
assert posted[0][0] == "http://master/distributed/job_complete"
|
||||||
|
assert posted[0][2] == 600
|
||||||
|
assert posted[0][1] == {
|
||||||
|
"job_id": "audio-job",
|
||||||
|
"worker_id": "worker-a",
|
||||||
|
"batch_idx": 0,
|
||||||
|
"audio": {"encoded": True},
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_audio_only_master_combines_local_and_worker_audio():
|
||||||
|
module = _load_collector_module()
|
||||||
|
collector = module.DistributedCollectorNode()
|
||||||
|
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
|
||||||
|
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||||
|
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||||
|
queue = asyncio.Queue()
|
||||||
|
queue.put_nowait(
|
||||||
|
{
|
||||||
|
"worker_id": "worker-a",
|
||||||
|
"image_index": 0,
|
||||||
|
"tensor": None,
|
||||||
|
"audio": worker_audio,
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
|
||||||
|
|
||||||
|
images, combined_audio = asyncio.run(
|
||||||
|
collector.execute(
|
||||||
|
images=None,
|
||||||
|
audio=master_audio,
|
||||||
|
multi_job_id="audio-job",
|
||||||
|
enabled_worker_ids='["worker-a"]',
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert images is None
|
||||||
|
assert combined_audio["sample_rate"] == 48000
|
||||||
|
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
|
||||||
|
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
|
||||||
|
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
|
||||||
|
|
||||||
|
|
||||||
|
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
|
||||||
|
module = _load_collector_module()
|
||||||
|
collector = module.DistributedCollectorNode()
|
||||||
|
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||||
|
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||||
|
queue = asyncio.Queue()
|
||||||
|
queue.put_nowait(
|
||||||
|
{
|
||||||
|
"worker_id": "worker-a",
|
||||||
|
"image_index": 0,
|
||||||
|
"tensor": None,
|
||||||
|
"audio": worker_audio,
|
||||||
|
"is_last": True,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
|
||||||
|
|
||||||
|
images, combined_audio = asyncio.run(
|
||||||
|
collector.execute(
|
||||||
|
images=None,
|
||||||
|
audio=None,
|
||||||
|
multi_job_id="delegate-audio-job",
|
||||||
|
enabled_worker_ids='["worker-a"]',
|
||||||
|
delegate_only=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert images is None
|
||||||
|
assert combined_audio["sample_rate"] == 48000
|
||||||
|
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
|
||||||
|
|||||||
@@ -71,6 +71,15 @@ def _collector_only_prompt():
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _audio_only_prompt():
|
||||||
|
"""1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
|
||||||
|
return {
|
||||||
|
"1": {"class_type": "LoadAudio", "inputs": {}},
|
||||||
|
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
|
||||||
|
"3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def _delegate_prompt():
|
def _delegate_prompt():
|
||||||
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
|
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
|
||||||
return {
|
return {
|
||||||
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
|
|||||||
self.assertEqual(len(preview_nodes), 1)
|
self.assertEqual(len(preview_nodes), 1)
|
||||||
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
|
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
|
||||||
|
|
||||||
|
def test_injects_preview_audio_for_audio_only_collector(self):
|
||||||
|
result = pt.prune_prompt_for_worker(_audio_only_prompt())
|
||||||
|
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
|
||||||
|
self.assertEqual(len(preview_nodes), 1)
|
||||||
|
self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
|
||||||
|
self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
|
||||||
|
|
||||||
|
def test_prefers_image_preview_when_collector_has_images_and_audio(self):
|
||||||
|
prompt = _linear_prompt()
|
||||||
|
prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
|
||||||
|
prompt["4"]["inputs"]["audio"] = ["6", 0]
|
||||||
|
result = pt.prune_prompt_for_worker(prompt)
|
||||||
|
self.assertEqual(
|
||||||
|
len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
|
||||||
|
1,
|
||||||
|
)
|
||||||
|
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
|
||||||
|
|
||||||
|
def test_preserves_image_preview_for_distributed_upscale(self):
|
||||||
|
prompt = {
|
||||||
|
"1": {"class_type": "LoadImage", "inputs": {}},
|
||||||
|
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
|
||||||
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
||||||
|
}
|
||||||
|
result = pt.prune_prompt_for_worker(prompt)
|
||||||
|
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
||||||
|
self.assertEqual(len(preview_nodes), 1)
|
||||||
|
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
|
||||||
|
|
||||||
def test_no_preview_image_when_no_downstream(self):
|
def test_no_preview_image_when_no_downstream(self):
|
||||||
result = pt.prune_prompt_for_worker(_collector_only_prompt())
|
result = pt.prune_prompt_for_worker(_collector_only_prompt())
|
||||||
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
||||||
@@ -242,7 +280,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
|
|||||||
def test_upscale_node_is_treated_as_distributed(self):
|
def test_upscale_node_is_treated_as_distributed(self):
|
||||||
prompt = {
|
prompt = {
|
||||||
"1": {"class_type": "KSampler", "inputs": {}},
|
"1": {"class_type": "KSampler", "inputs": {}},
|
||||||
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
|
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
|
||||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
||||||
}
|
}
|
||||||
result = pt.prune_prompt_for_worker(prompt)
|
result = pt.prune_prompt_for_worker(prompt)
|
||||||
@@ -285,12 +323,22 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
|
|||||||
placeholder_id = empty_nodes[0][0]
|
placeholder_id = empty_nodes[0][0]
|
||||||
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
|
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
|
||||||
|
|
||||||
|
def test_audio_only_collector_does_not_get_image_placeholder(self):
|
||||||
|
prompt = _audio_only_prompt()
|
||||||
|
result = pt.prepare_delegate_master_prompt(prompt, ["2"])
|
||||||
|
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
|
||||||
|
self.assertEqual(empty_nodes, [])
|
||||||
|
self.assertNotIn("images", result["2"].get("inputs", {}))
|
||||||
|
self.assertNotIn("audio", result["2"].get("inputs", {}))
|
||||||
|
|
||||||
def test_one_placeholder_per_collector(self):
|
def test_one_placeholder_per_collector(self):
|
||||||
"""Two collectors → two placeholders."""
|
"""Two collectors → two placeholders."""
|
||||||
prompt = {
|
prompt = {
|
||||||
"1": {"class_type": "DistributedCollector", "inputs": {}},
|
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
|
||||||
"2": {"class_type": "DistributedCollector", "inputs": {}},
|
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
|
||||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
||||||
|
"10": {"class_type": "LoadImage", "inputs": {}},
|
||||||
|
"11": {"class_type": "LoadImage", "inputs": {}},
|
||||||
}
|
}
|
||||||
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
|
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
|
||||||
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
|
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
|
||||||
|
|||||||
@@ -0,0 +1,21 @@
|
|||||||
|
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def test_native_v3_runtime():
|
||||||
|
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||||
|
if not comfy_root:
|
||||||
|
pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance')
|
||||||
|
helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py'
|
||||||
|
result = subprocess.run(
|
||||||
|
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
|
||||||
|
capture_output=True, text=True, timeout=120,
|
||||||
|
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
|
||||||
|
)
|
||||||
|
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||||
|
assert 'V3_ACCEPTANCE_OK' in result.stdout
|
||||||
+9
-12
@@ -2,7 +2,6 @@ import { app } from "/scripts/app.js";
|
|||||||
import { ENDPOINTS } from "./constants.js";
|
import { ENDPOINTS } from "./constants.js";
|
||||||
|
|
||||||
const NODE_CLASS = "DistributedValue";
|
const NODE_CLASS = "DistributedValue";
|
||||||
const CONVERTED_WIDGET = "converted-widget";
|
|
||||||
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
|
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
|
||||||
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
|
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
|
||||||
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
|
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
|
||||||
@@ -42,15 +41,13 @@ function getDynamicWorkerWidgets(node) {
|
|||||||
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
|
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
|
||||||
}
|
}
|
||||||
|
|
||||||
function hideWidgetForGood(node, widget, suffix = "") {
|
function hideWidgetForGood(widget) {
|
||||||
if (!widget) return;
|
if (!widget) return;
|
||||||
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
|
// These are serialized backing fields, not widgets converted into sockets.
|
||||||
|
// Keep their native type/serializer and mutate the shared options in place
|
||||||
widget.origType = widget.type;
|
// so both the classic canvas and Nodes 2.0 suppress the whole widget row.
|
||||||
widget.origComputeSize = widget.computeSize;
|
widget.options ??= {};
|
||||||
widget.origSerializeValue = widget.serializeValue;
|
widget.options.hidden = true;
|
||||||
widget.computeSize = () => [0, -4];
|
|
||||||
widget.type = `${CONVERTED_WIDGET}${suffix}`;
|
|
||||||
|
|
||||||
// Hide any attached DOM element (multiline widgets).
|
// Hide any attached DOM element (multiline widgets).
|
||||||
if (widget.element) widget.element.style.display = "none";
|
if (widget.element) widget.element.style.display = "none";
|
||||||
@@ -58,14 +55,14 @@ function hideWidgetForGood(node, widget, suffix = "") {
|
|||||||
|
|
||||||
if (widget.linkedWidgets) {
|
if (widget.linkedWidgets) {
|
||||||
for (const linked of widget.linkedWidgets) {
|
for (const linked of widget.linkedWidgets) {
|
||||||
hideWidgetForGood(node, linked, `:${widget.name}`);
|
hideWidgetForGood(linked);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
function hideRawWidgets(node) {
|
function hideRawWidgets(node) {
|
||||||
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
|
hideWidgetForGood(getRawDefaultWidget(node));
|
||||||
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
|
hideWidgetForGood(getRawWorkerValuesWidget(node));
|
||||||
}
|
}
|
||||||
|
|
||||||
function removeDynamicDefaultWidget(node) {
|
function removeDynamicDefaultWidget(node) {
|
||||||
|
|||||||
@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
|
|||||||
createCheckboxSetting(
|
createCheckboxSetting(
|
||||||
"setting-debug",
|
"setting-debug",
|
||||||
"Debug Mode",
|
"Debug Mode",
|
||||||
"Enable verbose logging in the browser console.",
|
"Enable verbose logging in the browser console and ComfyUI server output.",
|
||||||
extension.config?.settings?.debug || false,
|
extension.config?.settings?.debug || false,
|
||||||
(event) => extension._updateSetting("debug", event.target.checked)
|
(event) => extension._updateSetting("debug", event.target.checked)
|
||||||
)
|
)
|
||||||
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
|
|||||||
createNumberSetting(
|
createNumberSetting(
|
||||||
"setting-worker-timeout",
|
"setting-worker-timeout",
|
||||||
"Worker Timeout",
|
"Worker Timeout",
|
||||||
"Seconds without a heartbeat before a worker is considered timed out. Default 60.",
|
"Maximum result-wait and heartbeat inactivity period before recovery begins. Busy workers may receive additional grace. Default: 60 seconds.",
|
||||||
extension.config?.settings?.worker_timeout_seconds ?? 60,
|
extension.config?.settings?.worker_timeout_seconds ?? 60,
|
||||||
10,
|
10,
|
||||||
1,
|
1,
|
||||||
|
|||||||
@@ -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");
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -814,7 +814,7 @@
|
|||||||
"outputs": [],
|
"outputs": [],
|
||||||
"properties": {},
|
"properties": {},
|
||||||
"widgets_values": [
|
"widgets_values": [
|
||||||
"If all your GPUs are the same/similar, set static_distribution to true\n"
|
"Workers pull tiles from a shared queue, so faster workers can process more tiles automatically.\n"
|
||||||
],
|
],
|
||||||
"color": "#432",
|
"color": "#432",
|
||||||
"bgcolor": "#653"
|
"bgcolor": "#653"
|
||||||
|
|||||||
Reference in New Issue
Block a user