robertvoy-ComfyUI-Distributed/api/orchestration/prompt_transform.py

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