From 3de44fc580681f785065565fbe63f46c5a872d43 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 12 Sep 2022 20:47:46 +0300
Subject: [PATCH] Include the model name (or the SHA256 of the file) in the
 metadata #271

---
 modules/images.py     |  4 +++-
 modules/processing.py |  1 +
 modules/shared.py     |  4 +++-
 webui.py              | 10 +++++++++-
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/modules/images.py b/modules/images.py
index d742ed98..a08b549a 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -9,7 +9,7 @@ from fonts.ttf import Roboto
 import string
 
 import modules.shared
-from modules import sd_samplers
+from modules import sd_samplers, shared
 from modules.shared import opts
 
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
@@ -278,6 +278,8 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
         file_decoration = file_decoration.replace("[height]", str(p.height))
         file_decoration = file_decoration.replace("[sampler]", sd_samplers.samplers[p.sampler_index].name)
 
+    file_decoration = file_decoration.replace("[model_hash]", shared.sd_model_hash)
+
     if extension == 'png' and opts.enable_pnginfo and info is not None:
         pnginfo = PngImagePlugin.PngInfo()
 
diff --git a/modules/processing.py b/modules/processing.py
index 23b0c08f..17b4a37b 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -188,6 +188,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             "Seed": all_seeds[index],
             "Face restoration": (opts.face_restoration_model if p.restore_faces else None),
             "Size": f"{p.width}x{p.height}",
+            "Model hash": (None if not opts.add_model_hash_to_info or not shared.sd_model_hash else shared.sd_model_hash),
             "Batch size": (None if p.batch_size < 2 else p.batch_size),
             "Batch pos": (None if p.batch_size < 2 else position_in_batch),
             "Variation seed": (None if p.subseed_strength == 0 else all_subseeds[index]),
diff --git a/modules/shared.py b/modules/shared.py
index ea1c879b..8e07c9b1 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -97,7 +97,7 @@ class Options:
     data = None
     hide_dirs = {"visible": False} if cmd_opts.hide_ui_dir_config else None
     data_labels = {
-        "samples_filename_format": OptionInfo("", "Samples filename format using following tags: [steps],[cfg],[prompt],[prompt_spaces],[width],[height],[sampler],[seed]. Leave blank for default."),
+        "samples_filename_format": OptionInfo("", "Samples filename format using following tags: [steps], [cfg], [prompt], [prompt_spaces], [width], [height], [sampler], [seed], [model_hash]. Leave blank for default."),
         "outdir_samples": OptionInfo("", "Output directory for images; if empty, defaults to two directories below", component_args=hide_dirs),
         "outdir_txt2img_samples": OptionInfo("outputs/txt2img-images", 'Output directory for txt2img images', component_args=hide_dirs),
         "outdir_img2img_samples": OptionInfo("outputs/img2img-images", 'Output directory for img2img images', component_args=hide_dirs),
@@ -120,6 +120,7 @@ class Options:
         "jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
         "export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),
         "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),
+        "add_model_hash_to_info": OptionInfo(False, "Add model hash to generation information"),
         "font": OptionInfo("", "Font for image grids that have text"),
         "enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text text and [text] to make it pay less attention"),
         "save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),
@@ -178,6 +179,7 @@ if os.path.exists(config_filename):
 sd_upscalers = []
 
 sd_model = None
+sd_model_hash = ''
 
 progress_print_out = sys.stdout
 
diff --git a/webui.py b/webui.py
index 953e6620..add72123 100644
--- a/webui.py
+++ b/webui.py
@@ -35,7 +35,7 @@ realesrgan.setup_realesrgan()
 
 
 def load_model_from_config(config, ckpt, verbose=False):
-    print(f"Loading model from {ckpt}")
+    print(f"Loading model [{shared.sd_model_hash}] from {ckpt}")
     pl_sd = torch.load(ckpt, map_location="cpu")
     if "global_step" in pl_sd:
         print(f"Global Step: {pl_sd['global_step']}")
@@ -89,6 +89,14 @@ try:
 except Exception:
     pass
 
+with open(cmd_opts.ckpt, "rb") as file:
+    import hashlib
+    m = hashlib.sha256()
+
+    file.seek(0x100000)
+    m.update(file.read(0x10000))
+    shared.sd_model_hash = m.hexdigest()[0:8]
+
 sd_config = OmegaConf.load(cmd_opts.config)
 shared.sd_model = load_model_from_config(sd_config, cmd_opts.ckpt)
 shared.sd_model = (shared.sd_model if cmd_opts.no_half else shared.sd_model.half())