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.
|
||||
|
||||
* **Endpoint:** `POST /distributed/queue`
|
||||
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
|
||||
* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
|
||||
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
|
||||
|
||||
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
||||
@@ -175,7 +175,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
|
||||
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
|
||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||
|
||||
|
||||
+15
-22
@@ -1,29 +1,22 @@
|
||||
# Import everything needed from the main module
|
||||
from .distributed import (
|
||||
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
"""ComfyUI-Distributed's native V3 extension entrypoint."""
|
||||
from comfy_api.v0_0_2 import ComfyExtension, io
|
||||
|
||||
# Import utilities
|
||||
from .utils.config import ensure_config_exists, CONFIG_FILE
|
||||
from .utils.logging import debug_log
|
||||
from .nodes.v3 import NODES
|
||||
from .runtime.bootstrap import initialize
|
||||
|
||||
# Import distributed upscale nodes
|
||||
from .nodes.distributed_upscale import (
|
||||
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
WEB_DIRECTORY = './web'
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
ensure_config_exists()
|
||||
class DistributedExtension(ComfyExtension):
|
||||
async def on_load(self) -> None:
|
||||
initialize()
|
||||
|
||||
# Merge node mappings
|
||||
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return list(NODES)
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
async def comfy_entrypoint() -> DistributedExtension:
|
||||
return DistributedExtension()
|
||||
|
||||
|
||||
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
|
||||
|
||||
@@ -3,4 +3,3 @@ from . import tunnel_routes # noqa: F401
|
||||
from . import worker_routes # noqa: F401
|
||||
from . import job_routes # noqa: F401
|
||||
from . import usdu_routes # noqa: F401
|
||||
from . import ltx_routes # noqa: F401
|
||||
|
||||
+11
-4
@@ -295,17 +295,24 @@ async def job_complete_endpoint(request):
|
||||
errors.append("worker_id: expected non-empty string")
|
||||
if not isinstance(batch_idx, int) or batch_idx < 0:
|
||||
errors.append("batch_idx: expected non-negative integer")
|
||||
if not isinstance(image_payload, str) or not image_payload.strip():
|
||||
image_provided = image_payload is not None
|
||||
audio_provided = audio_payload is not None
|
||||
if not image_provided and not audio_provided:
|
||||
errors.append("expected at least one of image or audio")
|
||||
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
|
||||
errors.append("image: expected non-empty base64 PNG string")
|
||||
if audio_payload is not None and not isinstance(audio_payload, dict):
|
||||
if audio_provided and not isinstance(audio_payload, dict):
|
||||
errors.append("audio: expected object when provided")
|
||||
if not isinstance(is_last, bool):
|
||||
errors.append("is_last: expected boolean")
|
||||
if errors:
|
||||
return await handle_api_error(request, errors, 400)
|
||||
|
||||
tensor = _decode_canonical_png_tensor(image_payload)
|
||||
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
|
||||
try:
|
||||
tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
|
||||
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
multi_job_id = job_id.strip()
|
||||
worker_id = worker_id.strip()
|
||||
|
||||
|
||||
@@ -1,76 +0,0 @@
|
||||
"""HTTP routes for distributed LTX latent tile results."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..upscale.job_models import TileJobState
|
||||
from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
|
||||
from ..utils.logging import debug_log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/ltx/submit_tile")
|
||||
async def submit_ltx_tile_endpoint(request):
|
||||
"""Accept one processed LTX latent tile from a worker.
|
||||
|
||||
Body is JSON:
|
||||
- multi_job_id: distributed job id
|
||||
- worker_id: worker id
|
||||
- tile_index: spatial tile id, or -1 for completion-only messages
|
||||
- payload: base64 safetensors payload produced by nodes.ltx_tiled_sampler
|
||||
- is_last: optional worker completion flag
|
||||
"""
|
||||
try:
|
||||
content_length = request.headers.get("content-length")
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.json()
|
||||
multi_job_id = data.get("multi_job_id")
|
||||
worker_id = str(data.get("worker_id") or "")
|
||||
is_last = bool(data.get("is_last", False))
|
||||
payload = data.get("payload")
|
||||
tile_index = data.get("tile_index")
|
||||
|
||||
if not multi_job_id or not worker_id:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 404)
|
||||
|
||||
job_data.worker_status[worker_id] = time.time()
|
||||
|
||||
if is_last and (payload is None or int(tile_index or -1) < 0):
|
||||
await job_data.queue.put({
|
||||
"worker_id": worker_id,
|
||||
"is_last": True,
|
||||
"tiles": [],
|
||||
})
|
||||
debug_log(f"Received LTX completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
if payload is None or tile_index is None:
|
||||
return await handle_api_error(request, "Missing tile_index or payload", 400)
|
||||
|
||||
tile_index = int(tile_index)
|
||||
await job_data.queue.put({
|
||||
"worker_id": worker_id,
|
||||
"is_last": is_last,
|
||||
"tiles": [{
|
||||
"tile_idx": tile_index,
|
||||
"global_idx": tile_index,
|
||||
"payload": payload,
|
||||
"worker_id": worker_id,
|
||||
}],
|
||||
})
|
||||
debug_log(f"Received LTX tile {tile_index} from worker {worker_id} (is_last={is_last})")
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, exc, 500)
|
||||
@@ -351,16 +351,33 @@ def prune_prompt_for_worker(prompt_obj):
|
||||
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
||||
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
|
||||
if has_removed_downstream:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {
|
||||
"images": [dist_id, 0],
|
||||
},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {
|
||||
"title": "Preview Image (auto-added)",
|
||||
},
|
||||
}
|
||||
original_node = prompt_obj.get(str(dist_id), {})
|
||||
class_type = original_node.get("class_type")
|
||||
inputs = original_node.get("inputs", {})
|
||||
image_connected = class_type != "DistributedCollector" or (
|
||||
isinstance(inputs.get("images"), list)
|
||||
and len(inputs["images"]) == 2
|
||||
)
|
||||
audio_connected = (
|
||||
class_type == "DistributedCollector"
|
||||
and isinstance(inputs.get("audio"), list)
|
||||
and len(inputs["audio"]) == 2
|
||||
)
|
||||
|
||||
if image_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"images": [dist_id, 0]},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {"title": "Preview Image (auto-added)"},
|
||||
}
|
||||
elif audio_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"audio": [dist_id, 1]},
|
||||
"class_type": "PreviewAudio",
|
||||
"_meta": {"title": "Preview Audio (auto-added)"},
|
||||
}
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
@@ -400,6 +417,10 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
continue
|
||||
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
|
||||
original_images = original_inputs.get("images")
|
||||
if not (isinstance(original_images, list) and len(original_images) == 2):
|
||||
continue
|
||||
placeholder_id = next_id()
|
||||
pruned_prompt[placeholder_id] = {
|
||||
"class_type": "DistributedEmptyImage",
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
# conftest.py — project-level pytest configuration.
|
||||
#
|
||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||
# (from .distributed import ...) that fail when pytest tries to import it as a
|
||||
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
|
||||
# standalone module during Package.setup() for the root package node.
|
||||
#
|
||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
"""
|
||||
ComfyUI-Distributed: thin entry point.
|
||||
All implementation lives in workers/, nodes/, api/.
|
||||
"""
|
||||
import atexit
|
||||
import os
|
||||
|
||||
import server
|
||||
|
||||
from .utils.config import ensure_config_exists
|
||||
from .utils.logging import debug_log
|
||||
from .utils.network import cleanup_client_session
|
||||
from .workers import get_worker_manager
|
||||
from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
|
||||
from .upscale.job_store import ensure_tile_jobs_initialized
|
||||
from .nodes import (
|
||||
NODE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS,
|
||||
ImageBatchDivider,
|
||||
DistributedCollectorNode,
|
||||
DistributedSeed,
|
||||
DistributedModelName,
|
||||
DistributedValue,
|
||||
AudioBatchDivider,
|
||||
DistributedEmptyImage,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
)
|
||||
from . import api # noqa: F401 - triggers all @routes.* registrations
|
||||
from .api.queue_orchestration import ensure_distributed_state
|
||||
|
||||
ensure_config_exists()
|
||||
|
||||
# Aiohttp session cleanup
|
||||
async def _cleanup_session():
|
||||
await cleanup_client_session()
|
||||
|
||||
|
||||
atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup
|
||||
|
||||
# Initialize distributed job state on prompt_server
|
||||
prompt_server = server.PromptServer.instance
|
||||
ensure_distributed_state(prompt_server)
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
# Worker startup
|
||||
if not os.environ.get('COMFYUI_IS_WORKER'):
|
||||
atexit.register(sync_cleanup)
|
||||
delayed_auto_launch()
|
||||
register_async_signals()
|
||||
@@ -19,8 +19,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
|
||||
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
|
||||
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
|
||||
- Resolves enabled/selected workers.
|
||||
- Pings workers (`GET /prompt`) to include only reachable ones.
|
||||
- Dispatches the workflow to workers (`POST /prompt`).
|
||||
- Probes and dispatches workers through `/distributed/worker_ws` by default.
|
||||
- If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
|
||||
- Queues the master workflow in ComfyUI’s prompt queue.
|
||||
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
|
||||
|
||||
@@ -61,7 +61,8 @@ Queue a workflow for distributed execution.
|
||||
#### Fields
|
||||
|
||||
- `prompt` (required unless `workflow.prompt` is present, object)
|
||||
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
|
||||
- A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
|
||||
- This is not the normal visual workflow export from the ComfyUI editor.
|
||||
- `workflow` (optional, object)
|
||||
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
|
||||
- If you don’t care about UI metadata, you can omit it.
|
||||
@@ -123,10 +124,14 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
||||
|
||||
## Worker requirements (important)
|
||||
|
||||
For a worker to participate, it must be reachable from the master:
|
||||
For a worker to participate, it must be reachable from the master. By default:
|
||||
|
||||
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
|
||||
|
||||
If `settings.websocket_orchestration=false`:
|
||||
|
||||
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
|
||||
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
|
||||
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
|
||||
|
||||
Also, for collector-based flows:
|
||||
|
||||
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
|
||||
|
||||
## Examples
|
||||
|
||||
### 1) Minimal `curl`
|
||||
### 1) Minimal `curl` request envelope
|
||||
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||
@@ -209,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
Where `payload.json` contains at least:
|
||||
`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": {
|
||||
"1": {"class_type": "KSampler", "inputs": {} }
|
||||
"<node_id>": {
|
||||
"class_type": "<node class>",
|
||||
"inputs": {"<required_input>": "<value or connection>"}
|
||||
}
|
||||
},
|
||||
"enabled_worker_ids": [],
|
||||
"client_id": "external-client"
|
||||
}
|
||||
```
|
||||
|
||||
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
|
||||
|
||||
### 2) Python (`requests`)
|
||||
|
||||
```python
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
# Native V3 node API
|
||||
|
||||
The package registers one `ComfyExtension` through `comfy_entrypoint()` and
|
||||
imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
|
||||
that API; there is no V1 registration fallback.
|
||||
|
||||
## Compatibility
|
||||
|
||||
- All eight node IDs, display names, categories, visible input order, defaults
|
||||
and output order are retained. Existing execution algorithms remain in the
|
||||
private collector, utilities and upscale modules.
|
||||
- The image/audio dividers explicitly declare ten typed outputs. Their existing
|
||||
frontend extensions still show only the selected number of outputs. This
|
||||
replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
|
||||
IMAGE/AUDIO socket indices or the ten returned values.
|
||||
- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
|
||||
its existing prompt-input names and Python defaults via `accept_all_inputs`;
|
||||
it does not become visible widgets.
|
||||
- The collector retains list-input handling. Upscale retains its always-changing
|
||||
fingerprint and creates a private runtime helper for each V3 execution, so
|
||||
mutable helper state is not attached to sanitized V3 class clones.
|
||||
- Routes, distributed state and the existing worker startup/shutdown hooks are
|
||||
initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
|
||||
Worker mode still suppresses automatic worker launch. `distributed.py` is no
|
||||
longer an entrypoint.
|
||||
|
||||
## Verification
|
||||
|
||||
Run ordinary unit tests from this repository:
|
||||
|
||||
```bash
|
||||
python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
|
||||
/path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
|
||||
parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
|
||||
all node schemas with the V1 fixture, checks all five bundled workflows' node
|
||||
IDs and socket/link contracts, exercises injected worker metadata, image/audio
|
||||
lists, collector aggregation and divider outputs, rejects invalid upscale enums,
|
||||
and decodes an actual preview PNG referenced by executor history.
|
||||
|
||||
The upscale GPU/model boundary is mocked to check argument forwarding and
|
||||
per-execution helper isolation. Full checkpoint inference, browser canvas
|
||||
acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
|
||||
workers or model downloads are started; preview files use a temporary scratch
|
||||
directory. The test does not install this branch into a live custom-node folder.
|
||||
@@ -79,7 +79,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
||||
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
|
||||
|
||||
**On Runpod:**
|
||||
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
|
||||
> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
|
||||
|
||||
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
||||
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
|
||||
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
||||
- SAGE_ATTENTION: optional optimisation (set to true/false)
|
||||
5. Deploy your pod.
|
||||
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
|
||||
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
```
|
||||
# Download from CivitAI
|
||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||
|
||||
+1
-37
@@ -1,37 +1 @@
|
||||
from .utilities import (
|
||||
DistributedSeed,
|
||||
DistributedModelName,
|
||||
DistributedValue,
|
||||
ImageBatchDivider,
|
||||
AudioBatchDivider,
|
||||
DistributedEmptyImage,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
)
|
||||
from .collector import DistributedCollectorNode
|
||||
from .ltx_tiled_sampler import (
|
||||
NODE_CLASS_MAPPINGS as LTX_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as LTX_DISPLAY_NAME_MAPPINGS,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DistributedCollector": DistributedCollectorNode,
|
||||
"DistributedSeed": DistributedSeed,
|
||||
"DistributedModelName": DistributedModelName,
|
||||
"DistributedValue": DistributedValue,
|
||||
"ImageBatchDivider": ImageBatchDivider,
|
||||
"AudioBatchDivider": AudioBatchDivider,
|
||||
"DistributedEmptyImage": DistributedEmptyImage,
|
||||
**LTX_CLASS_MAPPINGS,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DistributedCollector": "Distributed Collector",
|
||||
"DistributedSeed": "Distributed Seed",
|
||||
"DistributedModelName": "Distributed Model Name",
|
||||
"DistributedValue": "Distributed Value",
|
||||
"ImageBatchDivider": "Image Batch Divider",
|
||||
"AudioBatchDivider": "Audio Batch Divider",
|
||||
"DistributedEmptyImage": "Distributed Empty Image",
|
||||
**LTX_DISPLAY_NAME_MAPPINGS,
|
||||
}
|
||||
"""Private execution helpers; public registration lives in nodes.v3."""
|
||||
|
||||
+73
-45
@@ -29,7 +29,6 @@ class DistributedCollectorNode:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"load_balance": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
@@ -38,7 +37,10 @@ class DistributedCollectorNode:
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": { "audio": ("AUDIO",) },
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
"audio": ("AUDIO",),
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
@@ -102,8 +104,9 @@ class DistributedCollectorNode:
|
||||
return None
|
||||
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
|
||||
|
||||
def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
|
||||
images = self._normalize_images_input(images)
|
||||
def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
|
||||
if images is not None:
|
||||
images = self._normalize_images_input(images)
|
||||
audio = self._normalize_audio_input(audio)
|
||||
load_balance = self._unwrap_list_input(load_balance)
|
||||
multi_job_id = self._unwrap_list_input(multi_job_id)
|
||||
@@ -115,6 +118,14 @@ class DistributedCollectorNode:
|
||||
pass_through = self._unwrap_list_input(pass_through)
|
||||
delegate_only = self._unwrap_list_input(delegate_only)
|
||||
|
||||
remote_only_master = (
|
||||
bool(multi_job_id)
|
||||
and not is_worker
|
||||
and (delegate_only or is_master_delegate_only())
|
||||
)
|
||||
if images is None and audio is None and not remote_only_master:
|
||||
raise ValueError("DistributedCollector requires at least one image or audio input")
|
||||
|
||||
# Create empty audio if not provided
|
||||
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
|
||||
@@ -141,41 +152,56 @@ class DistributedCollectorNode:
|
||||
return result
|
||||
|
||||
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
||||
"""Send image batch to master via canonical JSON envelopes."""
|
||||
batch_size = image_batch.shape[0]
|
||||
if batch_size == 0:
|
||||
return
|
||||
|
||||
"""Send an image batch, optionally with audio, or an audio-only completion."""
|
||||
encoded_audio = encode_audio_payload(audio)
|
||||
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_complete"
|
||||
for batch_idx in range(batch_size):
|
||||
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
|
||||
byte_io = io.BytesIO()
|
||||
img.save(byte_io, format='PNG', compress_level=0)
|
||||
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
|
||||
payload = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": int(batch_idx),
|
||||
"image": f"data:image/png;base64,{encoded_image}",
|
||||
"is_last": bool(batch_idx == batch_size - 1),
|
||||
}
|
||||
if payload["is_last"] and encoded_audio is not None:
|
||||
payload["audio"] = encoded_audio
|
||||
|
||||
payloads = []
|
||||
batch_size = 0 if image_batch is None else image_batch.shape[0]
|
||||
if batch_size == 0:
|
||||
if encoded_audio is None:
|
||||
raise ValueError("Worker completion requires image or audio data")
|
||||
payloads.append(
|
||||
{
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": 0,
|
||||
"audio": encoded_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
else:
|
||||
for batch_idx in range(batch_size):
|
||||
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
|
||||
byte_io = io.BytesIO()
|
||||
img.save(byte_io, format='PNG', compress_level=0)
|
||||
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
|
||||
payload = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": int(batch_idx),
|
||||
"image": f"data:image/png;base64,{encoded_image}",
|
||||
"is_last": bool(batch_idx == batch_size - 1),
|
||||
}
|
||||
if payload["is_last"] and encoded_audio is not None:
|
||||
payload["audio"] = encoded_audio
|
||||
payloads.append(payload)
|
||||
|
||||
for payload in payloads:
|
||||
timeout_seconds = 60 if "image" in payload else 600
|
||||
try:
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
log(f"Worker - Failed to send canonical image envelope to master: {e}")
|
||||
media_type = "image/audio" if "image" in payload else "audio-only"
|
||||
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
|
||||
debug_log(f"Worker - Full error details: URL={url}")
|
||||
raise # Re-raise to handle at caller level
|
||||
raise
|
||||
|
||||
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
|
||||
"""Combine audio from master and workers into a single audio output.
|
||||
@@ -257,8 +283,8 @@ class DistributedCollectorNode:
|
||||
images_on_cpu,
|
||||
delegate_mode: bool,
|
||||
fallback_images,
|
||||
) -> torch.Tensor:
|
||||
"""Assemble final tensor: master first, then workers in enabled order."""
|
||||
):
|
||||
"""Assemble final tensor, or return None when the job contains only audio."""
|
||||
ordered_tensors = []
|
||||
if not delegate_mode and images_on_cpu is not None:
|
||||
for i in range(master_batch_size):
|
||||
@@ -289,15 +315,15 @@ class DistributedCollectorNode:
|
||||
|
||||
if cpu_tensors:
|
||||
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
||||
elif fallback_images is not None:
|
||||
if fallback_images is not None:
|
||||
return ensure_contiguous(fallback_images)
|
||||
else:
|
||||
raise ValueError("No image data collected from master or workers")
|
||||
return None
|
||||
|
||||
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
|
||||
if is_worker:
|
||||
# Worker mode: send images and audio to master in a single batch
|
||||
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
|
||||
image_count = 0 if images is None else images.shape[0]
|
||||
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
|
||||
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
else:
|
||||
@@ -332,14 +358,15 @@ class DistributedCollectorNode:
|
||||
master_audio = None
|
||||
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
|
||||
else:
|
||||
images_on_cpu = images.cpu()
|
||||
master_batch_size = images.shape[0]
|
||||
if images is None:
|
||||
images_on_cpu = None
|
||||
master_batch_size = 0
|
||||
else:
|
||||
images_on_cpu = ensure_contiguous(images.cpu())
|
||||
master_batch_size = images.shape[0]
|
||||
master_audio = audio # Keep master's audio for later
|
||||
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
|
||||
|
||||
# Ensure master images are contiguous
|
||||
images_on_cpu = ensure_contiguous(images_on_cpu)
|
||||
|
||||
|
||||
# Initialize storage for collected images and audio
|
||||
worker_images = {} # Dict to store images by worker_id and index
|
||||
@@ -511,18 +538,19 @@ class DistributedCollectorNode:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
del prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
|
||||
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
||||
try:
|
||||
combined = self._reorder_and_combine_tensors(
|
||||
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
|
||||
)
|
||||
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||
|
||||
# Combine audio from master and workers
|
||||
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
||||
if combined is None:
|
||||
debug_log(f"Master - Job {multi_job_id} complete with audio only")
|
||||
else:
|
||||
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||
|
||||
return (combined, combined_audio)
|
||||
except Exception as e:
|
||||
log(f"Master - Error combining images: {e}")
|
||||
# Return just the master images as fallback
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
# Preserve collected audio even when image assembly fails.
|
||||
return (images, combined_audio)
|
||||
|
||||
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
|
||||
"""
|
||||
Distributed version of Ultimate SD Upscale (No Upscale).
|
||||
|
||||
Supports three processing modes:
|
||||
Supports two currently selected processing modes:
|
||||
1. Single GPU: No workers available, process everything locally
|
||||
2. Static Mode: Small batches, distributes tiles across workers (flattened)
|
||||
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
|
||||
2. Distributed tile queue: Workers pull tiles from a shared queue
|
||||
|
||||
Features:
|
||||
- Multi-mode batch handling for efficient video/image upscaling
|
||||
- Tile-based batch handling for video/image upscaling
|
||||
- Tiled VAE support for memory efficiency
|
||||
- Dynamic load balancing for large batches
|
||||
- Shared work queue so faster workers can process more tiles
|
||||
- Backward compatible with single-image workflows
|
||||
|
||||
Environment Variables:
|
||||
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
|
||||
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
|
||||
|
||||
Threshold: dynamic_threshold input controls mode switch (default 8)
|
||||
The hidden dynamic_threshold input is retained for workflow compatibility but
|
||||
does not affect the current mode-selection policy.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
|
||||
|
||||
# Ensure initialization before registering routes
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
# Node registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
|
||||
}
|
||||
|
||||
@@ -1,668 +0,0 @@
|
||||
"""Distributed LTX tiled sampler node.
|
||||
|
||||
This node adapts the upstream LTX tiled sampler interface to the existing
|
||||
ComfyUI-Distributed job queue. It distributes spatial latent tiles across the
|
||||
same worker/master infrastructure used by the distributed upscaler instead of
|
||||
creating a parallel queue implementation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import math
|
||||
from typing import Callable, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
import aiohttp
|
||||
import torch
|
||||
from safetensors.torch import load as safetensors_load
|
||||
from safetensors.torch import save as safetensors_save
|
||||
|
||||
from ..upscale.job_models import TileJobState
|
||||
from ..upscale.job_store import (
|
||||
_cleanup_job,
|
||||
_drain_results_queue,
|
||||
_get_completed_count,
|
||||
_mark_task_completed,
|
||||
ensure_tile_jobs_initialized,
|
||||
init_static_job_batched,
|
||||
)
|
||||
from ..upscale.job_timeout import _check_and_requeue_timed_out_workers
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..utils.constants import HEARTBEAT_INTERVAL, TILE_WAIT_TIMEOUT
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import get_client_session
|
||||
from ..utils.usdu_managment import _send_heartbeat_to_master
|
||||
|
||||
try: # Optional external LTX tiled sampler dependency.
|
||||
from latent_tiled_sampler import LTXTiledSampler # type: ignore
|
||||
except Exception: # pragma: no cover - depends on optional custom node install.
|
||||
LTXTiledSampler = None # type: ignore
|
||||
|
||||
|
||||
def _default_ltx_inputs() -> dict:
|
||||
"""Fallback schema used when the source LTX sampler is not importable."""
|
||||
return {
|
||||
"required": {
|
||||
"noise": ("NOISE",),
|
||||
"guider": ("GUIDER",),
|
||||
"sampler": ("SAMPLER",),
|
||||
"sigmas": ("SIGMAS",),
|
||||
"latent_image": ("LATENT",),
|
||||
"h_tiles": ("INT", {"default": 1, "min": 1, "max": 64}),
|
||||
"w_tiles": ("INT", {"default": 1, "min": 1, "max": 64}),
|
||||
"overlap": ("INT", {"default": 0, "min": 0, "max": 4096}),
|
||||
},
|
||||
"optional": {
|
||||
"denoise_mask": ("MASK",),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _copy_base_inputs() -> dict:
|
||||
if LTXTiledSampler is None:
|
||||
return _default_ltx_inputs()
|
||||
try:
|
||||
return copy.deepcopy(LTXTiledSampler.INPUT_TYPES())
|
||||
except Exception:
|
||||
return _default_ltx_inputs()
|
||||
|
||||
|
||||
def _compute_tile_starts(total_size: int, n_tiles: int, overlap: int) -> Tuple[List[int], int]:
|
||||
"""Compute evenly spaced tile starts and tile size for one dimension."""
|
||||
total_size = int(total_size)
|
||||
n_tiles = max(1, int(n_tiles))
|
||||
overlap = max(0, int(overlap))
|
||||
if n_tiles <= 1 or total_size <= 1:
|
||||
return [0], total_size
|
||||
tile_size = math.ceil((total_size + (n_tiles - 1) * overlap) / n_tiles)
|
||||
tile_size = min(tile_size, total_size)
|
||||
if n_tiles == 2:
|
||||
return [0, max(0, total_size - tile_size)], tile_size
|
||||
stride = (total_size - tile_size) / (n_tiles - 1)
|
||||
return [int(round(i * stride)) for i in range(n_tiles)], tile_size
|
||||
|
||||
|
||||
def _make_window_1d(size: int, fade_left: int, fade_right: int, dtype, device):
|
||||
"""Create a one-dimensional cosine edge window for tile blending."""
|
||||
win = torch.ones(int(size), dtype=dtype, device=device)
|
||||
if fade_left > 0:
|
||||
fl = min(int(fade_left), int(size))
|
||||
i = torch.arange(fl, dtype=dtype, device=device)
|
||||
# Use fl rather than fl - 1 to preserve a nonzero center plateau for small fades.
|
||||
win[:fl] = 0.5 * (1.0 - torch.cos(math.pi * i / max(1, fl)))
|
||||
if fade_right > 0:
|
||||
fr = min(int(fade_right), int(size))
|
||||
i = torch.arange(fr, dtype=dtype, device=device)
|
||||
win[int(size) - fr:] = 0.5 * (1.0 + torch.cos(math.pi * i / max(1, fr)))
|
||||
return win
|
||||
|
||||
|
||||
def _pack_tensors(tensors: Dict[str, torch.Tensor]) -> str:
|
||||
"""Serialize tensors to a base64 safetensors payload for JSON transport."""
|
||||
cpu_tensors = {
|
||||
name: tensor.detach().contiguous().to("cpu")
|
||||
for name, tensor in tensors.items()
|
||||
if tensor is not None
|
||||
}
|
||||
payload = safetensors_save(cpu_tensors)
|
||||
return base64.b64encode(payload).decode("ascii")
|
||||
|
||||
|
||||
def _unpack_tensors(payload_b64: str, device) -> Dict[str, torch.Tensor]:
|
||||
payload = base64.b64decode(payload_b64.encode("ascii"))
|
||||
return {name: tensor.to(device) for name, tensor in safetensors_load(payload).items()}
|
||||
|
||||
|
||||
def _extract_tensor_and_rebuilder(value) -> Tuple[torch.Tensor, Callable[[torch.Tensor], object]]:
|
||||
"""Extract the tensor from common Comfy latent containers and preserve shape."""
|
||||
if isinstance(value, dict) and "samples" in value:
|
||||
original = value
|
||||
|
||||
def rebuild(tensor):
|
||||
updated = dict(original)
|
||||
updated["samples"] = tensor
|
||||
return updated
|
||||
|
||||
return value["samples"], rebuild
|
||||
|
||||
if isinstance(value, tuple) and value and torch.is_tensor(value[0]):
|
||||
original = value
|
||||
|
||||
def rebuild(tensor):
|
||||
return (tensor, *original[1:])
|
||||
|
||||
return value[0], rebuild
|
||||
|
||||
if isinstance(value, list) and value and torch.is_tensor(value[0]):
|
||||
original = value
|
||||
|
||||
def rebuild(tensor):
|
||||
updated = list(original)
|
||||
updated[0] = tensor
|
||||
return updated
|
||||
|
||||
return value[0], rebuild
|
||||
|
||||
if torch.is_tensor(value):
|
||||
return value, lambda tensor: tensor
|
||||
|
||||
raise TypeError("Expected latent_image to be a tensor or latent container with a tensor samples field")
|
||||
|
||||
|
||||
def _extract_optional_mask(mask):
|
||||
if mask is None:
|
||||
return None
|
||||
if isinstance(mask, dict) and "samples" in mask:
|
||||
return mask["samples"]
|
||||
return mask
|
||||
|
||||
|
||||
def _tile_latent(value: torch.Tensor, tile: dict) -> torch.Tensor:
|
||||
return value[:, :, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
|
||||
|
||||
|
||||
def _tile_mask(mask, tile: dict):
|
||||
if mask is None:
|
||||
return None
|
||||
if torch.is_tensor(mask) and mask.ndim >= 5:
|
||||
return mask[:, :, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
|
||||
if torch.is_tensor(mask) and mask.ndim >= 4:
|
||||
return mask[:, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
|
||||
return mask
|
||||
|
||||
|
||||
def _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask):
|
||||
result = guider.sample(
|
||||
tile_noise,
|
||||
tile_latent,
|
||||
sampler,
|
||||
sigmas,
|
||||
denoise_mask=tile_mask,
|
||||
disable_pbar=True,
|
||||
)
|
||||
if isinstance(result, tuple):
|
||||
tile_samples = result[0]
|
||||
tile_denoised = result[1] if len(result) > 1 else None
|
||||
else:
|
||||
tile_samples = result
|
||||
tile_denoised = None
|
||||
return tile_samples, tile_denoised
|
||||
|
||||
|
||||
def _build_tile_queue(latent_shape: Iterable[int], n_h_tiles: int, n_w_tiles: int, overlap: int) -> List[dict]:
|
||||
_batch, _channels, _frames, height, width = [int(v) for v in latent_shape]
|
||||
starts_y, tile_h = _compute_tile_starts(height, n_h_tiles, overlap)
|
||||
starts_x, tile_w = _compute_tile_starts(width, n_w_tiles, overlap)
|
||||
queue = []
|
||||
tile_idx = 0
|
||||
for grid_y, y_start in enumerate(starts_y):
|
||||
y_end = min(y_start + tile_h, height)
|
||||
for grid_x, x_start in enumerate(starts_x):
|
||||
x_end = min(x_start + tile_w, width)
|
||||
queue.append(
|
||||
{
|
||||
"tile_index": tile_idx,
|
||||
"h_start": y_start,
|
||||
"h_end": y_end,
|
||||
"w_start": x_start,
|
||||
"w_end": x_end,
|
||||
"grid_y": grid_y,
|
||||
"grid_x": grid_x,
|
||||
}
|
||||
)
|
||||
tile_idx += 1
|
||||
return queue
|
||||
|
||||
|
||||
def _blend_tile_into_sums(
|
||||
tile: dict,
|
||||
samples: torch.Tensor,
|
||||
denoised: Optional[torch.Tensor],
|
||||
output_sum: torch.Tensor,
|
||||
weight_sum: torch.Tensor,
|
||||
denoised_sum: Optional[torch.Tensor],
|
||||
n_h_tiles: int,
|
||||
n_w_tiles: int,
|
||||
overlap: int,
|
||||
):
|
||||
y_start = tile["h_start"]
|
||||
y_end = tile["h_end"]
|
||||
x_start = tile["w_start"]
|
||||
x_end = tile["w_end"]
|
||||
grid_y = tile["grid_y"]
|
||||
grid_x = tile["grid_x"]
|
||||
fade_top = overlap if grid_y > 0 else 0
|
||||
fade_bottom = overlap if grid_y < n_h_tiles - 1 else 0
|
||||
fade_left = overlap if grid_x > 0 else 0
|
||||
fade_right = overlap if grid_x < n_w_tiles - 1 else 0
|
||||
dtype = output_sum.dtype
|
||||
device = output_sum.device
|
||||
wy = _make_window_1d(y_end - y_start, fade_top, fade_bottom, dtype, device)
|
||||
wx = _make_window_1d(x_end - x_start, fade_left, fade_right, dtype, device)
|
||||
weights = wy.view(1, 1, 1, y_end - y_start, 1) * wx.view(1, 1, 1, 1, x_end - x_start)
|
||||
output_sum[:, :, :, y_start:y_end, x_start:x_end] += samples.to(device=device, dtype=dtype) * weights
|
||||
weight_sum[:, :, :, y_start:y_end, x_start:x_end] += weights
|
||||
if denoised_sum is not None and denoised is not None:
|
||||
denoised_sum[:, :, :, y_start:y_end, x_start:x_end] += denoised.to(device=device, dtype=dtype) * weights
|
||||
|
||||
|
||||
async def _get_next_ltx_tile_index(multi_job_id: str) -> Optional[int]:
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
pending_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
if pending_queue is None:
|
||||
return None
|
||||
try:
|
||||
return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
|
||||
async def _get_ltx_completed_tasks(multi_job_id: str) -> dict:
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
return dict(job_data.completed_tasks)
|
||||
return {}
|
||||
|
||||
|
||||
async def _check_ltx_job_status(multi_job_id: str, master_url: str) -> bool:
|
||||
session = await get_client_session()
|
||||
async with session.get(f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}") as response:
|
||||
if response.status != 200:
|
||||
return False
|
||||
data = await response.json()
|
||||
return bool(data.get("ready", False))
|
||||
|
||||
|
||||
async def _request_tile_from_master(multi_job_id: str, master_url: str, worker_id: str):
|
||||
session = await get_client_session()
|
||||
async with session.post(
|
||||
f"{master_url}/distributed/request_image",
|
||||
json={"multi_job_id": multi_job_id, "worker_id": str(worker_id)},
|
||||
timeout=aiohttp.ClientTimeout(total=TILE_WAIT_TIMEOUT),
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
return await response.json()
|
||||
if response.status == 404:
|
||||
return None
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
async def _submit_ltx_tile_to_master(
|
||||
multi_job_id: str,
|
||||
master_url: str,
|
||||
worker_id: str,
|
||||
tile_index: int,
|
||||
payload_b64: Optional[str] = None,
|
||||
is_last: bool = False,
|
||||
):
|
||||
session = await get_client_session()
|
||||
body = {
|
||||
"multi_job_id": multi_job_id,
|
||||
"worker_id": str(worker_id),
|
||||
"tile_index": int(tile_index) if tile_index is not None else None,
|
||||
"payload": payload_b64,
|
||||
"is_last": bool(is_last),
|
||||
}
|
||||
async with session.post(f"{master_url}/distributed/ltx/submit_tile", json=body) as response:
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def _ltx_worker_process(
|
||||
node,
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_tensor,
|
||||
denoise_mask,
|
||||
n_h_tiles,
|
||||
n_w_tiles,
|
||||
tile_overlap,
|
||||
):
|
||||
job_id = node.multi_job_id
|
||||
master_url = node.master_url
|
||||
worker_id = node.worker_id or "ltx-worker"
|
||||
tiles_by_index = {
|
||||
tile["tile_index"]: tile
|
||||
for tile in _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
|
||||
}
|
||||
for attempt in range(120):
|
||||
if await _check_ltx_job_status(job_id, master_url):
|
||||
break
|
||||
await asyncio.sleep(1.0)
|
||||
else:
|
||||
log(f"LTX Dist Worker[{worker_id[:8]}]: job {job_id} was not ready; aborting")
|
||||
return
|
||||
|
||||
full_noise = noise.generate_noise({"samples": latent_tensor})
|
||||
|
||||
last_heartbeat = 0.0
|
||||
while True:
|
||||
data = await _request_tile_from_master(job_id, master_url, worker_id)
|
||||
if not data:
|
||||
log(f"LTX Dist Worker[{worker_id[:8]}]: master job unavailable; stopping")
|
||||
break
|
||||
tile_index = data.get("tile_idx")
|
||||
if tile_index is None:
|
||||
await _submit_ltx_tile_to_master(job_id, master_url, worker_id, -1, is_last=True)
|
||||
break
|
||||
tile = tiles_by_index.get(int(tile_index))
|
||||
if tile is None:
|
||||
raise RuntimeError(f"Master assigned unknown LTX tile index {tile_index}")
|
||||
|
||||
tile_latent = _tile_latent(latent_tensor, tile)
|
||||
tile_noise = _tile_latent(full_noise, tile)
|
||||
tile_mask = _tile_mask(denoise_mask, tile)
|
||||
tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
|
||||
tensors = {"samples": tile_samples}
|
||||
if tile_denoised is not None:
|
||||
tensors["denoised"] = tile_denoised
|
||||
await _submit_ltx_tile_to_master(
|
||||
job_id,
|
||||
master_url,
|
||||
worker_id,
|
||||
int(tile_index),
|
||||
_pack_tensors(tensors),
|
||||
)
|
||||
now = asyncio.get_running_loop().time()
|
||||
if now - last_heartbeat >= HEARTBEAT_INTERVAL:
|
||||
await _send_heartbeat_to_master(job_id, master_url, worker_id)
|
||||
last_heartbeat = now
|
||||
|
||||
|
||||
async def _ltx_master_process(
|
||||
node,
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_tensor,
|
||||
denoise_mask,
|
||||
n_h_tiles,
|
||||
n_w_tiles,
|
||||
tile_overlap,
|
||||
):
|
||||
job_id = node.multi_job_id
|
||||
enabled_workers = node.enabled_worker_ids
|
||||
tiles = _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
|
||||
tiles_by_index = {tile["tile_index"]: tile for tile in tiles}
|
||||
total_tiles = len(tiles)
|
||||
await init_static_job_batched(job_id, batch_size=1, num_tiles_per_image=total_tiles, enabled_workers=enabled_workers)
|
||||
|
||||
full_noise = noise.generate_noise({"samples": latent_tensor})
|
||||
output_sum = torch.zeros_like(latent_tensor)
|
||||
weight_sum = torch.zeros_like(latent_tensor)
|
||||
denoised_sum = torch.zeros_like(latent_tensor) if denoise_mask is not None else None
|
||||
integrated = set()
|
||||
last_timeout_check = 0.0
|
||||
|
||||
try:
|
||||
while len(integrated) < total_tiles:
|
||||
await _drain_results_queue(job_id)
|
||||
completed = await _get_ltx_completed_tasks(job_id)
|
||||
for tile_index, result in list(completed.items()):
|
||||
if tile_index in integrated:
|
||||
continue
|
||||
if not isinstance(result, dict) or "payload" not in result:
|
||||
# Locally completed tiles are blended immediately and then marked complete.
|
||||
continue
|
||||
tile = tiles_by_index.get(int(tile_index))
|
||||
if tile is None:
|
||||
continue
|
||||
tensors = _unpack_tensors(result["payload"], latent_tensor.device)
|
||||
_blend_tile_into_sums(
|
||||
tile,
|
||||
tensors["samples"],
|
||||
tensors.get("denoised"),
|
||||
output_sum,
|
||||
weight_sum,
|
||||
denoised_sum,
|
||||
n_h_tiles,
|
||||
n_w_tiles,
|
||||
tile_overlap,
|
||||
)
|
||||
integrated.add(int(tile_index))
|
||||
|
||||
if len(integrated) >= total_tiles:
|
||||
break
|
||||
|
||||
tile_index = await _get_next_ltx_tile_index(job_id)
|
||||
if tile_index is not None and int(tile_index) not in integrated:
|
||||
tile = tiles_by_index[int(tile_index)]
|
||||
tile_latent = _tile_latent(latent_tensor, tile)
|
||||
tile_noise = _tile_latent(full_noise, tile)
|
||||
tile_mask = _tile_mask(denoise_mask, tile)
|
||||
tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
|
||||
_blend_tile_into_sums(
|
||||
tile,
|
||||
tile_samples,
|
||||
tile_denoised,
|
||||
output_sum,
|
||||
weight_sum,
|
||||
denoised_sum,
|
||||
n_h_tiles,
|
||||
n_w_tiles,
|
||||
tile_overlap,
|
||||
)
|
||||
integrated.add(int(tile_index))
|
||||
await _mark_task_completed(job_id, int(tile_index), {"tile_idx": int(tile_index), "source": "master"})
|
||||
continue
|
||||
|
||||
now = asyncio.get_running_loop().time()
|
||||
if now - last_timeout_check >= 10.0:
|
||||
requeued = await _check_and_requeue_timed_out_workers(job_id, total_tiles)
|
||||
if requeued:
|
||||
log(f"LTX Dist: Requeued {requeued} timed-out tile(s)")
|
||||
last_timeout_check = now
|
||||
|
||||
completed_count = await _get_completed_count(job_id)
|
||||
debug_log(f"LTX Dist: Progress {completed_count}/{total_tiles} completed, {len(integrated)} integrated")
|
||||
await asyncio.sleep(0.1)
|
||||
finally:
|
||||
await _cleanup_job(job_id)
|
||||
|
||||
weight_sum = torch.where(weight_sum == 0, torch.ones_like(weight_sum), weight_sum)
|
||||
samples = output_sum / weight_sum
|
||||
denoised = denoised_sum / weight_sum if denoised_sum is not None else None
|
||||
return samples, denoised
|
||||
|
||||
|
||||
def _ltx_local_tiled_process(
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_tensor,
|
||||
denoise_mask,
|
||||
n_h_tiles,
|
||||
n_w_tiles,
|
||||
tile_overlap,
|
||||
):
|
||||
"""Process all LTX latent tiles locally in the current ComfyUI execution context.
|
||||
|
||||
This path is used when no workers are enabled. Keeping it synchronous avoids
|
||||
crossing ComfyUI's execution/inference context boundary while still using the
|
||||
same tile extraction and blend logic as the distributed master/worker path.
|
||||
"""
|
||||
tiles = _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
|
||||
full_noise = noise.generate_noise({"samples": latent_tensor})
|
||||
output_sum = torch.zeros_like(latent_tensor)
|
||||
weight_sum = torch.zeros_like(latent_tensor)
|
||||
denoised_sum = None
|
||||
|
||||
for tile in tiles:
|
||||
tile_latent = _tile_latent(latent_tensor, tile)
|
||||
tile_noise = _tile_latent(full_noise, tile)
|
||||
tile_mask = _tile_mask(denoise_mask, tile)
|
||||
tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
|
||||
if denoised_sum is None and tile_denoised is not None:
|
||||
denoised_sum = torch.zeros_like(latent_tensor)
|
||||
_blend_tile_into_sums(
|
||||
tile,
|
||||
tile_samples,
|
||||
tile_denoised,
|
||||
output_sum,
|
||||
weight_sum,
|
||||
denoised_sum,
|
||||
n_h_tiles,
|
||||
n_w_tiles,
|
||||
tile_overlap,
|
||||
)
|
||||
|
||||
weight_sum = torch.where(weight_sum == 0, torch.ones_like(weight_sum), weight_sum)
|
||||
samples = output_sum / weight_sum
|
||||
denoised = denoised_sum / weight_sum if denoised_sum is not None else None
|
||||
return samples, denoised
|
||||
|
||||
|
||||
class LTXTiledSamplerDistributed:
|
||||
"""Distributed wrapper around the LTX tiled sampler."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
base_inputs = _copy_base_inputs()
|
||||
base_inputs.setdefault("hidden", {})
|
||||
base_inputs["hidden"].update(
|
||||
{
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
}
|
||||
)
|
||||
return base_inputs
|
||||
|
||||
RETURN_TYPES = getattr(LTXTiledSampler, "RETURN_TYPES", ("LATENT", "LATENT", "LATENT"))
|
||||
RETURN_NAMES = getattr(LTXTiledSampler, "RETURN_NAMES", ("samples", "denoised", "preview"))
|
||||
FUNCTION = "sample_tiled_distributed"
|
||||
CATEGORY = "latent/diffusion"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
|
||||
def __init__(self):
|
||||
self.base = LTXTiledSampler() if LTXTiledSampler is not None else None
|
||||
|
||||
def _run_base(self, noise, guider, sampler, sigmas, latent_image, denoise_mask, h_tiles, w_tiles, overlap):
|
||||
if self.base is None:
|
||||
raise RuntimeError(
|
||||
"LTXTiledSamplerDistributed requires the base LTXTiledSampler node for local fallback. "
|
||||
"Install the LTX tiled sampler dependency or run with distributed multi_job_id."
|
||||
)
|
||||
return self.base.sample_tiled(
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_image,
|
||||
denoise_mask,
|
||||
h_tiles,
|
||||
w_tiles,
|
||||
overlap,
|
||||
)
|
||||
|
||||
def sample_tiled_distributed(
|
||||
self,
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_image,
|
||||
denoise_mask=None,
|
||||
h_tiles: int = 1,
|
||||
w_tiles: int = 1,
|
||||
overlap: int = 0,
|
||||
multi_job_id: str = "",
|
||||
is_worker: bool = False,
|
||||
master_url: str = "",
|
||||
worker_id: str = "",
|
||||
enabled_worker_ids: str = "[]",
|
||||
):
|
||||
if not multi_job_id:
|
||||
return self._run_base(noise, guider, sampler, sigmas, latent_image, denoise_mask, h_tiles, w_tiles, overlap)
|
||||
|
||||
latent_tensor, rebuild_latent = _extract_tensor_and_rebuilder(latent_image)
|
||||
if latent_tensor.ndim != 5:
|
||||
raise ValueError(
|
||||
f"LTXTiledSamplerDistributed expects a 5D latent tensor [B,C,F,H,W]; got shape {tuple(latent_tensor.shape)}"
|
||||
)
|
||||
mask_tensor = _extract_optional_mask(denoise_mask)
|
||||
self.multi_job_id = str(multi_job_id)
|
||||
self.master_url = str(master_url or "").rstrip("/")
|
||||
self.worker_id = str(worker_id or "")
|
||||
if isinstance(enabled_worker_ids, str):
|
||||
try:
|
||||
import json
|
||||
|
||||
parsed_workers = json.loads(enabled_worker_ids) if enabled_worker_ids else []
|
||||
except Exception:
|
||||
parsed_workers = []
|
||||
else:
|
||||
parsed_workers = enabled_worker_ids or []
|
||||
self.enabled_worker_ids = [str(worker) for worker in parsed_workers]
|
||||
|
||||
if is_worker:
|
||||
run_async_in_server_loop(
|
||||
_ltx_worker_process(
|
||||
self,
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_tensor,
|
||||
mask_tensor,
|
||||
int(h_tiles),
|
||||
int(w_tiles),
|
||||
int(overlap),
|
||||
),
|
||||
timeout=None,
|
||||
)
|
||||
placeholder = rebuild_latent(latent_tensor)
|
||||
return (placeholder, placeholder, placeholder)
|
||||
|
||||
if not self.enabled_worker_ids:
|
||||
samples, denoised = _ltx_local_tiled_process(
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_tensor,
|
||||
mask_tensor,
|
||||
int(h_tiles),
|
||||
int(w_tiles),
|
||||
int(overlap),
|
||||
)
|
||||
else:
|
||||
samples, denoised = run_async_in_server_loop(
|
||||
_ltx_master_process(
|
||||
self,
|
||||
noise,
|
||||
guider,
|
||||
sampler,
|
||||
sigmas,
|
||||
latent_tensor,
|
||||
mask_tensor,
|
||||
int(h_tiles),
|
||||
int(w_tiles),
|
||||
int(overlap),
|
||||
),
|
||||
timeout=None,
|
||||
)
|
||||
samples_latent = rebuild_latent(samples)
|
||||
denoised_latent = rebuild_latent(denoised if denoised is not None else samples)
|
||||
return (samples_latent, denoised_latent, samples_latent)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"LTXTiledSamplerDistributed": LTXTiledSamplerDistributed}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"LTXTiledSamplerDistributed": "LTX Tiled Sampler (Distributed)"}
|
||||
+2
-2
@@ -269,7 +269,7 @@ class ImageBatchDivider:
|
||||
|
||||
|
||||
class AudioBatchDivider:
|
||||
"""Divides an audio waveform into multiple parts along the time/samples dimension."""
|
||||
"""Divides an audio waveform into sequential segments along the time dimension."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -282,7 +282,7 @@ class AudioBatchDivider:
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the audio into"
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
+181
@@ -0,0 +1,181 @@
|
||||
"""Native V3 schemas with the existing execution algorithms kept intact.
|
||||
|
||||
Runtime objects are private, per-execution helpers, not registered V1 nodes.
|
||||
This avoids sharing mutable instance state through V3's sanitized class clones.
|
||||
"""
|
||||
import comfy.samplers
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
|
||||
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
|
||||
max=1125899906842624, force_input=False)],
|
||||
outputs=[io.Int.Output(display_name='seed')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seed, is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedValue(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedValue', display_name='Distributed Value', category='utils',
|
||||
inputs=[io.String.Input('default_value', default=''),
|
||||
io.String.Input('worker_values', default='{}')],
|
||||
outputs=[io.AnyType.Output(display_name='value')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedValue().distribute(
|
||||
default_value, worker_values, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedModelName(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
|
||||
inputs=[io.String.Input('text', default='')],
|
||||
outputs=[io.AnyType.Output(display_name='output')],
|
||||
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, text):
|
||||
result = _utilities.DistributedModelName().log_input(
|
||||
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
|
||||
return io.NodeOutput(*result['result'], ui=result['ui'])
|
||||
|
||||
|
||||
class ImageBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
|
||||
inputs=[io.Image.Input('images'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of parts to divide the batch into')],
|
||||
# The existing frontend still displays only divide_by sockets.
|
||||
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, divide_by):
|
||||
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
|
||||
|
||||
|
||||
class AudioBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
|
||||
inputs=[io.Audio.Input('audio'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of sequential time segments to create')],
|
||||
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, divide_by):
|
||||
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
|
||||
|
||||
|
||||
class DistributedEmptyImage(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
|
||||
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('width', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
|
||||
outputs=[io.Image.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, height, width, channels):
|
||||
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
|
||||
|
||||
|
||||
class DistributedCollector(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
|
||||
inputs=[io.Boolean.Input('load_balance', default=False,
|
||||
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
|
||||
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)],
|
||||
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
|
||||
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
|
||||
worker_id='', pass_through=False, delegate_only=False):
|
||||
return io.NodeOutput(*_CollectorRuntime().run(
|
||||
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
|
||||
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
|
||||
worker_batch_size=worker_batch_size, worker_id=worker_id,
|
||||
pass_through=pass_through, delegate_only=delegate_only))
|
||||
|
||||
|
||||
class UltimateSDUpscaleDistributed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='UltimateSDUpscaleDistributed',
|
||||
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
|
||||
inputs=[
|
||||
io.Image.Input('upscaled_image'), io.Model.Input('model'),
|
||||
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
|
||||
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
|
||||
io.Int.Input('steps', default=20, min=1, max=10000),
|
||||
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
|
||||
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
|
||||
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
|
||||
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
|
||||
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('padding', default=32, min=0, max=256, step=8),
|
||||
io.Int.Input('mask_blur', default=8, min=0, max=256),
|
||||
io.Boolean.Input('force_uniform_tiles', default=True),
|
||||
io.Boolean.Input('tiled_decode', default=False),
|
||||
], outputs=[io.Image.Output()], accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, **kwargs):
|
||||
return _UpscaleRuntime.IS_CHANGED(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
|
||||
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
|
||||
return io.NodeOutput(*_UpscaleRuntime().run(
|
||||
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
|
||||
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
|
||||
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
|
||||
+2
-2
@@ -1,9 +1,9 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.4.6"
|
||||
version = "1.4.7"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["safetensors"]
|
||||
dependencies = []
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
|
||||
|
||||
@@ -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(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_job_complete_accepts_audio_without_image(self):
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "audio-only-job",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
decode_image.assert_not_called()
|
||||
queued = await queue.get()
|
||||
self.assertIsNone(queued["tensor"])
|
||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||
|
||||
async def test_job_complete_rejects_payload_without_image_or_audio(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("image or audio", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": 123,
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("image", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_job_complete_returns_400_for_malformed_audio(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("audio.shape", response.payload.get("message", "").lower())
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_shape(self):
|
||||
bad = {
|
||||
"sample_rate": 44100,
|
||||
|
||||
+691
@@ -0,0 +1,691 @@
|
||||
{
|
||||
"DistributedCollector": {
|
||||
"input": {
|
||||
"required": {
|
||||
"load_balance": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false,
|
||||
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
|
||||
}
|
||||
]
|
||||
},
|
||||
"optional": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_batch_size": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1024
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"pass_through": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"delegate_only": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"load_balance"
|
||||
],
|
||||
"optional": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_batch_size",
|
||||
"worker_id",
|
||||
"pass_through",
|
||||
"delegate_only"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE",
|
||||
"AUDIO"
|
||||
],
|
||||
"output_name": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": true,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Collector"
|
||||
},
|
||||
"DistributedSeed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1125899906842,
|
||||
"min": 0,
|
||||
"max": 1125899906842624,
|
||||
"forceInput": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"seed"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"INT"
|
||||
],
|
||||
"output_name": [
|
||||
"seed"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Seed"
|
||||
},
|
||||
"DistributedModelName": {
|
||||
"input": {
|
||||
"required": {
|
||||
"text": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO"
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"text"
|
||||
],
|
||||
"hidden": [
|
||||
"unique_id",
|
||||
"extra_pnginfo"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"output"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Model Name"
|
||||
},
|
||||
"DistributedValue": {
|
||||
"input": {
|
||||
"required": {
|
||||
"default_value": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"worker_values": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "{}"
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"default_value",
|
||||
"worker_values"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"value"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Value"
|
||||
},
|
||||
"ImageBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the batch into"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"images",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"batch_1",
|
||||
"batch_2",
|
||||
"batch_3",
|
||||
"batch_4",
|
||||
"batch_5",
|
||||
"batch_6",
|
||||
"batch_7",
|
||||
"batch_8",
|
||||
"batch_9",
|
||||
"batch_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "image",
|
||||
"display_name": "Image Batch Divider"
|
||||
},
|
||||
"AudioBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"audio",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"audio_1",
|
||||
"audio_2",
|
||||
"audio_3",
|
||||
"audio_4",
|
||||
"audio_5",
|
||||
"audio_6",
|
||||
"audio_7",
|
||||
"audio_8",
|
||||
"audio_9",
|
||||
"audio_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "audio",
|
||||
"display_name": "Audio Segment Divider"
|
||||
},
|
||||
"DistributedEmptyImage": {
|
||||
"input": {
|
||||
"required": {
|
||||
"height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"channels": [
|
||||
"INT",
|
||||
{
|
||||
"default": 3,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"step": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"height",
|
||||
"width",
|
||||
"channels"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Empty Image"
|
||||
},
|
||||
"UltimateSDUpscaleDistributed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"upscaled_image": [
|
||||
"IMAGE"
|
||||
],
|
||||
"model": [
|
||||
"MODEL"
|
||||
],
|
||||
"positive": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"negative": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"vae": [
|
||||
"VAE"
|
||||
],
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 18446744073709551615
|
||||
}
|
||||
],
|
||||
"steps": [
|
||||
"INT",
|
||||
{
|
||||
"default": 20,
|
||||
"min": 1,
|
||||
"max": 10000
|
||||
}
|
||||
],
|
||||
"cfg": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 8.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0
|
||||
}
|
||||
],
|
||||
"sampler_name": [
|
||||
[
|
||||
"euler",
|
||||
"euler_cfg_pp",
|
||||
"euler_ancestral",
|
||||
"euler_ancestral_cfg_pp",
|
||||
"heun",
|
||||
"heunpp2",
|
||||
"exp_heun_2_x0",
|
||||
"exp_heun_2_x0_sde",
|
||||
"dpm_2",
|
||||
"dpm_2_ancestral",
|
||||
"lms",
|
||||
"dpm_fast",
|
||||
"dpm_adaptive",
|
||||
"dpmpp_2s_ancestral",
|
||||
"dpmpp_2s_ancestral_cfg_pp",
|
||||
"dpmpp_sde",
|
||||
"dpmpp_sde_gpu",
|
||||
"dpmpp_2m",
|
||||
"dpmpp_2m_cfg_pp",
|
||||
"dpmpp_2m_sde",
|
||||
"dpmpp_2m_sde_gpu",
|
||||
"dpmpp_2m_sde_heun",
|
||||
"dpmpp_2m_sde_heun_gpu",
|
||||
"dpmpp_3m_sde",
|
||||
"dpmpp_3m_sde_gpu",
|
||||
"ddpm",
|
||||
"lcm",
|
||||
"ipndm",
|
||||
"ipndm_v",
|
||||
"deis",
|
||||
"cfgpp_ud10_ab",
|
||||
"res_multistep",
|
||||
"res_multistep_cfg_pp",
|
||||
"res_multistep_ancestral",
|
||||
"res_multistep_ancestral_cfg_pp",
|
||||
"gradient_estimation",
|
||||
"gradient_estimation_cfg_pp",
|
||||
"er_sde",
|
||||
"seeds_2",
|
||||
"seeds_3",
|
||||
"sa_solver",
|
||||
"sa_solver_pece",
|
||||
"ddim",
|
||||
"uni_pc",
|
||||
"uni_pc_bh2"
|
||||
]
|
||||
],
|
||||
"scheduler": [
|
||||
[
|
||||
"simple",
|
||||
"sgm_uniform",
|
||||
"karras",
|
||||
"exponential",
|
||||
"ddim_uniform",
|
||||
"beta",
|
||||
"normal",
|
||||
"linear_quadratic",
|
||||
"kl_optimal"
|
||||
]
|
||||
],
|
||||
"denoise": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}
|
||||
],
|
||||
"tile_width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"tile_height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"padding": [
|
||||
"INT",
|
||||
{
|
||||
"default": 32,
|
||||
"min": 0,
|
||||
"max": 256,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"mask_blur": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 0,
|
||||
"max": 256
|
||||
}
|
||||
],
|
||||
"force_uniform_tiles": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": true
|
||||
}
|
||||
],
|
||||
"tiled_decode": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"tile_indices": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"dynamic_threshold": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 64
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"upscaled_image",
|
||||
"model",
|
||||
"positive",
|
||||
"negative",
|
||||
"vae",
|
||||
"seed",
|
||||
"steps",
|
||||
"cfg",
|
||||
"sampler_name",
|
||||
"scheduler",
|
||||
"denoise",
|
||||
"tile_width",
|
||||
"tile_height",
|
||||
"padding",
|
||||
"mask_blur",
|
||||
"force_uniform_tiles",
|
||||
"tiled_decode"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_id",
|
||||
"tile_indices",
|
||||
"dynamic_threshold"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image/upscaling",
|
||||
"display_name": "Ultimate SD Upscale Distributed (No Upscale)"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
|
||||
|
||||
No HTTP listener, workers, model downloads or live installation changes.
|
||||
The original contracts were captured from 32ac027 using the same core.
|
||||
"""
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
||||
sys.path.insert(0, str(COMFY_ROOT))
|
||||
os.environ['COMFYUI_IS_WORKER'] = '1'
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
from app.assets.manager import default_asset_manager
|
||||
from comfy_api.v0_0_2 import io
|
||||
import execution
|
||||
import nodes
|
||||
import server
|
||||
import torch
|
||||
|
||||
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
|
||||
|
||||
|
||||
def normalized_input(value):
|
||||
kind = value[0]
|
||||
opts = dict(value[1]) if len(value) > 1 else {}
|
||||
if kind == 'STRING':
|
||||
# V3 serializes the same default single-line widget explicitly.
|
||||
opts.setdefault('multiline', False)
|
||||
if kind == 'COMBO':
|
||||
kind = opts['options']
|
||||
opts = {key: val for key, val in opts.items() if key != 'options'}
|
||||
if isinstance(kind, list):
|
||||
# Single-selection is the V1 dropdown default as well.
|
||||
opts.setdefault('multiselect', False)
|
||||
return [kind, opts]
|
||||
|
||||
|
||||
async def map_node(cls, values, extra=None):
|
||||
"""Use the actual input parser and sanitized V3 executor clones."""
|
||||
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
|
||||
assert not missing, missing
|
||||
# Use a worker thread for synchronous nodes: collector/upscale bridge back
|
||||
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
|
||||
def run():
|
||||
return asyncio.run(execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
|
||||
return await asyncio.to_thread(run)
|
||||
|
||||
|
||||
def check_saved_workflows(mapping):
|
||||
checked = 0
|
||||
occurrences = 0
|
||||
for path in sorted((ROOT / 'workflows').glob('*.json')):
|
||||
workflow = json.loads(path.read_text())
|
||||
local_nodes = {node['id']: node for node in workflow['nodes']}
|
||||
for saved in workflow['nodes']:
|
||||
if saved['type'] not in mapping:
|
||||
continue
|
||||
occurrences += 1
|
||||
info = mapping[saved['type']].GET_NODE_INFO_V1()
|
||||
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
|
||||
for socket in saved.get('inputs', []):
|
||||
assert socket['name'] in inputs, (path.name, saved['id'], socket)
|
||||
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
|
||||
for index, output in enumerate(saved.get('outputs', [])):
|
||||
assert output['type'] == info['output'][index], (path.name, index)
|
||||
for link in workflow['links']:
|
||||
if link[1] == saved['id']:
|
||||
assert link[2] < len(info['output']), (path.name, link)
|
||||
target = local_nodes[link[3]]['inputs'][link[4]]
|
||||
assert target['type'] == info['output'][link[2]], (path.name, link)
|
||||
checked += 1
|
||||
assert checked == 5 and occurrences == 10, (checked, occurrences)
|
||||
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
|
||||
|
||||
|
||||
async def check_execution(mapping, module, prompt_server, asset_manager):
|
||||
result = await map_node(mapping['DistributedSeed'],
|
||||
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (126,)
|
||||
values = json.dumps({'_type': 'INT', '3': '17'})
|
||||
result = await map_node(mapping['DistributedValue'],
|
||||
{'default_value': '4', 'worker_values': values,
|
||||
'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (17,)
|
||||
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
|
||||
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
|
||||
{'extra_pnginfo': metadata})
|
||||
assert result[0].result == ('model.ckpt',)
|
||||
assert result[0].ui == {'text': ['model.ckpt']}
|
||||
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
|
||||
result = await map_node(mapping['DistributedEmptyImage'],
|
||||
{'height': 8, 'width': 8, 'channels': 3})
|
||||
empty = result[0].result[0]
|
||||
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
|
||||
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
|
||||
for count in (1, 3, 10):
|
||||
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat(result[0].result[:count]), images)
|
||||
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
|
||||
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
|
||||
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
|
||||
collector = mapping['DistributedCollector']
|
||||
prepared, missing, hidden = execution.get_input_data(
|
||||
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
|
||||
assert not missing
|
||||
prepared['images'] = [images[:2], images[2:]]
|
||||
prepared['audio'] = [audio, audio]
|
||||
result = await execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
|
||||
assert torch.equal(result[0].result[0], images)
|
||||
assert result[0].result[1]['waveform'].shape[-1] == 66
|
||||
# Exercise real async aggregation without an HTTP endpoint or worker.
|
||||
queue = asyncio.Queue()
|
||||
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
|
||||
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
|
||||
result = await map_node(collector,
|
||||
{'images': images[:2], 'load_balance': False,
|
||||
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
|
||||
assert torch.equal(result[0].result[0], images[:3])
|
||||
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
|
||||
|
||||
# Prove V3's per-execution class clones do not leak helper instance state;
|
||||
# replace only the GPU/model boundary, not the input parser or executor.
|
||||
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
|
||||
seen = []
|
||||
def fake_upscale(self, *args):
|
||||
assert not hasattr(self, 'acceptance_marker')
|
||||
self.acceptance_marker = True
|
||||
seen.append(args)
|
||||
return (args[0],)
|
||||
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
|
||||
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
|
||||
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
|
||||
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
|
||||
'force_uniform_tiles': True, 'tiled_decode': False,
|
||||
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
|
||||
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
|
||||
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
|
||||
with patch.object(upscale_runtime, 'run', fake_upscale):
|
||||
for _ in range(2):
|
||||
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
|
||||
assert result[0].result[0] is images
|
||||
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
|
||||
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
|
||||
'tile_indices', 'dynamic_threshold')), seen[0]
|
||||
import math
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
|
||||
|
||||
graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
|
||||
'pass_through': True}},
|
||||
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
|
||||
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('v3-graph', graph, None)
|
||||
assert valid[0], valid
|
||||
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
|
||||
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
|
||||
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
|
||||
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
|
||||
assert not rejected[0]
|
||||
errors = [error for entry in rejected[3].values() for error in entry['errors']]
|
||||
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
|
||||
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
|
||||
import folder_paths
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
history = executor.history_result
|
||||
record = history['outputs']['4']['images'][0]
|
||||
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
||||
with Image.open(preview) as image:
|
||||
assert image.size == (8, 8) and image.mode == 'RGB'
|
||||
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
|
||||
print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
|
||||
|
||||
|
||||
async def main():
|
||||
asset_manager = default_asset_manager()
|
||||
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
|
||||
assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
|
||||
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
||||
assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
|
||||
extension = await module.comfy_entrypoint()
|
||||
classes = await extension.get_node_list()
|
||||
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
|
||||
assert len(classes) == len(mapping) == len(BASELINE) == 8
|
||||
assert set(mapping) == set(BASELINE)
|
||||
for node_id, cls in mapping.items():
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
|
||||
old = dict(BASELINE[node_id])
|
||||
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
|
||||
# V1's ByPassTypeTuple advertises '*' when indexed, while its
|
||||
# underlying tuple and existing frontend declare IMAGE/AUDIO.
|
||||
# Native V3 declares all ten existing typed sockets explicitly.
|
||||
assert old['output'] == ['*'] * 10
|
||||
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
|
||||
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
|
||||
for group in ('required', 'optional'):
|
||||
old_inputs = old['input'].get(group, {})
|
||||
new_inputs = new['input'].get(group, {})
|
||||
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
|
||||
for name, original in old_inputs.items():
|
||||
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
|
||||
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
|
||||
assert old[key] == new[key], (node_id, key, old[key], new[key])
|
||||
# Standard context lives in cls.hidden; orchestrator metadata remains
|
||||
# accepted by its original kwarg name, without creating new widgets.
|
||||
signature = inspect.signature(cls.execute)
|
||||
for name, field in old['input'].get('hidden', {}).items():
|
||||
if isinstance(field, list):
|
||||
assert cls.GET_SCHEMA().accept_all_inputs, node_id
|
||||
assert name in signature.parameters, (node_id, name)
|
||||
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
|
||||
else:
|
||||
assert name in new['input']['hidden'], (node_id, name)
|
||||
expected_hidden = []
|
||||
if node_id == 'DistributedModelName':
|
||||
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
|
||||
if old['output_node']:
|
||||
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
||||
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
||||
print('SCHEMA_PARITY_OK', len(mapping))
|
||||
check_saved_workflows(mapping)
|
||||
await check_execution(mapping, module, prompt_server, asset_manager)
|
||||
print('V3_ACCEPTANCE_OK')
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -124,6 +124,49 @@ def test_collector_opts_into_comfyui_list_inputs():
|
||||
assert collector.INPUT_IS_LIST is True
|
||||
|
||||
|
||||
def test_collector_exposes_images_as_optional_input():
|
||||
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
|
||||
|
||||
assert "images" not in input_types["required"]
|
||||
assert input_types["optional"]["images"] == ("IMAGE",)
|
||||
|
||||
|
||||
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||
|
||||
images, returned_audio = collector.run(images=None, audio=[audio])
|
||||
|
||||
assert images is None
|
||||
assert returned_audio is audio
|
||||
|
||||
|
||||
def test_collector_rejects_missing_images_and_audio():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
|
||||
try:
|
||||
collector.run(images=None, audio=None)
|
||||
except ValueError as exc:
|
||||
assert "image or audio" in str(exc).lower()
|
||||
else:
|
||||
raise AssertionError("Expected collector to reject a run with no media input")
|
||||
|
||||
|
||||
def test_delegate_only_master_allows_no_local_media_input():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
|
||||
images, audio = collector.run(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=["delegate-audio-job"],
|
||||
delegate_only=[True],
|
||||
enabled_worker_ids=["[]"],
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
|
||||
|
||||
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
first = torch.zeros(1, 2, 2, 3)
|
||||
@@ -203,3 +246,119 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
|
||||
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
|
||||
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
|
||||
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
|
||||
|
||||
|
||||
def test_audio_only_worker_sends_one_completion_without_image():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||
posted = []
|
||||
|
||||
class _FakeResponse:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
class _FakeSession:
|
||||
def post(self, url, json, timeout):
|
||||
posted.append((url, json, timeout.total))
|
||||
return _FakeResponse()
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return _FakeSession()
|
||||
|
||||
module.get_client_session = _fake_get_client_session
|
||||
module.encode_audio_payload = lambda value: {"encoded": value is audio}
|
||||
|
||||
images, returned_audio = collector.run(
|
||||
images=None,
|
||||
audio=[audio],
|
||||
multi_job_id=["audio-job"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
worker_id=["worker-a"],
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert returned_audio is audio
|
||||
assert len(posted) == 1
|
||||
assert posted[0][0] == "http://master/distributed/job_complete"
|
||||
assert posted[0][2] == 600
|
||||
assert posted[0][1] == {
|
||||
"job_id": "audio-job",
|
||||
"worker_id": "worker-a",
|
||||
"batch_idx": 0,
|
||||
"audio": {"encoded": True},
|
||||
"is_last": True,
|
||||
}
|
||||
|
||||
|
||||
def test_audio_only_master_combines_local_and_worker_audio():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
|
||||
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait(
|
||||
{
|
||||
"worker_id": "worker-a",
|
||||
"image_index": 0,
|
||||
"tensor": None,
|
||||
"audio": worker_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
|
||||
|
||||
images, combined_audio = asyncio.run(
|
||||
collector.execute(
|
||||
images=None,
|
||||
audio=master_audio,
|
||||
multi_job_id="audio-job",
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
)
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert combined_audio["sample_rate"] == 48000
|
||||
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
|
||||
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
|
||||
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
|
||||
|
||||
|
||||
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait(
|
||||
{
|
||||
"worker_id": "worker-a",
|
||||
"image_index": 0,
|
||||
"tensor": None,
|
||||
"audio": worker_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
|
||||
|
||||
images, combined_audio = asyncio.run(
|
||||
collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id="delegate-audio-job",
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert combined_audio["sample_rate"] == 48000
|
||||
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
|
||||
|
||||
@@ -1,157 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_ltx_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "ltx_tiled_sampler.py"
|
||||
package_name = "dist_ltx_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
for subpkg in ["nodes", "upscale", "utils"]:
|
||||
pkg = types.ModuleType(f"{package_name}.{subpkg}")
|
||||
pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.{subpkg}"] = pkg
|
||||
|
||||
job_models = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class TileJobState:
|
||||
pass
|
||||
|
||||
job_models.TileJobState = TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models
|
||||
|
||||
async def _noop(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
job_store = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store._cleanup_job = _noop
|
||||
job_store._drain_results_queue = _noop
|
||||
job_store._get_completed_count = _noop
|
||||
job_store._mark_task_completed = _noop
|
||||
job_store.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
job_store.init_static_job_batched = _noop
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store
|
||||
|
||||
job_timeout = types.ModuleType(f"{package_name}.upscale.job_timeout")
|
||||
job_timeout._check_and_requeue_timed_out_workers = _noop
|
||||
sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout
|
||||
|
||||
async_helpers = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers.run_async_in_server_loop = lambda coro, timeout=None: asyncio.run(coro)
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers
|
||||
|
||||
constants = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants.HEARTBEAT_INTERVAL = 10.0
|
||||
constants.TILE_WAIT_TIMEOUT = 1.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants
|
||||
|
||||
logging = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging.debug_log = lambda *_args, **_kwargs: None
|
||||
logging.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging
|
||||
|
||||
network = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _missing_session():
|
||||
raise RuntimeError("network not available in unit test")
|
||||
|
||||
network.get_client_session = _missing_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network
|
||||
|
||||
usdu = types.ModuleType(f"{package_name}.utils.usdu_managment")
|
||||
usdu._send_heartbeat_to_master = _noop
|
||||
sys.modules[f"{package_name}.utils.usdu_managment"] = usdu
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.ltx_tiled_sampler", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
ltx = _load_ltx_module()
|
||||
|
||||
|
||||
class LtxTiledSamplerTests(unittest.TestCase):
|
||||
def test_compute_tile_starts_with_overlap_covers_edges(self):
|
||||
starts, tile_size = ltx._compute_tile_starts(total_size=10, n_tiles=3, overlap=2)
|
||||
|
||||
self.assertEqual(starts[0], 0)
|
||||
self.assertEqual(starts[-1] + tile_size, 10)
|
||||
self.assertGreaterEqual(tile_size, 4)
|
||||
|
||||
def test_pack_unpack_tensors_round_trips_with_safetensors(self):
|
||||
tensor = torch.arange(12, dtype=torch.float32).reshape(1, 1, 3, 2, 2)
|
||||
|
||||
payload = ltx._pack_tensors({"samples": tensor})
|
||||
restored = ltx._unpack_tensors(payload, device="cpu")
|
||||
|
||||
self.assertTrue(torch.equal(restored["samples"], tensor))
|
||||
|
||||
def test_latent_dict_rebuilder_preserves_metadata(self):
|
||||
samples = torch.zeros((1, 2, 3, 4, 5))
|
||||
latent = {"samples": samples, "noise_mask": "keep-me"}
|
||||
|
||||
extracted, rebuild = ltx._extract_tensor_and_rebuilder(latent)
|
||||
rebuilt = rebuild(torch.ones_like(extracted))
|
||||
|
||||
self.assertIs(extracted, samples)
|
||||
self.assertEqual(rebuilt["noise_mask"], "keep-me")
|
||||
self.assertTrue(torch.equal(rebuilt["samples"], torch.ones_like(samples)))
|
||||
|
||||
def test_node_schema_includes_distributed_hidden_inputs(self):
|
||||
schema = ltx.LTXTiledSamplerDistributed.INPUT_TYPES()
|
||||
|
||||
self.assertIn("multi_job_id", schema["hidden"])
|
||||
self.assertIn("is_worker", schema["hidden"])
|
||||
self.assertIn("master_url", schema["hidden"])
|
||||
self.assertIn("enabled_worker_ids", schema["hidden"])
|
||||
self.assertIn("LTXTiledSamplerDistributed", ltx.NODE_CLASS_MAPPINGS)
|
||||
|
||||
|
||||
def test_local_tiled_process_runs_without_async_master_path(self):
|
||||
class _Noise:
|
||||
def generate_noise(self, latent):
|
||||
return torch.ones_like(latent['samples']) * 0.25
|
||||
|
||||
class _Guider:
|
||||
def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, disable_pbar=True):
|
||||
return latent_image + noise
|
||||
|
||||
latent = torch.zeros((1, 1, 2, 4, 4), dtype=torch.float32)
|
||||
|
||||
samples, denoised = ltx._ltx_local_tiled_process(
|
||||
_Noise(),
|
||||
_Guider(),
|
||||
sampler=object(),
|
||||
sigmas=torch.tensor([1.0, 0.0]),
|
||||
latent_tensor=latent,
|
||||
denoise_mask=None,
|
||||
n_h_tiles=2,
|
||||
n_w_tiles=2,
|
||||
tile_overlap=0,
|
||||
)
|
||||
|
||||
self.assertIsNone(denoised)
|
||||
self.assertEqual(tuple(samples.shape), tuple(latent.shape))
|
||||
self.assertTrue(torch.allclose(samples, torch.full_like(latent, 0.25)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -71,6 +71,15 @@ def _collector_only_prompt():
|
||||
}
|
||||
|
||||
|
||||
def _audio_only_prompt():
|
||||
"""1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
|
||||
return {
|
||||
"1": {"class_type": "LoadAudio", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
|
||||
"3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
|
||||
}
|
||||
|
||||
|
||||
def _delegate_prompt():
|
||||
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
|
||||
return {
|
||||
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
|
||||
self.assertEqual(len(preview_nodes), 1)
|
||||
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
|
||||
|
||||
def test_injects_preview_audio_for_audio_only_collector(self):
|
||||
result = pt.prune_prompt_for_worker(_audio_only_prompt())
|
||||
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
|
||||
self.assertEqual(len(preview_nodes), 1)
|
||||
self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
|
||||
self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
|
||||
|
||||
def test_prefers_image_preview_when_collector_has_images_and_audio(self):
|
||||
prompt = _linear_prompt()
|
||||
prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
|
||||
prompt["4"]["inputs"]["audio"] = ["6", 0]
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
self.assertEqual(
|
||||
len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
|
||||
1,
|
||||
)
|
||||
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
|
||||
|
||||
def test_preserves_image_preview_for_distributed_upscale(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
||||
self.assertEqual(len(preview_nodes), 1)
|
||||
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
|
||||
|
||||
def test_no_preview_image_when_no_downstream(self):
|
||||
result = pt.prune_prompt_for_worker(_collector_only_prompt())
|
||||
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
|
||||
@@ -242,7 +280,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
|
||||
def test_upscale_node_is_treated_as_distributed(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "KSampler", "inputs": {}},
|
||||
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
|
||||
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
result = pt.prune_prompt_for_worker(prompt)
|
||||
@@ -285,12 +323,22 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
|
||||
placeholder_id = empty_nodes[0][0]
|
||||
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
|
||||
|
||||
def test_audio_only_collector_does_not_get_image_placeholder(self):
|
||||
prompt = _audio_only_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["2"])
|
||||
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
|
||||
self.assertEqual(empty_nodes, [])
|
||||
self.assertNotIn("images", result["2"].get("inputs", {}))
|
||||
self.assertNotIn("audio", result["2"].get("inputs", {}))
|
||||
|
||||
def test_one_placeholder_per_collector(self):
|
||||
"""Two collectors → two placeholders."""
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {}},
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
||||
"10": {"class_type": "LoadImage", "inputs": {}},
|
||||
"11": {"class_type": "LoadImage", "inputs": {}},
|
||||
}
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
|
||||
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
|
||||
|
||||
@@ -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";
|
||||
|
||||
const NODE_CLASS = "DistributedValue";
|
||||
const CONVERTED_WIDGET = "converted-widget";
|
||||
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
|
||||
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
|
||||
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
|
||||
@@ -42,15 +41,13 @@ function getDynamicWorkerWidgets(node) {
|
||||
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
|
||||
}
|
||||
|
||||
function hideWidgetForGood(node, widget, suffix = "") {
|
||||
function hideWidgetForGood(widget) {
|
||||
if (!widget) return;
|
||||
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
|
||||
|
||||
widget.origType = widget.type;
|
||||
widget.origComputeSize = widget.computeSize;
|
||||
widget.origSerializeValue = widget.serializeValue;
|
||||
widget.computeSize = () => [0, -4];
|
||||
widget.type = `${CONVERTED_WIDGET}${suffix}`;
|
||||
// These are serialized backing fields, not widgets converted into sockets.
|
||||
// Keep their native type/serializer and mutate the shared options in place
|
||||
// so both the classic canvas and Nodes 2.0 suppress the whole widget row.
|
||||
widget.options ??= {};
|
||||
widget.options.hidden = true;
|
||||
|
||||
// Hide any attached DOM element (multiline widgets).
|
||||
if (widget.element) widget.element.style.display = "none";
|
||||
@@ -58,14 +55,14 @@ function hideWidgetForGood(node, widget, suffix = "") {
|
||||
|
||||
if (widget.linkedWidgets) {
|
||||
for (const linked of widget.linkedWidgets) {
|
||||
hideWidgetForGood(node, linked, `:${widget.name}`);
|
||||
hideWidgetForGood(linked);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function hideRawWidgets(node) {
|
||||
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
|
||||
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
|
||||
hideWidgetForGood(getRawDefaultWidget(node));
|
||||
hideWidgetForGood(getRawWorkerValuesWidget(node));
|
||||
}
|
||||
|
||||
function removeDynamicDefaultWidget(node) {
|
||||
|
||||
@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
|
||||
createCheckboxSetting(
|
||||
"setting-debug",
|
||||
"Debug Mode",
|
||||
"Enable verbose logging in the browser console.",
|
||||
"Enable verbose logging in the browser console and ComfyUI server output.",
|
||||
extension.config?.settings?.debug || false,
|
||||
(event) => extension._updateSetting("debug", event.target.checked)
|
||||
)
|
||||
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
|
||||
createNumberSetting(
|
||||
"setting-worker-timeout",
|
||||
"Worker Timeout",
|
||||
"Seconds without a heartbeat before a worker is considered timed out. Default 60.",
|
||||
"Maximum result-wait and heartbeat inactivity period before recovery begins. Busy workers may receive additional grace. Default: 60 seconds.",
|
||||
extension.config?.settings?.worker_timeout_seconds ?? 60,
|
||||
10,
|
||||
1,
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
import { readFileSync } from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
// Evaluate this repository's browser entrypoint without needing a ComfyUI server.
|
||||
// The real-renderer acceptance additionally checks the input rows in Chromium.
|
||||
const source = readFileSync(new URL("../distributedValue.js", import.meta.url), "utf8")
|
||||
.replace(/^import .*;\r?\n/gm, "");
|
||||
|
||||
function setup({ connected = false, targetType = "STRING" } = {}) {
|
||||
let extension;
|
||||
const listeners = new Map();
|
||||
const tasks = [];
|
||||
const workers = [
|
||||
{ id: "gpu-a", name: "GPU A", enabled: true },
|
||||
{ id: "gpu-b", name: "GPU B", enabled: true },
|
||||
{ id: "off", name: "Disabled", enabled: false },
|
||||
];
|
||||
const target = {
|
||||
inputs: [{ name: "value" }],
|
||||
widgets: [{ name: "value", type: targetType === "INT" ? "number" : "string", options: { step: 1, precision: 0 } }],
|
||||
};
|
||||
const graph = { links: { 1: { target_id: 2, target_slot: 0 } }, getNodeById: () => target };
|
||||
const raw = ["default_value", "worker_values"].map((name, index) => ({
|
||||
name,
|
||||
type: "text",
|
||||
value: index === 0 ? "saved default" : '{"1":"saved worker"}',
|
||||
options: { multiline: false },
|
||||
computeSize: vi.fn(() => [200, 20]),
|
||||
serializeValue: vi.fn(function () { return this.value; }),
|
||||
}));
|
||||
const node = {
|
||||
comfyClass: "DistributedValue", graph, widgets: [...raw], size: [200, 100],
|
||||
inputs: raw.map(w => ({ name: w.name, widget: { name: w.name }, link: null })),
|
||||
outputs: [{ links: connected ? [1] : [] }],
|
||||
computeSize: () => [200, 100], setSize: vi.fn(), setDirtyCanvas: vi.fn(),
|
||||
addWidget(type, name, value, callback, options) {
|
||||
const widget = { type, name, value, callback, options };
|
||||
this.widgets.push(widget);
|
||||
return widget;
|
||||
},
|
||||
configure(data) {
|
||||
data.widgets_values.forEach((value, i) => { this.widgets[i].value = value; });
|
||||
return "configured";
|
||||
},
|
||||
};
|
||||
vm.runInNewContext(source, {
|
||||
app: { graph, registerExtension: ext => { extension = ext; } },
|
||||
ENDPOINTS: { CONFIG: "/distributed/config" },
|
||||
fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
|
||||
window: { addEventListener: (name, handler) => listeners.set(name, handler) },
|
||||
setTimeout: callback => { tasks.push(callback); }, console,
|
||||
}, { filename: "distributedValue.js" });
|
||||
return { node, raw, extension, listeners, tasks, flush: () => { while (tasks.length) tasks.shift()(); } };
|
||||
}
|
||||
|
||||
describe("Distributed Value internal widget visibility", () => {
|
||||
it("hides raw fields natively without converting their types or replacing their options", async () => {
|
||||
const h = setup();
|
||||
const options = h.raw.map(w => w.options);
|
||||
const sizes = h.raw.map(w => w.computeSize);
|
||||
const serializers = h.raw.map(w => w.serializeValue);
|
||||
const slots = h.node.inputs;
|
||||
await h.extension.nodeCreated(h.node);
|
||||
|
||||
for (const [i, widget] of h.raw.entries()) {
|
||||
expect(widget.type).toBe("text");
|
||||
expect(widget.options).toBe(options[i]);
|
||||
expect(widget.options.hidden).toBe(true);
|
||||
expect(widget.computeSize).toBe(sizes[i]);
|
||||
expect(widget.serializeValue).toBe(serializers[i]);
|
||||
expect(h.node.widgets[i]).toBe(widget);
|
||||
}
|
||||
expect(h.node.inputs).toBe(slots);
|
||||
expect(h.node.widgets.slice(2).map(w => w.label)).toEqual(["default_value", "GPU A", "GPU B"]);
|
||||
expect(h.node.widgets.slice(2).every(w => !w.options.hidden)).toBe(true);
|
||||
});
|
||||
|
||||
it("keeps typed edits in the original serialized fields across rebuild/configure", async () => {
|
||||
const h = setup({ connected: true, targetType: "INT" });
|
||||
await h.extension.nodeCreated(h.node);
|
||||
h.node.widgets.find(w => w.name === "_dv_default").callback(42);
|
||||
h.node.widgets.find(w => w.name === "_dv_worker_1").callback(7);
|
||||
h.node.widgets.find(w => w.name === "_dv_worker_2").callback(9);
|
||||
const values = h.raw.map(w => w.serializeValue());
|
||||
expect(values[0]).toBe(42);
|
||||
expect(JSON.parse(values[1])).toMatchObject({
|
||||
_type: "INT", 1: "7", 2: "9", _by_worker_id: { "gpu-a": "7", "gpu-b": "9" },
|
||||
});
|
||||
expect(h.node.configure({ widgets_values: values })).toBe("configured");
|
||||
h.flush();
|
||||
expect(h.raw.map(w => w.serializeValue())).toEqual(values);
|
||||
expect(h.raw.every(w => w.options.hidden && w.type === "text")).toBe(true);
|
||||
expect(h.node.widgets.slice(2).map(w => w.value)).toEqual([42, 7, 9]);
|
||||
});
|
||||
|
||||
it("reapplies visibility on worker refresh without duplicating raw or dynamic fields", async () => {
|
||||
const h = setup();
|
||||
await h.extension.nodeCreated(h.node);
|
||||
h.raw[0].options.hidden = false;
|
||||
await h.listeners.get("distributed:workers-changed")({ detail: { workers: [
|
||||
{ id: "gpu-b", name: "Renamed B", enabled: true },
|
||||
] } });
|
||||
expect(h.raw.every(w => w.options.hidden)).toBe(true);
|
||||
expect(h.node.widgets.map(w => w.name)).toEqual(["default_value", "worker_values", "_dv_default", "_dv_worker_1"]);
|
||||
expect(h.node.widgets[3].label).toBe("Renamed B");
|
||||
});
|
||||
|
||||
it("hides linked backing widgets and handles absent options without altering serializers", async () => {
|
||||
const h = setup();
|
||||
delete h.raw[1].options;
|
||||
const linked = { type: "text", options: {}, serializeValue: vi.fn(() => "linked") };
|
||||
h.raw[0].linkedWidgets = [linked];
|
||||
await h.extension.nodeCreated(h.node);
|
||||
expect(h.raw[1].options.hidden).toBe(true);
|
||||
expect(linked.options.hidden).toBe(true);
|
||||
expect(linked.type).toBe("text");
|
||||
expect(linked.serializeValue()).toBe("linked");
|
||||
});
|
||||
});
|
||||
@@ -814,7 +814,7 @@
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"If all your GPUs are the same/similar, set static_distribution to true\n"
|
||||
"Workers pull tiles from a shared queue, so faster workers can process more tiles automatically.\n"
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
|
||||
Reference in New Issue
Block a user