From 33b1c5cb87bfeccc51d08cda4558721112861e63 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Mon, 1 Jul 2024 17:20:20 +0200 Subject: [PATCH 1/4] feat: add hash_cache --- .gitignore | 1 + modules/config.py | 4 ++++ modules/hash_cache.py | 41 +++++++++++++++++++++++++++++++++++++++++ modules/meta_parser.py | 20 ++++++-------------- 4 files changed, 52 insertions(+), 14 deletions(-) create mode 100644 modules/hash_cache.py diff --git a/.gitignore b/.gitignore index 85914986..8408db59 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ __pycache__ *.partial *.onnx sorted_styles.json +hash_cache.json /input /cache /language/default.json diff --git a/modules/config.py b/modules/config.py index a6767c37..1fd4412b 100644 --- a/modules/config.py +++ b/modules/config.py @@ -7,6 +7,7 @@ import args_manager import tempfile import modules.flags import modules.sdxl_styles +from modules.hash_cache import load_cache_from_file, save_cache_to_file from modules.model_loader import load_file_from_url from modules.extra_utils import makedirs_with_log, get_files_from_folder, try_eval_env_var @@ -755,3 +756,6 @@ def downloading_safety_checker_model(): update_files() +load_cache_from_file() +# write cache to file again for cleanup of invalid cache entries +save_cache_to_file() diff --git a/modules/hash_cache.py b/modules/hash_cache.py new file mode 100644 index 00000000..bc9f479b --- /dev/null +++ b/modules/hash_cache.py @@ -0,0 +1,41 @@ +from modules.util import sha256, HASH_SHA256_LENGTH +import os +import json +hash_cache_filename = 'hash_cache.json' +hash_cache = {} + + +def sha256_from_cache(filepath): + global hash_cache + if filepath not in hash_cache: + hash_cache[filepath] = sha256(filepath) + save_cache_to_file() + + return hash_cache[filepath] + + +def load_cache_from_file(): + global hash_cache + + try: + if os.path.exists(hash_cache_filename): + with open(hash_cache_filename, 'rt', encoding='utf-8') as fp: + for filepath, hash in json.load(fp).items(): + if not os.path.exists(filepath) or not isinstance(hash, str) and len(hash) != HASH_SHA256_LENGTH: + print(f'[Cache] Skipping invalid cache entry: {filepath}') + continue + hash_cache[filepath] = hash + print(f'[Cache] Warmed cache from file') + except Exception as e: + print(f'[Cache] Warming failed: {e}') + + +def save_cache_to_file(): + global hash_cache + + try: + with open(hash_cache_filename, 'wt', encoding='utf-8') as fp: + json.dump(hash_cache, fp, indent=4) + print(f'[Cache] Updated cache file') + except Exception as e: + print(f'[Cache] Saving failed: {e}') diff --git a/modules/meta_parser.py b/modules/meta_parser.py index ff930cc0..307b656d 100644 --- a/modules/meta_parser.py +++ b/modules/meta_parser.py @@ -9,16 +9,16 @@ from PIL import Image import fooocus_version import modules.config import modules.sdxl_styles +from modules import hash_cache from modules.flags import MetadataScheme, Performance, Steps from modules.flags import SAMPLERS, CIVITAI_NO_KARRAS -from modules.util import quote, unquote, extract_styles_from_prompt, is_json, get_file_from_folder_list, sha256 +from modules.hash_cache import sha256_from_cache +from modules.util import quote, unquote, extract_styles_from_prompt, is_json, get_file_from_folder_list re_param_code = r'\s*(\w[\w \-/]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)' re_param = re.compile(re_param_code) re_imagesize = re.compile(r"^(\d+)x(\d+)$") -hash_cache = {} - def load_parameter_button_click(raw_metadata: dict | str, is_generating: bool): loaded_parameter_dict = raw_metadata @@ -215,14 +215,6 @@ def get_lora(key: str, fallback: str | None, source_dict: dict, results: list, p results.append(1) -def get_sha256(filepath): - global hash_cache - if filepath not in hash_cache: - hash_cache[filepath] = sha256(filepath) - - return hash_cache[filepath] - - def parse_meta_from_preset(preset_content): assert isinstance(preset_content, dict) preset_prepared = {} @@ -290,18 +282,18 @@ class MetadataParser(ABC): self.base_model_name = Path(base_model_name).stem base_model_path = get_file_from_folder_list(base_model_name, modules.config.paths_checkpoints) - self.base_model_hash = get_sha256(base_model_path) + self.base_model_hash = sha256_from_cache(base_model_path) if refiner_model_name not in ['', 'None']: self.refiner_model_name = Path(refiner_model_name).stem refiner_model_path = get_file_from_folder_list(refiner_model_name, modules.config.paths_checkpoints) - self.refiner_model_hash = get_sha256(refiner_model_path) + self.refiner_model_hash = sha256_from_cache(refiner_model_path) self.loras = [] for (lora_name, lora_weight) in loras: if lora_name != 'None': lora_path = get_file_from_folder_list(lora_name, modules.config.paths_loras) - lora_hash = get_sha256(lora_path) + lora_hash = sha256_from_cache(lora_path) self.loras.append((Path(lora_name).stem, lora_weight, lora_hash)) self.vae_name = Path(vae_name).stem From 1fa7eeac84b261f76481f93d77c80e6f9ae5a4c0 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Mon, 1 Jul 2024 17:21:30 +0200 Subject: [PATCH 2/4] feat: change type of hash_cache file to txt, append to file if single line is provided --- .gitignore | 2 +- modules/hash_cache.py | 40 +++++++++++++++++++++++++++------------- 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/.gitignore b/.gitignore index 8408db59..5bf633a8 100644 --- a/.gitignore +++ b/.gitignore @@ -10,7 +10,7 @@ __pycache__ *.partial *.onnx sorted_styles.json -hash_cache.json +hash_cache.txt /input /cache /language/default.json diff --git a/modules/hash_cache.py b/modules/hash_cache.py index bc9f479b..b2117ae9 100644 --- a/modules/hash_cache.py +++ b/modules/hash_cache.py @@ -1,15 +1,18 @@ -from modules.util import sha256, HASH_SHA256_LENGTH -import os import json -hash_cache_filename = 'hash_cache.json' +import os + +from modules.util import sha256, HASH_SHA256_LENGTH + +hash_cache_filename = 'hash_cache.txt' hash_cache = {} def sha256_from_cache(filepath): global hash_cache if filepath not in hash_cache: - hash_cache[filepath] = sha256(filepath) - save_cache_to_file() + hash_value = sha256(filepath) + hash_cache[filepath] = hash_value + save_cache_to_file(filepath, hash_value) return hash_cache[filepath] @@ -20,22 +23,33 @@ def load_cache_from_file(): try: if os.path.exists(hash_cache_filename): with open(hash_cache_filename, 'rt', encoding='utf-8') as fp: - for filepath, hash in json.load(fp).items(): - if not os.path.exists(filepath) or not isinstance(hash, str) and len(hash) != HASH_SHA256_LENGTH: - print(f'[Cache] Skipping invalid cache entry: {filepath}') - continue - hash_cache[filepath] = hash + for line in fp: + entry = json.loads(line) + for filepath, hash_value in entry.items(): + if not os.path.exists(filepath) or not isinstance(hash_value, str) and len(hash_value) != HASH_SHA256_LENGTH: + print(f'[Cache] Skipping invalid cache entry: {filepath}') + continue + hash_cache[filepath] = hash_value print(f'[Cache] Warmed cache from file') except Exception as e: print(f'[Cache] Warming failed: {e}') -def save_cache_to_file(): +def save_cache_to_file(filename=None, hash_value=None): global hash_cache + if filename is not None and hash_value is not None: + items = [(filename, hash_value)] + mode = 'at' + else: + items = sorted(hash_cache.items()) + mode = 'wt' + try: - with open(hash_cache_filename, 'wt', encoding='utf-8') as fp: - json.dump(hash_cache, fp, indent=4) + with open(hash_cache_filename, mode, encoding='utf-8') as fp: + for filepath, hash_value in items: + json.dump({filepath: hash_value}, fp) + fp.write('\n') print(f'[Cache] Updated cache file') except Exception as e: print(f'[Cache] Saving failed: {e}') From b19ecf410b6cc7e69311385eb7750dbe1c9bf4e7 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Mon, 1 Jul 2024 17:44:52 +0200 Subject: [PATCH 3/4] feat: add attribute --rebuild-hash-cache, add handling --- args_manager.py | 5 ++++- modules/config.py | 14 +++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/args_manager.py b/args_manager.py index 5a2b37c9..27fe14a5 100644 --- a/args_manager.py +++ b/args_manager.py @@ -32,7 +32,10 @@ args_parser.parser.add_argument("--enable-describe-uov-image", action='store_tru help="Disables automatic description of uov images when prompt is empty", default=False) args_parser.parser.add_argument("--always-download-new-model", action='store_true', - help="Always download newer models ", default=False) + help="Always download newer models", default=False) + +args_parser.parser.add_argument("--rebuild-hash-cache", action='store_true', + help="Generates missing model and LoRA hashes.", default=False) args_parser.parser.set_defaults( disable_cuda_malloc=True, diff --git a/modules/config.py b/modules/config.py index 1fd4412b..2f958ff6 100644 --- a/modules/config.py +++ b/modules/config.py @@ -757,5 +757,17 @@ def downloading_safety_checker_model(): update_files() load_cache_from_file() -# write cache to file again for cleanup of invalid cache entries + +if args_manager.args.rebuild_hash_cache: + from modules.hash_cache import sha256_from_cache + from modules.util import get_file_from_folder_list + + for filename in model_filenames: + filepath = get_file_from_folder_list(filename, paths_checkpoints) + sha256_from_cache(filepath) + for filename in lora_filenames: + filepath = get_file_from_folder_list(filename, paths_loras) + sha256_from_cache(filepath) + +# write cache to file again for sorting and cleanup of invalid cache entries save_cache_to_file() From df2dd194cca115f92274cb1c4bbb66f0ef2282e9 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Mon, 1 Jul 2024 17:54:20 +0200 Subject: [PATCH 4/4] feat: only use hash cache logs for exceptions --- modules/config.py | 2 ++ modules/hash_cache.py | 4 +--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/config.py b/modules/config.py index 2f958ff6..e9e6b1cf 100644 --- a/modules/config.py +++ b/modules/config.py @@ -762,12 +762,14 @@ if args_manager.args.rebuild_hash_cache: from modules.hash_cache import sha256_from_cache from modules.util import get_file_from_folder_list + print('[Cache] Rebuilding hash cache') for filename in model_filenames: filepath = get_file_from_folder_list(filename, paths_checkpoints) sha256_from_cache(filepath) for filename in lora_filenames: filepath = get_file_from_folder_list(filename, paths_loras) sha256_from_cache(filepath) + print('[Cache] Done') # write cache to file again for sorting and cleanup of invalid cache entries save_cache_to_file() diff --git a/modules/hash_cache.py b/modules/hash_cache.py index b2117ae9..10566560 100644 --- a/modules/hash_cache.py +++ b/modules/hash_cache.py @@ -30,9 +30,8 @@ def load_cache_from_file(): print(f'[Cache] Skipping invalid cache entry: {filepath}') continue hash_cache[filepath] = hash_value - print(f'[Cache] Warmed cache from file') except Exception as e: - print(f'[Cache] Warming failed: {e}') + print(f'[Cache] Loading failed: {e}') def save_cache_to_file(filename=None, hash_value=None): @@ -50,6 +49,5 @@ def save_cache_to_file(filename=None, hash_value=None): for filepath, hash_value in items: json.dump({filepath: hash_value}, fp) fp.write('\n') - print(f'[Cache] Updated cache file') except Exception as e: print(f'[Cache] Saving failed: {e}')