Robert Wojciechowski
a9ef36f951
Use the existing images/audio Collector for typed serverless pools. Preserve saved legacy_png defaults, map them to serverless png, and retain every transfer format with exact opt-in. Remove the separate Media Collector. Send direct exact frames through bounded raw tensor multipart transport without 8-bit fallback. Reject incompatible legacy pools, invalid payloads and incomplete exact results. Cover real CPU HTTP round trips, native ComfyUI execution, mixed participant ordering and failure paths.
668 lines
24 KiB
Python
668 lines
24 KiB
Python
import json
|
|
from collections import deque
|
|
|
|
from ...utils.logging import debug_log
|
|
|
|
|
|
class PromptIndex:
|
|
"""Cache prompt metadata for faster worker/master prompt preparation."""
|
|
|
|
def __init__(self, prompt_obj):
|
|
self._prompt_json = json.dumps(prompt_obj)
|
|
self.nodes_by_class = {}
|
|
self.class_by_node = {}
|
|
self.inputs_by_node = {}
|
|
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
|
class_type = node.get("class_type")
|
|
node_id_str = str(node_id)
|
|
if class_type:
|
|
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
|
|
self.class_by_node[node_id_str] = class_type
|
|
self.inputs_by_node[node_id_str] = node.get("inputs", {})
|
|
self._upstream_cache = {}
|
|
|
|
def copy_prompt(self):
|
|
return json.loads(self._prompt_json)
|
|
|
|
def nodes_for_class(self, class_name):
|
|
return self.nodes_by_class.get(class_name, [])
|
|
|
|
def has_upstream(self, start_node_id, target_class):
|
|
cache_key = (str(start_node_id), target_class)
|
|
if cache_key in self._upstream_cache:
|
|
return self._upstream_cache[cache_key]
|
|
|
|
visited = set()
|
|
stack = [str(start_node_id)]
|
|
while stack:
|
|
node_id = stack.pop()
|
|
if node_id in visited:
|
|
continue
|
|
visited.add(node_id)
|
|
inputs = self.inputs_by_node.get(node_id, {})
|
|
for value in inputs.values():
|
|
if isinstance(value, list) and len(value) == 2:
|
|
upstream_id = str(value[0])
|
|
if self.class_by_node.get(upstream_id) == target_class:
|
|
self._upstream_cache[cache_key] = True
|
|
return True
|
|
if upstream_id in self.inputs_by_node:
|
|
stack.append(upstream_id)
|
|
|
|
self._upstream_cache[cache_key] = False
|
|
return False
|
|
|
|
|
|
def _iter_prompt_nodes(prompt_obj):
|
|
for node_id, node in prompt_obj.items():
|
|
if isinstance(node, dict):
|
|
yield str(node_id), node
|
|
|
|
|
|
def find_nodes_by_class(prompt_obj, class_name):
|
|
nodes = []
|
|
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
|
if node.get("class_type") == class_name:
|
|
nodes.append(node_id)
|
|
return nodes
|
|
|
|
|
|
def _find_downstream_nodes(prompt_obj, start_ids):
|
|
"""Return all nodes reachable downstream from the provided IDs."""
|
|
adjacency = {}
|
|
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
|
inputs = node.get("inputs", {})
|
|
for value in inputs.values():
|
|
if isinstance(value, list) and len(value) == 2:
|
|
source_id = str(value[0])
|
|
adjacency.setdefault(source_id, set()).add(str(node_id))
|
|
|
|
connected = set(start_ids)
|
|
queue = deque(start_ids)
|
|
while queue:
|
|
current = queue.popleft()
|
|
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
|
|
if dependent not in connected:
|
|
connected.add(dependent)
|
|
queue.append(dependent)
|
|
return connected
|
|
|
|
|
|
def _create_numeric_id_generator(prompt_obj):
|
|
"""Return a closure that yields new numeric string IDs."""
|
|
max_id = 0
|
|
for node_id in prompt_obj.keys():
|
|
try:
|
|
numeric = int(node_id)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
max_id = max(max_id, numeric)
|
|
|
|
counter = max_id
|
|
|
|
def _next_id():
|
|
nonlocal counter
|
|
counter += 1
|
|
return str(counter)
|
|
|
|
return _next_id
|
|
|
|
|
|
def _find_upstream_nodes(prompt_obj, start_ids):
|
|
"""Return all nodes reachable upstream from start_ids, including start nodes."""
|
|
connected = set(str(node_id) for node_id in start_ids)
|
|
queue = deque(connected)
|
|
while queue:
|
|
node_id = queue.popleft()
|
|
node = prompt_obj.get(node_id) or {}
|
|
inputs = node.get("inputs", {})
|
|
for value in inputs.values():
|
|
if isinstance(value, list) and len(value) == 2:
|
|
source_id = str(value[0])
|
|
if source_id in prompt_obj and source_id not in connected:
|
|
connected.add(source_id)
|
|
queue.append(source_id)
|
|
return connected
|
|
|
|
|
|
_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
|
|
"PrimitiveBoolean",
|
|
"PrimitiveFloat",
|
|
"PrimitiveInt",
|
|
"PrimitiveNode",
|
|
"PrimitiveString",
|
|
}
|
|
|
|
_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
|
|
"LoadImage",
|
|
}
|
|
|
|
_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
|
|
_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
|
|
|
|
# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
|
|
# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
|
|
# only after its connected inputs recursively prove safe.
|
|
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
|
|
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
|
|
"CreateList": ("inputs.",),
|
|
}
|
|
|
|
# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
|
|
_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
|
|
|
|
|
|
def _get_delegate_master_node_class_mappings():
|
|
"""Return ComfyUI node-class mappings when available."""
|
|
if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
|
|
return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
|
|
try:
|
|
import nodes as comfy_nodes # type: ignore
|
|
except Exception: # pragma: no cover - depends on ComfyUI runtime imports
|
|
return {}
|
|
return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
|
|
|
|
|
|
def _get_delegate_master_node_class(class_type):
|
|
mappings = _get_delegate_master_node_class_mappings()
|
|
return mappings.get(class_type) if isinstance(mappings, dict) else None
|
|
|
|
|
|
def _normalize_delegate_master_return_type(return_type):
|
|
if return_type is None:
|
|
return ""
|
|
return str(return_type).strip().upper()
|
|
|
|
|
|
def _delegate_master_type_is_safe_scalar(type_name):
|
|
return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
|
|
|
|
|
|
def _delegate_master_type_is_safe_config(type_name):
|
|
return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
|
|
|
|
|
|
def _delegate_master_output_is_safe_scalar(class_type, output_index):
|
|
"""Return True when a registered node output is lightweight config data."""
|
|
if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
|
|
return True
|
|
node_class = _get_delegate_master_node_class(class_type)
|
|
return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
|
|
try:
|
|
output_type = return_types[int(output_index)]
|
|
except (IndexError, TypeError, ValueError):
|
|
return False
|
|
return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
|
|
|
|
|
|
def _get_delegate_master_input_types(class_type):
|
|
node_class = _get_delegate_master_node_class(class_type)
|
|
input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
|
|
if callable(input_types):
|
|
try:
|
|
input_types = input_types()
|
|
except TypeError:
|
|
return {}
|
|
return input_types if isinstance(input_types, dict) else {}
|
|
|
|
|
|
def _normalize_delegate_master_input_type(input_spec):
|
|
if isinstance(input_spec, (list, tuple)) and input_spec:
|
|
return _normalize_delegate_master_return_type(input_spec[0])
|
|
return _normalize_delegate_master_return_type(input_spec)
|
|
|
|
|
|
def _delegate_master_input_is_safe_scalar(class_type, input_name):
|
|
"""Return True when a registered downstream input expects config data."""
|
|
for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
|
|
if input_name.startswith(prefix):
|
|
return True
|
|
input_types = _get_delegate_master_input_types(class_type)
|
|
for section_name in ("required", "optional"):
|
|
section = input_types.get(section_name, {})
|
|
if isinstance(section, dict) and input_name in section:
|
|
input_type = _normalize_delegate_master_input_type(section[input_name])
|
|
return _delegate_master_type_is_safe_config(input_type)
|
|
return False
|
|
|
|
|
|
def _is_delegate_master_always_retained_upstream_node(node):
|
|
if not isinstance(node, dict):
|
|
return False
|
|
class_type = node.get("class_type")
|
|
return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
|
|
|
|
|
|
def _is_delegate_master_retained_upstream_node(node, output_index=0):
|
|
"""Return True for lightweight upstream nodes safe to keep on the master."""
|
|
if not isinstance(node, dict):
|
|
return False
|
|
class_type = node.get("class_type")
|
|
if not isinstance(class_type, str):
|
|
return False
|
|
return (
|
|
class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
|
|
or class_type.startswith("Primitive")
|
|
or _delegate_master_output_is_safe_scalar(class_type, output_index)
|
|
)
|
|
|
|
|
|
def _collect_delegate_master_retained_upstream_branch(
|
|
prompt_obj,
|
|
node_id,
|
|
output_index,
|
|
memo,
|
|
visiting,
|
|
):
|
|
"""Return safe retained branch nodes, or None when the branch is not safe."""
|
|
node_id = str(node_id)
|
|
cache_key = (node_id, output_index)
|
|
if cache_key in memo:
|
|
cached = memo[cache_key]
|
|
return None if cached is None else set(cached)
|
|
if cache_key in visiting:
|
|
memo[cache_key] = None
|
|
return None
|
|
|
|
node = prompt_obj.get(node_id)
|
|
if not _is_delegate_master_retained_upstream_node(node, output_index):
|
|
memo[cache_key] = None
|
|
return None
|
|
|
|
visiting.add(cache_key)
|
|
retained = {node_id}
|
|
inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
|
|
class_type = node.get("class_type") if isinstance(node, dict) else None
|
|
for input_name, value in inputs.items():
|
|
if not (isinstance(value, list) and len(value) == 2):
|
|
continue
|
|
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
|
|
visiting.remove(cache_key)
|
|
memo[cache_key] = None
|
|
return None
|
|
source_id = str(value[0])
|
|
branch = _collect_delegate_master_retained_upstream_branch(
|
|
prompt_obj,
|
|
source_id,
|
|
value[1],
|
|
memo,
|
|
visiting,
|
|
)
|
|
if branch is None:
|
|
visiting.remove(cache_key)
|
|
memo[cache_key] = None
|
|
return None
|
|
retained.update(branch)
|
|
|
|
visiting.remove(cache_key)
|
|
memo[cache_key] = frozenset(retained)
|
|
return retained
|
|
|
|
|
|
def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
|
|
"""Return lightweight upstream nodes needed by kept delegate-master nodes."""
|
|
connected = set()
|
|
memo = {}
|
|
for node_id in start_ids:
|
|
node = prompt_obj.get(str(node_id)) or {}
|
|
inputs = node.get("inputs", {})
|
|
class_type = node.get("class_type") if isinstance(node, dict) else None
|
|
for input_name, value in inputs.items():
|
|
if not (isinstance(value, list) and len(value) == 2):
|
|
continue
|
|
source_node = prompt_obj.get(str(value[0]))
|
|
if _is_delegate_master_always_retained_upstream_node(source_node):
|
|
connected.add(str(value[0]))
|
|
continue
|
|
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
|
|
continue
|
|
branch = _collect_delegate_master_retained_upstream_branch(
|
|
prompt_obj,
|
|
value[0],
|
|
value[1],
|
|
memo,
|
|
set(),
|
|
)
|
|
if branch is not None:
|
|
connected.update(branch)
|
|
return connected
|
|
|
|
|
|
WORKER_SINK_CLASS = "DistributedWorkerSink"
|
|
|
|
|
|
def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
|
|
"""Prune worker prompt to distributed nodes and their upstream dependencies.
|
|
|
|
Removing downstream nodes can leave the worker prompt without an output
|
|
node. One is re-added after each such distributed node: the hidden
|
|
``DistributedWorkerSink`` when ``use_worker_sink`` is set (the worker must
|
|
have it registered), otherwise PreviewImage/PreviewAudio for workers
|
|
running older versions. The preview nodes write every frame to the
|
|
worker's temp directory, which the sink avoids.
|
|
"""
|
|
collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
|
|
upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
|
|
distributed_ids = collector_ids + upscale_ids
|
|
if not distributed_ids:
|
|
return prompt_obj
|
|
|
|
connected = _find_upstream_nodes(prompt_obj, distributed_ids)
|
|
pruned_prompt = {}
|
|
for node_id in connected:
|
|
node = prompt_obj.get(node_id)
|
|
if node is not None:
|
|
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
|
|
|
# Generate IDs from the original prompt so we never reuse IDs from pruned downstream nodes.
|
|
next_id = _create_numeric_id_generator(prompt_obj)
|
|
for dist_id in distributed_ids:
|
|
if dist_id not in pruned_prompt:
|
|
continue
|
|
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
|
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
|
|
if has_removed_downstream:
|
|
original_node = prompt_obj.get(str(dist_id), {})
|
|
class_type = original_node.get("class_type")
|
|
inputs = original_node.get("inputs", {})
|
|
image_connected = class_type != "DistributedCollector" or (
|
|
isinstance(inputs.get("images"), list)
|
|
and len(inputs["images"]) == 2
|
|
)
|
|
audio_connected = (
|
|
class_type == "DistributedCollector"
|
|
and isinstance(inputs.get("audio"), list)
|
|
and len(inputs["audio"]) == 2
|
|
)
|
|
|
|
if use_worker_sink and (image_connected or audio_connected):
|
|
pruned_prompt[next_id()] = {
|
|
"inputs": {"value": [dist_id, 0 if image_connected else 1]},
|
|
"class_type": WORKER_SINK_CLASS,
|
|
"_meta": {"title": "Distributed Worker Sink (auto-added)"},
|
|
}
|
|
elif image_connected:
|
|
preview_id = next_id()
|
|
pruned_prompt[preview_id] = {
|
|
"inputs": {"images": [dist_id, 0]},
|
|
"class_type": "PreviewImage",
|
|
"_meta": {"title": "Preview Image (auto-added)"},
|
|
}
|
|
elif audio_connected:
|
|
preview_id = next_id()
|
|
pruned_prompt[preview_id] = {
|
|
"inputs": {"audio": [dist_id, 1]},
|
|
"class_type": "PreviewAudio",
|
|
"_meta": {"title": "Preview Audio (auto-added)"},
|
|
}
|
|
|
|
return pruned_prompt
|
|
|
|
|
|
def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
|
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
|
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
|
nodes_to_keep = set(collector_ids)
|
|
nodes_to_keep.update(downstream)
|
|
nodes_to_keep.update(
|
|
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
|
|
)
|
|
|
|
pruned_prompt = {}
|
|
for node_id in nodes_to_keep:
|
|
node = prompt_obj.get(node_id)
|
|
if node is not None:
|
|
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
|
|
|
pruned_ids = set(pruned_prompt.keys())
|
|
for node_id, node in pruned_prompt.items():
|
|
inputs = node.get("inputs")
|
|
if not inputs:
|
|
continue
|
|
for input_name, input_value in list(inputs.items()):
|
|
if isinstance(input_value, list) and len(input_value) == 2:
|
|
source_id = str(input_value[0])
|
|
if source_id not in pruned_ids:
|
|
inputs.pop(input_name, None)
|
|
debug_log(
|
|
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
|
|
)
|
|
|
|
# Generate IDs from the original prompt to avoid ID collisions with pruned nodes.
|
|
next_id = _create_numeric_id_generator(prompt_obj)
|
|
for collector_id in collector_ids:
|
|
collector_entry = pruned_prompt.get(collector_id)
|
|
if not collector_entry:
|
|
continue
|
|
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
|
|
original_images = original_inputs.get("images")
|
|
if not (isinstance(original_images, list) and len(original_images) == 2):
|
|
continue
|
|
placeholder_id = next_id()
|
|
pruned_prompt[placeholder_id] = {
|
|
"class_type": "DistributedEmptyImage",
|
|
"inputs": {
|
|
"height": 64,
|
|
"width": 64,
|
|
"channels": 3,
|
|
},
|
|
"_meta": {
|
|
"title": "Distributed Empty Image (auto-added)",
|
|
},
|
|
}
|
|
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
|
|
debug_log(
|
|
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
|
|
)
|
|
|
|
return pruned_prompt
|
|
|
|
|
|
def _resolve_direct_workers_for_image_list(config, requested_ids):
|
|
"""Resolve enabled workers by default, or exact explicitly selected direct workers."""
|
|
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
|
|
configured_workers = (config or {}).get("workers", [])
|
|
configured_ids = {
|
|
str(worker.get("id") or "").strip()
|
|
for worker in configured_workers
|
|
if str(worker.get("id") or "").strip()
|
|
}
|
|
if requested is not None:
|
|
unknown_ids = requested - configured_ids
|
|
if unknown_ids:
|
|
raise RuntimeError(
|
|
"Unknown explicitly selected image-list worker ID(s): "
|
|
+ ", ".join(sorted(unknown_ids))
|
|
)
|
|
direct = []
|
|
for worker in configured_workers:
|
|
worker_id = str(worker.get("id") or "").strip()
|
|
if not worker_id:
|
|
continue
|
|
if requested is None:
|
|
if not worker.get("enabled", False):
|
|
continue
|
|
elif 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")
|
|
+ prompt_index.nodes_for_class("DistributedImageListStart")
|
|
)
|
|
for node_id in distributed_nodes:
|
|
job_map[node_id] = f"{prefix}_{node_id}"
|
|
return job_map
|
|
|
|
|
|
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
|
"""Configure DistributedSeed nodes for master or worker role."""
|
|
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
|
|
node = prompt_copy.get(node_id)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.setdefault("inputs", {})
|
|
inputs["is_worker"] = not is_master
|
|
if is_master:
|
|
inputs["worker_id"] = ""
|
|
else:
|
|
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
|
|
|
|
|
def _override_collector_nodes(
|
|
prompt_copy,
|
|
prompt_index,
|
|
is_master,
|
|
participant_id,
|
|
job_id_map,
|
|
master_url,
|
|
enabled_json,
|
|
delegate_master,
|
|
):
|
|
"""Configure DistributedCollector nodes for master or worker role."""
|
|
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
|
node = prompt_copy.get(node_id)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
|
|
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
|
|
node.setdefault("inputs", {})["pass_through"] = True
|
|
continue
|
|
|
|
inputs = node.setdefault("inputs", {})
|
|
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
|
inputs["is_worker"] = not is_master
|
|
inputs["enabled_worker_ids"] = enabled_json
|
|
if is_master:
|
|
inputs["delegate_only"] = bool(delegate_master)
|
|
inputs.pop("master_url", None)
|
|
inputs.pop("worker_id", None)
|
|
else:
|
|
inputs["master_url"] = master_url
|
|
inputs["worker_id"] = participant_id
|
|
inputs["delegate_only"] = False
|
|
|
|
|
|
def _override_upscale_nodes(
|
|
prompt_copy,
|
|
prompt_index,
|
|
is_master,
|
|
participant_id,
|
|
job_id_map,
|
|
master_url,
|
|
enabled_json,
|
|
):
|
|
"""Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
|
|
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
|
|
node = prompt_copy.get(node_id)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.setdefault("inputs", {})
|
|
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
|
inputs["is_worker"] = not is_master
|
|
inputs["enabled_worker_ids"] = enabled_json
|
|
if is_master:
|
|
inputs.pop("master_url", None)
|
|
inputs.pop("worker_id", None)
|
|
else:
|
|
inputs["master_url"] = master_url
|
|
inputs["worker_id"] = participant_id
|
|
|
|
|
|
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
|
"""Configure DistributedValue nodes for master or worker role."""
|
|
for node_id in prompt_index.nodes_for_class("DistributedValue"):
|
|
node = prompt_copy.get(node_id)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.setdefault("inputs", {})
|
|
inputs["is_worker"] = not is_master
|
|
if is_master:
|
|
inputs["worker_id"] = ""
|
|
else:
|
|
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, delegate_master, 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_participates"] = participant_id == "master" and not delegate_master
|
|
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,
|
|
enabled_worker_ids,
|
|
job_id_map,
|
|
master_url,
|
|
delegate_master,
|
|
prompt_index,
|
|
workflow_meta=None,
|
|
):
|
|
"""Return a prompt copy with hidden inputs configured for master/worker."""
|
|
is_master = participant_id == "master"
|
|
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
|
|
enabled_json = json.dumps(enabled_worker_ids)
|
|
|
|
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
|
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
|
_override_collector_nodes(
|
|
prompt_copy,
|
|
prompt_index,
|
|
is_master,
|
|
participant_id,
|
|
job_id_map,
|
|
master_url,
|
|
enabled_json,
|
|
delegate_master,
|
|
)
|
|
_override_upscale_nodes(
|
|
prompt_copy,
|
|
prompt_index,
|
|
is_master,
|
|
participant_id,
|
|
job_id_map,
|
|
master_url,
|
|
enabled_json,
|
|
)
|
|
_override_image_list_nodes(
|
|
prompt_copy,
|
|
prompt_index,
|
|
participant_id,
|
|
enabled_worker_ids,
|
|
job_id_map,
|
|
master_url,
|
|
delegate_master,
|
|
workflow_meta=workflow_meta,
|
|
)
|
|
|
|
return prompt_copy
|