Merge pull request #4762 from awsr/hashcache

Update hash cache

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)

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
 

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
 
 

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 = {}

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")