update model_manager.py

- read files in chunks when calculating sha
  - windows runner is crashing without
This commit is contained in:
mauwii 2023-02-05 07:13:42 +01:00 committed by Matthias Wild
parent fc2670b4d6
commit a40bdef29f

View File

@ -753,7 +753,7 @@ class ModelManager(object):
return search_folder, found_models
def _choose_diffusers_vae(self, model_name:str, vae:str=None)->Union[dict,str]:
# In the event that the original entry is using a custom ckpt VAE, we try to
# map that VAE onto a diffuser VAE using a hard-coded dictionary.
# I would prefer to do this differently: We load the ckpt model into memory, swap the
@ -954,7 +954,7 @@ class ModelManager(object):
def _has_cuda(self) -> bool:
return self.device.type == 'cuda'
def _diffuser_sha256(self,name_or_path:Union[str, Path])->Union[str,bytes]:
def _diffuser_sha256(self,name_or_path:Union[str, Path],chunksize=4096)->Union[str,bytes]:
path = None
if isinstance(name_or_path,Path):
path = name_or_path
@ -976,7 +976,8 @@ class ModelManager(object):
for name in files:
count += 1
with open(os.path.join(root,name),'rb') as f:
sha.update(f.read())
while chunk := f.read(chunksize):
sha.update(chunk)
hash = sha.hexdigest()
toc = time.time()
print(f' | sha256 = {hash} ({count} files hashed in','%4.2fs)' % (toc - tic))