diff --git a/api/image_list_routes.py b/api/image_list_routes.py index 3a43db2..8940c8b 100644 --- a/api/image_list_routes.py +++ b/api/image_list_routes.py @@ -6,6 +6,7 @@ import server from ..image_list.runtime import ( MAX_REQUEST_BYTES, + PROTOCOL_VERSION, decode_result, delete_worker_item, set_master_result, @@ -34,7 +35,7 @@ async def _json_body(request): async def image_list_capabilities_endpoint(_request): return web.json_response({ "supported": True, - "protocol_version": 1, + "protocol_version": PROTOCOL_VERSION, "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"], "transport": "typed-float-tensor-json-v1", }) diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 0aa8251..39a17b4 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -6,6 +6,9 @@ import aiohttp from ...utils.logging import debug_log, log from ...utils.network import build_worker_url, get_client_session, probe_worker +# Must match image_list.runtime.PROTOCOL_VERSION (kept local so this module +# stays importable without the image-list runtime; a test pins them equal). +IMAGE_LIST_PROTOCOL_VERSION = 2 try: from ...utils.trace_logger import trace_debug, trace_info except ImportError: @@ -215,11 +218,12 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data") if ( supported is not True - or protocol_version != 1 + or protocol_version != IMAGE_LIST_PROTOCOL_VERSION or not required_nodes.issubset(nodes) ): raise RuntimeError( - f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes " + f"Image-list worker {worker_id} is incompatible: requires protocol " + f"v{IMAGE_LIST_PROTOCOL_VERSION} and nodes " f"{sorted(required_nodes)}; no work was submitted" ) diff --git a/docs/image-list-queue.md b/docs/image-list-queue.md index 88a610f..d17ace1 100644 --- a/docs/image-list-queue.md +++ b/docs/image-list-queue.md @@ -23,7 +23,7 @@ Distributed Image List Start -------- context ---------+ 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. +1. Install the same feature-capable version of Distributed on the master and every selected worker. Workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation. 2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. 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**. @@ -36,13 +36,21 @@ ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a 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. +- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames; each frame becomes one job. +- Optional `mask`. +- Optional `value_1`, `value_2`, … inputs. Only `value_1` is shown at first; connecting it adds `value_2`, and so on, up to `value_10`. Each accepts an IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN. -Reference images and masks default to `per_item`. Text, strength and seed default to `shared`; strength defaults to `1.0`, seed to `0`. +`mask` and every value follow one rule: -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. +- **1 value** is shared by every image. +- **One value per image** pairs them in order. +- **Any other count** is an error. Values are never padded or repeated silently. + +Only `images` is split into frames. A batch connected to a value input stays a single value, so a 2-image batch can travel with each item (for example, two reference images per job). To pass a different batch to each image, connect a list with one batch per image. + +Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`, then one `value_N` output per connected `value_N` input, with the same type. Unused value outputs are hidden. Seeds and strengths are ordinary values: connect a single number to share it, a list for one per image, or derive it from `item_index`. + +Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below. 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. @@ -63,7 +71,7 @@ File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations wh - Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt. - 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. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported 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. In-process master items use tensors directly and do not incur those HTTP limits. +- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and do not incur those HTTP limits. - 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 fail the region. The worker timeout measures loss of evidence for the **owned prompt ID**, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue. - User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are **ambiguous**, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network. @@ -75,7 +83,7 @@ File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations wh Uploads, media preparation and submission acknowledgements are independently bounded. WebSocket acknowledgement waiting includes connect/send/receive and is not extended by heartbeats or unrelated messages. A user interrupt allows up to one second for a late acknowledgement; cleanup requests use five-second bounds each. -**4K float32 RGB outputs cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/reference/mask combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits. +**4K float32 RGB outputs cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/mask/value combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits. 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. diff --git a/image_list/runtime.py b/image_list/runtime.py index af505de..c8f6848 100644 --- a/image_list/runtime.py +++ b/image_list/runtime.py @@ -15,13 +15,20 @@ import time from dataclasses import dataclass, field from typing import Any -PROTOCOL_VERSION = 1 +# v2: Start carries generic value_N slots instead of fixed reference/text/strength/seed fields. +PROTOCOL_VERSION = 2 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) +# Start output indexes are a stable internal prompt protocol. Fixed outputs come +# first; value_1..value_N follow so the frontend can hide unused trailing slots +# without renumbering any connected output. +IMAGE_OUT, MASK_OUT, INDEX_OUT, CONTEXT_OUT = range(4) +VALUE_OUT_BASE = 4 +MAX_VALUE_SLOTS = 10 +VALUE_NAMES = tuple(f"value_{i}" for i in range(1, MAX_VALUE_SLOTS + 1)) +VALUES_INPUT = "values" MAX_PAYLOAD_BYTES = 50 * 1024 * 1024 MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024 MAX_ITEMS = 4096 @@ -29,6 +36,7 @@ SUBMISSION_TIMEOUT_SECONDS = 190.0 SUBMISSION_CANCEL_GRACE_SECONDS = 1.0 CLEANUP_TIMEOUT_SECONDS = 16.0 _ALLOWED_FLOAT_DTYPES = {"float16", "bfloat16", "float32", "float64"} +_VALUE_KINDS = {"TENSOR", "STRING", "INT", "FLOAT", "BOOLEAN"} _LOCK = threading.RLock() _MASTER_JOBS: dict[str, "ImageListJob"] = {} @@ -116,53 +124,75 @@ def _scalar_items(value: Any, name: str) -> list[Any]: return [value] -def _normalize_aligned(value: Any, mode: str, count: int, name: str, kind: str) -> list[Any]: +def _align(items: list[Any], count: int, name: str) -> list[Any]: + """One value is shared by every image; exactly one value per image pairs in order.""" + if len(items) == 1: + return [items[0] for _ in range(count)] + if len(items) == count: + return items + raise ValueError(f"{name} has {len(items)} values; connect exactly 1 (shared by every image) " + f"or {count} (one per image, in order)") + + +def value_kind(value: Any, name: str) -> str: + """Return the transport kind for one generic value, or raise for unsupported types.""" + torch = _torch() + if isinstance(value, torch.Tensor): + if not value.is_floating_point(): + raise TypeError(f"{name} tensor must have a floating-point dtype") + if value.ndim < 1 or value.ndim > 4 or min(value.shape) < 1: + raise ValueError(f"{name} tensor must have 1 to 4 positive dimensions") + return "TENSOR" + if isinstance(value, bool): + return "BOOLEAN" + if isinstance(value, int): + return "INT" + if isinstance(value, float): + if not math.isfinite(value): + raise TypeError(f"{name} numbers must be finite") + return "FLOAT" + if isinstance(value, str): + return "STRING" 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 + raise ValueError(f"{name} contains an empty value") + raise TypeError(f"{name} has unsupported type {type(value).__name__}; " + "use IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN") -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.""" +def _normalize_values(values: Any, count: int) -> dict[str, list[Any]]: + if values is None: + return {} + if not isinstance(values, dict): + raise TypeError("Start values must be a mapping of value_N slots") + aligned = {} + for name, raw in values.items(): + if name not in VALUE_NAMES: + raise ValueError(f"Unknown Start value slot {name!r}; expected value_1..value_{MAX_VALUE_SLOTS}") + if raw is None or (isinstance(raw, (list, tuple)) and all(item is None for item in raw)): + continue # Unconnected Autogrow slot. + # Each list element is one value. A batch tensor stays one value; only + # `images` is split into frames. + items = list(raw) if isinstance(raw, (list, tuple)) else [raw] + for item in items: + value_kind(item, name) + aligned[name] = _align(items, count, name) + return aligned + + +def normalize_image_list_inputs(images: Any, mask: Any = None, values: Any = None) -> list[dict[str, Any]]: + """Split images into items and pair masks/generic values by the 1-or-N rule.""" 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") + mask_frames = _image_frames(mask, "mask", mask=True) + masks = _align(mask_frames, count, "mask") if mask_frames else [None] * count + aligned = _normalize_values(values, count) return [ - {"image": frames[i], "reference_image": refs[i], "mask": masks[i], - "text": texts[i], "strength": strengths[i], "seed": seeds[i], "index": i} + {"image": frames[i], "mask": masks[i], + "values": {name: column[i] for name, column in aligned.items()}, "index": i} for i in range(count) ] @@ -175,14 +205,17 @@ def _tensor_dtype_name(tensor: Any) -> str: def encode_value(value: Any, kind: str) -> dict[str, Any] | None: - """Encode only the explicit IMAGE/MASK/STRING/FLOAT/INT protocol types.""" + """Encode only the explicit IMAGE/MASK/TENSOR/STRING/FLOAT/INT/BOOLEAN protocol types.""" if value is None: return None - if kind in ("IMAGE", "MASK"): + if kind in ("IMAGE", "MASK", "TENSOR"): 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 kind == "TENSOR": + if value.ndim < 1 or value.ndim > 4 or min(value.shape) < 1: + raise ValueError("TENSOR value must have 1 to 4 positive dimensions") + elif 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: @@ -208,6 +241,10 @@ def encode_value(value: Any, kind: str) -> dict[str, Any] | None: 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 == "BOOLEAN": + if not isinstance(value, bool): + raise TypeError("BOOLEAN transport value must be a boolean") + 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") @@ -220,7 +257,7 @@ def decode_value(data: Any, expected_kind: str) -> Any: 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"): + if expected_kind in ("IMAGE", "MASK", "TENSOR"): torch = _torch() dtype_name = data.get("dtype") if dtype_name not in _ALLOWED_FLOAT_DTYPES: @@ -256,18 +293,21 @@ def encode_request_payload(payload: dict[str, Any]) -> bytes: if len(body) > MAX_REQUEST_BYTES: raise ValueError( f"Image-list complete JSON/base64 request is {len(body)} bytes; " - f"limit is {MAX_REQUEST_BYTES} bytes (image/reference/mask combined)" + f"limit is {MAX_REQUEST_BYTES} bytes (image/mask/values combined)" ) return body def encode_item(item: dict[str, Any]) -> dict[str, Any]: + values = item.get("values") or {} + encoded_values = {} + for name, value in values.items(): + if name not in VALUE_NAMES: + raise ValueError(f"Unknown Start value slot {name!r}") + encoded_values[name] = encode_value(value, value_kind(value, name)) 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"), + "values": encoded_values, "index": int(item["index"])} @@ -280,18 +320,22 @@ def decode_item(data: Any) -> dict[str, Any]: 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} + raw_values = data.get("values", {}) + if not isinstance(raw_values, dict): + raise ValueError("item.values must be an object") + values = {} + for name, encoded in raw_values.items(): + if name not in VALUE_NAMES: + raise ValueError(f"item.values has unknown slot {name!r}") + if not isinstance(encoded, dict) or encoded.get("type") not in _VALUE_KINDS: + raise ValueError(f"item.values.{name} has an unsupported type") + values[name] = decode_value(encoded, encoded["type"]) + if values[name] is None: + raise ValueError(f"item.values.{name} is empty") + return {"image": image, "mask": mask, "values": values, "index": index} def _sweep_worker_items() -> None: @@ -533,8 +577,10 @@ def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_tok _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) + # Per-item data reaches workers only through the staged item token. + for name in list(start_inputs): + if name in ("images", "mask") or name.startswith(VALUES_INPUT + "."): + start_inputs.pop(name) 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, diff --git a/nodes/image_list.py b/nodes/image_list.py index 6ed9fa6..3c0bb62 100644 --- a/nodes/image_list.py +++ b/nodes/image_list.py @@ -7,6 +7,9 @@ from comfy_api.v0_0_2 import io from ..utils.async_helpers import run_async_in_server_loop from ..image_list.runtime import ( + MAX_VALUE_SLOTS, + VALUE_NAMES, + VALUES_INPUT, create_master_job, normalize_image_list_inputs, normalize_worker_result, @@ -17,6 +20,12 @@ from ..image_list.runtime import ( ) +# Types the transport can carry exactly. Anything else (models, latents, ...) +# must be loaded inside the region or shared as a model dependency instead. +_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean] +assert len(VALUE_NAMES) == MAX_VALUE_SLOTS + + class DistributedImageListStart(io.ComfyNode): """Materialize a batch once, then expose ordered per-item typed values.""" @@ -26,30 +35,28 @@ class DistributedImageListStart(io.ComfyNode): 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."), + description=("Starts one per-image processing region. Each image becomes one job. " + "mask and value_N inputs take 1 value (shared by every image) or exactly " + "one value per image, paired in order."), 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"), + io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one job."), + io.Mask.Input("mask", optional=True, + tooltip="1 mask shared by every image, or one mask per image."), + io.Autogrow.Input(VALUES_INPUT, optional=True, template=io.Autogrow.TemplateNames( + io.MultiType.Input("value", types=_VALUE_TYPES, + tooltip=("Any image, mask, text, number or boolean. 1 value is shared by " + "every image; one value per image pairs them in order. " + "A batch stays one value.")), + names=list(VALUE_NAMES), min=0)), ], 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), + # value_N outputs mirror value_N inputs. The web extension hides + # unused ones and shows the connected input's type. + *[io.AnyType.Output(id=name, display_name=name, is_output_list=True) for name in VALUE_NAMES], ], is_input_list=True, accept_all_inputs=True, @@ -58,9 +65,7 @@ class DistributedImageListStart(io.ComfyNode): ) @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", + def execute(cls, images=None, mask=None, values=None, worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]", image_list_master_callback="", image_list_execution_id="", image_list_workflow_meta=None, image_list_master_participates=None, @@ -77,16 +82,13 @@ class DistributedImageListStart(io.ComfyNode): 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") + item_values = item.get("values") or {} 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], + [item["image"]], [item.get("mask")], [item["index"]], [context_id], + *[[item_values.get(name)] for name in VALUE_NAMES], ) - items = normalize_image_list_inputs( - images, reference_image, mask, text, strength, seed, - reference_image_mode, mask_mode, text_mode, strength_mode, seed_mode, - ) + items = normalize_image_list_inputs(images, mask, values) prompt = getattr(cls.hidden, "prompt", None) start_id = str(getattr(cls.hidden, "unique_id", "")) if not prompt or not start_id: @@ -119,13 +121,10 @@ class DistributedImageListStart(io.ComfyNode): ) 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], + *[[item["values"].get(name) for item in items] for name in VALUE_NAMES], ) diff --git a/tests/helpers/image_list_network_check.py b/tests/helpers/image_list_network_check.py index fcdcff9..95fb384 100644 --- a/tests/helpers/image_list_network_check.py +++ b/tests/helpers/image_list_network_check.py @@ -63,6 +63,17 @@ class TestImageSource: ('image', 'reference_image', 'mask', 'text', 'strength', 'seed')) +class TestSharedFlag: + @classmethod + def INPUT_TYPES(cls): + return {'required': {}} + RETURN_TYPES = ('BOOLEAN',) + FUNCTION = 'run' + def run(self): + assert not WORKER, 'shared value source leaked into worker prompt' + return (True,) + + class TestCachedModel: @classmethod def INPUT_TYPES(cls): @@ -81,10 +92,10 @@ class TestImageOperation: 'text': ('STRING', {'forceInput': True}), 'strength': ('FLOAT', {'forceInput': True}), 'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}), - 'model': ('MODEL',)}} + 'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}} RETURN_TYPES = ('IMAGE', 'MASK', 'STRING') FUNCTION = 'run' - def run(self, image, reference_image, mask, text, strength, seed, index, model): + def run(self, image, reference_image, mask, text, strength, seed, index, flag, model): assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode' BRANCH_CALLS.append(index) # Sleep in short chunks so the real Comfy interrupt signal can be tested. @@ -93,6 +104,7 @@ class TestImageOperation: comfy.model_management.throw_exception_if_processing_interrupted() time.sleep(0.01) assert seed == 100 + index + assert flag is True assert text == f'item-{index}' assert strength == (index + 1) / 16 assert torch.all(reference_image == strength) @@ -119,7 +131,7 @@ async def setup(): module = sys.modules[str(ROOT).replace('.', '_x_')] nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource, TestImageOperation=TestImageOperation, TestCapture=TestCapture, - TestCachedModel=TestCachedModel) + TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag) # Match the CORS-enabled direct-worker launch contract for canvas probes. app = web.Application(middlewares=[server.create_cors_middleware('*')]) # Register the actual feature routes, not test replacements. @@ -206,19 +218,22 @@ async def worker_main(): def graph(callback, worker_ids): return { '1': {'class_type': 'TestImageSource', 'inputs': {}}, + # Generic Autogrow value slots carry per-item reference images, text, + # strength and seed (one value per image), plus a shared boolean. '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', + 'images': ['1', 0], 'mask': ['1', 2], + 'values.value_1': ['1', 1], 'values.value_2': ['1', 3], + 'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0], '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)]} | {'model': ['6', 0]}}, + key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2), + ('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]} + | {'model': ['6', 0]}}, '4': {'class_type': 'DistributedImageListEnd', 'inputs': { - 'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 7]}}, + 'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}}, '5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}}, '6': {'class_type': 'TestCachedModel', 'inputs': {}}, + '7': {'class_type': 'TestSharedFlag', 'inputs': {}}, } @@ -325,10 +340,9 @@ async def coordinator(): 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'}} + 'images': [previous, 0]}} chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': { - 'images': [start, 0], 'context': [start, 7]}} + 'images': [start, 0], 'context': [start, 3]}} previous = end chained['8'] = {'class_type': 'TestCapture', 'inputs': { 'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}} diff --git a/tests/helpers/v3_runtime_check.py b/tests/helpers/v3_runtime_check.py index 3f64e02..be4d552 100644 --- a/tests/helpers/v3_runtime_check.py +++ b/tests/helpers/v3_runtime_check.py @@ -228,15 +228,13 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m 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"]', + 'images': ['1', 0], 'image_list_enabled_worker_ids': '["test-worker"]', 'image_list_master_callback': 'http://127.0.0.1:1', 'image_list_workflow_meta': workflow_meta, }}, '3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}}, '4': {'class_type': 'DistributedImageListEnd', 'inputs': { - 'images': ['3', 0], 'context': ['2', 7], + 'images': ['3', 0], 'context': ['2', 3], }}, '5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}}, } @@ -247,7 +245,7 @@ async def check_image_list_lazy_executor(mapping, module, prompt_server, asset_m assert len(validator(graph)) == 1 connected_text_graph = json.loads(json.dumps(graph)) connected_text_graph['6'] = { - 'class_type': 'ImageListTextProbe', 'inputs': {'text': ['2', 3]}, + 'class_type': 'ImageListTextProbe', 'inputs': {'text': ['2', 4]}, } connected_text_graph['4']['inputs']['texts'] = ['6', 0] assert len(validator(connected_text_graph)) == 1 @@ -359,9 +357,44 @@ async def main(): assert texts_schema.get('forceInput') is True, texts_schema assert 'default' not in texts_schema, texts_schema assert [out.io_type for out in start_schema.outputs] == [ - 'IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT'] + 'IMAGE', 'MASK', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT'] + ['*'] * 10 + assert [out.display_name for out in start_schema.outputs][4:] == [f'value_{i}' for i in range(1, 11)] assert all(out.is_output_list for out in start_schema.outputs) assert start_schema.is_input_list and start_schema.accept_all_inputs + start_info = json.loads(json.dumps(mapping['DistributedImageListStart'].GET_NODE_INFO_V1())) + assert list(start_info['input']['optional']) == ['images', 'mask', 'values'], start_info['input'] + values_spec = start_info['input']['optional']['values'] + assert values_spec[0] == 'COMFY_AUTOGROW_V3', values_spec + template = values_spec[1]['template'] + assert template['names'] == [f'value_{i}' for i in range(1, 11)] and template['min'] == 0, template + assert template['input']['required']['value'][0] == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN', template + # Autogrow + INPUT_IS_LIST through the real parser/list mapper, values + # still split 1-or-N, unconnected slots ignored. + # get_input_data must accept Autogrow names; constants are wrapped once. + parsed, missing, hidden = execution.get_input_data( + {'values.value_1': 'shared', 'values.value_3': 4}, + mapping['DistributedImageListStart'], unique_id='autogrow-probe', extra_data={}) + assert not missing and parsed['values.value_1'] == ['shared'], parsed + # Linked inputs arrive as the upstream output list (one element per value). + prepared = {'images': [torch.zeros((3, 2, 2, 3))], 'values.value_1': ['shared'], + 'values.value_3': [4, 5, 6]} + seen = {} + runtime_mod = sys.modules[module.__name__ + '.image_list.runtime'] + def capture(_images, _mask, values): + seen['values'] = values + raise RuntimeError('captured') + nodes_mod = sys.modules[module.__name__ + '.nodes.image_list'] + with patch.object(nodes_mod, 'normalize_image_list_inputs', capture): + try: + await asyncio.to_thread(lambda: asyncio.run(execution._async_map_node_over_list( + 'autogrow-probe', 'autogrow-probe', mapping['DistributedImageListStart'], prepared, + 'execute', v3_data=hidden))) + except RuntimeError as exc: + assert 'captured' in str(exc), exc + assert seen['values'] == {'value_1': ['shared'], 'value_3': [4, 5, 6]}, seen + rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values']) + assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)] + print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item') 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 diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py index fedf562..401121b 100644 --- a/tests/test_dispatch_selection.py +++ b/tests/test_dispatch_selection.py @@ -168,7 +168,7 @@ class _FakeSession: def _valid_capability(): return { "supported": True, - "protocol_version": 1, + "protocol_version": 2, "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"], } @@ -418,7 +418,7 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): await dispatch.preflight_image_list_workers(workers, require_idle=True) async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self): - malformed = {"supported": True, "protocol_version": 1, "node_ids": None} + malformed = {"supported": True, "protocol_version": 2, "node_ids": None} session = _FakeSession({("worker", "capabilities"): (200, malformed)}) async def fake_session(): @@ -432,7 +432,7 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self): incompatible = _valid_capability() - incompatible["protocol_version"] = 2 + incompatible["protocol_version"] = 1 session = _FakeSession({("worker", "capabilities"): (200, incompatible)}) async def fake_session(): diff --git a/tests/test_image_list.py b/tests/test_image_list.py index ddaad1e..c0c7f0f 100644 --- a/tests/test_image_list.py +++ b/tests/test_image_list.py @@ -9,6 +9,7 @@ from image_list.runtime import ( IMAGE_OUT, CONTEXT_OUT, START_CLASS, + VALUE_OUT_BASE, ImageListCancelledError, build_worker_prompt, create_master_job, @@ -26,11 +27,7 @@ from image_list.runtime import ( 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", - }} + return {"class_type": START_CLASS, "inputs": {"images": image_link}} def graph_one_region(): @@ -39,7 +36,7 @@ def graph_one_region(): "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], + "image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", VALUE_OUT_BASE], }}, "5": {"class_type": END_CLASS, "inputs": { "images": ["4", 0], "context": ["3", CONTEXT_OUT], @@ -70,53 +67,85 @@ class ImageListNormalizationTests(unittest.TestCase): 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", - ) + rows = normalize_image_list_inputs(images, masks, { + "value_1": [refs[:1], refs[1:]], "value_2": ["first", "second"], + "value_3": [0.2], "value_4": [37], + }) 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.assertTrue(torch.equal(rows[0]["values"]["value_1"], 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]) + self.assertEqual([row["values"]["value_2"] for row in rows], ["first", "second"]) + self.assertEqual([row["values"]["value_3"] for row in rows], [0.2, 0.2]) + self.assertEqual([row["values"]["value_4"] for row in rows], [37, 37]) - def test_explicit_per_item_and_shared_scalar_contracts(self): + def test_one_value_is_shared_and_n_values_pair_in_order(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") + rows = normalize_image_list_inputs(image, values={ + "value_1": ["same"], "value_2": [0.4], "value_3": [9], "value_4": [True]}) + self.assertEqual([row["values"]["value_1"] for row in rows], ["same"] * 3) + self.assertEqual([row["values"]["value_2"] for row in rows], [0.4] * 3) + self.assertEqual([row["values"]["value_3"] for row in rows], [9] * 3) + self.assertEqual([row["values"]["value_4"] for row in rows], [True] * 3) + rows = normalize_image_list_inputs(image, values={"value_2": [1, 2, 3]}) + self.assertEqual([row["values"] for row in rows], [{"value_2": 1}, {"value_2": 2}, {"value_2": 3}]) + with self.assertRaisesRegex(ValueError, "value_1 has 2 values; connect exactly 1.*or 3"): + normalize_image_list_inputs(image, values={"value_1": ["one", "two"]}) + with self.assertRaisesRegex(ValueError, "mask has 2 values"): + normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4))) + shared_mask = normalize_image_list_inputs(image, mask=torch.ones((1, 2, 4))) + self.assertTrue(all(torch.equal(row["mask"], torch.ones((1, 2, 4))) for row in shared_mask)) - def test_no_implicit_tail_repetition_and_no_empty_primary_batch(self): + def test_value_batches_stay_whole_and_only_images_split(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") + pair = torch.rand((2, 4, 4, 3)) # e.g. two references for one item + rows = normalize_image_list_inputs(images, values={"value_1": [pair]}) + self.assertTrue(all(torch.equal(row["values"]["value_1"], pair) for row in rows)) + per_item = [torch.rand((2, 4, 4, 3)), torch.rand((2, 4, 4, 3))] + rows = normalize_image_list_inputs(images, values={"value_1": per_item}) + self.assertTrue(torch.equal(rows[1]["values"]["value_1"], per_item[1])) + + def test_rejects_unsupported_values_and_counts(self): + images = torch.zeros((2, 2, 2, 3), dtype=torch.float32) + with self.assertRaisesRegex(ValueError, "value_1 has 3 values"): + normalize_image_list_inputs(images, values={"value_1": [1, 2, 3]}) + with self.assertRaisesRegex(TypeError, "unsupported type dict"): + normalize_image_list_inputs(images, values={"value_1": [{"samples": 1}]}) + with self.assertRaisesRegex(TypeError, "floating-point"): + normalize_image_list_inputs(images, values={"value_1": [torch.zeros(2, dtype=torch.int64)]}) + with self.assertRaisesRegex(TypeError, "finite"): + normalize_image_list_inputs(images, values={"value_1": [float("nan")]}) + with self.assertRaisesRegex(ValueError, "Unknown Start value slot"): + normalize_image_list_inputs(images, values={"value_11": [1]}) with self.assertRaisesRegex(ValueError, "at least one image"): normalize_image_list_inputs(None) + # Unconnected Autogrow slots arrive as None and are ignored. + rows = normalize_image_list_inputs(images, values={"value_1": [None], "value_2": None}) + self.assertEqual([row["values"] for row in rows], [{}, {}]) 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} + "values": {"value_1": torch.ones((2, 1, 3, 4), dtype=torch.float16), + "value_2": "safe text", "value_3": 0.625, "value_4": 42, + "value_5": False}, + "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["values"]["value_1"], item["values"]["value_1"])) + self.assertEqual(restored["values"]["value_1"].dtype, torch.float16) self.assertTrue(torch.equal(restored["mask"], item["mask"])) - self.assertEqual(restored["text"], "safe text") - self.assertEqual(restored["seed"], 42) + self.assertEqual(restored["values"]["value_2"], "safe text") + self.assertEqual(restored["values"]["value_3"], 0.625) + self.assertIs(restored["values"]["value_4"], 42) + self.assertIs(restored["values"]["value_5"], False) + with self.assertRaisesRegex(ValueError, "unknown slot"): + decode_item({**encode_item(item), "values": {"strength": {"type": "FLOAT", "value": 1.0}}}) + with self.assertRaisesRegex(ValueError, "unsupported type"): + decode_item({**encode_item(item), "values": {"value_1": {"type": "PATH", "value": "/x"}}}) 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"): @@ -148,6 +177,19 @@ class ImageListGraphBoundaryTests(unittest.TestCase): self.assertEqual(set(worker), {"1", "3", "4", "5"}) self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012") self.assertNotIn("images", worker["3"]["inputs"]) + + def test_worker_prompt_drops_per_item_value_links(self): + prompt = graph_one_region() + prompt["8"] = {"class_type": "StringList", "inputs": {}} + prompt["3"]["inputs"].update({"values.value_1": ["8", 0], "values.value_2": ["8", 0]}) + import sys, types + from unittest.mock import patch + registry = {"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL",)})} + with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}): + region = validate_image_list_regions(prompt)[0] + worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012", "ctx", 0, "http://m") + self.assertFalse([name for name in worker["3"]["inputs"] if name.startswith("values.")]) + self.assertNotIn("8", worker) self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0]) self.assertNotIn("6", worker) diff --git a/tests/test_image_list_integration_gaps.py b/tests/test_image_list_integration_gaps.py index 9175ca5..b11ea66 100644 --- a/tests/test_image_list_integration_gaps.py +++ b/tests/test_image_list_integration_gaps.py @@ -154,11 +154,9 @@ class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase): 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'}): + for value in (float('nan'), object(), torch.zeros(2, dtype=torch.int32)): with self.assertRaises(TypeError): - runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), **kwargs) + runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), values={'value_1': [value]}) def test_three_sequential_regions_can_share_worker_loader(self): prompt = graph_two_sequential_regions() diff --git a/tests/test_image_list_master_runtime.py b/tests/test_image_list_master_runtime.py index d232bdd..bc317a9 100644 --- a/tests/test_image_list_master_runtime.py +++ b/tests/test_image_list_master_runtime.py @@ -21,7 +21,7 @@ def graph(): 'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}}, 'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}}, 'end': {'class_type': 'DistributedImageListEnd', 'inputs': { - 'images': ['op', 0], 'context': ['start', 7]}}, + 'images': ['op', 0], 'context': ['start', 3]}}, } @@ -90,7 +90,8 @@ class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase): raise fixture.local_error item = fixture.rt.consume_worker_item(token) fixture.rt.set_master_result(fixture.context, index, token, { - 'image': item['image'], 'mask': item.get('mask'), 'text': item.get('text')}) + 'image': item['image'], 'mask': item.get('mask'), + 'text': item['values'].get('value_1')}) finally: fixture.local_finished.set() task = asyncio.create_task(execute()) @@ -134,7 +135,8 @@ class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase): def create(self, participates=True, count=5): self.items = [{'image': torch.full((1, i + 1, 2, 3), float(i)), - 'mask': torch.full((1, i + 1, 2), float(i)), 'text': str(i), 'index': i} + 'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)}, + 'index': i} for i in range(count)] self.context = self.rt.create_master_job('start', 'end', graph(), self.items, [w['id'] for w in self.workers], '', master_participates=participates) @@ -162,7 +164,7 @@ class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase): # A remote callback is ready immediately; the local slot is slower. item = self.rt.consume_worker_item(token) self.rt.set_master_result(self.context, index, token, { - 'image': item['image'], 'mask': item['mask'], 'text': item['text']}) + 'image': item['image'], 'mask': item['mask'], 'text': item['values']['value_1']}) return f'remote-{index}' async def test_local_only_preserves_order_and_uses_no_http(self): @@ -328,10 +330,12 @@ def real_framework_probe(comfy_root): nodes.NODE_CLASS_MAPPINGS.update(Loader=Loader, Operation=Operation, Expanded=Expanded) prompt = graph() prompt['loader'] = {'class_type': 'Loader', 'inputs': {}} - prompt['op']['inputs'].update(model=['loader', 0], mask=['start', 2], text=['start', 3]) + prompt['op']['inputs'].update(model=['loader', 0], mask=['start', rt.MASK_OUT], + text=['start', rt.VALUE_OUT_BASE]) prompt['end']['inputs'].update(masks=['op', 1], texts=['op', 2]) items = [{'image': torch.full((1, i + 1, 2, 3), float(i)), - 'mask': torch.full((1, i + 1, 2), float(i)), 'text': str(i), 'index': i} for i in range(3)] + 'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)}, 'index': i} + for i in range(3)] region = rt.validate_image_list_regions(prompt)[0] parent_progress = progress.ProgressRegistry('outer', execution.DynamicPrompt({})) parent_progress.nodes['outer'] = {'state': progress.NodeState.Running, 'value': 3, 'max': 7} @@ -364,7 +368,7 @@ def real_framework_probe(comfy_root): data['images'] = [torch.zeros((3, 2, 2, 3))] result = await execution._async_map_node_over_list('outer', 'start', cls, data, cls.FUNCTION, v3_data=v3) result = await execution.resolve_map_node_over_list_results(result) - context = result[0].result[7][0] + context = result[0].result[rt.CONTEXT_OUT][0] assert rt.get_master_job(context).master_participates is expected rt.remove_master_job(context) diff --git a/tests/test_image_list_master_selection.py b/tests/test_image_list_master_selection.py index 49838a6..58e2e5a 100644 --- a/tests/test_image_list_master_selection.py +++ b/tests/test_image_list_master_selection.py @@ -25,7 +25,7 @@ class MasterSelectionTests(unittest.IsolatedAsyncioTestCase): self.prompt = { "1": {"class_type": "EmptyImage", "inputs": {}}, "2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}}, - "3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 7]}}, + "3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 3]}}, } async def queue(self, **kwargs): diff --git a/tests/test_image_list_runtime_boundaries.py b/tests/test_image_list_runtime_boundaries.py index 70d87be..341a97f 100644 --- a/tests/test_image_list_runtime_boundaries.py +++ b/tests/test_image_list_runtime_boundaries.py @@ -237,7 +237,7 @@ class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32') self.rt.MAX_REQUEST_BYTES = len(single) self.assertEqual(self.rt.encode_request_payload(payload), single) - item['reference_image'] = item['image'] + item['values'] = {'value_1': item['image']} item['mask'] = torch.zeros((1, 4, 4)) payload['item'] = self.rt.encode_item(item) with self.assertRaisesRegex(ValueError, 'complete.*request'): diff --git a/web/imageListStart.js b/web/imageListStart.js new file mode 100644 index 0000000..dd1fbba --- /dev/null +++ b/web/imageListStart.js @@ -0,0 +1,170 @@ +import { app } from "/scripts/app.js"; + +// Distributed Image List Start declares value_1..value_10 outputs in Python so +// saved links and prompt output indexes stay valid. ComfyUI grows the matching +// `values.value_N` inputs natively (Autogrow) but has no growable outputs, so +// this extension only presents them: it shows outputs for connected inputs, +// copies each input's type onto its output, and moves output links when +// Autogrow compacts inputs after a middle input is disconnected. +// If this script fails, the node still runs; all outputs simply stay visible. + +const START_CLASS = "DistributedImageListStart"; +const FIXED_OUTPUTS = 4; // image, mask, item_index, context +const MAX_VALUES = 10; +const INPUT_PREFIX = "values.value_"; + +function valueOrdinal(name) { + if (typeof name !== "string" || !name.startsWith(INPUT_PREFIX)) return null; + const ordinal = Number.parseInt(name.slice(INPUT_PREFIX.length), 10); + return ordinal >= 1 && ordinal <= MAX_VALUES ? ordinal : null; +} + +function getLink(graph, id) { + if (id == null || !graph) return undefined; + if (typeof graph.getLink === "function") return graph.getLink(id) ?? undefined; + const links = graph.links; + return (typeof links?.get === "function" ? links.get(id) : links?.[id]) ?? undefined; +} + +function connectedValues(node) { + const result = new Map(); + for (const input of node.inputs ?? []) { + const ordinal = valueOrdinal(input.name); + if (!ordinal) continue; + const link = getLink(node.graph, input.link); + if (!link) continue; + const origin = node.graph.getNodeById?.(link.origin_id); + const originType = origin?.outputs?.[link.origin_slot]?.type; + const type = originType && originType !== "*" ? originType : (link.type ?? "*"); + result.set(ordinal, { type, key: `${link.origin_id}:${link.origin_slot}` }); + } + return result; +} + +function outputLinks(node, ordinal) { + return node.outputs?.[FIXED_OUTPUTS + ordinal - 1]?.links ?? []; +} + +function typesCompatible(outputType, inputType) { + if (!outputType || !inputType || outputType === "*" || inputType === "*") return true; + const accepted = String(inputType).split(","); + return String(outputType).split(",").some(type => accepted.includes(type)); +} + +function moveOutputLinks(node, from, to) { + const graph = node.graph; + const targets = outputLinks(node, from) + .map(id => getLink(graph, id)) + .filter(Boolean) + .map(link => [graph.getNodeById?.(link.target_id), link.target_slot]) + .filter(([target]) => target); + // The moved links carry the same value, so give the destination output its + // type first; LiteGraph refuses connections to a mismatched output type. + node.outputs[FIXED_OUTPUTS + to - 1].type = node.outputs[FIXED_OUTPUTS + from - 1].type; + // Connecting a target input replaces its existing link from `from`. + for (const [target, slot] of targets) node.connect(FIXED_OUTPUTS + to - 1, target, slot); +} + +function disconnectValueOutput(node, ordinal) { + if (outputLinks(node, ordinal).length) node.disconnectOutput(FIXED_OUTPUTS + ordinal - 1); +} + +// Autogrow removes a disconnected middle input and shifts later input links +// down by one. Shift the matching outputs' links the same way, and drop the +// removed value's downstream links, so no consumer silently receives a +// different value. +function followCompactedInputs(node, current) { + const previous = node._distributedValueSources; + if (!previous) return; + const now = []; + for (const [ordinal, info] of current) now[ordinal - 1] = info.key; + const before = previous; + if (now.length !== before.length - 1) return; + let removed = 0; + while (removed < now.length && now[removed] === before[removed]) removed++; + // Removing the last value shifts nothing; keep its links so a temporary + // rewire does not silently drop downstream connections. + if (removed >= now.length) return; + for (let i = removed; i < now.length; i++) { + if (now[i] !== before[i + 1]) return; // Not a single Autogrow compaction. + } + disconnectValueOutput(node, removed + 1); + for (let i = removed + 1; i < before.length; i++) moveOutputLinks(node, i + 1, i); +} + +function syncImageListStartOutputs(node) { + if (!node?.graph || !Array.isArray(node.outputs) || node.outputs.length < FIXED_OUTPUTS) return; + const current = connectedValues(node); + followCompactedInputs(node, current); + const sources = []; + for (const [ordinal, info] of current) sources[ordinal - 1] = info.key; + node._distributedValueSources = Array.from(sources); + + let visible = Math.max(0, ...current.keys()); + for (let ordinal = MAX_VALUES; ordinal > visible; ordinal--) { + // Never drop a wired output; its links keep it visible until removed. + if (outputLinks(node, ordinal).length) visible = ordinal; + } + const listShape = node.outputs[0]?.shape; + while (node.outputs.length < FIXED_OUTPUTS + visible) { + const ordinal = node.outputs.length - FIXED_OUTPUTS + 1; + node.addOutput(`value_${ordinal}`, "*", listShape === undefined ? undefined : { shape: listShape }); + } + while (node.outputs.length > FIXED_OUTPUTS + visible) node.removeOutput(node.outputs.length - 1); + + for (let ordinal = 1; ordinal <= visible; ordinal++) { + const output = node.outputs[FIXED_OUTPUTS + ordinal - 1]; + const type = current.get(ordinal)?.type ?? "*"; + output.name = `value_${ordinal}`; + if (output.type !== type) { + output.type = type; + for (const id of [...(output.links ?? [])]) { + const link = getLink(node.graph, id); + if (!link) continue; + const target = node.graph.getNodeById?.(link.target_id); + const targetType = target?.inputs?.[link.target_slot]?.type; + if (target && !typesCompatible(type, targetType)) target.disconnectInput(link.target_slot); + else link.type = type; + } + } + } + node.setDirtyCanvas?.(true, true); +} + +function scheduleSync(node) { + if (node._distributedValueSyncPending) return; + node._distributedValueSyncPending = true; + const run = () => { + node._distributedValueSyncPending = false; + try { + syncImageListStartOutputs(node); + } catch (error) { + console.error("[Distributed] Image List Start output sync failed:", error); + } + }; + // Run after Autogrow's own animation-frame input compaction. + if (typeof requestAnimationFrame === "function") requestAnimationFrame(() => setTimeout(run, 0)); + else setTimeout(run, 0); +} + +app.registerExtension({ + name: "Distributed.ImageListStartValues", + nodeCreated(node) { + if (node.comfyClass !== START_CLASS) return; + const onConnectionsChange = node.onConnectionsChange; + node.onConnectionsChange = function (...args) { + const result = onConnectionsChange?.apply(this, args); + scheduleSync(this); + return result; + }; + const onConfigure = node.onConfigure; + node.onConfigure = function (...args) { + const result = onConfigure?.apply(this, args); + // A loaded graph is the source of truth for prior wiring. + this._distributedValueSources = undefined; + scheduleSync(this); + return result; + }; + scheduleSync(node); + }, +}); diff --git a/web/tests/imageListStart.test.js b/web/tests/imageListStart.test.js new file mode 100644 index 0000000..6abcd92 --- /dev/null +++ b/web/tests/imageListStart.test.js @@ -0,0 +1,174 @@ +import { readFileSync } from "node:fs"; +import vm from "node:vm"; +import { describe, expect, it } from "vitest"; + +const source = readFileSync(new URL("../imageListStart.js", import.meta.url), "utf8") + .replace(/^import .*;\r?\n/gm, ""); + +// Minimal LiteGraph-like graph: enough link bookkeeping to exercise output +// add/remove, retyping, and link moves the extension performs. +function setup() { + let extension; + const tasks = []; + const links = new Map(); + let nextLink = 1; + const nodes = new Map(); + const graph = { + links, + getLink: id => links.get(id), + getNodeById: id => nodes.get(id), + }; + const makeSource = (id, type) => { + const node = { id, outputs: [{ type, links: [] }] }; + nodes.set(id, node); + return node; + }; + const makeTarget = (id, type) => { + const node = { + id, inputs: [{ type, link: null }], + disconnectInput(slot) { + const link = links.get(this.inputs[slot].link); + if (!link) return; + const origin = nodes.get(link.origin_id); + origin.outputs[link.origin_slot].links = + origin.outputs[link.origin_slot].links.filter(l => l !== link.id); + links.delete(link.id); + this.inputs[slot].link = null; + }, + }; + nodes.set(id, node); + return node; + }; + const fixed = ["image", "mask", "item_index", "context"].map(name => ({ name, type: "IMAGE", links: [], shape: 6 })); + const start = { + id: 1, comfyClass: "DistributedImageListStart", graph, + inputs: [{ name: "images", link: null }, { name: "mask", link: null }, { name: "values.value_1", link: null }], + outputs: [...fixed, ...Array.from({ length: 10 }, (_, i) => ({ name: `value_${i + 1}`, type: "*", links: [] }))], + addOutput(name, type, extra) { this.outputs.push({ name, type, links: [], ...(extra ?? {}) }); }, + removeOutput(slot) { + if (this.outputs[slot].links.length) throw new Error("removed a wired output"); + this.outputs.splice(slot, 1); + }, + connect(slot, target, targetSlot) { + // LiteGraph refuses a link whose output type does not fit the input. + const outType = this.outputs[slot].type; + const inType = target.inputs[targetSlot].type; + if (outType !== "*" && inType !== "*" && !String(inType).split(",").includes(outType)) return null; + target.disconnectInput(targetSlot); + const id = nextLink++; + links.set(id, { id, origin_id: this.id, origin_slot: slot, target_id: target.id, + target_slot: targetSlot, type: this.outputs[slot].type }); + this.outputs[slot].links.push(id); + target.inputs[targetSlot].link = id; + return links.get(id); + }, + disconnectOutput(slot) { + for (const id of [...this.outputs[slot].links]) { + const link = links.get(id); + nodes.get(link.target_id).disconnectInput(link.target_slot); + } + }, + setDirtyCanvas() {}, + }; + nodes.set(1, start); + vm.runInNewContext(source, { + app: { registerExtension: ext => { extension = ext; } }, + setTimeout: callback => { tasks.push(callback); }, + console, + }, { filename: "imageListStart.js" }); + extension.nodeCreated(start); + const flush = () => { while (tasks.length) tasks.shift()(); }; + // Wire a source output into Start's value input, as Autogrow would. + const connectValue = (ordinal, sourceNode) => { + const name = `values.value_${ordinal}`; + let input = start.inputs.find(i => i.name === name); + if (!input) { input = { name, link: null }; start.inputs.push(input); } + const id = nextLink++; + links.set(id, { id, origin_id: sourceNode.id, origin_slot: 0, target_id: 1, + target_slot: start.inputs.indexOf(input), type: sourceNode.outputs[0].type }); + sourceNode.outputs[0].links.push(id); + input.link = id; + if (!start.inputs.some(i => i.name === `values.value_${ordinal + 1}`)) { + start.inputs.push({ name: `values.value_${ordinal + 1}`, link: null }); + } + start.onConnectionsChange(1, start.inputs.indexOf(input), true); + flush(); + }; + // Mimic Autogrow's compaction after disconnecting a middle input. + const disconnectValueAndCompact = ordinal => { + const valueInputs = start.inputs.filter(i => i.name.startsWith("values.")); + const linksBefore = valueInputs.map(i => i.link); + links.delete(linksBefore[ordinal - 1]); + const remaining = linksBefore.filter((l, i) => i !== ordinal - 1 && l != null); + start.inputs = start.inputs.filter(i => !i.name.startsWith("values.")); + remaining.forEach((l, i) => start.inputs.push({ name: `values.value_${i + 1}`, link: l })); + start.inputs.push({ name: `values.value_${remaining.length + 1}`, link: null }); + start.onConnectionsChange(1, 0, false); + flush(); + }; + flush(); + return { start, makeSource, makeTarget, connectValue, disconnectValueAndCompact, links, flush }; +} + +const valueOutputs = start => start.outputs.slice(4).map(o => `${o.name}:${o.type}`); + +describe("Distributed Image List Start value outputs", () => { + it("shows no value outputs until a value input is connected", () => { + const h = setup(); + expect(h.start.outputs.map(o => o.name)).toEqual(["image", "mask", "item_index", "context"]); + }); + + it("adds one typed output per connected value input", () => { + const h = setup(); + h.connectValue(1, h.makeSource(10, "IMAGE")); + h.connectValue(2, h.makeSource(11, "STRING")); + expect(valueOutputs(h.start)).toEqual(["value_1:IMAGE", "value_2:STRING"]); + expect(h.start.outputs[4].shape).toBe(6); + }); + + it("keeps downstream wiring on the same value when a middle input is removed", () => { + const h = setup(); + h.connectValue(1, h.makeSource(10, "IMAGE")); + h.connectValue(2, h.makeSource(11, "STRING")); + h.connectValue(3, h.makeSource(12, "INT")); + const imageUser = h.makeTarget(20, "IMAGE"); + const stringUser = h.makeTarget(21, "STRING"); + const intUser = h.makeTarget(22, "INT"); + h.start.connect(4, imageUser, 0); + h.start.connect(5, stringUser, 0); + h.start.connect(6, intUser, 0); + + h.disconnectValueAndCompact(2); + + expect(valueOutputs(h.start)).toEqual(["value_1:IMAGE", "value_2:INT"]); + expect(h.links.get(imageUser.inputs[0].link).origin_slot).toBe(4); + expect(h.links.get(intUser.inputs[0].link).origin_slot).toBe(5); + expect(stringUser.inputs[0].link).toBe(null); + }); + + it("disconnects a consumer whose type no longer matches the input", () => { + const h = setup(); + h.connectValue(1, h.makeSource(10, "IMAGE")); + const imageUser = h.makeTarget(20, "IMAGE"); + h.start.connect(4, imageUser, 0); + // Replace value_1's source with a STRING source. + const input = h.start.inputs.find(i => i.name === "values.value_1"); + h.links.delete(input.link); + input.link = null; + h.connectValue(1, h.makeSource(11, "STRING")); + expect(valueOutputs(h.start)).toEqual(["value_1:STRING"]); + expect(imageUser.inputs[0].link).toBe(null); + }); + + it("never removes an output that still has links", () => { + const h = setup(); + h.connectValue(1, h.makeSource(10, "IMAGE")); + h.start.connect(4, h.makeTarget(20, "*"), 0); + const input = h.start.inputs.find(i => i.name === "values.value_1"); + h.links.delete(input.link); + input.link = null; + h.start.onConnectionsChange(1, 2, false); + h.flush(); + expect(valueOutputs(h.start)).toEqual(["value_1:*"]); + }); +});