feat: update sha256 generation functions

29be1da7cf/modules/hashes.py

modules/meta_parser.py

@@ -12,7 +12,7 @@ import modules.config
 import modules.sdxl_styles
 from modules.flags import MetadataScheme, Performance, Steps
 from modules.flags import SAMPLERS, CIVITAI_NO_KARRAS
-from modules.util import quote, unquote, extract_styles_from_prompt, is_json, get_file_from_folder_list, calculate_sha256
+from modules.util import quote, unquote, extract_styles_from_prompt, is_json, get_file_from_folder_list, sha256
 
 re_param_code = r'\s*(\w[\w \-/]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)'
 re_param = re.compile(re_param_code)
@@ -204,7 +204,8 @@ def get_lora(key: str, fallback: str | None, source_dict: dict, results: list):
 def get_sha256(filepath):
     global hash_cache
     if filepath not in hash_cache:
-        hash_cache[filepath] = calculate_sha256(filepath)
+        is_safetensors = os.path.splitext(filepath)[1].lower() == '.safetensors'
+        hash_cache[filepath] = sha256(filepath, is_safetensors)
 
     return hash_cache[filepath]
 

modules/util.py

@@ -7,9 +7,9 @@ import math
 import os
 import cv2
 import json
+import hashlib
 
 from PIL import Image
-from hashlib import sha256
 
 import modules.sdxl_styles
 
@@ -182,16 +182,44 @@ def get_files_from_folder(folder_path, extensions=None, name_filter=None):
     return filenames
 
 
-def calculate_sha256(filename, length=HASH_SHA256_LENGTH) -> str:
-    hash_sha256 = sha256()
+def sha256(filename, use_addnet_hash=False, length=HASH_SHA256_LENGTH):
+    print(f"Calculating sha256 for {filename}: ", end='')
+    if use_addnet_hash:
+        with open(filename, "rb") as file:
+            sha256_value = addnet_hash_safetensors(file)
+    else:
+        sha256_value = calculate_sha256(filename)
+    print(f"{sha256_value}")
+
+    return sha256_value[:length] if length is not None else sha256_value
+
+
+def addnet_hash_safetensors(b):
+    """kohya-ss hash for safetensors from https://github.com/kohya-ss/sd-scripts/blob/main/library/train_util.py"""
+    hash_sha256 = hashlib.sha256()
+    blksize = 1024 * 1024
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+    for chunk in iter(lambda: b.read(blksize), b""):
+        hash_sha256.update(chunk)
+
+    return hash_sha256.hexdigest()
+
+
+def calculate_sha256(filename) -> str:
+    hash_sha256 = hashlib.sha256()
     blksize = 1024 * 1024
 
     with open(filename, "rb") as f:
         for chunk in iter(lambda: f.read(blksize), b""):
             hash_sha256.update(chunk)
 
-    res = hash_sha256.hexdigest()
-    return res[:length] if length else res
+    return hash_sha256.hexdigest()
 
 
 def quote(text):