From d9f72b066f612e0220f61d142cb5e3106df6afb8 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Wed, 31 May 2023 09:14:15 -0400
Subject: [PATCH] precalc hashes

---
 modules/hashes.py    |  6 +-----
 modules/lora         |  2 +-
 modules/sd_models.py | 19 +++++++++++++++++++
 modules/ui.py        |  2 ++
 4 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/modules/hashes.py b/modules/hashes.py
index fb7ce62fa..3dd9bdb6c 100644
--- a/modules/hashes.py
+++ b/modules/hashes.py
@@ -17,8 +17,7 @@ def dump_cache():
 
 
 def cache(subsection):
-    global cache_data
-
+    global cache_data # pylint: disable=global-statement
     if cache_data is None:
         with filelock.FileLock(f"{cache_filename}.lock"):
             if not os.path.isfile(cache_filename):
@@ -26,10 +25,8 @@ def cache(subsection):
             else:
                 with open(cache_filename, "r", encoding="utf8") as file:
                     cache_data = json.load(file)
-
     s = cache_data.get(subsection, {})
     cache_data[subsection] = s
-
     return s
 
 
@@ -86,4 +83,3 @@ def addnet_hash_safetensors(b):
     for chunk in iter(lambda: b.read(blksize), b""):
         hash_sha256.update(chunk)
     return hash_sha256.hexdigest()
-
diff --git a/modules/lora b/modules/lora
index 16e5981d3..8a5e3904a 160000
--- a/modules/lora
+++ b/modules/lora
@@ -1 +1 @@
-Subproject commit 16e5981d3153ba02c34445089b998c5002a60abc
+Subproject commit 8a5e3904a07362bf380b27c65241849b57502f91
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 7b1444e6a..5a5737bec 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -136,6 +136,25 @@ def list_models():
                     checkpoint_info.register()
 
 
+def update_model_hashes():
+    txt = []
+    lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None]
+    shared.log.info(f'Models list: short hash missing for {len(lst)} out of {len(checkpoints_list)} models')
+    for ckpt in lst:
+        ckpt.hash = model_hash(ckpt.filename)
+        txt.append(f'Calculated short hash: <b>{ckpt.title}</b> {ckpt.hash}')
+    txt.append(f'Updated short hashes for <b>{len(lst)}</b> out of <b>{len(checkpoints_list)}</b> models')
+    lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.sha256 is None or ckpt.shorthash is None]
+    shared.log.info(f'Models list: full hash missing for {len(lst)} out of {len(checkpoints_list)} models')
+    for ckpt in lst:
+        ckpt.sha256 = hashes.sha256(ckpt.filename, f"checkpoint/{ckpt.name}")
+        ckpt.shorthash = ckpt.sha256[0:10]
+        txt.append(f'Calculated full hash: <b>{ckpt.title}</b> {ckpt.shorthash}')
+    txt.append(f'Updated full hashes for <b>{len(lst)}</b> out of <b>{len(checkpoints_list)}</b> models')
+    txt = '<br>'.join(txt)
+    return txt
+
+
 def get_closet_checkpoint_match(search_string):
     checkpoint_info = checkpoint_aliases.get(search_string, None)
     if checkpoint_info is not None:
diff --git a/modules/ui.py b/modules/ui.py
index e6dfc2fb3..4dd6bf4ed 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -934,6 +934,7 @@ def create_ui():
                                 discard_weights = gr.Textbox(value="", label="Discard weights with matching name", elem_id="modelmerger_discard_weights")
                             with gr.Row():
                                 modelmerger_merge = gr.Button(elem_id="modelmerger_merge", value="Merge", variant='primary')
+                                model_checkhash = gr.Button(elem_id="modelmerger_hash", value="Calculate hash for all models (may take a long time)", variant='primary')
 
                         with gr.Column(variant='compact', elem_id="modelmerger_results_container"):
                             with gr.Group(elem_id="modelmerger_results_panel"):
@@ -1519,6 +1520,7 @@ def create_ui():
                 modelmerger_result,
             ]
         )
+        model_checkhash.click(fn=sd_models.update_model_hashes, inputs=[], outputs=[modelmerger_result])
 
     ui_config_file = cmd_opts.ui_config
     ui_settings = {}