diff --git a/modules/civitai/api_civitai.py b/modules/civitai/api_civitai.py index 7612797d5..4a8302757 100644 --- a/modules/civitai/api_civitai.py +++ b/modules/civitai/api_civitai.py @@ -477,9 +477,8 @@ def post_check_local(request: dict): return {"found": {}} # Build reverse lookup: lowercase sha256 -> {filename, type} found = {} - hash_cache = hash_module.cache("hashes") - for title, entry in hash_cache.items(): - sha = entry.get("sha256") + for title, entry in hash_module.cache().items(): + sha = entry["sha256"] if not sha: continue parts = title.split("/", 1) diff --git a/modules/civitai/download_civitai.py b/modules/civitai/download_civitai.py index 402f2ac28..154b93448 100644 --- a/modules/civitai/download_civitai.py +++ b/modules/civitai/download_civitai.py @@ -285,12 +285,8 @@ class DownloadManager: prefix = model_type_map.get(item.model_type, item.model_type.lower()) name = os.path.splitext(item.filename)[0] title = f"{prefix}/{name}" - hash_cache = hashes.cache("hashes") - hash_cache[title] = { - "mtime": os.path.getmtime(final_file), - "sha256": item.expected_hash.lower(), - } - hashes.dump_cache() + hashes.cache().add_hash(title, os.path.getmtime(final_file), item.expected_hash.lower()) + hashes.save_cache() except Exception: pass diff --git a/modules/hashes.py b/modules/hashes.py index 64741253d..487271e52 100644 --- a/modules/hashes.py +++ b/modules/hashes.py @@ -1,5 +1,7 @@ import hashlib import os.path +from collections import defaultdict +from typing import Literal, TypeAlias, TypedDict from rich import progress, errors from modules.logger import console from modules.logger import log @@ -7,28 +9,44 @@ from modules.json_helpers import readfile, writefile from modules.paths import data_path -cache_filename = os.path.join(data_path, 'data', 'cache.json') -cache_data = None +class HashEntry(TypedDict): + mtime: float + sha256: str + + +class HashStore(dict[str, HashEntry]): + def __init__(self, *args): + super().__init__(*args) + + def add_hash(self, key: str, mtime: float = 0, sha256: str | None = None): + self.__setitem__(key, {"mtime": mtime, "sha256": sha256 or ""}) + + +default_hash_store = "hashes" +KnownHashStores: TypeAlias = Literal["hashes", "hashes-addnet"] # For autocomplete in IDE + +cache_filename = os.path.join(data_path, "data", "cache.json") progress_ok = True +# defaultdict allows for easily using new stores without needing to define them ahead of time +_data: defaultdict[str, HashStore] = defaultdict(HashStore) -def init_cache(): - global cache_data # pylint: disable=global-statement - if cache_data is None: - cache_data = {} if not os.path.isfile(cache_filename) else readfile(cache_filename, lock=True, as_type="dict") +def load_cache(): + if os.path.isfile(cache_filename): + for store, data in readfile(cache_filename, lock=True, as_type="dict").items(): + _data[store] = HashStore(data) -def dump_cache(): - writefile(cache_data, cache_filename) +def save_cache(): + # Don't include empty hash stores + filtered = filter(lambda item: len(item[1]) > 0, _data.items()) + writefile(dict(filtered), cache_filename) -def cache(subsection): - global cache_data # pylint: disable=global-statement - if cache_data is None: - cache_data = {} if not os.path.isfile(cache_filename) else readfile(cache_filename, lock=True, as_type="dict") - s = cache_data.get(subsection, {}) - cache_data[subsection] = s - return s +def cache(store: KnownHashStores | str | None = None) -> HashStore: + if store is None: + return _data[default_hash_store] + return _data[store] def calculate_sha256(filename, quiet=False): @@ -55,23 +73,21 @@ def calculate_sha256(filename, quiet=False): return hash_sha256.hexdigest() -def sha256_from_cache(filename, title, use_addnet_hash=False): - hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes") +def sha256_from_cache(filename: str, title: str, *, store: KnownHashStores | str | None = None): + hashes = cache(store) if title not in hashes: return None - cached_sha256 = hashes[title].get("sha256", None) - cached_mtime = hashes[title].get("mtime", 0) + cached = hashes[title] ondisk_mtime = os.path.getmtime(filename) if os.path.isfile(filename) else 0 - if ondisk_mtime > cached_mtime or cached_sha256 is None: + if ondisk_mtime > cached["mtime"] or not cached["sha256"]: return None - return cached_sha256 + return cached["sha256"] -def sha256(filename, title, use_addnet_hash=False): +def sha256(filename: str, title: str, *, store: KnownHashStores | str | None = None): from modules import shared global progress_ok # pylint: disable=global-statement - hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes") - sha256_value = sha256_from_cache(filename, title, use_addnet_hash) + sha256_value = sha256_from_cache(filename, title, store=store) if sha256_value is not None: return sha256_value if shared.cmd_opts.no_hashing: @@ -79,7 +95,7 @@ def sha256(filename, title, use_addnet_hash=False): if not os.path.isfile(filename): return None jobid = shared.state.begin("Hash") - if use_addnet_hash: + if store == "hashes-addnet": if progress_ok: try: with progress.open(filename, 'rb', description=f'[cyan]Calculating hash: [yellow]{filename}', auto_refresh=True, console=console) as f: @@ -92,12 +108,9 @@ def sha256(filename, title, use_addnet_hash=False): sha256_value = addnet_hash_safetensors(f) else: sha256_value = calculate_sha256(filename) - hashes[title] = { - "mtime": os.path.getmtime(filename), - "sha256": sha256_value - } + cache(store).add_hash(title, os.path.getmtime(filename), sha256_value) shared.state.end(jobid) - dump_cache() + save_cache() return sha256_value diff --git a/modules/lora/network.py b/modules/lora/network.py index 5d3fd9fc7..a6753a500 100644 --- a/modules/lora/network.py +++ b/modules/lora/network.py @@ -40,7 +40,7 @@ class NetworkOnDisk: m[k] = v self.metadata = m self.alias = self.metadata.get('ss_output_name', self.name) - sha256 = hashes.sha256_from_cache(self.filename, "lora/" + self.name) or hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=True) or self.metadata.get('sshs_model_hash') + sha256 = hashes.sha256_from_cache(self.filename, "lora/" + self.name) or hashes.sha256_from_cache(self.filename, "lora/" + self.name, store='hashes-addnet') or self.metadata.get('sshs_model_hash') self.set_hash(sha256) self.sd_version = self.detect_version() @@ -107,7 +107,7 @@ class NetworkOnDisk: def read_hash(self): if not self.hash: - self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '') + self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, store='hashes-addnet' if self.is_safetensors else None) or '') def get_info(self): data = {} diff --git a/webui.py b/webui.py index 11c48e8a5..b8d06697f 100644 --- a/webui.py +++ b/webui.py @@ -77,7 +77,7 @@ def initialize(): from concurrent.futures import ThreadPoolExecutor, as_completed modules.sd_checkpoint.init_metadata() - modules.hashes.init_cache() + modules.hashes.load_cache() modules.sd_samplers.list_samplers() timer.startup.record("samplers")