mirrors/stable-diffusion-webui
1
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2026-02-08 08:42:08 -08:00
Code Issues Projects Releases Packages Wiki Activity

Attempting to solve slow loads for safetensors.

Browse source 
Fixes #5893
This commit is contained in:
Nicolas Patry 2022-12-21 13:45:58 +01:00
parent 4af3ca5393
commit 5ba04f9ec0
No known key found for this signature in database
GPG key ID: 6AE76ACC68FFBAF9
1 changed files with 4 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

5
modules/sd_models.py
View file

@ -168,7 +168,10 @@ def get_state_dict_from_checkpoint(pl_sd):
def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):
_, extension = os.path.splitext(checkpoint_file)
if extension.lower() == ".safetensors":
pl_sd = safetensors.torch.load_file(checkpoint_file, device=map_location or shared.weight_load_location)
device = map_location or shared.weight_load_location
if device is None:
device = "cuda:0" if torch.cuda.is_available() else "cpu"
pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
else:
pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)