Skip to content

Commit

Permalink
very simple strong-cache on model list (comfyanonymous#4969)
Browse files Browse the repository at this point in the history
* very simple strong-cache on model list

* store the cache after validation too

* only cache object_info for now

* use a 'with' context
  • Loading branch information
mcmonkey4eva authored Sep 19, 2024
1 parent 0bfc7cc commit a1e71cf
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 8 deletions.
35 changes: 35 additions & 0 deletions folder_paths.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,36 @@

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]] = {}
self.active = False

def get(self, key: str, default=None) -> tuple[list[str], dict[str, float], float]:
if not self.active:
return default
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

def clear(self):
self.cache.clear()

def __enter__(self):
self.active = True
return self

def __exit__(self, exc_type, exc_value, traceback):
self.active = False
self.clear()

cache_helper = CacheHelper()

extension_mimetypes_cache = {
"webp" : "image",
}
Expand Down Expand Up @@ -257,6 +287,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)
Expand Down Expand Up @@ -285,6 +319,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]:
Expand Down
17 changes: 9 additions & 8 deletions server.py
Original file line number Diff line number Diff line change
Expand Up @@ -552,14 +552,15 @@ def node_info(node_class):

@routes.get("/object_info")
async def get_object_info(request):
out = {}
for x in nodes.NODE_CLASS_MAPPINGS:
try:
out[x] = node_info(x)
except Exception as e:
logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
logging.error(traceback.format_exc())
return web.json_response(out)
with folder_paths.cache_helper:
out = {}
for x in nodes.NODE_CLASS_MAPPINGS:
try:
out[x] = node_info(x)
except Exception as e:
logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")
logging.error(traceback.format_exc())
return web.json_response(out)

@routes.get("/object_info/{node_class}")
async def get_object_info_node(request):
Expand Down

0 comments on commit a1e71cf

Please sign in to comment.