From 08c6f009a5ee92dd3218a942c08e8337c26352be Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sat, 14 Jan 2023 15:55:40 +0300
Subject: load hashes from cache for checkpoints that have them add checkpoint
 hash to footer

---
 modules/hashes.py    | 26 +++++++++++++++++++-------
 modules/sd_models.py |  9 ++++++---
 modules/shared.py    |  1 +
 modules/ui.py        |  2 ++
 4 files changed, 28 insertions(+), 10 deletions(-)

(limited to 'modules')

diff --git a/modules/hashes.py b/modules/hashes.py
index ebfbd90c..14231771 100644
--- a/modules/hashes.py
+++ b/modules/hashes.py
@@ -42,23 +42,35 @@ def calculate_sha256(filename):
     return hash_sha256.hexdigest()
 
 
-def sha256(filename, title):
+def sha256_from_cache(filename, title):
     hashes = cache("hashes")
     ondisk_mtime = os.path.getmtime(filename)
 
-    if title in hashes:
-        cached_sha256 = hashes[title].get("sha256", None)
-        cached_mtime = hashes[title].get("mtime", 0)
+    if title not in hashes:
+        return None
+
+    cached_sha256 = hashes[title].get("sha256", None)
+    cached_mtime = hashes[title].get("mtime", 0)
+
+    if ondisk_mtime > cached_mtime or cached_sha256 is None:
+        return None
+
+    return cached_sha256
+
+
+def sha256(filename, title):
+    hashes = cache("hashes")
 
-        if ondisk_mtime <= cached_mtime and cached_sha256 is not None:
-            return cached_sha256
+    sha256_value = sha256_from_cache(filename, title)
+    if sha256_value is not None:
+        return sha256_value
 
     print(f"Calculating sha256 for {filename}: ", end='')
     sha256_value = calculate_sha256(filename)
     print(f"{sha256_value}")
 
     hashes[title] = {
-        "mtime": ondisk_mtime,
+        "mtime": os.path.getmtime(filename),
         "sha256": sha256_value,
     }
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 1fe6d11b..e5a0bc63 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -44,9 +44,11 @@ class CheckpointInfo:
         self.title = name
         self.model_name = os.path.splitext(name.replace("/", "_").replace("\\", "_"))[0]
         self.hash = model_hash(filename)
-        self.ids = [self.hash, self.model_name, self.title, f'{name} [{self.hash}]']
-        self.shorthash = None
-        self.sha256 = None
+
+        self.sha256 = hashes.sha256_from_cache(self.filename, "checkpoint/" + self.title)
+        self.shorthash = self.sha256[0:10] if self.sha256 else None
+
+        self.ids = [self.hash, self.model_name, self.title, f'{name} [{self.hash}]'] + ([self.shorthash, self.sha256] if self.shorthash else [])
 
     def register(self):
         checkpoints_list[self.title] = self
@@ -269,6 +271,7 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, vae_file="auto"):
     model.sd_model_hash = sd_model_hash
     model.sd_model_checkpoint = checkpoint_info.filename
     model.sd_checkpoint_info = checkpoint_info
+    shared.opts.data["sd_checkpoint_hash"] = checkpoint_info.sha256
 
     model.logvar = model.logvar.to(devices.device)  # fix for training
 
diff --git a/modules/shared.py b/modules/shared.py
index a6c61db3..c9988d4d 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -458,6 +458,7 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
 
 options_templates.update(options_section((None, "Hidden options"), {
     "disabled_extensions": OptionInfo([], "Disable those extensions"),
+    "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"),
 }))
 
 options_templates.update()
diff --git a/modules/ui.py b/modules/ui.py
index e86a624b..2625ae32 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1841,4 +1841,6 @@ xformers: {xformers_version}
 gradio: {gr.__version__}
  • 
 commit: <a href="https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/{commit}">{short_commit}</a>
+ • 
+checkpoint: <a id="sd_checkpoint_hash">N/A</a>
 """
-- 
cgit v1.2.3