diff options
| author | dhwz <dhwz@gmx.net> | 2023-06-16 16:10:15 +0000 |
|---|---|---|
| committer | dhwz <dhwz@gmx.net> | 2023-06-16 16:10:15 +0000 |
| commit | 41363e0d27bbaa0c84eebe3c7c8451075390ec4e (patch) | |
| tree | 31047355fe5851188c63bcf57651f1b1a33002ab /modules/sd_models.py | |
| parent | 59419bd64a1581caccaac04dceb66c1c069a2db1 (diff) | |
| download | stable-diffusion-webui-gfx803-41363e0d27bbaa0c84eebe3c7c8451075390ec4e.tar.gz stable-diffusion-webui-gfx803-41363e0d27bbaa0c84eebe3c7c8451075390ec4e.tar.bz2 stable-diffusion-webui-gfx803-41363e0d27bbaa0c84eebe3c7c8451075390ec4e.zip | |
fix very slow loading speed of .safetensors files
Diffstat (limited to 'modules/sd_models.py')
| -rw-r--r-- | modules/sd_models.py | 7 |
1 files changed, 5 insertions, 2 deletions
diff --git a/modules/sd_models.py b/modules/sd_models.py index 918f6fd6..d9ac675b 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -247,8 +247,11 @@ def read_metadata_from_safetensors(filename): def read_state_dict(checkpoint_file, print_global_state=False, map_location=None):
_, extension = os.path.splitext(checkpoint_file)
if extension.lower() == ".safetensors":
- device = map_location or shared.weight_load_location or devices.get_optimal_device_name()
- pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
+ if not shared.opts.disable_mmap_load_safetensors:
+ device = map_location or shared.weight_load_location or devices.get_optimal_device_name()
+ pl_sd = safetensors.torch.load_file(checkpoint_file, device=device)
+ else:
+ pl_sd = safetensors.torch.load(open(checkpoint_file, 'rb').read())
else:
pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)
|