Skip to content

Commit

Permalink
update model_manager.py
Browse files Browse the repository at this point in the history
- read files in chunks when calculating sha
  - windows runner is crashing without
  • Loading branch information
mauwii committed Feb 6, 2023
1 parent fc2670b commit a40bdef
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions ldm/invoke/model_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -753,7 +753,7 @@ def search_models(self, search_folder):
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
Expand Down Expand Up @@ -954,7 +954,7 @@ def _push_newest_model(self,model_name:str) -> None:
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
Expand All @@ -976,7 +976,8 @@ def _diffuser_sha256(self,name_or_path:Union[str, Path])->Union[str,bytes]:
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))
Expand Down

0 comments on commit a40bdef

Please sign in to comment.