304 lines
11 KiB
Python
304 lines
11 KiB
Python
"""Saved-workflow traversal and queued-prompt node identifier mapping."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
import logging
|
|
from typing import Any, Iterator, Mapping
|
|
|
|
if __package__:
|
|
from .settings import ModalSyncSettings, get_settings
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from settings import ModalSyncSettings, get_settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _looks_like_workflow_node(fragment: dict[str, Any]) -> bool:
|
|
"""Return whether a JSON fragment resembles a saved ComfyUI workflow node."""
|
|
return "id" in fragment and "properties" in fragment
|
|
|
|
|
|
def _iter_workflow_nodes(
|
|
workflow_fragment: Any,
|
|
visited_object_ids: set[int] | None = None,
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""Yield workflow node dictionaries from a nested saved workflow fragment."""
|
|
if visited_object_ids is None:
|
|
visited_object_ids = set()
|
|
|
|
if isinstance(workflow_fragment, dict):
|
|
object_id = id(workflow_fragment)
|
|
if object_id in visited_object_ids:
|
|
return
|
|
visited_object_ids.add(object_id)
|
|
|
|
if _looks_like_workflow_node(workflow_fragment):
|
|
yield workflow_fragment
|
|
|
|
for value in workflow_fragment.values():
|
|
yield from _iter_workflow_nodes(value, visited_object_ids)
|
|
return
|
|
|
|
if isinstance(workflow_fragment, list):
|
|
object_id = id(workflow_fragment)
|
|
if object_id in visited_object_ids:
|
|
return
|
|
visited_object_ids.add(object_id)
|
|
|
|
for item in workflow_fragment:
|
|
yield from _iter_workflow_nodes(item, visited_object_ids)
|
|
|
|
|
|
def _workflow_subgraph_definitions(
|
|
workflow_fragment: Any,
|
|
) -> dict[str, dict[str, Any]]:
|
|
"""Return reusable ComfyUI subgraph definitions keyed by their node type id."""
|
|
if not isinstance(workflow_fragment, dict):
|
|
return {}
|
|
definitions = workflow_fragment.get("definitions")
|
|
if not isinstance(definitions, dict):
|
|
return {}
|
|
subgraphs = definitions.get("subgraphs")
|
|
if not isinstance(subgraphs, list):
|
|
return {}
|
|
|
|
return {
|
|
str(subgraph["id"]): subgraph
|
|
for subgraph in subgraphs
|
|
if isinstance(subgraph, dict) and subgraph.get("id") is not None
|
|
}
|
|
|
|
|
|
def _iter_workflow_nodes_with_ancestors(
|
|
workflow_fragment: Any,
|
|
ancestor_node_ids: tuple[str, ...] = (),
|
|
visited_fragments: set[tuple[int, tuple[str, ...], tuple[str, ...]]] | None = None,
|
|
subgraph_definitions: Mapping[str, dict[str, Any]] | None = None,
|
|
active_subgraph_definition_ids: tuple[str, ...] = (),
|
|
) -> Iterator[tuple[dict[str, Any], tuple[str, ...]]]:
|
|
"""Yield workflow nodes with instance paths across embedded and defined subgraphs."""
|
|
if visited_fragments is None:
|
|
visited_fragments = set()
|
|
if subgraph_definitions is None:
|
|
subgraph_definitions = _workflow_subgraph_definitions(workflow_fragment)
|
|
|
|
if isinstance(workflow_fragment, dict):
|
|
fragment_identity = (
|
|
id(workflow_fragment),
|
|
ancestor_node_ids,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
if fragment_identity in visited_fragments:
|
|
return
|
|
visited_fragments.add(fragment_identity)
|
|
|
|
next_ancestor_node_ids = ancestor_node_ids
|
|
if _looks_like_workflow_node(workflow_fragment):
|
|
node_id = str(workflow_fragment.get("id"))
|
|
yield workflow_fragment, ancestor_node_ids
|
|
next_ancestor_node_ids = ancestor_node_ids + (node_id,)
|
|
|
|
subgraph_definition_id = str(workflow_fragment.get("type") or "")
|
|
subgraph_definition = subgraph_definitions.get(subgraph_definition_id)
|
|
if (
|
|
subgraph_definition is not None
|
|
and subgraph_definition_id not in active_subgraph_definition_ids
|
|
):
|
|
yield from _iter_workflow_nodes_with_ancestors(
|
|
subgraph_definition.get("nodes", []),
|
|
next_ancestor_node_ids,
|
|
visited_fragments,
|
|
subgraph_definitions,
|
|
active_subgraph_definition_ids + (subgraph_definition_id,),
|
|
)
|
|
|
|
for key, value in workflow_fragment.items():
|
|
if key == "definitions":
|
|
continue
|
|
yield from _iter_workflow_nodes_with_ancestors(
|
|
value,
|
|
next_ancestor_node_ids,
|
|
visited_fragments,
|
|
subgraph_definitions,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
return
|
|
|
|
if isinstance(workflow_fragment, list):
|
|
fragment_identity = (
|
|
id(workflow_fragment),
|
|
ancestor_node_ids,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
if fragment_identity in visited_fragments:
|
|
return
|
|
visited_fragments.add(fragment_identity)
|
|
|
|
for item in workflow_fragment:
|
|
yield from _iter_workflow_nodes_with_ancestors(
|
|
item,
|
|
ancestor_node_ids,
|
|
visited_fragments,
|
|
subgraph_definitions,
|
|
active_subgraph_definition_ids,
|
|
)
|
|
|
|
|
|
def _resolve_prompt_node_ids_for_workflow_node(
|
|
workflow_node_id: str,
|
|
ancestor_node_ids: tuple[str, ...],
|
|
prompt_node_ids: set[str],
|
|
) -> set[str]:
|
|
"""Resolve one saved workflow node id to matching queued prompt node ids."""
|
|
path_segments = ancestor_node_ids + (workflow_node_id,)
|
|
|
|
for index in range(len(path_segments)):
|
|
candidate = ":".join(path_segments[index:])
|
|
resolved_prompt_node_ids: set[str] = set()
|
|
if candidate in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(candidate)
|
|
descendant_prefix = f"{candidate}:"
|
|
resolved_prompt_node_ids.update(
|
|
prompt_node_id
|
|
for prompt_node_id in prompt_node_ids
|
|
if prompt_node_id.startswith(descendant_prefix)
|
|
)
|
|
if resolved_prompt_node_ids:
|
|
return resolved_prompt_node_ids
|
|
|
|
resolved_prompt_node_ids: set[str] = set()
|
|
|
|
for ancestor_node_id in reversed(ancestor_node_ids):
|
|
if ancestor_node_id in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(ancestor_node_id)
|
|
descendant_prefix = f"{ancestor_node_id}:"
|
|
resolved_prompt_node_ids.update(
|
|
prompt_node_id
|
|
for prompt_node_id in prompt_node_ids
|
|
if prompt_node_id.startswith(descendant_prefix)
|
|
)
|
|
break
|
|
|
|
return resolved_prompt_node_ids
|
|
|
|
|
|
def _workflow_node_path(
|
|
workflow_node_id: str, ancestor_node_ids: tuple[str, ...]
|
|
) -> str:
|
|
"""Return one workflow node's composed path, including subgraph ancestors."""
|
|
if not ancestor_node_ids:
|
|
return workflow_node_id
|
|
return ":".join((*ancestor_node_ids, workflow_node_id))
|
|
|
|
|
|
def _extract_marked_workflow_node_paths(
|
|
workflow: dict[str, Any] | None,
|
|
settings: ModalSyncSettings | None = None,
|
|
) -> set[str]:
|
|
"""Return composed workflow paths for nodes explicitly marked remote in metadata."""
|
|
if workflow is None:
|
|
return set()
|
|
|
|
marker = (settings or get_settings()).marker_property
|
|
marked_workflow_node_paths: set[str] = set()
|
|
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
|
properties = node.get("properties") or {}
|
|
if not properties.get(marker):
|
|
continue
|
|
marked_workflow_node_paths.add(
|
|
_workflow_node_path(str(node.get("id")), ancestor_node_ids)
|
|
)
|
|
return marked_workflow_node_paths
|
|
|
|
|
|
def _build_workflow_prompt_resolution_maps(
|
|
workflow: dict[str, Any] | None,
|
|
prompt_node_ids: set[str],
|
|
) -> tuple[dict[str, set[str]], dict[str, set[str]]]:
|
|
"""Return bidirectional mappings between workflow paths and queued prompt ids."""
|
|
workflow_path_to_prompt_node_ids: dict[str, set[str]] = {}
|
|
prompt_node_id_to_workflow_paths: dict[str, set[str]] = defaultdict(set)
|
|
|
|
if workflow is None:
|
|
for prompt_node_id in prompt_node_ids:
|
|
workflow_path_to_prompt_node_ids[prompt_node_id] = {prompt_node_id}
|
|
prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id)
|
|
return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths
|
|
|
|
for node, ancestor_node_ids in _iter_workflow_nodes_with_ancestors(workflow):
|
|
workflow_node_id = str(node.get("id"))
|
|
workflow_node_path = _workflow_node_path(workflow_node_id, ancestor_node_ids)
|
|
resolved_prompt_node_ids = _resolve_prompt_node_ids_for_workflow_node(
|
|
workflow_node_id,
|
|
ancestor_node_ids,
|
|
prompt_node_ids,
|
|
)
|
|
if not ancestor_node_ids and workflow_node_id in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(workflow_node_id)
|
|
if workflow_node_path in prompt_node_ids:
|
|
resolved_prompt_node_ids.add(workflow_node_path)
|
|
if not resolved_prompt_node_ids:
|
|
continue
|
|
|
|
workflow_path_to_prompt_node_ids[workflow_node_path] = resolved_prompt_node_ids
|
|
for prompt_node_id in resolved_prompt_node_ids:
|
|
prompt_node_id_to_workflow_paths[prompt_node_id].add(workflow_node_path)
|
|
|
|
for prompt_node_id in prompt_node_ids:
|
|
workflow_path_to_prompt_node_ids.setdefault(prompt_node_id, {prompt_node_id})
|
|
prompt_node_id_to_workflow_paths[prompt_node_id].add(prompt_node_id)
|
|
|
|
return workflow_path_to_prompt_node_ids, prompt_node_id_to_workflow_paths
|
|
|
|
|
|
def _resolve_requested_prompt_node_ids(
|
|
requested_workflow_node_paths: set[str],
|
|
prompt_node_ids: set[str],
|
|
workflow_path_to_prompt_node_ids: dict[str, set[str]],
|
|
) -> set[str]:
|
|
"""Resolve requested workflow-node paths to the prompt ids that queue-time rewrite sees."""
|
|
requested_prompt_node_ids: set[str] = set()
|
|
for requested_workflow_node_path in requested_workflow_node_paths:
|
|
if requested_workflow_node_path in workflow_path_to_prompt_node_ids:
|
|
requested_prompt_node_ids.update(
|
|
workflow_path_to_prompt_node_ids[requested_workflow_node_path]
|
|
)
|
|
continue
|
|
if requested_workflow_node_path in prompt_node_ids:
|
|
requested_prompt_node_ids.add(requested_workflow_node_path)
|
|
return requested_prompt_node_ids
|
|
|
|
|
|
def _best_workflow_path_for_prompt_node(
|
|
prompt_node_id: str,
|
|
prompt_node_id_to_workflow_paths: dict[str, set[str]],
|
|
) -> str:
|
|
"""Choose the most specific workflow path for one queued prompt node id."""
|
|
workflow_node_paths = prompt_node_id_to_workflow_paths.get(
|
|
prompt_node_id, {prompt_node_id}
|
|
)
|
|
return max(
|
|
workflow_node_paths,
|
|
key=lambda workflow_node_path: (
|
|
workflow_node_path.count(":"),
|
|
len(workflow_node_path),
|
|
workflow_node_path,
|
|
),
|
|
)
|
|
|
|
|
|
def _resolve_workflow_node_paths_for_prompt_nodes(
|
|
prompt_node_ids: set[str],
|
|
prompt_node_id_to_workflow_paths: dict[str, set[str]],
|
|
) -> set[str]:
|
|
"""Map queued prompt ids back to the workflow node paths the UI can mark remote."""
|
|
return {
|
|
_best_workflow_path_for_prompt_node(
|
|
prompt_node_id, prompt_node_id_to_workflow_paths
|
|
)
|
|
for prompt_node_id in prompt_node_ids
|
|
}
|
|
|
|
|