Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
Vladimir Mandic
2023-05-20 13:36:27 -04:00
parent 2ce9852cff
commit b6289d56c7
2 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sd_models.py
+4 -4
View File
@@ -223,19 +223,19 @@ def read_state_dict(checkpoint_file, map_location=None): # pylint: disable=unuse
_, extension = os.path.splitext(checkpoint_file)
if shared.opts.stream_load:
if extension.lower() == ".safetensors":
shared.log.debug('Model weights loading: type=safetensors mode=buffered')
# shared.log.debug('Model weights loading: type=safetensors mode=buffered')
buffer = f.read()
pl_sd = safetensors.torch.load(buffer)
else:
shared.log.debug('Model weights loading: type=checkpoint mode=buffered')
# shared.log.debug('Model weights loading: type=checkpoint mode=buffered')
buffer = io.BytesIO(f.read())
pl_sd = torch.load(buffer, map_location='cpu')
else:
if extension.lower() == ".safetensors":
shared.log.debug('Model weights loading: type=safetensors mode=mmap')
# shared.log.debug('Model weights loading: type=safetensors mode=mmap')
pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu')
else:
shared.log.debug('Model weights loading: type=checkpoint mode=direct')
# shared.log.debug('Model weights loading: type=checkpoint mode=direct')
pl_sd = torch.load(f, map_location='cpu')
sd = get_state_dict_from_checkpoint(pl_sd)
del pl_sd