feat: preserve required output utilities in image-list execution

api/image_list_routes.py

@@ -35,6 +35,7 @@ async def _json_body(request):
 async def image_list_capabilities_endpoint(_request):
     return web.json_response({
         "supported": True,
+        "execution_targets": True,
         "protocol_version": PROTOCOL_VERSION,
         "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
         "transport": "typed-float-tensor-json-v1",

api/orchestration/dispatch.py

@@ -77,6 +77,8 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
         "workflow": payload.get("workflow"),
         "client_id": client_id,
     }
+    if "partial_execution_targets" in payload:
+        ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
     ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
     ws_url = f"{ws_url}/distributed/worker_ws"
     session = await get_client_session()
@@ -128,11 +130,17 @@ async def dispatch_worker_prompt(
     use_websocket=False,
     trace_execution_id=None,
     return_prompt_id=False,
+    *,
+    partial_execution_targets=None,
 ):
     """Send the prepared prompt to a worker ComfyUI instance."""
     worker_url = build_worker_url(worker)
     url = build_worker_url(worker, "/prompt")
     payload = {"prompt": prompt_obj}
+    target_payload = {}
+    if partial_execution_targets is not None:
+        target_payload["partial_execution_targets"] = list(partial_execution_targets)
+        payload.update(target_payload)
     extra_data = {}
     if workflow_meta:
         extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
@@ -146,6 +154,7 @@ async def dispatch_worker_prompt(
                 {
                     "prompt": prompt_obj,
                     "workflow": workflow_meta,
+                    **target_payload,
                 },
                 client_id,
             )
@@ -227,6 +236,12 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
                         f"{sorted(required_nodes)}; no work was submitted"
                     )
 
+                if capability.get("execution_targets") is not True:
+                    raise RuntimeError(
+                        f"Image-list worker {worker_id} lacks explicit execution_targets support; "
+                        "upgrade that worker before running this feature; no work was submitted"
+                    )
+
                 check = "queue"
                 queue_url = build_worker_url(worker, "/queue")
                 async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:

api/queue_orchestration.py

@@ -38,6 +38,7 @@ from .orchestration.prompt_transform import (
     _resolve_direct_workers_for_image_list,
 )
 from ..image_list.runtime import validate_image_list_regions
+from ..image_list.execution_plan import prepare_image_list_execution
 
 
 prompt_server = server.PromptServer.instance
@@ -265,13 +266,17 @@ 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
+    image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart"
+                          for node in prompt_obj.values())
+    target_kwargs = {}
+    if image_list_mode:
+        prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
+        target_kwargs["partial_execution_targets"] = outer_targets
     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)
+        validate_image_list_regions(prompt_obj, explicit_targets=True)
         other_distributed = (
             prompt_index.nodes_for_class("DistributedCollector")
             + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
@@ -409,6 +414,7 @@ async def orchestrate_distributed_execution(
             workflow_meta,
             client_id,
             include_queue_metadata=True,
+            **target_kwargs,
         )
         return (
             queue_result["prompt_id"],
@@ -506,6 +512,7 @@ async def orchestrate_distributed_execution(
             workflow_meta,
             client_id,
             include_queue_metadata=True,
+            **target_kwargs,
         )
     except Exception:
         if serverless_jobs:

api/worker_routes.py

@@ -61,6 +61,13 @@ async def worker_ws_endpoint(request):
                 })
                 continue
 
+            if not isinstance(data, dict):
+                await ws.send_json({
+                    "type": "dispatch_ack", "request_id": None, "ok": False,
+                    "error": "Expected a websocket message object.",
+                })
+                continue
+
             if data.get("type") != "dispatch_prompt":
                 await ws.send_json({
                     "type": "dispatch_ack",
@@ -81,10 +88,16 @@ async def worker_ws_endpoint(request):
                 continue
 
             try:
+                target_kwargs = {}
+                if "partial_execution_targets" in data:
+                    from ..utils.async_helpers import validate_execution_targets
+                    target_kwargs["partial_execution_targets"] = validate_execution_targets(
+                        prompt, data["partial_execution_targets"])
                 prompt_id = await queue_prompt_payload(
                     prompt,
                     workflow_meta=data.get("workflow"),
                     client_id=data.get("client_id"),
+                    **target_kwargs,
                 )
                 await ws.send_json({
                     "type": "dispatch_ack",