From 3049afcfd089fa70fd7461d1f1c1cfc23d0406fe Mon Sep 17 00:00:00 2001
From: "Alex \"mcmonkey\" Goodwin" <git_commits@alexgoodwin.dev>
Date: Thu, 19 Sep 2024 01:45:13 +0900
Subject: [PATCH] very simple strong-cache on model list

---
 folder_paths.py | 23 +++++++++++++++++++++++
 server.py       |  1 +
 2 files changed, 24 insertions(+)

diff --git a/folder_paths.py b/folder_paths.py
index 38fad6238..b573fa281 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -46,6 +46,24 @@ user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user
 
 filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {}
 
+class CacheHelper:
+    """
+    Helper class for managing file list cache data.
+    """
+    def __init__(self):
+        self.cache: dict[str, tuple[list[str], dict[str, float], float]] = {}
+
+    def get(self, key: str, default=None) -> tuple[list[str], dict[str, float], float]:
+        return self.cache.get(key, default)
+    
+    def set(self, key: str, value: tuple[list[str], dict[str, float], float]) -> None:
+        self.cache[key] = value
+
+    def clear(self):
+        self.cache.clear()
+
+cache_helper = CacheHelper()
+
 extension_mimetypes_cache = {
     "webp" : "image",
 }
@@ -257,6 +275,10 @@ def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], f
     return sorted(list(output_list)), output_folders, time.perf_counter()
 
 def cached_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float] | None:
+    strong_cache = cache_helper.get(folder_name)
+    if strong_cache is not None:
+        return strong_cache
+    
     global filename_list_cache
     global folder_names_and_paths
     folder_name = map_legacy(folder_name)
@@ -285,6 +307,7 @@ def get_filename_list(folder_name: str) -> list[str]:
         out = get_filename_list_(folder_name)
         global filename_list_cache
         filename_list_cache[folder_name] = out
+        cache_helper.set(folder_name, out)
     return list(out[0])
 
 def get_save_image_path(filename_prefix: str, output_dir: str, image_width=0, image_height=0) -> tuple[str, str, int, str, str]:
diff --git a/server.py b/server.py
index e370e2467..7bf13b4a6 100644
--- a/server.py
+++ b/server.py
@@ -552,6 +552,7 @@ class PromptServer():
 
         @routes.get("/object_info")
         async def get_object_info(request):
+            folder_paths.cache_helper.clear()
             out = {}
             for x in nodes.NODE_CLASS_MAPPINGS:
                 try: