From 3f466f97bd760b1372f8d689e2d71bbe87f210e4 Mon Sep 17 00:00:00 2001
From: JettHu <jett.hux@gmail.com>
Date: Thu, 27 Mar 2025 11:16:49 +0800
Subject: [PATCH] Enable cache_helper in more places where filename list is not
 actually needed

---
 comfy_execution/caching.py |  8 +++++---
 comfy_execution/graph.py   |  4 +++-
 execution.py               |  4 +++-
 folder_paths.py            | 12 +++++++++---
 server.py                  |  2 +-
 5 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py
index 630f280fc..f255ca1b3 100644
--- a/comfy_execution/caching.py
+++ b/comfy_execution/caching.py
@@ -1,9 +1,10 @@
 import itertools
-from typing import Sequence, Mapping, Dict
-from comfy_execution.graph import DynamicPrompt
+from typing import Dict, Mapping, Sequence
 
+import folder_paths
 import nodes
 
+from comfy_execution.graph import DynamicPrompt
 from comfy_execution.graph_utils import is_link
 
 NODE_CLASS_CONTAINS_UNIQUE_ID: Dict[str, bool] = {}
@@ -13,7 +14,8 @@ 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()
+    with folder_paths.cache_helper:  # Because we don't care about other except UNIQUE_ID
+        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:
diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py
index 59b42b746..64769898d 100644
--- a/comfy_execution/graph.py
+++ b/comfy_execution/graph.py
@@ -1,3 +1,4 @@
+import folder_paths
 import nodes
 
 from comfy_execution.graph_utils import is_link
@@ -126,7 +127,8 @@ class TopologicalSort:
                     from_node_id, from_socket = value
                     if subgraph_nodes is not None and from_node_id not in subgraph_nodes:
                         continue
-                    input_type, input_category, input_info = self.get_input_info(unique_id, input_name)
+                    with folder_paths.cache_helper:  # Because we don't care about other except lazy
+                        input_type, input_category, 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)
diff --git a/execution.py b/execution.py
index fcb4f6f40..47be48cc4 100644
--- a/execution.py
+++ b/execution.py
@@ -10,6 +10,7 @@ import inspect
 from typing import List, Literal, NamedTuple, Optional
 
 import torch
+import folder_paths
 import nodes
 
 import comfy.model_management
@@ -88,7 +89,8 @@ class CacheSet:
         return result
 
 def get_input_data(inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}):
-    valid_inputs = class_def.INPUT_TYPES()
+    with folder_paths.cache_helper:  # Only `rawLink` has been used so far
+        valid_inputs = class_def.INPUT_TYPES()
     input_data_all = {}
     missing_keys = {}
     for x in inputs:
diff --git a/folder_paths.py b/folder_paths.py
index 72c70f594..33e61143d 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -60,6 +60,7 @@ class CacheHelper:
     def __init__(self):
         self.cache: dict[str, tuple[list[str], dict[str, float], float]] = {}
         self.active = False
+        self.refresh = False
 
     def get(self, key: str, default=None) -> tuple[list[str], dict[str, float], float]:
         if not self.active:
@@ -67,19 +68,24 @@ class CacheHelper:
         return self.cache.get(key, default)
 
     def set(self, key: str, value: tuple[list[str], dict[str, float], float]) -> None:
-        if self.active:
-            self.cache[key] = value
+        self.cache[key] = value
 
     def clear(self):
         self.cache.clear()
 
     def __enter__(self):
+        if self.refresh:
+            self.clear()
+            self.refresh = False
         self.active = True
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.active = False
-        self.clear()
+
+    def __call__(self, refresh=False):
+        self.refresh = refresh
+        return self
 
 cache_helper = CacheHelper()
 
diff --git a/server.py b/server.py
index 76a99167d..a00bdfd50 100644
--- a/server.py
+++ b/server.py
@@ -584,7 +584,7 @@ class PromptServer():
 
         @routes.get("/object_info")
         async def get_object_info(request):
-            with folder_paths.cache_helper:
+            with folder_paths.cache_helper(refresh=True):
                 out = {}
                 for x in nodes.NODE_CLASS_MAPPINGS:
                     try: