diff --git a/README.md b/README.md index 8397368..6b442ce 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/api/__init__.py b/api/__init__.py index a2c9ad8..061bac1 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -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 diff --git a/api/image_list_routes.py b/api/image_list_routes.py new file mode 100644 index 0000000..2e97b7e --- /dev/null +++ b/api/image_list_routes.py @@ -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) diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index e2c02a4..3ef6da4 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -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( diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index ef25de8..2bc79db 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index d66d382..28b8082 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -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 diff --git a/docs/image-list-queue.md b/docs/image-list-queue.md new file mode 100644 index 0000000..75dbc24 --- /dev/null +++ b/docs/image-list-queue.md @@ -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. diff --git a/image_list/__init__.py b/image_list/__init__.py new file mode 100644 index 0000000..34f6346 --- /dev/null +++ b/image_list/__init__.py @@ -0,0 +1 @@ +"""Dynamic per-image image-list execution internals.""" diff --git a/image_list/runtime.py b/image_list/runtime.py new file mode 100644 index 0000000..4c6a989 --- /dev/null +++ b/image_list/runtime.py @@ -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) diff --git a/nodes/image_list.py b/nodes/image_list.py new file mode 100644 index 0000000..ecfe571 --- /dev/null +++ b/nodes/image_list.py @@ -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}") diff --git a/nodes/v3.py b/nodes/v3.py index 0bce348..1b2752f 100644 --- a/nodes/v3.py +++ b/nodes/v3.py @@ -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] diff --git a/tests/helpers/image_list_network_check.py b/tests/helpers/image_list_network_check.py new file mode 100644 index 0000000..b0fbc32 --- /dev/null +++ b/tests/helpers/image_list_network_check.py @@ -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()) diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py index dc148f2..ef64ce3 100644 --- a/tests/helpers/v3_runtime_check.py +++ b/tests/helpers/v3_runtime_check.py @@ -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') diff --git a/tests/test_image_list.py b/tests/test_image_list.py new file mode 100644 index 0000000..2af0eb9 --- /dev/null +++ b/tests/test_image_list.py @@ -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() diff --git a/tests/test_image_list_integration_gaps.py b/tests/test_image_list_integration_gaps.py new file mode 100644 index 0000000..291bc8b --- /dev/null +++ b/tests/test_image_list_integration_gaps.py @@ -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']) diff --git a/tests/test_v3_migration.py b/tests/test_v3_migration.py index f9ef1d7..df7e245 100644 --- a/tests/test_v3_migration.py +++ b/tests/test_v3_migration.py @@ -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 diff --git a/web/constants.js b/web/constants.js index 6a5a09c..b7a604d 100644 --- a/web/constants.js +++ b/web/constants.js @@ -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', diff --git a/web/executionUtils.js b/web/executionUtils.js index 8102bc6..3beceb4 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -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."); diff --git a/web/tests/imageListExecution.test.js b/web/tests/imageListExecution.test.js new file mode 100644 index 0000000..018874b --- /dev/null +++ b/web/tests/imageListExecution.test.js @@ -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(); + }); +});