From 6a6519b1e04519c9aa08509be0f3c69fb8e2252d Mon Sep 17 00:00:00 2001 From: rlfetchr Date: Fri, 6 Sep 2024 18:07:51 +0100 Subject: [PATCH] use memory mapping to accelerate hashing of large models. --- py/model_info.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/py/model_info.py b/py/model_info.py index b360baf..ee7d120 100644 --- a/py/model_info.py +++ b/py/model_info.py @@ -1,11 +1,21 @@ import hashlib import json +import mmap from aiohttp import web from server import PromptServer import folder_paths import os +async def generate_sha256(path): + hasher = hashlib.sha256() + with open(path, "rb") as f: + with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm: + hasher.update(mm) + + return hasher.hexdigest() + + def get_metadata(filepath): with open(filepath, "rb") as file: # https://github.com/huggingface/safetensors#format @@ -28,7 +38,7 @@ async def save_notes(request): name = request.match_info["name"] pos = name.index("/") type = name[0:pos] - name = name[pos+1:] + name = name[pos + 1:] file_path = None if type == "embeddings" or type == "loras": @@ -64,7 +74,7 @@ async def load_metadata(request): name = request.match_info["name"] pos = name.index("/") type = name[0:pos] - name = name[pos+1:] + name = name[pos + 1:] file_path = None if type == "embeddings" or type == "loras": @@ -107,8 +117,7 @@ async def load_metadata(request): with open(hash_file, "rt") as f: meta["pysssss.sha256"] = f.read() else: - with open(file_path, "rb") as f: - meta["pysssss.sha256"] = hashlib.sha256(f.read()).hexdigest() + meta["pysssss.sha256"] = await generate_sha256(file_path) with open(hash_file, "wt") as f: f.write(meta["pysssss.sha256"])