From f6b7194f64ab4b32018eea0da9d9d89a30b582aa Mon Sep 17 00:00:00 2001
From: JettHu <35261585+JettHu@users.noreply.github.com>
Date: Fri, 13 Sep 2024 11:02:52 +0800
Subject: [PATCH] Reduce repeated calls of get_immediate_node_signature for
 ancestors in cache (#4871)

---
 comfy_execution/caching.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py
index 62311ed7..f163bc98 100644
--- a/comfy_execution/caching.py
+++ b/comfy_execution/caching.py
@@ -67,6 +67,7 @@ class CacheKeySetInputSignature(CacheKeySet):
         super().__init__(dynprompt, node_ids, is_changed_cache)
         self.dynprompt = dynprompt
         self.is_changed_cache = is_changed_cache
+        self.immediate_node_signature = {}
         self.add_keys(node_ids)
 
     def include_node_id_in_input(self) -> bool:
@@ -94,6 +95,8 @@ class CacheKeySetInputSignature(CacheKeySet):
         if not dynprompt.has_node(node_id):
             # This node doesn't exist -- we can't cache it.
             return [float("NaN")]
+        if node_id in self.immediate_node_signature:    # reduce repeated calls of ancestors
+            return self.immediate_node_signature[node_id]
         node = dynprompt.get_node(node_id)
         class_type = node["class_type"]
         class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
@@ -108,6 +111,7 @@ class CacheKeySetInputSignature(CacheKeySet):
                 signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
             else:
                 signature.append((key, inputs[key]))
+        self.immediate_node_signature[node_id] = signature
         return signature
 
     # This function returns a list of all ancestors of the given node. The order of the list is