diff --git a/.gitignore b/.gitignore index 6790e9ee7..519b4a53d 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,4 @@ notification.mp3 /package-lock.json /.coverage* /test/test_outputs +/cache diff --git a/modules/cache.py b/modules/cache.py index a9822a0eb..9df248d7b 100644 --- a/modules/cache.py +++ b/modules/cache.py @@ -2,48 +2,47 @@ import json import os import os.path import threading -import time + +import diskcache +import tqdm from modules.paths import data_path, script_path cache_filename = os.environ.get('SD_WEBUI_CACHE_FILE', os.path.join(data_path, "cache.json")) -cache_data = None +cache_dir = os.environ.get('SD_WEBUI_CACHE_DIR', os.path.join(data_path, "cache")) +caches = {} cache_lock = threading.Lock() -dump_cache_after = None -dump_cache_thread = None - def dump_cache(): - """ - Marks cache for writing to disk. 5 seconds after no one else flags the cache for writing, it is written. - """ + """old function for dumping cache to disk; does nothing since diskcache.""" - global dump_cache_after - global dump_cache_thread + pass - def thread_func(): - global dump_cache_after - global dump_cache_thread - while dump_cache_after is not None and time.time() < dump_cache_after: - time.sleep(1) +def convert_old_cached_data(): + try: + with open(cache_filename, "r", encoding="utf8") as file: + data = json.load(file) + except FileNotFoundError: + return + except Exception: + os.replace(cache_filename, os.path.join(script_path, "tmp", "cache.json")) + print('[ERROR] issue occurred while trying to read cache.json; old cache has been moved to tmp/cache.json') + return - with cache_lock: - cache_filename_tmp = cache_filename + "-" - with open(cache_filename_tmp, "w", encoding="utf8") as file: - json.dump(cache_data, file, indent=4, ensure_ascii=False) + total_count = sum(len(keyvalues) for keyvalues in data.values()) - os.replace(cache_filename_tmp, cache_filename) + with tqdm.tqdm(total=total_count, desc="converting cache") as progress: + for subsection, keyvalues in data.items(): + cache_obj = caches.get(subsection) + if cache_obj is None: + cache_obj = diskcache.Cache(os.path.join(cache_dir, subsection)) + caches[subsection] = cache_obj - dump_cache_after = None - dump_cache_thread = None - - with cache_lock: - dump_cache_after = time.time() + 5 - if dump_cache_thread is None: - dump_cache_thread = threading.Thread(name='cache-writer', target=thread_func) - dump_cache_thread.start() + for key, value in keyvalues.items(): + cache_obj[key] = value + progress.update(1) def cache(subsection): @@ -54,28 +53,21 @@ def cache(subsection): subsection (str): The subsection identifier for the cache. Returns: - dict: The cache data for the specified subsection. + diskcache.Cache: The cache data for the specified subsection. """ - global cache_data - - if cache_data is None: + cache_obj = caches.get(subsection) + if not cache_obj: with cache_lock: - if cache_data is None: - try: - with open(cache_filename, "r", encoding="utf8") as file: - cache_data = json.load(file) - except FileNotFoundError: - cache_data = {} - except Exception: - os.replace(cache_filename, os.path.join(script_path, "tmp", "cache.json")) - print('[ERROR] issue occurred while trying to read cache.json, move current cache to tmp/cache.json and create new cache') - cache_data = {} + if not os.path.exists(cache_dir) and os.path.isfile(cache_filename): + convert_old_cached_data() - s = cache_data.get(subsection, {}) - cache_data[subsection] = s + cache_obj = caches.get(subsection) + if not cache_obj: + cache_obj = diskcache.Cache(os.path.join(cache_dir, subsection)) + caches[subsection] = cache_obj - return s + return cache_obj def cached_data_for_file(subsection, title, filename, func): diff --git a/requirements.txt b/requirements.txt index 731a1be7d..8699c02be 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,6 +4,7 @@ accelerate blendmodes clean-fid +diskcache einops facexlib fastapi>=0.90.1 diff --git a/requirements_versions.txt b/requirements_versions.txt index 5e30b5ea1..87aae9136 100644 --- a/requirements_versions.txt +++ b/requirements_versions.txt @@ -3,6 +3,7 @@ Pillow==9.5.0 accelerate==0.21.0 blendmodes==2022 clean-fid==0.1.35 +diskcache==5.6.3 einops==0.4.1 facexlib==0.3.0 fastapi==0.94.0