Delete comfy_execution

comfy_execution/caching.py

@@ -1,472 +0,0 @@
-import itertools
-from typing import Sequence, Mapping, Dict
-from comfy_execution.graph import DynamicPrompt
-from abc import ABC, abstractmethod
-
-import nodes
-
-from comfy_execution.graph_utils import is_link
-
-NODE_CLASS_CONTAINS_UNIQUE_ID: Dict[str, bool] = {}
-
-
-def include_unique_id_in_input(class_type: str) -> bool:
-    if class_type in NODE_CLASS_CONTAINS_UNIQUE_ID:
-        return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type]
-    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-    NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] = "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values()
-    return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type]
-
-class CacheKeySet(ABC):
-    def __init__(self, dynprompt, node_ids, is_changed_cache):
-        self.keys = {}
-        self.subcache_keys = {}
-
-    @abstractmethod
-    async def add_keys(self, node_ids):
-        raise NotImplementedError()
-
-    def all_node_ids(self):
-        return set(self.keys.keys())
-
-    def get_used_keys(self):
-        return self.keys.values()
-
-    def get_used_subcache_keys(self):
-        return self.subcache_keys.values()
-
-    def get_data_key(self, node_id):
-        return self.keys.get(node_id, None)
-
-    def get_subcache_key(self, node_id):
-        return self.subcache_keys.get(node_id, None)
-
-class Unhashable:
-    def __init__(self):
-        self.value = float("NaN")
-
-def to_hashable(obj):
-    # So that we don't infinitely recurse since frozenset and tuples
-    # are Sequences.
-    if isinstance(obj, (int, float, str, bool, type(None))):
-        return obj
-    elif isinstance(obj, Mapping):
-        return frozenset([(to_hashable(k), to_hashable(v)) for k, v in sorted(obj.items())])
-    elif isinstance(obj, Sequence):
-        return frozenset(zip(itertools.count(), [to_hashable(i) for i in obj]))
-    else:
-        # TODO - Support other objects like tensors?
-        return Unhashable()
-
-class CacheKeySetID(CacheKeySet):
-    def __init__(self, dynprompt, node_ids, is_changed_cache):
-        super().__init__(dynprompt, node_ids, is_changed_cache)
-        self.dynprompt = dynprompt
-
-    async def add_keys(self, node_ids):
-        for node_id in node_ids:
-            if node_id in self.keys:
-                continue
-            if not self.dynprompt.has_node(node_id):
-                continue
-            node = self.dynprompt.get_node(node_id)
-            self.keys[node_id] = (node_id, node["class_type"])
-            self.subcache_keys[node_id] = (node_id, node["class_type"])
-
-class CacheKeySetInputSignature(CacheKeySet):
-    def __init__(self, dynprompt, node_ids, is_changed_cache):
-        super().__init__(dynprompt, node_ids, is_changed_cache)
-        self.dynprompt = dynprompt
-        self.is_changed_cache = is_changed_cache
-
-    def include_node_id_in_input(self) -> bool:
-        return False
-
-    async def add_keys(self, node_ids):
-        for node_id in node_ids:
-            if node_id in self.keys:
-                continue
-            if not self.dynprompt.has_node(node_id):
-                continue
-            node = self.dynprompt.get_node(node_id)
-            self.keys[node_id] = await self.get_node_signature(self.dynprompt, node_id)
-            self.subcache_keys[node_id] = (node_id, node["class_type"])
-
-    async def get_node_signature(self, dynprompt, node_id):
-        signature = []
-        ancestors, order_mapping = self.get_ordered_ancestry(dynprompt, node_id)
-        signature.append(await self.get_immediate_node_signature(dynprompt, node_id, order_mapping))
-        for ancestor_id in ancestors:
-            signature.append(await self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping))
-        return to_hashable(signature)
-
-    async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
-        if not dynprompt.has_node(node_id):
-            # This node doesn't exist -- we can't cache it.
-            return [float("NaN")]
-        node = dynprompt.get_node(node_id)
-        class_type = node["class_type"]
-        class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-        signature = [class_type, await self.is_changed_cache.get(node_id)]
-        if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or include_unique_id_in_input(class_type):
-            signature.append(node_id)
-        inputs = node["inputs"]
-        for key in sorted(inputs.keys()):
-            if is_link(inputs[key]):
-                (ancestor_id, ancestor_socket) = inputs[key]
-                ancestor_index = ancestor_order_mapping[ancestor_id]
-                signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
-            else:
-                signature.append((key, inputs[key]))
-        return signature
-
-    # This function returns a list of all ancestors of the given node. The order of the list is
-    # deterministic based on which specific inputs the ancestor is connected by.
-    def get_ordered_ancestry(self, dynprompt, node_id):
-        ancestors = []
-        order_mapping = {}
-        self.get_ordered_ancestry_internal(dynprompt, node_id, ancestors, order_mapping)
-        return ancestors, order_mapping
-
-    def get_ordered_ancestry_internal(self, dynprompt, node_id, ancestors, order_mapping):
-        if not dynprompt.has_node(node_id):
-            return
-        inputs = dynprompt.get_node(node_id)["inputs"]
-        input_keys = sorted(inputs.keys())
-        for key in input_keys:
-            if is_link(inputs[key]):
-                ancestor_id = inputs[key][0]
-                if ancestor_id not in order_mapping:
-                    ancestors.append(ancestor_id)
-                    order_mapping[ancestor_id] = len(ancestors) - 1
-                    self.get_ordered_ancestry_internal(dynprompt, ancestor_id, ancestors, order_mapping)
-
-class BasicCache:
-    def __init__(self, key_class):
-        self.key_class = key_class
-        self.initialized = False
-        self.dynprompt: DynamicPrompt
-        self.cache_key_set: CacheKeySet
-        self.cache = {}
-        self.subcaches = {}
-
-    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
-        self.dynprompt = dynprompt
-        self.cache_key_set = self.key_class(dynprompt, node_ids, is_changed_cache)
-        await self.cache_key_set.add_keys(node_ids)
-        self.is_changed_cache = is_changed_cache
-        self.initialized = True
-
-    def all_node_ids(self):
-        assert self.initialized
-        node_ids = self.cache_key_set.all_node_ids()
-        for subcache in self.subcaches.values():
-            node_ids = node_ids.union(subcache.all_node_ids())
-        return node_ids
-
-    def _clean_cache(self):
-        preserve_keys = set(self.cache_key_set.get_used_keys())
-        to_remove = []
-        for key in self.cache:
-            if key not in preserve_keys:
-                to_remove.append(key)
-        for key in to_remove:
-            del self.cache[key]
-
-    def _clean_subcaches(self):
-        preserve_subcaches = set(self.cache_key_set.get_used_subcache_keys())
-
-        to_remove = []
-        for key in self.subcaches:
-            if key not in preserve_subcaches:
-                to_remove.append(key)
-        for key in to_remove:
-            del self.subcaches[key]
-
-    def clean_unused(self):
-        assert self.initialized
-        self._clean_cache()
-        self._clean_subcaches()
-
-    def _set_immediate(self, node_id, value):
-        assert self.initialized
-        cache_key = self.cache_key_set.get_data_key(node_id)
-        self.cache[cache_key] = value
-
-    def _get_immediate(self, node_id):
-        if not self.initialized:
-            return None
-        cache_key = self.cache_key_set.get_data_key(node_id)
-        if cache_key in self.cache:
-            return self.cache[cache_key]
-        else:
-            return None
-
-    async def _ensure_subcache(self, node_id, children_ids):
-        subcache_key = self.cache_key_set.get_subcache_key(node_id)
-        subcache = self.subcaches.get(subcache_key, None)
-        if subcache is None:
-            subcache = BasicCache(self.key_class)
-            self.subcaches[subcache_key] = subcache
-        await subcache.set_prompt(self.dynprompt, children_ids, self.is_changed_cache)
-        return subcache
-
-    def _get_subcache(self, node_id):
-        assert self.initialized
-        subcache_key = self.cache_key_set.get_subcache_key(node_id)
-        if subcache_key in self.subcaches:
-            return self.subcaches[subcache_key]
-        else:
-            return None
-
-    def recursive_debug_dump(self):
-        result = []
-        for key in self.cache:
-            result.append({"key": key, "value": self.cache[key]})
-        for key in self.subcaches:
-            result.append({"subcache_key": key, "subcache": self.subcaches[key].recursive_debug_dump()})
-        return result
-
-class HierarchicalCache(BasicCache):
-    def __init__(self, key_class):
-        super().__init__(key_class)
-
-    def _get_cache_for(self, node_id):
-        assert self.dynprompt is not None
-        parent_id = self.dynprompt.get_parent_node_id(node_id)
-        if parent_id is None:
-            return self
-
-        hierarchy = []
-        while parent_id is not None:
-            hierarchy.append(parent_id)
-            parent_id = self.dynprompt.get_parent_node_id(parent_id)
-
-        cache = self
-        for parent_id in reversed(hierarchy):
-            cache = cache._get_subcache(parent_id)
-            if cache is None:
-                return None
-        return cache
-
-    def get(self, node_id):
-        cache = self._get_cache_for(node_id)
-        if cache is None:
-            return None
-        return cache._get_immediate(node_id)
-
-    def set(self, node_id, value):
-        cache = self._get_cache_for(node_id)
-        assert cache is not None
-        cache._set_immediate(node_id, value)
-
-    async def ensure_subcache_for(self, node_id, children_ids):
-        cache = self._get_cache_for(node_id)
-        assert cache is not None
-        return await cache._ensure_subcache(node_id, children_ids)
-
-class LRUCache(BasicCache):
-    def __init__(self, key_class, max_size=100):
-        super().__init__(key_class)
-        self.max_size = max_size
-        self.min_generation = 0
-        self.generation = 0
-        self.used_generation = {}
-        self.children = {}
-
-    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
-        await super().set_prompt(dynprompt, node_ids, is_changed_cache)
-        self.generation += 1
-        for node_id in node_ids:
-            self._mark_used(node_id)
-
-    def clean_unused(self):
-        while len(self.cache) > self.max_size and self.min_generation < self.generation:
-            self.min_generation += 1
-            to_remove = [key for key in self.cache if self.used_generation[key] < self.min_generation]
-            for key in to_remove:
-                del self.cache[key]
-                del self.used_generation[key]
-                if key in self.children:
-                    del self.children[key]
-        self._clean_subcaches()
-
-    def get(self, node_id):
-        self._mark_used(node_id)
-        return self._get_immediate(node_id)
-
-    def _mark_used(self, node_id):
-        cache_key = self.cache_key_set.get_data_key(node_id)
-        if cache_key is not None:
-            self.used_generation[cache_key] = self.generation
-
-    def set(self, node_id, value):
-        self._mark_used(node_id)
-        return self._set_immediate(node_id, value)
-
-    async def ensure_subcache_for(self, node_id, children_ids):
-        # Just uses subcaches for tracking 'live' nodes
-        await super()._ensure_subcache(node_id, children_ids)
-
-        await self.cache_key_set.add_keys(children_ids)
-        self._mark_used(node_id)
-        cache_key = self.cache_key_set.get_data_key(node_id)
-        self.children[cache_key] = []
-        for child_id in children_ids:
-            self._mark_used(child_id)
-            self.children[cache_key].append(self.cache_key_set.get_data_key(child_id))
-        return self
-
-
-class DependencyAwareCache(BasicCache):
-    """
-    A cache implementation that tracks dependencies between nodes and manages
-    their execution and caching accordingly. It extends the BasicCache class.
-    Nodes are removed from this cache once all of their descendants have been
-    executed.
-    """
-
-    def __init__(self, key_class):
-        """
-        Initialize the DependencyAwareCache.
-
-        Args:
-            key_class: The class used for generating cache keys.
-        """
-        super().__init__(key_class)
-        self.descendants = {}  # Maps node_id -> set of descendant node_ids
-        self.ancestors = {}    # Maps node_id -> set of ancestor node_ids
-        self.executed_nodes = set()  # Tracks nodes that have been executed
-
-    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
-        """
-        Clear the entire cache and rebuild the dependency graph.
-
-        Args:
-            dynprompt: The dynamic prompt object containing node information.
-            node_ids: List of node IDs to initialize the cache for.
-            is_changed_cache: Flag indicating if the cache has changed.
-        """
-        # Clear all existing cache data
-        self.cache.clear()
-        self.subcaches.clear()
-        self.descendants.clear()
-        self.ancestors.clear()
-        self.executed_nodes.clear()
-
-        # Call the parent method to initialize the cache with the new prompt
-        await super().set_prompt(dynprompt, node_ids, is_changed_cache)
-
-        # Rebuild the dependency graph
-        self._build_dependency_graph(dynprompt, node_ids)
-
-    def _build_dependency_graph(self, dynprompt, node_ids):
-        """
-        Build the dependency graph for all nodes.
-
-        Args:
-            dynprompt: The dynamic prompt object containing node information.
-            node_ids: List of node IDs to build the graph for.
-        """
-        self.descendants.clear()
-        self.ancestors.clear()
-        for node_id in node_ids:
-            self.descendants[node_id] = set()
-            self.ancestors[node_id] = set()
-
-        for node_id in node_ids:
-            inputs = dynprompt.get_node(node_id)["inputs"]
-            for input_data in inputs.values():
-                if is_link(input_data):  # Check if the input is a link to another node
-                    ancestor_id = input_data[0]
-                    self.descendants[ancestor_id].add(node_id)
-                    self.ancestors[node_id].add(ancestor_id)
-
-    def set(self, node_id, value):
-        """
-        Mark a node as executed and store its value in the cache.
-
-        Args:
-            node_id: The ID of the node to store.
-            value: The value to store for the node.
-        """
-        self._set_immediate(node_id, value)
-        self.executed_nodes.add(node_id)
-        self._cleanup_ancestors(node_id)
-
-    def get(self, node_id):
-        """
-        Retrieve the cached value for a node.
-
-        Args:
-            node_id: The ID of the node to retrieve.
-
-        Returns:
-            The cached value for the node.
-        """
-        return self._get_immediate(node_id)
-
-    async def ensure_subcache_for(self, node_id, children_ids):
-        """
-        Ensure a subcache exists for a node and update dependencies.
-
-        Args:
-            node_id: The ID of the parent node.
-            children_ids: List of child node IDs to associate with the parent node.
-
-        Returns:
-            The subcache object for the node.
-        """
-        subcache = await super()._ensure_subcache(node_id, children_ids)
-        for child_id in children_ids:
-            self.descendants[node_id].add(child_id)
-            self.ancestors[child_id].add(node_id)
-        return subcache
-
-    def _cleanup_ancestors(self, node_id):
-        """
-        Check if ancestors of a node can be removed from the cache.
-
-        Args:
-            node_id: The ID of the node whose ancestors are to be checked.
-        """
-        for ancestor_id in self.ancestors.get(node_id, []):
-            if ancestor_id in self.executed_nodes:
-                # Remove ancestor if all its descendants have been executed
-                if all(descendant in self.executed_nodes for descendant in self.descendants[ancestor_id]):
-                    self._remove_node(ancestor_id)
-
-    def _remove_node(self, node_id):
-        """
-        Remove a node from the cache.
-
-        Args:
-            node_id: The ID of the node to remove.
-        """
-        cache_key = self.cache_key_set.get_data_key(node_id)
-        if cache_key in self.cache:
-            del self.cache[cache_key]
-        subcache_key = self.cache_key_set.get_subcache_key(node_id)
-        if subcache_key in self.subcaches:
-            del self.subcaches[subcache_key]
-
-    def clean_unused(self):
-        """
-        Clean up unused nodes. This is a no-op for this cache implementation.
-        """
-        pass
-
-    def recursive_debug_dump(self):
-        """
-        Dump the cache and dependency graph for debugging.
-
-        Returns:
-            A list containing the cache state and dependency graph.
-        """
-        result = super().recursive_debug_dump()
-        result.append({
-            "descendants": self.descendants,
-            "ancestors": self.ancestors,
-            "executed_nodes": list(self.executed_nodes),
-        })
-        return result

comfy_execution/graph.py

@@ -1,314 +0,0 @@
-from __future__ import annotations
-from typing import Type, Literal
-
-import nodes
-import asyncio
-import inspect
-from comfy_execution.graph_utils import is_link
-from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, InputTypeOptions
-
-class DependencyCycleError(Exception):
-    pass
-
-class NodeInputError(Exception):
-    pass
-
-class NodeNotFoundError(Exception):
-    pass
-
-class DynamicPrompt:
-    def __init__(self, original_prompt):
-        # The original prompt provided by the user
-        self.original_prompt = original_prompt
-        # Any extra pieces of the graph created during execution
-        self.ephemeral_prompt = {}
-        self.ephemeral_parents = {}
-        self.ephemeral_display = {}
-
-    def get_node(self, node_id):
-        if node_id in self.ephemeral_prompt:
-            return self.ephemeral_prompt[node_id]
-        if node_id in self.original_prompt:
-            return self.original_prompt[node_id]
-        raise NodeNotFoundError(f"Node {node_id} not found")
-
-    def has_node(self, node_id):
-        return node_id in self.original_prompt or node_id in self.ephemeral_prompt
-
-    def add_ephemeral_node(self, node_id, node_info, parent_id, display_id):
-        self.ephemeral_prompt[node_id] = node_info
-        self.ephemeral_parents[node_id] = parent_id
-        self.ephemeral_display[node_id] = display_id
-
-    def get_real_node_id(self, node_id):
-        while node_id in self.ephemeral_parents:
-            node_id = self.ephemeral_parents[node_id]
-        return node_id
-
-    def get_parent_node_id(self, node_id):
-        return self.ephemeral_parents.get(node_id, None)
-
-    def get_display_node_id(self, node_id):
-        while node_id in self.ephemeral_display:
-            node_id = self.ephemeral_display[node_id]
-        return node_id
-
-    def all_node_ids(self):
-        return set(self.original_prompt.keys()).union(set(self.ephemeral_prompt.keys()))
-
-    def get_original_prompt(self):
-        return self.original_prompt
-
-def get_input_info(
-    class_def: Type[ComfyNodeABC],
-    input_name: str,
-    valid_inputs: InputTypeDict | None = None
-) -> tuple[str, Literal["required", "optional", "hidden"], InputTypeOptions] | tuple[None, None, None]:
-    """Get the input type, category, and extra info for a given input name.
-
-    Arguments:
-        class_def: The class definition of the node.
-        input_name: The name of the input to get info for.
-        valid_inputs: The valid inputs for the node, or None to use the class_def.INPUT_TYPES().
-
-    Returns:
-        tuple[str, str, dict] | tuple[None, None, None]: The input type, category, and extra info for the input name.
-    """
-
-    valid_inputs = valid_inputs or class_def.INPUT_TYPES()
-    input_info = None
-    input_category = None
-    if "required" in valid_inputs and input_name in valid_inputs["required"]:
-        input_category = "required"
-        input_info = valid_inputs["required"][input_name]
-    elif "optional" in valid_inputs and input_name in valid_inputs["optional"]:
-        input_category = "optional"
-        input_info = valid_inputs["optional"][input_name]
-    elif "hidden" in valid_inputs and input_name in valid_inputs["hidden"]:
-        input_category = "hidden"
-        input_info = valid_inputs["hidden"][input_name]
-    if input_info is None:
-        return None, None, None
-    input_type = input_info[0]
-    if len(input_info) > 1:
-        extra_info = input_info[1]
-    else:
-        extra_info = {}
-    return input_type, input_category, extra_info
-
-class TopologicalSort:
-    def __init__(self, dynprompt):
-        self.dynprompt = dynprompt
-        self.pendingNodes = {}
-        self.blockCount = {} # Number of nodes this node is directly blocked by
-        self.blocking = {} # Which nodes are blocked by this node
-        self.externalBlocks = 0
-        self.unblockedEvent = asyncio.Event()
-
-    def get_input_info(self, unique_id, input_name):
-        class_type = self.dynprompt.get_node(unique_id)["class_type"]
-        class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-        return get_input_info(class_def, input_name)
-
-    def make_input_strong_link(self, to_node_id, to_input):
-        inputs = self.dynprompt.get_node(to_node_id)["inputs"]
-        if to_input not in inputs:
-            raise NodeInputError(f"Node {to_node_id} says it needs input {to_input}, but there is no input to that node at all")
-        value = inputs[to_input]
-        if not is_link(value):
-            raise NodeInputError(f"Node {to_node_id} says it needs input {to_input}, but that value is a constant")
-        from_node_id, from_socket = value
-        self.add_strong_link(from_node_id, from_socket, to_node_id)
-
-    def add_strong_link(self, from_node_id, from_socket, to_node_id):
-        if not self.is_cached(from_node_id):
-            self.add_node(from_node_id)
-            if to_node_id not in self.blocking[from_node_id]:
-                self.blocking[from_node_id][to_node_id] = {}
-                self.blockCount[to_node_id] += 1
-            self.blocking[from_node_id][to_node_id][from_socket] = True
-
-    def add_node(self, node_unique_id, include_lazy=False, subgraph_nodes=None):
-        node_ids = [node_unique_id]
-        links = []
-
-        while len(node_ids) > 0:
-            unique_id = node_ids.pop()
-            if unique_id in self.pendingNodes:
-                continue
-
-            self.pendingNodes[unique_id] = True
-            self.blockCount[unique_id] = 0
-            self.blocking[unique_id] = {}
-
-            inputs = self.dynprompt.get_node(unique_id)["inputs"]
-            for input_name in inputs:
-                value = inputs[input_name]
-                if is_link(value):
-                    from_node_id, from_socket = value
-                    if subgraph_nodes is not None and from_node_id not in subgraph_nodes:
-                        continue
-                    _, _, input_info = self.get_input_info(unique_id, input_name)
-                    is_lazy = input_info is not None and "lazy" in input_info and input_info["lazy"]
-                    if (include_lazy or not is_lazy) and not self.is_cached(from_node_id):
-                        node_ids.append(from_node_id)
-                        links.append((from_node_id, from_socket, unique_id))
-
-        for link in links:
-            self.add_strong_link(*link)
-
-    def add_external_block(self, node_id):
-        assert node_id in self.blockCount, "Can't add external block to a node that isn't pending"
-        self.externalBlocks += 1
-        self.blockCount[node_id] += 1
-        def unblock():
-            self.externalBlocks -= 1
-            self.blockCount[node_id] -= 1
-            self.unblockedEvent.set()
-        return unblock
-
-    def is_cached(self, node_id):
-        return False
-
-    def get_ready_nodes(self):
-        return [node_id for node_id in self.pendingNodes if self.blockCount[node_id] == 0]
-
-    def pop_node(self, unique_id):
-        del self.pendingNodes[unique_id]
-        for blocked_node_id in self.blocking[unique_id]:
-            self.blockCount[blocked_node_id] -= 1
-        del self.blocking[unique_id]
-
-    def is_empty(self):
-        return len(self.pendingNodes) == 0
-
-class ExecutionList(TopologicalSort):
-    """
-    ExecutionList implements a topological dissolve of the graph. After a node is staged for execution,
-    it can still be returned to the graph after having further dependencies added.
-    """
-    def __init__(self, dynprompt, output_cache):
-        super().__init__(dynprompt)
-        self.output_cache = output_cache
-        self.staged_node_id = None
-
-    def is_cached(self, node_id):
-        return self.output_cache.get(node_id) is not None
-
-    async def stage_node_execution(self):
-        assert self.staged_node_id is None
-        if self.is_empty():
-            return None, None, None
-        available = self.get_ready_nodes()
-        while len(available) == 0 and self.externalBlocks > 0:
-            # Wait for an external block to be released
-            await self.unblockedEvent.wait()
-            self.unblockedEvent.clear()
-            available = self.get_ready_nodes()
-        if len(available) == 0:
-            cycled_nodes = self.get_nodes_in_cycle()
-            # Because cycles composed entirely of static nodes are caught during initial validation,
-            # we will 'blame' the first node in the cycle that is not a static node.
-            blamed_node = cycled_nodes[0]
-            for node_id in cycled_nodes:
-                display_node_id = self.dynprompt.get_display_node_id(node_id)
-                if display_node_id != node_id:
-                    blamed_node = display_node_id
-                    break
-            ex = DependencyCycleError("Dependency cycle detected")
-            error_details = {
-                "node_id": blamed_node,
-                "exception_message": str(ex),
-                "exception_type": "graph.DependencyCycleError",
-                "traceback": [],
-                "current_inputs": []
-            }
-            return None, error_details, ex
-
-        self.staged_node_id = self.ux_friendly_pick_node(available)
-        return self.staged_node_id, None, None
-
-    def ux_friendly_pick_node(self, node_list):
-        # If an output node is available, do that first.
-        # Technically this has no effect on the overall length of execution, but it feels better as a user
-        # for a PreviewImage to display a result as soon as it can
-        # Some other heuristics could probably be used here to improve the UX further.
-        def is_output(node_id):
-            class_type = self.dynprompt.get_node(node_id)["class_type"]
-            class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-            if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
-                return True
-            return False
-
-        # If an available node is async, do that first.
-        # This will execute the asynchronous function earlier, reducing the overall time.
-        def is_async(node_id):
-            class_type = self.dynprompt.get_node(node_id)["class_type"]
-            class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-            return inspect.iscoroutinefunction(getattr(class_def, class_def.FUNCTION))
-
-        for node_id in node_list:
-            if is_output(node_id) or is_async(node_id):
-                return node_id
-
-        #This should handle the VAEDecode -> preview case
-        for node_id in node_list:
-            for blocked_node_id in self.blocking[node_id]:
-                if is_output(blocked_node_id):
-                    return node_id
-
-        #This should handle the VAELoader -> VAEDecode -> preview case
-        for node_id in node_list:
-            for blocked_node_id in self.blocking[node_id]:
-                for blocked_node_id1 in self.blocking[blocked_node_id]:
-                    if is_output(blocked_node_id1):
-                        return node_id
-
-        #TODO: this function should be improved
-        return node_list[0]
-
-    def unstage_node_execution(self):
-        assert self.staged_node_id is not None
-        self.staged_node_id = None
-
-    def complete_node_execution(self):
-        node_id = self.staged_node_id
-        self.pop_node(node_id)
-        self.staged_node_id = None
-
-    def get_nodes_in_cycle(self):
-        # We'll dissolve the graph in reverse topological order to leave only the nodes in the cycle.
-        # We're skipping some of the performance optimizations from the original TopologicalSort to keep
-        # the code simple (and because having a cycle in the first place is a catastrophic error)
-        blocked_by = { node_id: {} for node_id in self.pendingNodes }
-        for from_node_id in self.blocking:
-            for to_node_id in self.blocking[from_node_id]:
-                if True in self.blocking[from_node_id][to_node_id].values():
-                    blocked_by[to_node_id][from_node_id] = True
-        to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
-        while len(to_remove) > 0:
-            for node_id in to_remove:
-                for to_node_id in blocked_by:
-                    if node_id in blocked_by[to_node_id]:
-                        del blocked_by[to_node_id][node_id]
-                del blocked_by[node_id]
-            to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
-        return list(blocked_by.keys())
-
-class ExecutionBlocker:
-    """
-    Return this from a node and any users will be blocked with the given error message.
-    If the message is None, execution will be blocked silently instead.
-    Generally, you should avoid using this functionality unless absolutely necessary. Whenever it's
-    possible, a lazy input will be more efficient and have a better user experience.
-    This functionality is useful in two cases:
-    1. You want to conditionally prevent an output node from executing. (Particularly a built-in node
-       like SaveImage. For your own output nodes, I would recommend just adding a BOOL input and using
-       lazy evaluation to let it conditionally disable itself.)
-    2. You have a node with multiple possible outputs, some of which are invalid and should not be used.
-       (I would recommend not making nodes like this in the future -- instead, make multiple nodes with
-       different outputs. Unfortunately, there are several popular existing nodes using this pattern.)
-    """
-    def __init__(self, message):
-        self.message = message
-

comfy_execution/graph_utils.py

@@ -1,139 +0,0 @@
-def is_link(obj):
-    if not isinstance(obj, list):
-        return False
-    if len(obj) != 2:
-        return False
-    if not isinstance(obj[0], str):
-        return False
-    if not isinstance(obj[1], int) and not isinstance(obj[1], float):
-        return False
-    return True
-
-# The GraphBuilder is just a utility class that outputs graphs in the form expected by the ComfyUI back-end
-class GraphBuilder:
-    _default_prefix_root = ""
-    _default_prefix_call_index = 0
-    _default_prefix_graph_index = 0
-
-    def __init__(self, prefix = None):
-        if prefix is None:
-            self.prefix = GraphBuilder.alloc_prefix()
-        else:
-            self.prefix = prefix
-        self.nodes = {}
-        self.id_gen = 1
-
-    @classmethod
-    def set_default_prefix(cls, prefix_root, call_index, graph_index = 0):
-        cls._default_prefix_root = prefix_root
-        cls._default_prefix_call_index = call_index
-        cls._default_prefix_graph_index = graph_index
-
-    @classmethod
-    def alloc_prefix(cls, root=None, call_index=None, graph_index=None):
-        if root is None:
-            root = GraphBuilder._default_prefix_root
-        if call_index is None:
-            call_index = GraphBuilder._default_prefix_call_index
-        if graph_index is None:
-            graph_index = GraphBuilder._default_prefix_graph_index
-        result = f"{root}.{call_index}.{graph_index}."
-        GraphBuilder._default_prefix_graph_index += 1
-        return result
-
-    def node(self, class_type, id=None, **kwargs):
-        if id is None:
-            id = str(self.id_gen)
-            self.id_gen += 1
-        id = self.prefix + id
-        if id in self.nodes:
-            return self.nodes[id]
-
-        node = Node(id, class_type, kwargs)
-        self.nodes[id] = node
-        return node
-
-    def lookup_node(self, id):
-        id = self.prefix + id
-        return self.nodes.get(id)
-
-    def finalize(self):
-        output = {}
-        for node_id, node in self.nodes.items():
-            output[node_id] = node.serialize()
-        return output
-
-    def replace_node_output(self, node_id, index, new_value):
-        node_id = self.prefix + node_id
-        to_remove = []
-        for node in self.nodes.values():
-            for key, value in node.inputs.items():
-                if is_link(value) and value[0] == node_id and value[1] == index:
-                    if new_value is None:
-                        to_remove.append((node, key))
-                    else:
-                        node.inputs[key] = new_value
-        for node, key in to_remove:
-            del node.inputs[key]
-
-    def remove_node(self, id):
-        id = self.prefix + id
-        del self.nodes[id]
-
-class Node:
-    def __init__(self, id, class_type, inputs):
-        self.id = id
-        self.class_type = class_type
-        self.inputs = inputs
-        self.override_display_id = None
-
-    def out(self, index):
-        return [self.id, index]
-
-    def set_input(self, key, value):
-        if value is None:
-            if key in self.inputs:
-                del self.inputs[key]
-        else:
-            self.inputs[key] = value
-
-    def get_input(self, key):
-        return self.inputs.get(key)
-
-    def set_override_display_id(self, override_display_id):
-        self.override_display_id = override_display_id
-
-    def serialize(self):
-        serialized = {
-            "class_type": self.class_type,
-            "inputs": self.inputs
-        }
-        if self.override_display_id is not None:
-            serialized["override_display_id"] = self.override_display_id
-        return serialized
-
-def add_graph_prefix(graph, outputs, prefix):
-    # Change the node IDs and any internal links
-    new_graph = {}
-    for node_id, node_info in graph.items():
-        # Make sure the added nodes have unique IDs
-        new_node_id = prefix + node_id
-        new_node = { "class_type": node_info["class_type"], "inputs": {} }
-        for input_name, input_value in node_info.get("inputs", {}).items():
-            if is_link(input_value):
-                new_node["inputs"][input_name] = [prefix + input_value[0], input_value[1]]
-            else:
-                new_node["inputs"][input_name] = input_value
-        new_graph[new_node_id] = new_node
-
-    # Change the node IDs in the outputs
-    new_outputs = []
-    for n in range(len(outputs)):
-        output = outputs[n]
-        if is_link(output):
-            new_outputs.append([prefix + output[0], output[1]])
-        else:
-            new_outputs.append(output)
-
-    return new_graph, tuple(new_outputs)
-

comfy_execution/progress.py

@@ -1,347 +0,0 @@
-from typing import TypedDict, Dict, Optional
-from typing_extensions import override
-from PIL import Image
-from enum import Enum
-from abc import ABC
-from tqdm import tqdm
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from comfy_execution.graph import DynamicPrompt
-from protocol import BinaryEventTypes
-from comfy_api import feature_flags
-
-
-class NodeState(Enum):
-    Pending = "pending"
-    Running = "running"
-    Finished = "finished"
-    Error = "error"
-
-
-class NodeProgressState(TypedDict):
-    """
-    A class to represent the state of a node's progress.
-    """
-
-    state: NodeState
-    value: float
-    max: float
-
-
-class ProgressHandler(ABC):
-    """
-    Abstract base class for progress handlers.
-    Progress handlers receive progress updates and display them in various ways.
-    """
-
-    def __init__(self, name: str):
-        self.name = name
-        self.enabled = True
-
-    def set_registry(self, registry: "ProgressRegistry"):
-        pass
-
-    def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        """Called when a node starts processing"""
-        pass
-
-    def update_handler(
-        self,
-        node_id: str,
-        value: float,
-        max_value: float,
-        state: NodeProgressState,
-        prompt_id: str,
-        image: Optional[Image.Image] = None,
-    ):
-        """Called when a node's progress is updated"""
-        pass
-
-    def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        """Called when a node finishes processing"""
-        pass
-
-    def reset(self):
-        """Called when the progress registry is reset"""
-        pass
-
-    def enable(self):
-        """Enable this handler"""
-        self.enabled = True
-
-    def disable(self):
-        """Disable this handler"""
-        self.enabled = False
-
-
-class CLIProgressHandler(ProgressHandler):
-    """
-    Handler that displays progress using tqdm progress bars in the CLI.
-    """
-
-    def __init__(self):
-        super().__init__("cli")
-        self.progress_bars: Dict[str, tqdm] = {}
-
-    @override
-    def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Create a new tqdm progress bar
-        if node_id not in self.progress_bars:
-            self.progress_bars[node_id] = tqdm(
-                total=state["max"],
-                desc=f"Node {node_id}",
-                unit="steps",
-                leave=True,
-                position=len(self.progress_bars),
-            )
-
-    @override
-    def update_handler(
-        self,
-        node_id: str,
-        value: float,
-        max_value: float,
-        state: NodeProgressState,
-        prompt_id: str,
-        image: Optional[Image.Image] = None,
-    ):
-        # Handle case where start_handler wasn't called
-        if node_id not in self.progress_bars:
-            self.progress_bars[node_id] = tqdm(
-                total=max_value,
-                desc=f"Node {node_id}",
-                unit="steps",
-                leave=True,
-                position=len(self.progress_bars),
-            )
-            self.progress_bars[node_id].update(value)
-        else:
-            # Update existing progress bar
-            if max_value != self.progress_bars[node_id].total:
-                self.progress_bars[node_id].total = max_value
-            # Calculate the update amount (difference from current position)
-            current_position = self.progress_bars[node_id].n
-            update_amount = value - current_position
-            if update_amount > 0:
-                self.progress_bars[node_id].update(update_amount)
-
-    @override
-    def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Complete and close the progress bar if it exists
-        if node_id in self.progress_bars:
-            # Ensure the bar shows 100% completion
-            remaining = state["max"] - self.progress_bars[node_id].n
-            if remaining > 0:
-                self.progress_bars[node_id].update(remaining)
-            self.progress_bars[node_id].close()
-            del self.progress_bars[node_id]
-
-    @override
-    def reset(self):
-        # Close all progress bars
-        for bar in self.progress_bars.values():
-            bar.close()
-        self.progress_bars.clear()
-
-
-class WebUIProgressHandler(ProgressHandler):
-    """
-    Handler that sends progress updates to the WebUI via WebSockets.
-    """
-
-    def __init__(self, server_instance):
-        super().__init__("webui")
-        self.server_instance = server_instance
-
-    def set_registry(self, registry: "ProgressRegistry"):
-        self.registry = registry
-
-    def _send_progress_state(self, prompt_id: str, nodes: Dict[str, NodeProgressState]):
-        """Send the current progress state to the client"""
-        if self.server_instance is None:
-            return
-
-        # Only send info for non-pending nodes
-        active_nodes = {
-            node_id: {
-                "value": state["value"],
-                "max": state["max"],
-                "state": state["state"].value,
-                "node_id": node_id,
-                "prompt_id": prompt_id,
-                "display_node_id": self.registry.dynprompt.get_display_node_id(node_id),
-                "parent_node_id": self.registry.dynprompt.get_parent_node_id(node_id),
-                "real_node_id": self.registry.dynprompt.get_real_node_id(node_id),
-            }
-            for node_id, state in nodes.items()
-            if state["state"] != NodeState.Pending
-        }
-
-        # Send a combined progress_state message with all node states
-        self.server_instance.send_sync(
-            "progress_state", {"prompt_id": prompt_id, "nodes": active_nodes}
-        )
-
-    @override
-    def start_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Send progress state of all nodes
-        if self.registry:
-            self._send_progress_state(prompt_id, self.registry.nodes)
-
-    @override
-    def update_handler(
-        self,
-        node_id: str,
-        value: float,
-        max_value: float,
-        state: NodeProgressState,
-        prompt_id: str,
-        image: Optional[Image.Image] = None,
-    ):
-        # Send progress state of all nodes
-        if self.registry:
-            self._send_progress_state(prompt_id, self.registry.nodes)
-        if image:
-            # Only send new format if client supports it
-            if feature_flags.supports_feature(
-                self.server_instance.sockets_metadata,
-                self.server_instance.client_id,
-                "supports_preview_metadata",
-            ):
-                metadata = {
-                    "node_id": node_id,
-                    "prompt_id": prompt_id,
-                    "display_node_id": self.registry.dynprompt.get_display_node_id(
-                        node_id
-                    ),
-                    "parent_node_id": self.registry.dynprompt.get_parent_node_id(
-                        node_id
-                    ),
-                    "real_node_id": self.registry.dynprompt.get_real_node_id(node_id),
-                }
-                self.server_instance.send_sync(
-                    BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA,
-                    (image, metadata),
-                    self.server_instance.client_id,
-                )
-
-    @override
-    def finish_handler(self, node_id: str, state: NodeProgressState, prompt_id: str):
-        # Send progress state of all nodes
-        if self.registry:
-            self._send_progress_state(prompt_id, self.registry.nodes)
-
-
-class ProgressRegistry:
-    """
-    Registry that maintains node progress state and notifies registered handlers.
-    """
-
-    def __init__(self, prompt_id: str, dynprompt: "DynamicPrompt"):
-        self.prompt_id = prompt_id
-        self.dynprompt = dynprompt
-        self.nodes: Dict[str, NodeProgressState] = {}
-        self.handlers: Dict[str, ProgressHandler] = {}
-
-    def register_handler(self, handler: ProgressHandler) -> None:
-        """Register a progress handler"""
-        self.handlers[handler.name] = handler
-
-    def unregister_handler(self, handler_name: str) -> None:
-        """Unregister a progress handler"""
-        if handler_name in self.handlers:
-            # Allow handler to clean up resources
-            self.handlers[handler_name].reset()
-            del self.handlers[handler_name]
-
-    def enable_handler(self, handler_name: str) -> None:
-        """Enable a progress handler"""
-        if handler_name in self.handlers:
-            self.handlers[handler_name].enable()
-
-    def disable_handler(self, handler_name: str) -> None:
-        """Disable a progress handler"""
-        if handler_name in self.handlers:
-            self.handlers[handler_name].disable()
-
-    def ensure_entry(self, node_id: str) -> NodeProgressState:
-        """Ensure a node entry exists"""
-        if node_id not in self.nodes:
-            self.nodes[node_id] = NodeProgressState(
-                state=NodeState.Pending, value=0, max=1
-            )
-        return self.nodes[node_id]
-
-    def start_progress(self, node_id: str) -> None:
-        """Start progress tracking for a node"""
-        entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Running
-        entry["value"] = 0.0
-        entry["max"] = 1.0
-
-        # Notify all enabled handlers
-        for handler in self.handlers.values():
-            if handler.enabled:
-                handler.start_handler(node_id, entry, self.prompt_id)
-
-    def update_progress(
-        self, node_id: str, value: float, max_value: float, image: Optional[Image.Image]
-    ) -> None:
-        """Update progress for a node"""
-        entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Running
-        entry["value"] = value
-        entry["max"] = max_value
-
-        # Notify all enabled handlers
-        for handler in self.handlers.values():
-            if handler.enabled:
-                handler.update_handler(
-                    node_id, value, max_value, entry, self.prompt_id, image
-                )
-
-    def finish_progress(self, node_id: str) -> None:
-        """Finish progress tracking for a node"""
-        entry = self.ensure_entry(node_id)
-        entry["state"] = NodeState.Finished
-        entry["value"] = entry["max"]
-
-        # Notify all enabled handlers
-        for handler in self.handlers.values():
-            if handler.enabled:
-                handler.finish_handler(node_id, entry, self.prompt_id)
-
-    def reset_handlers(self) -> None:
-        """Reset all handlers"""
-        for handler in self.handlers.values():
-            handler.reset()
-
-# Global registry instance
-global_progress_registry: ProgressRegistry = None
-
-def reset_progress_state(prompt_id: str, dynprompt: "DynamicPrompt") -> None:
-    global global_progress_registry
-
-    # Reset existing handlers if registry exists
-    if global_progress_registry is not None:
-        global_progress_registry.reset_handlers()
-
-    # Create new registry
-    global_progress_registry = ProgressRegistry(prompt_id, dynprompt)
-
-
-def add_progress_handler(handler: ProgressHandler) -> None:
-    registry = get_progress_state()
-    handler.set_registry(registry)
-    registry.register_handler(handler)
-
-
-def get_progress_state() -> ProgressRegistry:
-    global global_progress_registry
-    if global_progress_registry is None:
-        from comfy_execution.graph import DynamicPrompt
-
-        global_progress_registry = ProgressRegistry(
-            prompt_id="", dynprompt=DynamicPrompt({})
-        )
-    return global_progress_registry

comfy_execution/utils.py

@@ -1,46 +0,0 @@
-import contextvars
-from typing import Optional, NamedTuple
-
-class ExecutionContext(NamedTuple):
-    """
-    Context information about the currently executing node.
-
-    Attributes:
-        node_id: The ID of the currently executing node
-        list_index: The index in a list being processed (for operations on batches/lists)
-    """
-    prompt_id: str
-    node_id: str
-    list_index: Optional[int]
-
-current_executing_context: contextvars.ContextVar[Optional[ExecutionContext]] = contextvars.ContextVar("current_executing_context", default=None)
-
-def get_executing_context() -> Optional[ExecutionContext]:
-    return current_executing_context.get(None)
-
-class CurrentNodeContext:
-    """
-    Context manager for setting the current executing node context.
-
-    Sets the current_executing_context on enter and resets it on exit.
-
-    Example:
-        with CurrentNodeContext(node_id="123", list_index=0):
-            # Code that should run with the current node context set
-            process_image()
-    """
-    def __init__(self, prompt_id: str, node_id: str, list_index: Optional[int] = None):
-        self.context = ExecutionContext(
-            prompt_id= prompt_id,
-            node_id= node_id,
-            list_index= list_index
-        )
-        self.token = None
-
-    def __enter__(self):
-        self.token = current_executing_context.set(self.context)
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if self.token is not None:
-            current_executing_context.reset(self.token)

comfy_execution/validation.py

@@ -1,39 +0,0 @@
-from __future__ import annotations
-
-
-def validate_node_input(
-    received_type: str, input_type: str, strict: bool = False
-) -> bool:
-    """
-    received_type and input_type are both strings of the form "T1,T2,...".
-
-    If strict is True, the input_type must contain the received_type.
-      For example, if received_type is "STRING" and input_type is "STRING,INT",
-      this will return True. But if received_type is "STRING,INT" and input_type is
-      "INT", this will return False.
-
-    If strict is False, the input_type must have overlap with the received_type.
-      For example, if received_type is "STRING,BOOLEAN" and input_type is "STRING,INT",
-      this will return True.
-
-    Supports pre-union type extension behaviour of ``__ne__`` overrides.
-    """
-    # If the types are exactly the same, we can return immediately
-    # Use pre-union behaviour: inverse of `__ne__`
-    if not received_type != input_type:
-        return True
-
-    # Not equal, and not strings
-    if not isinstance(received_type, str) or not isinstance(input_type, str):
-        return False
-
-    # Split the type strings into sets for comparison
-    received_types = set(t.strip() for t in received_type.split(","))
-    input_types = set(t.strip() for t in input_type.split(","))
-
-    if strict:
-        # In strict mode, all received types must be in the input types
-        return received_types.issubset(input_types)
-    else:
-        # In non-strict mode, there must be at least one type in common
-        return len(received_types.intersection(input_types)) > 0