From 5e68a4ce67450dd8d3aa1c51c6d06014ac46c03f Mon Sep 17 00:00:00 2001
From: JettHu <35261585+JettHu@users.noreply.github.com>
Date: Sun, 15 Sep 2024 13:03:09 +0800
Subject: [PATCH] Reduce repeated calls of INPUT_TYPES in cache (#4922)

---
 comfy_execution/caching.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py
index 62311ed7..630f280f 100644
--- a/comfy_execution/caching.py
+++ b/comfy_execution/caching.py
@@ -1,11 +1,21 @@
 import itertools
-from typing import Sequence, Mapping
+from typing import Sequence, Mapping, Dict
 from comfy_execution.graph import DynamicPrompt
 
 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:
     def __init__(self, dynprompt, node_ids, is_changed_cache):
         self.keys = {}
@@ -98,7 +108,7 @@ class CacheKeySetInputSignature(CacheKeySet):
         class_type = node["class_type"]
         class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
         signature = [class_type, 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 "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values():
+        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()):