feat: add dynamic per-image list queue with Start and End nodes

This commit is contained in:
Robert Wojciechowski
2026-10-11 08:49:23 +00:00
parent 87c0cda9f6
commit 3e280a0e29
19 changed files with 2367 additions and 17 deletions
+5
View File
@@ -27,6 +27,11 @@
- Intelligent distribution
- Handles single images and videos
#### Dynamic Image-List Processing
- Schedule one image at a time to the next available direct worker; collect results in input order
- Use **Distributed Image List Start / End** around a per-image processing region
- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
#### Ease of Use
- Auto-setup local workers; easily add remote/cloud ones
- Convert any workflow to distributed with 2 nodes
+1
View File
@@ -4,3 +4,4 @@ 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 image_list_routes # noqa: F401
+89
View File
@@ -0,0 +1,89 @@
"""Bounded, typed HTTP transport for dynamic image-list worker tasks."""
import json
from aiohttp import web
import server
from ..image_list.runtime import (
MAX_PAYLOAD_BYTES,
decode_result,
delete_worker_item,
set_master_result,
store_worker_item,
)
prompt_server = server.PromptServer.instance
MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024
async def _json_body(request):
length = request.content_length
if length is not None and length > MAX_REQUEST_BYTES:
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
raw = bytearray()
async for chunk in request.content.iter_chunked(64 * 1024):
raw.extend(chunk)
if len(raw) > MAX_REQUEST_BYTES:
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
try:
return json.loads(raw.decode("utf-8"))
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise ValueError(f"Invalid image-list JSON body: {exc}") from exc
@prompt_server.routes.get("/distributed/image-list/capabilities")
async def image_list_capabilities_endpoint(_request):
return web.json_response({
"supported": True,
"protocol_version": 1,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
"transport": "typed-float-tensor-json-v1",
})
@prompt_server.routes.post("/distributed/image-list/item")
async def image_list_item_endpoint(request):
try:
data = await _json_body(request)
if not isinstance(data, dict):
raise ValueError("Expected an image-list item object")
context_id = data.get("context_id")
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
item = data.get("item")
store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"]})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
@prompt_server.routes.delete("/distributed/image-list/item/{token}")
async def image_list_delete_item_endpoint(request):
token = request.match_info.get("token", "")
return web.json_response({"status": "deleted", "removed": delete_worker_item(token)})
@prompt_server.routes.post("/distributed/image-list/result")
async def image_list_result_endpoint(request):
try:
data = await _json_body(request)
if not isinstance(data, dict):
raise ValueError("Expected an image-list result object")
context_id = data.get("context_id")
token = data.get("token")
item_index = data.get("item_index")
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
if not isinstance(token, str) or not token:
raise ValueError("token must be a non-empty string")
if isinstance(item_index, bool) or not isinstance(item_index, int) or item_index < 0:
raise ValueError("item_index must be a non-negative integer")
result = decode_result(data.get("result"))
set_master_result(context_id, item_index, token, result)
return web.json_response({"status": "accepted", "item_index": item_index})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
+73 -3
View File
@@ -85,7 +85,7 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
data = json.loads(msg.data or "{}")
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
if data.get("ok"):
return
return data.get("prompt_id")
error_text = data.get("error") or "Worker rejected websocket dispatch."
validation_error = data.get("validation_error")
node_errors = data.get("node_errors")
@@ -107,6 +107,7 @@ async def dispatch_worker_prompt(
client_id=None,
use_websocket=False,
trace_execution_id=None,
return_prompt_id=False,
):
"""Send the prepared prompt to a worker ComfyUI instance."""
worker_url = build_worker_url(worker)
@@ -120,7 +121,7 @@ async def dispatch_worker_prompt(
if use_websocket:
try:
await _dispatch_via_websocket(
prompt_id = await _dispatch_via_websocket(
worker_url,
{
"prompt": prompt_obj,
@@ -128,7 +129,9 @@ async def dispatch_worker_prompt(
},
client_id,
)
return
if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
return prompt_id if return_prompt_id else None
except Exception as exc:
worker_id = worker.get("id")
if trace_execution_id:
@@ -144,6 +147,73 @@ async def dispatch_worker_prompt(
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
if not return_prompt_id:
return None
try:
response_data = await resp.json()
except Exception as exc:
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError("Worker /prompt response omitted prompt_id")
return prompt_id if return_prompt_id else None
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
"""Require protocol/node support and report which direct workers are idle.
This performs read-only GETs only; it never submits a prompt or contacts a
serverless provider. Unsupported worker releases fail before master inputs run.
"""
if not workers:
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
session = await get_client_session()
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
semaphore = asyncio.Semaphore(8)
async def inspect(worker):
worker_id = str(worker.get("id") or "unknown")
async with semaphore:
cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
try:
async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(
f"Image-list worker {worker_id} does not expose the capability endpoint "
f"(HTTP {response.status}); upgrade that worker before running this feature"
)
capability = await response.json()
except asyncio.TimeoutError as exc:
raise RuntimeError(f"Image-list worker {worker_id} capability check timed out") from exc
except aiohttp.ClientError as exc:
raise RuntimeError(f"Image-list worker {worker_id} capability check failed: {exc}") from exc
nodes = set(capability.get("node_ids", [])) if isinstance(capability, dict) else set()
if (not isinstance(capability, dict) or capability.get("supported") is not True
or capability.get("protocol_version") != 1 or not required_nodes.issubset(nodes)):
raise RuntimeError(
f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes "
f"{sorted(required_nodes)}; no work was submitted"
)
queue_url = build_worker_url(worker, "/queue")
try:
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
queue = await response.json()
except asyncio.TimeoutError as exc:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight timed out") from exc
except aiohttp.ClientError as exc:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight failed: {exc}") from exc
if not isinstance(queue, dict) or not isinstance(queue.get("queue_running"), list) or not isinstance(queue.get("queue_pending"), list):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
idle = not queue["queue_running"] and not queue["queue_pending"]
return worker_id, idle
statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
result = {worker_id: idle for worker_id, idle in statuses}
if require_idle and not any(result.values()):
raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
return {worker_id: True for worker_id, idle in result.items() if idle}
async def select_active_workers(
+53 -2
View File
@@ -441,11 +441,39 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
return pruned_prompt
def _resolve_direct_workers_for_image_list(config, requested_ids):
"""Resolve only enabled direct workers already selected by orchestration."""
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
direct = []
for worker in (config or {}).get("workers", []):
worker_id = str(worker.get("id") or "").strip()
if not worker_id or not worker.get("enabled", False):
continue
if requested is not None and worker_id not in requested:
continue
if str(worker.get("type") or "local").strip().lower() == "serverless":
continue
try:
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
direct.append({
"id": worker_id,
"name": str(worker.get("name") or worker_id),
"host": worker.get("host"),
"port": port,
"type": worker.get("type", "local"),
})
return direct
def generate_job_id_map(prompt_index, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
"UltimateSDUpscaleDistributed"
distributed_nodes = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+ prompt_index.nodes_for_class("DistributedImageListStart")
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
@@ -540,6 +568,21 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
job_id_map, master_url, workflow_meta=None):
"""Attach protocol context to Start without changing its visible V3 sockets."""
for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
inputs["image_list_master_callback"] = str(master_url or "")
inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
if workflow_meta is not None:
inputs["image_list_workflow_meta"] = workflow_meta
def apply_participant_overrides(
prompt_copy,
participant_id,
@@ -575,5 +618,13 @@ def apply_participant_overrides(
master_url,
enabled_json,
)
_override_image_list_nodes(
prompt_copy,
prompt_index,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
)
return prompt_copy
+58 -6
View File
@@ -29,7 +29,9 @@ from .orchestration.prompt_transform import (
generate_job_id_map,
prepare_delegate_master_prompt,
prune_prompt_for_worker,
_resolve_direct_workers_for_image_list,
)
from ..image_list.runtime import validate_image_list_regions
prompt_server = server.PromptServer.instance
@@ -246,8 +248,42 @@ async def orchestrate_distributed_execution(
media_sync_timeout_seconds,
) = _resolve_orchestration_limits(config)
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
workers = _resolve_enabled_workers(config, requested_ids)
prompt_index = PromptIndex(prompt_obj)
image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
image_list_mode = bool(image_list_start_ids)
if image_list_mode:
# Validate before probing or queueing anything so malformed boundaries and
# old worker releases fail before any upstream image is materialized.
validate_image_list_regions(prompt_obj)
other_distributed = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
)
if other_distributed:
raise ValueError(
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
"Ultimate SD Upscale nodes in one prompt"
)
configured_serverless = [
str(worker.get("id")) for worker in config.get("workers", [])
if worker.get("enabled", False)
and str(worker.get("type") or "local").lower() == "serverless"
and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
]
if configured_serverless:
trace_debug(
execution_trace_id,
"Image List v1 explicitly excludes unsupported serverless pools: "
+ ", ".join(configured_serverless),
)
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
if not workers:
raise RuntimeError(
"Distributed Image List v1 requires an enabled selected direct local/remote worker; "
"serverless pools and master-only execution are unsupported"
)
else:
workers = _resolve_enabled_workers(config, requested_ids)
load_balance_requested = _prompt_requests_load_balance(prompt_index)
trace_debug(
execution_trace_id,
@@ -266,6 +302,10 @@ async def orchestrate_distributed_execution(
# Respect master delegate-only configuration
if delegate_master is None:
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
if image_list_mode:
# The master must execute each upstream dependency once to obtain the
# full typed lists; only the enclosed branch is delegated.
delegate_master = False
if not workers and delegate_master:
trace_debug(
@@ -281,8 +321,19 @@ async def orchestrate_distributed_execution(
trace_execution_id=execution_trace_id,
probe_concurrency=worker_probe_concurrency,
)
if image_list_mode:
from .orchestration.dispatch import preflight_image_list_workers
if not active_workers:
raise RuntimeError(
"No selected direct worker is reachable for Distributed Image List; "
"the master will not duplicate per-image processing"
)
idle_ids = await preflight_image_list_workers(active_workers, require_idle=True)
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
if not active_workers:
raise RuntimeError("No feature-compatible direct worker is idle; no master job was queued")
if load_balance_requested:
if load_balance_requested and not image_list_mode:
candidate_workers = list(active_workers)
if not delegate_master:
# Include master in load balancing only when master participation is enabled.
@@ -385,7 +436,7 @@ async def orchestrate_distributed_execution(
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
)
worker_payloads = []
if active_workers:
if active_workers and not image_list_mode:
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
worker_payloads = await asyncio.gather(
@@ -416,8 +467,9 @@ async def orchestrate_distributed_execution(
manager.ensure_capacity(len(serverless_jobs))
# Do not leave collector queues behind when offline compilation/media validation fails.
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
if not image_list_mode:
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
if regular_payloads:
await asyncio.gather(
@@ -468,4 +520,4 @@ async def orchestrate_distributed_execution(
execution_trace_id,
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
)
return prompt_id, prompt_number, len(worker_payloads), node_errors
return prompt_id, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
+72
View File
@@ -0,0 +1,72 @@
# Dynamic image-list queue
Process a list of images across **direct local or remote ComfyUI workers** without assigning fixed chunks to each worker.
The master evaluates the input list once. Each selected idle worker receives one image job; when it returns a result, it gets the next unassigned image. Faster workers can therefore process more images. The master restores input order before running downstream nodes.
## Wiring
```text
Image list or IMAGE batch
|
Distributed Image List Start -------- context ---------+
| image |
v |
Per-image processing |
| IMAGE |
v v
Distributed Image List End
| image_list
v
Preview / Save / next stage
```
1. Install the same feature-capable version of Distributed on the master and every selected worker. Restart those instances and refresh the frontend after installation.
2. Enable Distributed and select at least one reachable, idle **direct** worker in the panel. The master must be reachable from remote workers at its configured callback address.
3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
6. Put preview/save nodes **after End**, then queue normally. API callers should submit through `/distributed/queue`, not directly through `/prompt`.
ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; on a worker, Start provides one frame and its aligned values. End returns ordered lists, not a concatenated tensor batch, so differently sized images can be preserved.
## Inputs and outputs
Start accepts:
- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames.
- Optional `reference_image`, `mask`, `text`, `strength`, and `seed`.
- A `per_item` / `shared` mode for each optional input. `per_item` requires exactly one value per image; `shared` requires exactly one value, reused for all images. Inputs are never silently padded by repeating the last element.
Reference images and masks default to `per_item`. Text, strength and seed default to `shared`; strength defaults to `1.0`, seed to `0`.
Start outputs, in order: `image`, `reference_image`, `mask`, `text`, `strength`, `seed`, `item_index` (zero-based), `context`. Connect only optional values that your processing region actually provides and uses. Seeds are tied to input values, not worker assignments; use a per-item seed list for deterministic variation.
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
Models are **loaded and cached on workers**, not serialized from the master. Connect loader outputs directly to the processing nodes. Each worker needs the region's node packs and model files. Per-image data must pass through Start; do not bypass it with a second image list connected directly into the processing region.
## Scope and failure behavior
- This version supports direct workers only, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
- Worker eligibility is determined before execution and checked again before dispatch. Busy workers are excluded. Workers joining later are not added to an already-running region.
- Only the enclosed processing is distributed. List creation and downstream saving run on the master, even when master-delegate-only is enabled. There is no master fallback when workers are unavailable.
- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End.
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/reference/mask combinations can hit the aggregate request limit first.
- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
- Worker errors and the configured worker timeout fail the region. User interruption stops further scheduling and issues targeted queue deletion/interruption for this region's acknowledged prompts. An in-progress submission is allowed to resolve so its prompt ID can be cleaned up. Unreachable workers or a lost submission acknowledgement can require checking that worker manually; cancellation is not a guarantee of remote termination through a broken network.
- State is in memory. Restarting the master does not resume an interrupted region. A region also has a one-hour overall deadline.
Use trusted workers and trusted workflows, as with the existing ComfyUI API. Capability tokens protect assignment/result matching; they are not a replacement for network access control.
## Verification
The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
```bash
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
bash scripts/test-web.sh
```
Without `COMFYUI_SOURCE_ROOT`, real-framework tests are explicitly skipped. This coverage is not a substitute for checking your model workflow and connectivity on the actual GPU worker machines.
+1
View File
@@ -0,0 +1 @@
"""Dynamic per-image image-list execution internals."""
+798
View File
@@ -0,0 +1,798 @@
"""Runtime primitives for the V3 dynamic image-list feature.
The network protocol transports typed float tensors as bounded JSON/base64 payloads;
no Python object deserialization or arbitrary-path access is used.
"""
from __future__ import annotations
import asyncio
import base64
import json
import math
import secrets
import threading
import time
from dataclasses import dataclass, field
from typing import Any
PROTOCOL_VERSION = 1
START_CLASS = "DistributedImageListStart"
END_CLASS = "DistributedImageListEnd"
START_NODE_ID = "DistributedImageListStart"
END_NODE_ID = "DistributedImageListEnd"
# Start output indexes are a stable internal prompt protocol.
IMAGE_OUT, REFERENCE_OUT, MASK_OUT, TEXT_OUT, STRENGTH_OUT, SEED_OUT, INDEX_OUT, CONTEXT_OUT = range(8)
MAX_PAYLOAD_BYTES = 50 * 1024 * 1024
MAX_ITEMS = 4096
_ALLOWED_FLOAT_DTYPES = {"float16", "bfloat16", "float32", "float64"}
_LOCK = threading.RLock()
_MASTER_JOBS: dict[str, "ImageListJob"] = {}
_WORKER_ITEMS: dict[str, tuple[float, dict[str, Any]]] = {}
_WORKER_ITEM_TTL_SECONDS = 3600.0
class ImageListCancelledError(RuntimeError):
"""Raised after owned worker prompts have been cleaned up on cancellation."""
@dataclass
class ImageListJob:
context_id: str
start_id: str
end_id: str
original_prompt: dict[str, Any]
items: list[dict[str, Any]]
worker_ids: list[str]
master_callback_url: str
workflow_meta: dict[str, Any] | None = None
results: dict[int, dict[str, Any]] = field(default_factory=dict)
received_indexes: set[int] = field(default_factory=set)
expected_tokens: dict[int, str] = field(default_factory=dict)
expected_workers: dict[int, str] = field(default_factory=dict)
event: asyncio.Event | None = None
created_at: float = field(default_factory=time.monotonic)
failed: str | None = None
def _torch():
import torch
return torch
def unwrap_single(value: Any) -> Any:
"""Unwrap the one-element wrappers V3/INPUT_IS_LIST uses for scalar inputs."""
while isinstance(value, (list, tuple)) and len(value) == 1:
value = value[0]
return value
def _image_frames(value: Any, name: str, *, mask: bool = False) -> list[Any]:
torch = _torch()
if value is None:
return []
if isinstance(value, (list, tuple)):
frames: list[Any] = []
for item in value:
frames.extend(_image_frames(item, name, mask=mask))
return frames
if not isinstance(value, torch.Tensor):
raise TypeError(f"{name} must be a ComfyUI tensor or a list of tensors")
if not value.is_floating_point():
raise TypeError(f"{name} tensor must have a floating-point dtype")
if mask:
if value.ndim == 2:
value = value.unsqueeze(0)
elif value.ndim == 4 and value.shape[-1] == 1:
value = value[..., 0]
if value.ndim != 3:
raise ValueError(f"{name} MASK must have shape [B,H,W] or [B,H,W,1]")
if min(value.shape) < 1:
raise ValueError(f"{name} MASK dimensions must be positive")
return [value[i:i + 1] for i in range(value.shape[0])]
if value.ndim == 3:
value = value.unsqueeze(0)
if value.ndim != 4:
raise ValueError(f"{name} IMAGE must have shape [B,H,W,C] or [H,W,C]")
if min(value.shape[:3]) < 1 or value.shape[-1] not in (3, 4):
raise ValueError(f"{name} IMAGE must have positive dimensions and 3 or 4 channels")
return [value[i:i + 1] for i in range(value.shape[0])]
def _scalar_items(value: Any, name: str) -> list[Any]:
if value is None:
return []
if isinstance(value, (list, tuple)):
return list(value)
return [value]
def _normalize_aligned(value: Any, mode: str, count: int, name: str, kind: str) -> list[Any]:
if value is None:
return [None] * count
if mode not in ("per_item", "shared"):
raise ValueError(f"{name}_mode must be 'per_item' or 'shared'")
if kind == "image":
items = _image_frames(value, name)
elif kind == "mask":
items = _image_frames(value, name, mask=True)
else:
items = _scalar_items(value, name)
if mode == "shared":
if len(items) != 1:
raise ValueError(f"{name} is shared but has {len(items)} values; shared inputs require exactly one value")
items = [items[0] for _ in range(count)]
if len(items) != count:
raise ValueError(f"{name} has {len(items)} per-item values; expected exactly {count}")
if kind == "text" and any(item is not None and not isinstance(item, str) for item in items):
raise TypeError(f"{name} per-item values must be strings")
if kind == "int" and any(item is not None and (isinstance(item, bool) or not isinstance(item, int)) for item in items):
raise TypeError(f"{name} per-item values must be integers")
if kind == "float" and any(item is not None and (isinstance(item, bool) or not isinstance(item, (int, float)) or not math.isfinite(float(item))) for item in items):
raise TypeError(f"{name} per-item values must be finite numbers")
return items
def normalize_image_list_inputs(images: Any, reference_image: Any = None, mask: Any = None,
text: Any = None, strength: Any = None, seed: Any = None,
reference_image_mode: Any = "per_item", mask_mode: Any = "per_item",
text_mode: Any = "per_item", strength_mode: Any = "per_item",
seed_mode: Any = "per_item") -> list[dict[str, Any]]:
"""Normalize batches and explicitly shared/per-item optional inputs."""
frames = _image_frames(images, "images")
if not frames:
raise ValueError("Distributed Image List Start requires at least one image")
if len(frames) > MAX_ITEMS:
raise ValueError(f"Image list contains {len(frames)} items; maximum is {MAX_ITEMS}")
count = len(frames)
refs = _normalize_aligned(reference_image, str(unwrap_single(reference_image_mode)), count,
"reference_image", "image")
masks = _normalize_aligned(mask, str(unwrap_single(mask_mode)), count, "mask", "mask")
texts = _normalize_aligned(text, str(unwrap_single(text_mode)), count, "text", "text")
strengths = _normalize_aligned(strength, str(unwrap_single(strength_mode)), count, "strength", "float")
seeds = _normalize_aligned(seed, str(unwrap_single(seed_mode)), count, "seed", "int")
return [
{"image": frames[i], "reference_image": refs[i], "mask": masks[i],
"text": texts[i], "strength": strengths[i], "seed": seeds[i], "index": i}
for i in range(count)
]
def _tensor_dtype_name(tensor: Any) -> str:
name = str(tensor.dtype).removeprefix("torch.")
if name not in _ALLOWED_FLOAT_DTYPES:
raise TypeError(f"Tensor dtype {name!r} is unsupported; use float16, bfloat16, float32, or float64")
return name
def encode_value(value: Any, kind: str) -> dict[str, Any] | None:
"""Encode only the explicit IMAGE/MASK/STRING/FLOAT/INT protocol types."""
if value is None:
return None
if kind in ("IMAGE", "MASK"):
torch = _torch()
if not isinstance(value, torch.Tensor) or not value.is_floating_point():
raise TypeError(f"{kind} result must be a floating-point tensor")
if kind == "IMAGE":
if value.ndim != 4 or value.shape[0] != 1 or value.shape[-1] not in (3, 4):
raise ValueError("IMAGE item must have shape [1,H,W,3|4]")
else:
if value.ndim == 2:
value = value.unsqueeze(0)
if value.ndim == 4 and value.shape[-1] == 1:
value = value[..., 0]
if value.ndim != 3 or value.shape[0] != 1:
raise ValueError("MASK item must have shape [1,H,W] or [1,H,W,1]")
if value.numel() * value.element_size() > MAX_PAYLOAD_BYTES:
raise ValueError(f"{kind} item exceeds the {MAX_PAYLOAD_BYTES}-byte transport limit")
cpu = value.detach().to("cpu").contiguous()
raw = cpu.view(torch.uint8).numpy().tobytes()
if len(raw) > MAX_PAYLOAD_BYTES:
raise ValueError(f"{kind} item exceeds the {MAX_PAYLOAD_BYTES}-byte transport limit")
return {"type": kind, "dtype": _tensor_dtype_name(cpu),
"shape": list(cpu.shape), "data": base64.b64encode(raw).decode("ascii")}
if kind == "STRING":
if not isinstance(value, str):
raise TypeError("STRING transport value must be a string")
return {"type": kind, "value": value}
if kind == "INT":
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError("INT transport value must be an integer")
return {"type": kind, "value": value}
if kind == "FLOAT":
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)):
raise TypeError("FLOAT transport value must be finite")
return {"type": kind, "value": float(value)}
raise ValueError(f"Unsupported image-list transport type {kind!r}")
def decode_value(data: Any, expected_kind: str) -> Any:
if data is None:
return None
if not isinstance(data, dict) or data.get("type") != expected_kind:
raise ValueError(f"Expected {expected_kind} transport value")
if expected_kind in ("IMAGE", "MASK"):
torch = _torch()
dtype_name = data.get("dtype")
if dtype_name not in _ALLOWED_FLOAT_DTYPES:
raise ValueError(f"Unsupported tensor dtype {dtype_name!r}")
shape = data.get("shape")
if not isinstance(shape, list) or not shape or len(shape) > 4:
raise ValueError("Tensor shape must be a list of 1 to 4 positive dimensions")
if any(isinstance(dim, bool) or not isinstance(dim, int) or dim < 1 for dim in shape):
raise ValueError("Tensor shape dimensions must be positive integers")
expected_count = math.prod(shape)
element_size = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}[dtype_name]
expected_bytes = expected_count * element_size
if expected_bytes > MAX_PAYLOAD_BYTES:
raise ValueError("Tensor exceeds the image-list transport limit")
encoded = data.get("data")
if not isinstance(encoded, str) or len(encoded) > ((MAX_PAYLOAD_BYTES + 2) // 3) * 4:
raise ValueError("Tensor data is missing or exceeds the transport limit")
try:
raw = base64.b64decode(encoded, validate=True)
except Exception as exc:
raise ValueError("Tensor data must be valid base64") from exc
if len(raw) != expected_bytes:
raise ValueError("Tensor byte length does not match shape and dtype")
dtype = getattr(torch, dtype_name)
return torch.frombuffer(bytearray(raw), dtype=torch.uint8).view(dtype).reshape(shape).clone()
value = data.get("value")
return encode_value(value, expected_kind) and value
def encode_item(item: dict[str, Any]) -> dict[str, Any]:
return {"image": encode_value(item["image"], "IMAGE"),
"reference_image": encode_value(item.get("reference_image"), "IMAGE"),
"mask": encode_value(item.get("mask"), "MASK"),
"text": encode_value(item.get("text"), "STRING"),
"strength": encode_value(item.get("strength"), "FLOAT"),
"seed": encode_value(item.get("seed"), "INT"),
"index": int(item["index"])}
def decode_item(data: Any) -> dict[str, Any]:
if not isinstance(data, dict):
raise ValueError("item must be an object")
index = data.get("index")
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise ValueError("item.index must be a non-negative integer")
image = decode_value(data.get("image"), "IMAGE")
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
raise ValueError("item.image must contain exactly one RGB/RGBA image frame")
reference = decode_value(data.get("reference_image"), "IMAGE")
if reference is not None and (reference.ndim != 4 or reference.shape[0] != 1 or reference.shape[-1] not in (3, 4)):
raise ValueError("item.reference_image must contain exactly one RGB/RGBA image frame")
mask = decode_value(data.get("mask"), "MASK")
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
raise ValueError("item.mask must contain exactly one MASK frame")
return {"image": image,
"reference_image": reference,
"mask": mask,
"text": decode_value(data.get("text"), "STRING"),
"strength": decode_value(data.get("strength"), "FLOAT"),
"seed": decode_value(data.get("seed"), "INT"), "index": index}
def _sweep_worker_items() -> None:
now = time.monotonic()
with _LOCK:
for token, (expires, _) in list(_WORKER_ITEMS.items()):
if expires <= now:
del _WORKER_ITEMS[token]
def store_worker_item(token: str, item: dict[str, Any]) -> None:
if not isinstance(token, str) or len(token) < 32 or len(token) > 128:
raise ValueError("token must be a 32-128 character capability")
clean = decode_item(item)
_sweep_worker_items()
with _LOCK:
if token in _WORKER_ITEMS:
raise ValueError("duplicate item token")
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, clean)
def consume_worker_item(token: str) -> dict[str, Any]:
if not isinstance(token, str) or not token:
raise ValueError("worker item token is required")
_sweep_worker_items()
with _LOCK:
entry = _WORKER_ITEMS.pop(token, None)
if entry is None:
raise RuntimeError("Image-list item is missing, expired, or already consumed")
return entry[1]
def delete_worker_item(token: str) -> bool:
with _LOCK:
return _WORKER_ITEMS.pop(str(token), None) is not None
def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: list[dict[str, Any]],
worker_ids: list[str], master_callback_url: str,
workflow_meta: dict[str, Any] | None = None) -> str:
if not isinstance(prompt, dict):
raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction")
if not worker_ids:
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
context_id = secrets.token_urlsafe(36)
state = ImageListJob(context_id, str(start_id), str(end_id), prompt, items,
[str(worker_id) for worker_id in worker_ids], master_callback_url,
workflow_meta)
with _LOCK:
_MASTER_JOBS[context_id] = state
return context_id
def get_master_job(context_id: str) -> ImageListJob | None:
with _LOCK:
return _MASTER_JOBS.get(context_id)
def remove_master_job(context_id: str) -> None:
with _LOCK:
_MASTER_JOBS.pop(context_id, None)
def _prompt_links(node: Any, prompt: dict[str, Any]):
if not isinstance(node, dict):
return
inputs = node.get("inputs", {})
if not isinstance(inputs, dict):
return
for input_name, value in inputs.items():
if isinstance(value, (list, tuple)) and len(value) == 2 and str(value[0]) in prompt:
yield input_name, str(value[0]), value[1]
def _direct_link(value: Any, prompt: dict[str, Any], node_id: str, input_name: str) -> tuple[str, int]:
if not (isinstance(value, (list, tuple)) and len(value) == 2 and str(value[0]) in prompt):
raise ValueError(f"{node_id}.{input_name} must be connected to an in-graph node")
try:
return str(value[0]), int(value[1])
except (TypeError, ValueError) as exc:
raise ValueError(f"{node_id}.{input_name} has an invalid link") from exc
def _ancestors_for_inputs(prompt: dict[str, Any], end_id: str, start_id: str,
input_names: tuple[str, ...]) -> set[str]:
result = {end_id}
pending = []
end_inputs = prompt[end_id].get("inputs", {})
for name in input_names:
if name in end_inputs:
source, _ = _direct_link(end_inputs[name], prompt, end_id, name)
pending.append(source)
while pending:
node_id = pending.pop()
if node_id in result:
continue
result.add(node_id)
if node_id == start_id:
continue
node = prompt[node_id]
for _, source, _ in _prompt_links(node, prompt):
if source not in result:
pending.append(source)
return result
def _downstream_until(prompt: dict[str, Any], start_id: str, end_id: str) -> set[str]:
downstream: dict[str, set[str]] = {str(node_id): set() for node_id in prompt}
for node_id, node in prompt.items():
for _, source, _ in _prompt_links(node, prompt):
downstream[source].add(str(node_id))
visited = set()
pending = [start_id]
while pending:
node_id = pending.pop()
for target in downstream.get(node_id, ()):
if target == end_id or target in visited:
continue
visited.add(target)
pending.append(target)
return visited
def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
"""Validate pairing, non-nesting, and closed per-image graph boundaries."""
if not isinstance(prompt, dict):
raise ValueError("Prompt must be an object")
starts = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == START_CLASS)
ends = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == END_CLASS)
if len(starts) != len(ends) or not starts:
raise ValueError("Distributed Image List needs one paired Start and End per region")
regions = []
for start_id in starts:
candidates = []
for end_id in ends:
context = prompt[end_id].get("inputs", {}).get("context")
if isinstance(context, (list, tuple)) and len(context) == 2 and str(context[0]) == start_id and context[1] == CONTEXT_OUT:
candidates.append(end_id)
if len(candidates) != 1:
raise ValueError(f"Start node {start_id} must connect its typed context output directly to exactly one End")
end_id = candidates[0]
if "images" not in prompt[end_id].get("inputs", {}):
raise ValueError(f"End node {end_id}.images must be connected to the per-item result branch")
end_inputs = prompt[end_id].get("inputs", {})
for name in ("images", "masks", "texts"):
if name in end_inputs:
_direct_link(end_inputs[name], prompt, end_id, name)
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, ("images", "masks", "texts"))
if start_id not in region_nodes:
raise ValueError(f"Start node {start_id} is not an ancestor of End node {end_id}")
forbidden = {n for n in region_nodes if n != start_id and n != end_id and
isinstance(prompt[n], dict) and prompt[n].get("class_type") in
{START_CLASS, END_CLASS, "DistributedCollector", "UltimateSDUpscaleDistributed"}}
if forbidden:
raise ValueError("Nested/ambiguous regions and other distributed nodes are unsupported in an image-list region: " + ", ".join(sorted(forbidden)))
# Output nodes execute independently of the lazy End, so they would
# evaluate the processing branch on the master as well as the worker.
import sys
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
for node_id in region_nodes - {start_id, end_id}:
node_class = registry.get(prompt[node_id].get("class_type"))
if node_class is None:
continue
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
if info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)):
raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End")
downstream = _downstream_until(prompt, start_id, end_id)
outside = sorted(downstream - region_nodes)
if outside:
raise ValueError(f"Image-list region {start_id}->{end_id} has unsafe cross-boundary nodes: {outside}")
# Model loaders and constants can be shared; item-dependent processing cannot.
regions.append({"start_id": start_id, "end_id": end_id,
"node_ids": region_nodes, "branch_ids": (region_nodes & downstream) - {start_id}})
regions.sort(key=lambda region: (0, int(region["start_id"])) if region["start_id"].isdigit()
else (1, region["start_id"]))
internal_sets = [r["branch_ids"] for r in regions]
for i in range(len(regions)):
for j in range(i + 1, len(regions)):
overlap = internal_sets[i] & internal_sets[j]
if overlap:
raise ValueError(f"Image-list regions share ambiguous processing nodes: {sorted(overlap)}")
return regions
def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_token: str,
context_id: str, item_index: int, callback_url: str) -> dict[str, Any]:
"""Copy only the per-image branch and its worker-side dependencies."""
start_id, end_id = region["start_id"], region["end_id"]
_direct_link(prompt[end_id].get("inputs", {}).get("context"), prompt, end_id, "context")
kept = {node_id: json.loads(json.dumps(prompt[node_id])) for node_id in region["node_ids"]}
start_inputs = kept[start_id].setdefault("inputs", {})
for name in ("images", "reference_image", "mask", "text", "strength", "seed"):
start_inputs.pop(name, None)
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
end_inputs = kept[end_id].setdefault("inputs", {})
end_inputs.update({"worker_result_token": item_token,
"worker_item_index": int(item_index),
"master_callback_url": callback_url})
return kept
def create_worker_task_token() -> str:
return secrets.token_urlsafe(36)
async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_next,
cleanup, interrupted=None):
"""One in-flight job per worker; completed workers immediately take next item."""
if item_count <= 0:
raise ValueError("Dynamic image-list queue has no items")
if not workers:
raise RuntimeError("No free, reachable, feature-compatible direct workers are available")
pending_index = 0
inflight: dict[str, Any] = {}
results: dict[int, Any] = {}
def check_interrupted():
if interrupted is not None and interrupted():
raise ImageListCancelledError("Image-list job interrupted by user")
async def next_result():
# Keep one waiter alive while checking ComfyUI's thread-safe interrupt flag.
# Restarting wait_next on a timeout can lose a result it already consumed.
waiter = asyncio.create_task(wait_next(inflight))
try:
while True:
check_interrupted()
done, _ = await asyncio.wait({waiter}, timeout=0.1)
check_interrupted()
if done:
return waiter.result()
finally:
if not waiter.done():
waiter.cancel()
await asyncio.gather(waiter, return_exceptions=True)
async def submit(worker, index):
# An outer timeout must not discard an acknowledgement for a job we own.
task = asyncio.create_task(dispatch(worker, index))
try:
return await asyncio.shield(task)
except asyncio.CancelledError:
try:
inflight[worker] = await task
except Exception:
pass
raise
try:
for worker in workers:
check_interrupted()
if pending_index >= item_count:
break
inflight[worker] = await submit(worker, pending_index)
pending_index += 1
while inflight:
check_interrupted()
worker, index, result = await next_result()
if worker not in inflight:
raise RuntimeError(f"Unexpected completion from unassigned worker {worker}")
if isinstance(index, bool) or not isinstance(index, int) or index != inflight[worker].get("index"):
raise RuntimeError(f"Worker {worker} returned an unexpected image index {index!r}")
if index in results:
raise RuntimeError(f"Duplicate image-list result for item {index}")
results[index] = result
del inflight[worker]
if pending_index < item_count:
inflight[worker] = await submit(worker, pending_index)
pending_index += 1
missing = sorted(set(range(item_count)) - set(results))
if missing:
raise RuntimeError(f"Image-list results are incomplete; missing indexes {missing}")
return [results[i] for i in range(item_count)]
except BaseException:
if inflight:
await asyncio.gather(*(cleanup(worker, handle) for worker, handle in inflight.items()), return_exceptions=True)
raise
def set_master_result(context_id: str, item_index: int, token: str, result: dict[str, Any]) -> None:
with _LOCK:
state = _MASTER_JOBS.get(context_id)
if state is None:
raise ValueError("Unknown or expired image-list context")
expected = state.expected_tokens.get(item_index)
if expected is None or not secrets.compare_digest(expected, token):
raise ValueError("Image-list result token does not match an assigned item")
if item_index in state.received_indexes:
raise ValueError(f"Duplicate image-list result for item {item_index}")
state.received_indexes.add(item_index)
state.results[item_index] = result
event = state.event
if event is not None:
event.set()
def register_assignment(context_id: str, item_index: int, worker_id: str, token: str) -> None:
state = get_master_job(context_id)
if state is None:
raise RuntimeError("Image-list context expired before worker dispatch")
with _LOCK:
if item_index in state.expected_tokens:
raise RuntimeError(f"Image-list item {item_index} was already assigned")
state.expected_tokens[item_index] = token
state.expected_workers[item_index] = worker_id
def decode_result(result: Any) -> dict[str, Any]:
if not isinstance(result, dict):
raise ValueError("result must be an object")
image = decode_value(result.get("image"), "IMAGE")
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
raise ValueError("result.image must contain exactly one RGB/RGBA image")
mask = decode_value(result.get("mask"), "MASK")
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
raise ValueError("result.mask must contain exactly one MASK frame")
text = decode_value(result.get("text"), "STRING")
return {"image": image, "mask": mask, "text": text}
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -> dict[str, Any]:
def one(value, name, kind):
if value is None:
return None
if kind in ("IMAGE", "MASK"):
values = _image_frames(value, name, mask=(kind == "MASK"))
if len(values) != 1:
raise ValueError(f"Worker End expects exactly one {kind} result, received {len(values)}")
return values[0]
values = _scalar_items(value, name)
values = [v for v in values if v is not None]
if not values:
return None
if len(values) != 1:
raise ValueError(f"Worker End expects at most one {kind} result, received {len(values)}")
return values[0]
image = one(images, "images", "IMAGE")
if image is None:
raise ValueError("Worker End received no image result")
mask = one(masks, "masks", "MASK")
text = one(texts, "texts", "STRING")
if text is not None and not isinstance(text, str):
raise TypeError("Worker End text result must be a string")
return {"image": image, "mask": mask, "text": text}
def encode_result(result: dict[str, Any]) -> dict[str, Any]:
normalized = normalize_worker_result(result.get("image"), result.get("mask"), result.get("text"))
return {"image": encode_value(normalized["image"], "IMAGE"),
"mask": encode_value(normalized["mask"], "MASK"),
"text": encode_value(normalized["text"], "STRING")}
def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]:
if not results:
raise RuntimeError("Image-list execution returned no results")
images, masks, texts = [], [], []
for expected_index, result in enumerate(results):
if not isinstance(result, dict) or "image" not in result:
raise RuntimeError(f"Image-list result {expected_index} is missing its image")
images.append(result["image"])
masks.append(result.get("mask"))
texts.append(result.get("text"))
return images, masks, texts
async def run_image_list_job(context_id: str) -> tuple[list[Any], list[Any], list[Any]]:
"""Resolve compatible idle direct workers and schedule one image at a time."""
state = get_master_job(context_id)
if state is None:
raise RuntimeError("Distributed Image List context is missing or expired")
try:
return await _run_image_list_job(state)
finally:
remove_master_job(context_id)
async def _run_image_list_job(state: ImageListJob):
context_id = state.context_id
state.event = asyncio.Event()
from ..api.orchestration.dispatch import (
preflight_image_list_workers, dispatch_worker_prompt,
)
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
from ..utils.config import load_config, get_worker_timeout_seconds
from ..utils.network import build_master_url, build_worker_url, get_client_session
from ..utils.logging import log
import aiohttp
config = load_config()
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
if not workers:
raise RuntimeError("Distributed Image List requires enabled direct workers; serverless pools are unsupported")
# Recheck worker capability and idle state immediately before the first item.
worker_statuses = await preflight_image_list_workers(workers, require_idle=True)
workers = [worker for worker in workers if worker["id"] in worker_statuses]
if not workers:
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
session = await get_client_session()
callback_url = state.master_callback_url or build_master_url(config=config)
regions = validate_image_list_regions(state.original_prompt)
region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None)
if region is None:
raise RuntimeError("Image-list region no longer matches the submitted prompt")
timeout_seconds = max(5.0, float(get_worker_timeout_seconds()))
use_websocket = bool((config.get("settings", {}) or {}).get("websocket_orchestration", False))
prepared_worker_ids: set[str] = set()
worker_path_separators: dict[str, str | None] = {}
async def dispatch(worker_id: str, index: int) -> dict[str, Any]:
worker = next((entry for entry in workers if entry["id"] == worker_id), None)
if worker is None:
raise RuntimeError(f"Worker {worker_id} was removed before dispatch")
token = create_worker_task_token()
register_assignment(context_id, index, worker_id, token)
item_payload = {"context_id": context_id, "item": encode_item(state.items[index])}
input_url = build_worker_url(worker, "/distributed/image-list/item")
async with session.post(input_url, json={"token": token, **item_payload},
timeout=aiohttp.ClientTimeout(total=60)) as response:
if response.status >= 400:
raise RuntimeError(f"Worker {worker_id} rejected image-list item {index}: HTTP {response.status}")
worker_prompt = build_worker_prompt(state.original_prompt, region, token, context_id,
index, callback_url)
if worker_id not in worker_path_separators:
worker_path_separators[worker_id] = None
if worker.get("host") and str(worker.get("type") or "local").lower() != "local":
from ..api.orchestration.media_sync import fetch_worker_path_separator
worker_path_separators[worker_id] = await fetch_worker_path_separator(worker)
separator = worker_path_separators[worker_id]
if separator:
from ..api.orchestration.media_sync import convert_paths_for_platform
worker_prompt = convert_paths_for_platform(worker_prompt, separator)
if worker_id not in prepared_worker_ids:
from ..api.orchestration.media_sync import sync_worker_media
await asyncio.wait_for(sync_worker_media(worker, worker_prompt), timeout=60)
prepared_worker_ids.add(worker_id)
prompt_id = await dispatch_worker_prompt(
worker, worker_prompt, state.workflow_meta,
client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket,
return_prompt_id=True,
)
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError(f"Worker {worker_id} did not return a prompt_id; submission status is ambiguous")
return {"prompt_id": prompt_id, "token": token, "index": index,
"started_at": time.monotonic(), "worker": worker}
async def wait_next(inflight: dict[str, Any]):
while True:
if state.failed:
raise RuntimeError(state.failed)
now = time.monotonic()
expired = [f"{wid}/item {handle['index']}" for wid, handle in inflight.items()
if now - handle["started_at"] > timeout_seconds]
if expired:
raise TimeoutError(f"Image-list worker timeout after {timeout_seconds:.1f}s: {', '.join(expired)}")
# Treat worker-side execution errors as failures even when End never posts a result.
async def history_status(worker_id, handle):
url = build_worker_url(handle["worker"], f"/history/{handle['prompt_id']}")
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as response:
if response.status >= 400:
return None
data = await response.json()
entry = data.get(handle["prompt_id"])
if not isinstance(entry, dict):
return None
status = entry.get("status", {}) or {}
if status.get("status_str") == "error":
return f"Worker {worker_id} failed image-list item {handle['index']}: {status.get('messages', status)}"
if status.get("completed") and status.get("status_str") not in (None, "success"):
return f"Worker {worker_id} failed image-list item {handle['index']}: {status}"
if status.get("completed"):
with _LOCK:
result = state.results.get(handle["index"])
if result is None:
return f"Worker {worker_id} completed item {handle['index']} without posting a result"
return worker_id, handle["index"], result
except (asyncio.TimeoutError, aiohttp.ClientError):
return None
return None
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()))
failure = next((error for error in errors if isinstance(error, str)), None)
if failure:
raise RuntimeError(failure)
completed = next((entry for entry in errors if isinstance(entry, tuple)), None)
if completed is not None:
# The callback precedes the worker prompt finishing. Wait for its
# successful history before releasing the slot / next region.
with _LOCK:
state.results.pop(completed[1], None)
return completed
await asyncio.sleep(0.25)
async def cleanup(worker_id: str, handle: dict[str, Any]) -> None:
worker = handle["worker"]
try:
prompt_id = handle["prompt_id"]
# Delete if still queued, then interrupt only this prompt if it began
# running meanwhile. Never send a global interrupt or clear a queue.
async with session.post(build_worker_url(worker, "/queue"), json={"delete": [prompt_id]},
timeout=aiohttp.ClientTimeout(total=5)) as response:
response.raise_for_status()
async with session.post(build_worker_url(worker, "/interrupt"), json={"prompt_id": prompt_id},
timeout=aiohttp.ClientTimeout(total=5)) as response:
response.raise_for_status()
except Exception as exc:
log(f"[Distributed Image List] Owned prompt cleanup failed for worker {worker_id}: {exc}")
try:
async with session.delete(build_worker_url(worker, f"/distributed/image-list/item/{handle['token']}"),
timeout=aiohttp.ClientTimeout(total=5)):
pass
except Exception:
pass
def interrupted() -> bool:
try:
import comfy.model_management
return bool(comfy.model_management.processing_interrupted())
except Exception:
return False
results = await dynamic_schedule(len(state.items), [worker["id"] for worker in workers],
dispatch, wait_next, cleanup, interrupted)
return _collect_output_lists(results)
+226
View File
@@ -0,0 +1,226 @@
"""Public V3 Start/End nodes for dynamic per-image image-list execution."""
from __future__ import annotations
import json
from comfy_api.v0_0_2 import io
from ..utils.async_helpers import run_async_in_server_loop
from ..image_list.runtime import (
create_master_job,
normalize_image_list_inputs,
normalize_worker_result,
run_image_list_job,
encode_result,
unwrap_single,
)
class DistributedImageListStart(io.ComfyNode):
"""Materialize a batch once, then expose ordered per-item typed values."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListStart",
display_name="Distributed Image List Start",
category="image/distributed",
description=("Starts one per-image processing region. Optional inputs explicitly choose shared "
"(exactly one value) or per-item (exactly one value per image) alignment."),
inputs=[
io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one worker job."),
io.Image.Input("reference_image", optional=True),
io.Combo.Input("reference_image_mode", options=["per_item", "shared"], default="per_item"),
io.Mask.Input("mask", optional=True),
io.Combo.Input("mask_mode", options=["per_item", "shared"], default="per_item"),
io.String.Input("text", optional=True),
io.Combo.Input("text_mode", options=["per_item", "shared"], default="shared"),
io.Float.Input("strength", optional=True, default=1.0),
io.Combo.Input("strength_mode", options=["per_item", "shared"], default="shared"),
io.Int.Input("seed", optional=True, default=0, min=0, max=0xffffffffffffffff),
io.Combo.Input("seed_mode", options=["per_item", "shared"], default="shared"),
],
outputs=[
io.Image.Output(display_name="image", is_output_list=True),
io.Image.Output(display_name="reference_image", is_output_list=True),
io.Mask.Output(display_name="mask", is_output_list=True),
io.String.Output(display_name="text", is_output_list=True),
io.Float.Output(display_name="strength", is_output_list=True),
io.Int.Output(display_name="seed", is_output_list=True),
io.Int.Output(display_name="item_index", is_output_list=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Output(display_name="context", is_output_list=True),
],
is_input_list=True,
accept_all_inputs=True,
not_idempotent=True,
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
)
@classmethod
def execute(cls, images=None, reference_image=None, reference_image_mode="per_item",
mask=None, mask_mode="per_item", text=None, text_mode="shared",
strength=None, strength_mode="shared", seed=None, seed_mode="shared",
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
image_list_master_callback="", image_list_execution_id="",
image_list_workflow_meta=None):
worker_item_token = unwrap_single(worker_item_token)
worker_context_id = unwrap_single(worker_context_id)
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
if worker_item_token:
item = run_async_in_server_loop(_consume_worker_item(worker_item_token), timeout=10)
context_id = worker_context_id
if not isinstance(context_id, str) or not context_id:
raise RuntimeError("Worker image-list prompt is missing its paired context")
return io.NodeOutput(
[item["image"]], [item.get("reference_image")], [item.get("mask")],
[item.get("text")], [item.get("strength")], [item.get("seed")],
[item["index"]], [context_id],
)
items = normalize_image_list_inputs(
images, reference_image, mask, text, strength, seed,
reference_image_mode, mask_mode, text_mode, strength_mode, seed_mode,
)
prompt = getattr(cls.hidden, "prompt", None)
start_id = str(getattr(cls.hidden, "unique_id", ""))
if not prompt or not start_id:
raise RuntimeError("Distributed Image List Start needs a submitted ComfyUI prompt context")
try:
worker_ids = json.loads(str(unwrap_single(image_list_enabled_worker_ids or "[]")))
except (TypeError, ValueError) as exc:
raise RuntimeError("Invalid preflighted image-list worker list") from exc
if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
raise RuntimeError("Invalid preflighted image-list worker list")
from ..image_list.runtime import validate_image_list_regions
regions = validate_image_list_regions(prompt)
region = next((row for row in regions if row["start_id"] == start_id), None)
if region is None:
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
context_id = create_master_job(
start_id, region["end_id"], prompt, items, worker_ids,
str(unwrap_single(image_list_master_callback or "")),
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
)
return io.NodeOutput(
[item["image"] for item in items],
[item.get("reference_image") for item in items],
[item.get("mask") for item in items],
[item.get("text") for item in items],
[item.get("strength") for item in items],
[item.get("seed") for item in items],
[item["index"] for item in items],
[context_id for _ in items],
)
async def _consume_worker_item(token):
# Item state is local to each worker process; source nodes consume it exactly once.
from ..image_list.runtime import consume_worker_item
return consume_worker_item(token)
class DistributedImageListEnd(io.ComfyNode):
"""Collect branch results in source order without evaluating the branch on master."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListEnd",
display_name="Distributed Image List End",
category="image/distributed",
description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
inputs=[
io.Image.Input("images", lazy=True),
io.Mask.Input("masks", optional=True, lazy=True),
io.String.Input("texts", optional=True, lazy=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
],
outputs=[
io.Image.Output(display_name="image_list", is_output_list=True),
io.Mask.Output(display_name="masks", is_output_list=True),
io.String.Output(display_name="texts", is_output_list=True),
],
is_input_list=True,
is_output_node=True,
accept_all_inputs=True,
not_idempotent=True,
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
)
@classmethod
def check_lazy_status(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", **kwargs):
worker_result_token = unwrap_single(worker_result_token)
if worker_result_token:
# Decide which optional sockets are connected from the actual worker prompt;
# absent optional inputs must not be demanded from Comfy's lazy executor.
prompt = getattr(cls.hidden, "prompt", {}) or {}
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
if isinstance(node, dict)
and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
required = []
for name, value in (("images", images), ("masks", masks), ("texts", texts)):
connected = name == "images" or name in end_inputs
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
if connected and missing:
required.append(name)
return required
# Master only needs Start's context. Never request branch outputs here.
if context is None or (isinstance(context, (tuple, list)) and all(item is None for item in context)):
return ["context"]
return []
@classmethod
def execute(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", worker_item_index=-1, master_callback_url=""):
worker_result_token = unwrap_single(worker_result_token)
master_callback_url = unwrap_single(master_callback_url)
if worker_result_token:
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Worker Image List End is missing its paired context")
result = normalize_worker_result(images, masks, texts)
payload = encode_result(result)
run_async_in_server_loop(
submit_worker_result(str(master_callback_url), str(worker_result_token),
context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
)
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Distributed Image List End requires the paired Start context wire")
context_id = context_values[0]
from ..image_list.runtime import ImageListCancelledError
try:
# Put the deadline inside the coroutine so timeout cancels owned remote work.
import asyncio
result_lists = run_async_in_server_loop(asyncio.wait_for(run_image_list_job(context_id), 3600))
except ImageListCancelledError:
import comfy.model_management
raise comfy.model_management.InterruptProcessingException()
images_out, masks_out, texts_out = result_lists
if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
return io.NodeOutput(images_out, masks_out, texts_out)
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
if not callback_url.startswith(("http://", "https://")):
raise RuntimeError("Worker image-list callback URL is missing or invalid")
import aiohttp
from ..utils.network import get_client_session
session = await get_client_session()
payload = {"context_id": context_id, "token": item_token,
"item_index": item_index, "result": result}
async with session.post(
f"{callback_url.rstrip('/')}/distributed/image-list/result",
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
) as response:
if response.status >= 400:
body = (await response.text())[:500]
raise RuntimeError(f"Master rejected image-list result (HTTP {response.status}): {body}")
+3 -1
View File
@@ -9,6 +9,7 @@ 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
from .image_list import DistributedImageListEnd, DistributedImageListStart
class DistributedSeed(io.ComfyNode):
@@ -178,4 +179,5 @@ class UltimateSDUpscaleDistributed(io.ComfyNode):
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
DistributedImageListStart, DistributedImageListEnd]
+395
View File
@@ -0,0 +1,395 @@
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
No managed ComfyUI service, user workflow, model, or remote worker is touched.
"""
import asyncio
from collections import Counter
import importlib
import json
import os
from pathlib import Path
import sys
import tempfile
import time
import uuid
from unittest.mock import patch
ROOT = Path(__file__).resolve().parents[2]
COMFY_ROOT = Path(sys.argv[1]).resolve()
WORKER = '--worker' in sys.argv
DELAY = float(sys.argv[3]) if WORKER else 0
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
import comfy.model_management
import execution
import nodes
import server
import torch
from aiohttp import web
import aiohttp
torch.set_num_threads(1)
SOURCE_CALLS, BRANCH_CALLS, CAPTURED = [], [], []
COUNT = 7
def source_rows():
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
for i in range(COUNT)]
class TestImageSource:
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT')
OUTPUT_IS_LIST = (True,) * 6
FUNCTION = 'run'
def run(self):
assert not WORKER, 'upstream image source leaked into worker prompt'
SOURCE_CALLS.append(True)
rows = source_rows()
return tuple([row[key] for row in rows] for key in
('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
class TestImageOperation:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'reference_image': ('IMAGE',), 'mask': ('MASK',),
'text': ('STRING', {'forceInput': True}),
'strength': ('FLOAT', {'forceInput': True}),
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True})}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, reference_image, mask, text, strength, seed, index):
assert WORKER, 'per-image branch executed on master'
BRANCH_CALLS.append(index)
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
until = time.monotonic() + DELAY
while time.monotonic() < until:
comfy.model_management.throw_exception_if_processing_interrupted()
time.sleep(0.01)
assert seed == 100 + index
assert text == f'item-{index}'
assert strength == (index + 1) / 16
assert torch.all(reference_image == strength)
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
class TestCapture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
INPUT_IS_LIST = True
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, images, masks, texts):
CAPTURED.append((images, masks, texts))
return ()
async def setup():
assets = default_asset_manager()
ps = server.PromptServer(asyncio.get_running_loop(), assets)
assert await nodes.load_custom_node(str(ROOT))
module = sys.modules[str(ROOT).replace('.', '_x_')]
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
TestImageOperation=TestImageOperation, TestCapture=TestCapture)
app = web.Application()
# Register the actual feature routes, not test replacements.
app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')])
return ps, assets, module, app
async def listen(app):
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
return runner, site._server.sockets[0].getsockname()[1]
async def worker_main():
ps, assets, module, app = await setup()
pending, running, history = [], {}, {}
wake = asyncio.Event()
interrupted_ids = []
async def prompt(request):
data = await request.json()
prompt_id = str(uuid.uuid4())
graph = data['prompt']
valid = await execution.validate_prompt(prompt_id, graph, None)
if not valid[0]:
return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400)
pending.append((prompt_id, graph, valid[2]))
wake.set()
return web.json_response({'prompt_id': prompt_id})
async def queue(request):
if request.method == 'POST':
deleted = (await request.json()).get('delete', [])
pending[:] = [entry for entry in pending if entry[0] not in deleted]
return web.json_response({})
return web.json_response({'queue_running': [[0, key] for key in running],
'queue_pending': [[0, entry[0]] for entry in pending]})
async def get_history(request):
pid = request.match_info['prompt_id']
return web.json_response({pid: history[pid]} if pid in history else {})
async def interrupt(request):
pid = (await request.json())['prompt_id']
interrupted_ids.append(pid)
if pid in running:
comfy.model_management.interrupt_current_processing(True)
return web.json_response({})
async def stats(request):
return web.json_response({'calls': BRANCH_CALLS, 'running': list(running),
'pending': [e[0] for e in pending], 'interrupted': interrupted_ids,
'history': history})
async def consume():
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
while True:
await wake.wait()
while pending:
pid, graph, outputs = pending.pop(0)
running[pid] = True
await asyncio.to_thread(executor.execute, graph, pid, {}, outputs)
history[pid] = {'status': {'completed': True,
'status_str': 'success' if executor.success else 'error',
'messages': str(executor.status_messages)}}
running.pop(pid, None)
wake.clear()
async def probe(_request):
return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}})
app.router.add_get('/prompt', probe)
app.router.add_post('/prompt', prompt)
app.router.add_get('/queue', queue)
app.router.add_post('/queue', queue)
app.router.add_get('/history/{prompt_id}', get_history)
app.router.add_post('/interrupt', interrupt)
app.router.add_get('/test-stats', stats)
runner, port = await listen(app)
consumer = asyncio.create_task(consume())
print(f'READY:{port}', flush=True)
try:
await asyncio.Event().wait()
finally:
consumer.cancel()
await asyncio.gather(consumer, return_exceptions=True)
await runner.cleanup()
def graph(callback, worker_ids):
return {
'1': {'class_type': 'TestImageSource', 'inputs': {}},
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
'images': ['1', 0], 'reference_image': ['1', 1], 'mask': ['1', 2], 'text': ['1', 3],
'strength': ['1', 4], 'seed': ['1', 5],
'reference_image_mode': 'per_item', 'mask_mode': 'per_item', 'text_mode': 'per_item',
'strength_mode': 'per_item', 'seed_mode': 'per_item',
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
'3': {'class_type': 'TestImageOperation', 'inputs': {
key: ['2', index] for key, index in [('image', 0), ('reference_image', 1), ('mask', 2),
('text', 3), ('strength', 4), ('seed', 5), ('index', 6)]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 7]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
}
async def coordinator():
ps, assets, module, app = await setup()
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
config_mod = importlib.import_module(module.__name__ + '.utils.config')
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
runner, port = await listen(app)
processes, stderr_files = [], []
# Old worker capability refusal is read-only and precedes input evaluation.
legacy_app = web.Application()
legacy_requests = []
async def legacy(request):
legacy_requests.append(request.method)
return web.json_response({'error': 'not found'}, status=404)
legacy_app.router.add_route('*', '/{path:.*}', legacy)
legacy_runner, legacy_port = await listen(legacy_app)
dispatch_module = importlib.import_module(module.__name__ + '.api.orchestration.dispatch')
try:
try:
await dispatch_module.preflight_image_list_workers([
{'id': 'legacy', 'type': 'local', 'host': '127.0.0.1', 'port': legacy_port}])
raise AssertionError('legacy worker accepted')
except RuntimeError as exc:
assert 'upgrade' in str(exc), str(exc)
assert legacy_requests == ['GET'] and not SOURCE_CALLS
finally:
await legacy_runner.cleanup()
print('IMAGE_LIST_PREFLIGHT_OK legacy worker rejected before materialization or writes')
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
scratch.mkdir(parents=True, exist_ok=True)
try:
with tempfile.TemporaryDirectory(prefix='image-list-network-', dir=scratch) as tmp:
workers = []
for name, delay in [('fast', 0.02), ('slow', 0.5)]:
err = open(Path(tmp) / f'{name}.log', 'w+')
stderr_files.append(err)
proc = await asyncio.create_subprocess_exec(sys.executable, str(Path(__file__).resolve()),
str(COMFY_ROOT), '--worker', str(delay), stdout=asyncio.subprocess.PIPE, stderr=err,
env={**os.environ, 'TMPDIR': tmp})
processes.append(proc)
async def ready():
while True:
line = await proc.stdout.readline()
if not line:
err.seek(0)
raise AssertionError(f'Worker exited: {err.read()}')
if line.startswith(b'READY:'):
return int(line.split(b':')[1])
worker_port = await asyncio.wait_for(ready(), 60)
workers.append({'id': name, 'type': 'local', 'host': '127.0.0.1',
'port': worker_port, 'enabled': True})
config = {'workers': workers, 'settings': {'websocket_orchestration': False},
'master': {'host': f'http://127.0.0.1:{port}'}}
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
valid = await execution.validate_prompt('network-acceptance', prompt, None)
assert valid[0], valid
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
# Exercise production queue preparation; only master enqueue is redirected
# to this isolated executor, not the managed server's live queue.
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
prompt['2']['inputs'].pop('image_list_master_callback')
async def enqueue(prepared, workflow, client_id, **kwargs):
checked = await execution.validate_prompt('network-acceptance', prepared, None)
assert checked[0], checked
await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
prompt, None, 'test', [w['id'] for w in workers]), 45)
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
if not executor.success:
for err in stderr_files:
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS
assert len(CAPTURED) == 1
images, masks, texts = CAPTURED[0]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
assert torch.equal(images[i], row['image'] + row['strength']), (i, images[i])
assert torch.equal(masks[i], row['mask'] * 0.5), i
assert texts[i] == f'item-{i}:seed-{100 + i}', texts
async with aiohttp.ClientSession() as session:
stats = []
for worker in workers:
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
stats.append(await response.json())
counts = [len(row['calls']) for row in stats]
assert counts[0] > counts[1] > 0, counts
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
assert not runtime._MASTER_JOBS
# Three sequential regions on one worker, including optional outputs
# left unconnected. Region N+1 must not race the prior worker End.
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
previous = '1'
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
'images': [previous, 0], 'reference_image_mode': 'per_item', 'mask_mode': 'per_item',
'text_mode': 'shared', 'strength_mode': 'shared', 'seed_mode': 'shared'}}
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': [start, 0], 'context': [start, 7]}}
previous = end
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
chained, None, 'test', ['slow']), 45)
assert executor.success, executor.status_messages
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
assert CAPTURED[1][1:] == ([None, None], [None, None])
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
# Real Comfy user-interrupt flag while a real worker executor is active.
slow = workers[1]
async with aiohttp.ClientSession() as session:
async def slow_stats():
async with session.get(f"http://127.0.0.1:{slow['port']}/test-stats") as response:
return await response.json()
async def until(predicate):
for _ in range(150):
current = await slow_stats()
if predicate(current):
return current
await asyncio.sleep(0.02)
raise AssertionError('worker state transition timed out')
before = await until(lambda state: not state['running'] and not state['pending'])
context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
task = asyncio.create_task(runtime.run_image_list_job(context))
active = await until(lambda state: bool(state['running']))
owned_id = active['running'][0]
comfy.model_management.interrupt_current_processing(True)
try:
await asyncio.wait_for(task, 5)
raise AssertionError('interrupt was ignored')
except runtime.ImageListCancelledError:
pass
finally:
comfy.model_management.interrupt_current_processing(False)
after = await until(lambda state: not state['running'] and not state['pending'])
assert after['interrupted'] == [owned_id], after
assert len(after['calls']) <= len(before['calls']) + 1
assert set(before['history']) <= set(after['history'])
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
# A worker exception before End must fail, not hang or return a partial list.
bad_rows = source_rows()
bad_rows[0]['seed'] = -1
context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
try:
await asyncio.wait_for(runtime.run_image_list_job(context), 5)
raise AssertionError('worker exception was ignored')
except RuntimeError as exc:
assert 'failed image-list item' in str(exc), str(exc)
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_FAILURE_OK worker exception surfaced; context removed')
print('IMAGE_LIST_NETWORK_OK ' + json.dumps({'workers': 2, 'items': COUNT,
'assignments': dict(zip(['fast', 'slow'], counts)), 'ordered_exact_outputs': True,
'master_branch_calls': len(BRANCH_CALLS), 'source_calls': len(SOURCE_CALLS)}))
finally:
for proc in processes:
if proc.returncode is None:
proc.terminate()
await asyncio.gather(*(proc.wait() for proc in processes))
for err in stderr_files:
err.close()
await runner.cleanup()
network = importlib.import_module(module.__name__ + '.utils.network')
session = await network.get_client_session()
await session.close()
if __name__ == '__main__':
asyncio.run(worker_main() if WORKER else coordinator())
+83 -3
View File
@@ -196,6 +196,70 @@ async def check_execution(mapping, module, prompt_server, asset_manager):
print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager):
"""Execute real native V3 list/lazy flow; master must skip the processing branch."""
branch_calls = []
class ImageListProbe:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',)}}
RETURN_TYPES = ('IMAGE',)
FUNCTION = 'run'
CATEGORY = 'tests'
def run(self, image):
branch_calls.append(image)
return (image,)
nodes.NODE_CLASS_MAPPINGS['ImageListTestProbe'] = ImageListProbe
runtime = sys.modules[module.__name__ + '.image_list.runtime']
original_contexts = set(runtime._MASTER_JOBS)
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
async def fake_worker_results(_context_id):
return (
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
[None, None, None], [None, None, None],
)
graph = {
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 3, 'color': 0}},
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
'images': ['1', 0], 'reference_image_mode': 'per_item', 'mask_mode': 'per_item',
'text_mode': 'shared', 'strength': 1.0, 'strength_mode': 'shared',
'seed': 17, 'seed_mode': 'shared', 'image_list_enabled_worker_ids': '["test-worker"]',
'image_list_master_callback': 'http://127.0.0.1:1',
}},
'3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'context': ['2', 7],
}},
'5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}},
}
valid = await execution.validate_prompt('image-list-lazy', graph, None)
assert valid[0], valid
import folder_paths
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
try:
with tempfile.TemporaryDirectory(prefix='image-list-lazy-', dir=scratch) as temp:
with patch.object(folder_paths, 'temp_directory', temp), \
patch.object(image_list_nodes, 'run_image_list_job', fake_worker_results):
executor = execution.PromptExecutor(
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, 'image-list-lazy', {}, valid[2])
assert executor.success, executor.status_messages
assert branch_calls == [], 'master evaluated the lazy per-image branch'
records = executor.history_result['outputs']['5']['images']
assert len(records) == 3, records
for record in records:
preview = Path(temp) / record.get('subfolder', '') / record['filename']
with Image.open(preview) as image:
assert image.size == (8, 8) and image.mode == 'RGB'
finally:
nodes.NODE_CLASS_MAPPINGS.pop('ImageListTestProbe', None)
for context_id in set(runtime._MASTER_JOBS) - original_contexts:
runtime.remove_master_job(context_id)
print('IMAGE_LIST_LAZY_EXECUTION_OK real PromptExecutor; branch skipped on master; 3 previews verified')
async def main():
asset_manager = default_asset_manager()
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
@@ -206,9 +270,10 @@ async def main():
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 len(classes) == len(mapping) == len(BASELINE) + 2 == 10
assert set(mapping) == set(BASELINE) | {'DistributedImageListStart', 'DistributedImageListEnd'}
for node_id in BASELINE:
cls = mapping[node_id]
assert issubclass(cls, io.ComfyNode)
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
old = dict(BASELINE[node_id])
@@ -243,9 +308,24 @@ async def main():
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'))
start_schema = mapping['DistributedImageListStart'].GET_SCHEMA()
end_schema = mapping['DistributedImageListEnd'].GET_SCHEMA()
assert start_schema.node_id == 'DistributedImageListStart'
assert start_schema.display_name == 'Distributed Image List Start'
assert end_schema.node_id == 'DistributedImageListEnd'
assert end_schema.display_name == 'Distributed Image List End'
assert [out.io_type for out in start_schema.outputs] == [
'IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT']
assert all(out.is_output_list for out in start_schema.outputs)
assert start_schema.is_input_list and start_schema.accept_all_inputs
assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
assert all(out.is_output_list for out in end_schema.outputs)
assert end_schema.is_input_list and end_schema.is_output_node
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
print('SCHEMA_PARITY_OK', len(mapping))
check_saved_workflows(mapping)
await check_execution(mapping, module, prompt_server, asset_manager)
await check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager)
print('V3_ACCEPTANCE_OK')
+248
View File
@@ -0,0 +1,248 @@
"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
import asyncio
import unittest
import torch
from image_list.runtime import (
END_CLASS,
IMAGE_OUT,
CONTEXT_OUT,
START_CLASS,
ImageListCancelledError,
build_worker_prompt,
create_master_job,
decode_item,
decode_result,
dynamic_schedule,
encode_item,
encode_result,
normalize_image_list_inputs,
normalize_worker_result,
remove_master_job,
set_master_result,
validate_image_list_regions,
)
def start_node(image_link):
return {"class_type": START_CLASS, "inputs": {
"images": image_link,
"reference_image_mode": "per_item", "mask_mode": "per_item",
"text_mode": "shared", "strength_mode": "shared", "seed_mode": "shared",
}}
def graph_one_region():
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
"3": start_node(["2", 0]),
"4": {"class_type": "GenericFrameOperation", "inputs": {
"image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", 5],
}},
"5": {"class_type": END_CLASS, "inputs": {
"images": ["4", 0], "context": ["3", CONTEXT_OUT],
}},
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
}
def graph_two_sequential_regions():
return {
"1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
"2": start_node(["1", 0]),
"3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
"4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
"5": start_node(["4", 0]),
"6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
"7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
}
class ImageListNormalizationTests(unittest.TestCase):
def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
images = [
torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
]
refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
rows = normalize_image_list_inputs(
images, refs, masks, ["first", "second"], [0.2], [37],
"per_item", "per_item", "per_item", "shared", "shared",
)
self.assertEqual(len(rows), 2)
self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
self.assertTrue(torch.equal(rows[0]["reference_image"], refs[:1]))
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
self.assertEqual([row["index"] for row in rows], [0, 1])
self.assertEqual([row["seed"] for row in rows], [37, 37])
self.assertEqual([row["strength"] for row in rows], [0.2, 0.2])
def test_explicit_per_item_and_shared_scalar_contracts(self):
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
rows = normalize_image_list_inputs(image, text="same", strength=0.4, seed=9,
text_mode="shared", strength_mode="shared", seed_mode="shared")
self.assertEqual([row["text"] for row in rows], ["same"] * 3)
self.assertEqual([row["strength"] for row in rows], [0.4] * 3)
self.assertEqual([row["seed"] for row in rows], [9] * 3)
with self.assertRaisesRegex(ValueError, "expected exactly 3"):
normalize_image_list_inputs(image, text=["one", "two"], text_mode="per_item")
with self.assertRaisesRegex(ValueError, "shared inputs require exactly one"):
normalize_image_list_inputs(image, seed=[1, 2], seed_mode="shared")
with self.assertRaisesRegex(ValueError, "expected exactly 3"):
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)), mask_mode="per_item")
def test_no_implicit_tail_repetition_and_no_empty_primary_batch(self):
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
with self.assertRaisesRegex(ValueError, "expected exactly 2"):
normalize_image_list_inputs(images, reference_image=torch.zeros((1, 2, 2, 3)),
reference_image_mode="per_item")
with self.assertRaisesRegex(ValueError, "at least one image"):
normalize_image_list_inputs(None)
def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
"reference_image": torch.ones((1, 1, 3, 4), dtype=torch.float16),
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
"text": "safe text", "strength": 0.625, "seed": 42, "index": 3}
restored = decode_item(encode_item(item))
self.assertTrue(torch.equal(restored["image"], item["image"]))
self.assertEqual(restored["image"].dtype, torch.float64)
self.assertTrue(torch.equal(restored["reference_image"], item["reference_image"]))
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
self.assertEqual(restored["text"], "safe text")
self.assertEqual(restored["seed"], 42)
with self.assertRaisesRegex(TypeError, "floating-point tensor"):
encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
from image_list.runtime import encode_value
encode_value("/untrusted/path", "PATH")
def test_result_transport_preserves_ragged_result_shapes(self):
result = normalize_worker_result(
[torch.ones((1, 3, 5, 4), dtype=torch.float32)],
[torch.ones((1, 3, 5), dtype=torch.float16)],
["done"],
)
encoded = encode_result(result)
decoded = decode_result(encoded)
self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
self.assertEqual(decoded["image"].dtype, torch.float32)
self.assertEqual(decoded["mask"].dtype, torch.float16)
self.assertEqual(decoded["text"], "done")
with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
normalize_worker_result(torch.zeros((2, 3, 4, 3)))
class ImageListGraphBoundaryTests(unittest.TestCase):
def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
prompt = graph_one_region()
region = validate_image_list_regions(prompt)[0]
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
"ctx", 1, "http://master.invalid")
self.assertEqual(set(worker), {"1", "3", "4", "5"})
self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
self.assertNotIn("images", worker["3"]["inputs"])
self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
self.assertNotIn("6", worker)
def test_two_sequential_regions_are_valid_and_isolated(self):
regions = validate_image_list_regions(graph_two_sequential_regions())
self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
p = graph_two_sequential_regions()
second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
"http://master.invalid")
self.assertEqual(set(second), {"5", "6", "7"})
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
prompt = graph_one_region()
prompt["5"]["inputs"]["context"] = ["2", 0]
with self.assertRaisesRegex(ValueError, "typed context output"):
validate_image_list_regions(prompt)
nested = graph_one_region()
nested["7"] = start_node(["4", 0])
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
nested["5"]["inputs"]["images"] = ["8", 0]
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
validate_image_list_regions(nested)
outside = graph_one_region()
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
with self.assertRaisesRegex(ValueError, "cross-boundary"):
validate_image_list_regions(outside)
def test_rejects_duplicate_or_missing_results(self):
prompt = graph_one_region()
prompt["5"]["inputs"].pop("images")
with self.assertRaisesRegex(ValueError, "images must be connected"):
validate_image_list_regions(prompt)
class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
async def test_fast_workers_get_more_items_and_order_is_stable(self):
assignments = {"fast": [], "slow": []}
async def dispatch(worker, index):
assignments[worker].append(index)
delay = 0.002 if worker == "fast" else 0.025
task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
return {"index": index, "task": task}
async def wait_next(inflight):
done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
task = next(iter(done))
worker = next(w for w, handle in inflight.items() if handle["task"] is task)
index, result = task.result()
return worker, index, result
async def cleanup(_worker, handle):
handle["task"].cancel()
ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
async def dispatch(worker, index):
return {"index": index, "prompt_id": f"{worker}-{index}"}
cancelled = []
async def cleanup(worker, handle):
cancelled.append((worker, handle["prompt_id"]))
async def timeout(_inflight):
raise TimeoutError("worker timed out")
with self.assertRaisesRegex(TimeoutError, "timed out"):
await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
cancelled.clear()
async def wrong_index(inflight):
worker, handle = next(iter(inflight.items()))
return worker, handle["index"] + 99, "bad"
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
self.assertEqual(len(cancelled), 1)
cancelled.clear()
with self.assertRaises(ImageListCancelledError):
await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
interrupted=lambda: True)
self.assertEqual(len(cancelled), 0)
async def test_no_workers_and_duplicate_results_fail_closed(self):
async def unused(*_args):
self.fail("must not dispatch without workers")
with self.assertRaisesRegex(RuntimeError, "No free"):
await dynamic_schedule(2, [], unused, unused, unused)
async def dispatch(worker, index):
return {"index": index}
async def duplicate(inflight):
worker, _handle = next(iter(inflight.items()))
return worker, 0, "same"
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
if __name__ == "__main__":
unittest.main()
+178
View File
@@ -0,0 +1,178 @@
"""Regression coverage at the scheduler and real HTTP transport boundaries."""
import asyncio
import importlib.util
import json
from pathlib import Path
import sys
import types
import unittest
from unittest.mock import patch
import torch
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
from image_list import runtime
from test_image_list import graph_two_sequential_regions, start_node
def load_routes():
root = Path(__file__).resolve().parents[1]
name = '_image_list_transport_tests'
routes = web.RouteTableDef()
modules = {
name: types.ModuleType(name),
name + '.api': types.ModuleType(name + '.api'),
name + '.image_list': types.ModuleType(name + '.image_list'),
name + '.image_list.runtime': runtime,
'server': types.SimpleNamespace(PromptServer=types.SimpleNamespace(
instance=types.SimpleNamespace(routes=routes))),
}
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', root / 'api/image_list_routes.py')
module = importlib.util.module_from_spec(spec)
with patch.dict(sys.modules, modules):
spec.loader.exec_module(module)
return module, routes
class TransportTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.module, routes = load_routes()
app = web.Application()
app.add_routes(routes)
self.client = TestClient(TestServer(app))
await self.client.start_server()
async def asyncTearDown(self):
await self.client.close()
runtime._WORKER_ITEMS.clear()
async def test_fragmented_chunked_item_upload_is_consumable(self):
token = runtime.create_worker_task_token()
item = {'image': torch.full((1, 2, 3, 3), 0.25), 'index': 2}
data = json.dumps({'token': token, 'context_id': 'context', 'item': runtime.encode_item(item)}).encode()
async def chunks():
yield data[:8]
await asyncio.sleep(0.02)
yield data[8:]
response = await self.client.post('/distributed/image-list/item', data=chunks())
self.assertEqual(response.status, 200, await response.text())
result = runtime.consume_worker_item(token)
self.assertTrue(torch.equal(result['image'], item['image']))
self.assertEqual(result['index'], 2)
async def test_complete_item_upload_and_replay_rejection(self):
token = runtime.create_worker_task_token()
data = {'token': token, 'context_id': 'context', 'item': runtime.encode_item(
{'image': torch.zeros((1, 2, 3, 3)), 'index': 0})}
response = await self.client.post('/distributed/image-list/item', json=data)
self.assertEqual(response.status, 200, await response.text())
response = await self.client.post('/distributed/image-list/item', json=data)
self.assertEqual(response.status, 400)
runtime.consume_worker_item(token)
with self.assertRaisesRegex(RuntimeError, 'already consumed'):
runtime.consume_worker_item(token)
async def test_chunked_body_limit_is_aggregate(self):
async def chunks():
yield b'{"context_id":"'
await asyncio.sleep(0.01)
yield b'x' * 60
yield b'"}'
with patch.object(self.module, 'MAX_REQUEST_BYTES', 64):
response = await self.client.post('/distributed/image-list/item', data=chunks())
self.assertEqual(response.status, 400)
self.assertIn('exceeds', await response.text())
async def test_invalid_result_shapes_are_rejected(self):
for kind, shape in [('image', (1, 2, 3)), ('mask', (1, 2, 3, 1))]:
result = runtime.encode_result({'image': torch.zeros((1, 2, 3, 3)),
'mask': torch.zeros((1, 2, 3))})
result[kind]['shape'] = list(shape)
with self.assertRaises(ValueError):
runtime.decode_result(result)
class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase):
async def test_user_interrupt_during_result_wait_cleans_up_promptly(self):
entered = asyncio.Event()
cancelled = asyncio.Event()
wait_cancelled = asyncio.Event()
dispatches, cleanups = [], []
async def dispatch(worker, index):
dispatches.append(index)
return {'index': index}
async def wait_next(_inflight):
entered.set()
try:
await asyncio.Event().wait()
finally:
wait_cancelled.set()
async def cleanup(worker, handle):
cleanups.append((worker, handle['index']))
task = asyncio.create_task(runtime.dynamic_schedule(
3, ['one'], dispatch, wait_next, cleanup, cancelled.is_set))
await entered.wait()
cancelled.set()
with self.assertRaises(runtime.ImageListCancelledError):
await asyncio.wait_for(task, timeout=1)
self.assertTrue(wait_cancelled.is_set())
self.assertEqual(dispatches, [0])
self.assertEqual(cleanups, [('one', 0)])
async def test_already_interrupted_does_not_submit_any_work(self):
async def unused(*args):
self.fail('pre-cancelled job must not dispatch or wait')
with self.assertRaises(runtime.ImageListCancelledError):
await runtime.dynamic_schedule(2, ['one'], unused, unused, unused, lambda: True)
async def test_task_cancel_during_submit_waits_for_owned_handle(self):
entered, release = asyncio.Event(), asyncio.Event()
cleanups = []
async def dispatch(worker, index):
entered.set()
await release.wait()
return {'index': index, 'prompt_id': 'owned'}
async def unused(*args):
self.fail('cancelled dispatch must not enter result wait')
async def cleanup(worker, handle):
cleanups.append(handle['prompt_id'])
task = asyncio.create_task(runtime.dynamic_schedule(2, ['one'], dispatch, unused, cleanup))
await entered.wait()
task.cancel()
release.set()
with self.assertRaises(asyncio.CancelledError):
await task
self.assertEqual(cleanups, ['owned'])
def test_output_nodes_inside_region_are_rejected(self):
prompt = graph_two_sequential_regions()
output_class = type('OutputOperation', (), {'OUTPUT_NODE': True})
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'FirstFrameOperation': output_class})
with patch.dict(sys.modules, {'nodes': registry}):
with self.assertRaisesRegex(ValueError, 'must be placed after'):
runtime.validate_image_list_regions(prompt)
def test_shared_scalars_are_validated_before_dispatch(self):
for kwargs in ({'text': 4, 'text_mode': 'shared'},
{'strength': float('nan'), 'strength_mode': 'shared'},
{'seed': True, 'seed_mode': 'shared'}):
with self.assertRaises(TypeError):
runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), **kwargs)
def test_three_sequential_regions_can_share_worker_loader(self):
prompt = graph_two_sequential_regions()
prompt.pop('8')
prompt.update({
'8': start_node(['7', 0]),
'9': {'class_type': 'ThirdOperation', 'inputs': {'image': ['8', 0]}},
'10': {'class_type': runtime.END_CLASS,
'inputs': {'images': ['9', 0], 'context': ['8', runtime.CONTEXT_OUT]}},
'11': {'class_type': 'GenericModelLoader', 'inputs': {}},
})
for node_id in ['3', '6', '9']:
prompt[node_id]['inputs']['model'] = ['11', 0]
regions = runtime.validate_image_list_regions(prompt)
self.assertEqual(len(regions), 3)
for region in regions:
self.assertIn('11', region['node_ids'])
+15
View File
@@ -19,3 +19,18 @@ def test_native_v3_runtime():
)
assert result.returncode == 0, result.stdout + '\n' + result.stderr
assert 'V3_ACCEPTANCE_OK' in result.stdout
def test_image_list_network_runtime():
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
if not comfy_root:
pytest.skip('Set COMFYUI_SOURCE_ROOT for model-free multi-process HTTP acceptance')
helper = Path(__file__).parent / 'helpers' / 'image_list_network_check.py'
result = subprocess.run(
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
capture_output=True, text=True, timeout=180,
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
)
assert result.returncode == 0, result.stdout + '\n' + result.stderr
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
+1
View File
@@ -225,6 +225,7 @@ export const ENDPOINTS = {
export const NODE_CLASSES = {
DISTRIBUTED_COLLECTOR: 'DistributedCollector',
IMAGE_LIST_START: 'DistributedImageListStart',
DISTRIBUTED_SEED: 'DistributedSeed',
DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',
+12 -2
View File
@@ -9,7 +9,9 @@ export function setupInterceptor(extension) {
const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
if (hasCollector || hasDistUpscale) {
const hasImageList = findNodesByClass(prompt.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
if (hasCollector || hasDistUpscale || hasImageList) {
const result = await executeParallelDistributed(extension, prompt);
// Immediate status check for instant feedback
checkAllWorkerStatuses(extension);
@@ -25,12 +27,20 @@ export function setupInterceptor(extension) {
export async function executeParallelDistributed(extension, promptWrapper) {
const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
try {
const enabledWorkers = extension.enabledWorkers;
const hasImageList = findNodesByClass(promptWrapper.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
// Image-list v1 uses direct workers only; do not contact paid provider pools.
const enabledWorkers = hasImageList
? extension.enabledWorkers.filter(worker => worker.type !== "serverless")
: extension.enabledWorkers;
extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
// Pre-flight health check on all enabled workers
const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
if (hasImageList && activeWorkers.length === 0) {
throw new Error("Distributed Image List requires an active direct worker; master fallback is disabled.");
}
// Case: Enabled workers but all offline
if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
extension.log("No active workers found. All enabled workers are offline.");
+56
View File
@@ -0,0 +1,56 @@
import { afterEach, describe, expect, it, vi } from 'vitest';
vi.mock('../../scripts/api.js', () => ({ api: { clientId: 'test-client' } }), { virtual: true });
vi.mock('../workerLifecycle.js', () => ({
checkAllWorkerStatuses: vi.fn(), getWorkerUrl: () => 'http://worker.invalid',
}));
vi.mock('../workerUtils.js', () => ({
applyProbeResultToWorkerDot: vi.fn(),
findNodesByClass: (prompt, cls) => Object.values(prompt).filter(node => node.class_type === cls),
}));
import { api } from '../../scripts/api.js';
import { setupInterceptor } from '../executionUtils.js';
const prompt = { output: { '1': { class_type: 'DistributedImageListStart', inputs: {} } } };
function extension(workers, online = true) {
return {
isEnabled: true, enabledWorkers: workers, log: vi.fn(), config: { settings: {} },
originalQueuePrompt: vi.fn().mockResolvedValue({ prompt_id: 'local' }),
api: { probeWorker: vi.fn().mockResolvedValue({ ok: online }),
getServerlessStatus: vi.fn(),
queueDistributed: vi.fn().mockResolvedValue({ prompt_id: 'distributed' }) },
ui: { showToast: vi.fn() },
};
}
afterEach(() => vi.useRealTimers());
describe('image-list queue interception', () => {
it('routes image-list graphs through backend orchestration', async () => {
vi.useFakeTimers();
const ext = extension([{ id: 'one', type: 'local' }]);
setupInterceptor(ext);
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: ['one'] }));
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
});
it('never falls back to master when all image-list workers are offline', async () => {
vi.useFakeTimers();
const ext = extension([{ id: 'one', type: 'local' }], false);
setupInterceptor(ext);
expect(await api.queuePrompt(0, prompt)).toBeNull();
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
});
it('does not probe serverless providers for an image-list job', async () => {
vi.useFakeTimers();
const ext = extension([{ id: 'pool', type: 'serverless' }]);
setupInterceptor(ext);
expect(await api.queuePrompt(0, prompt)).toBeNull();
expect(ext.api.getServerlessStatus).not.toHaveBeenCalled();
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
});
it('preserves ordinary prompt handling', async () => {
const ext = extension([]);
setupInterceptor(ext);
expect(await api.queuePrompt(0, { output: {} })).toEqual({ prompt_id: 'local' });
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
});
});