diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2023-06-27 09:11:08 +0300 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-06-27 09:11:08 +0300 |
| commit | 14196548c55dfe4775c96bdb939ce1a150933393 (patch) | |
| tree | 480bb6bbe68d30ee0fef88bdbeae5336b9cde976 | |
| parent | d35e24611172a49f7b3637c601dc0fb12c9d0326 (diff) | |
| parent | 41363e0d27bbaa0c84eebe3c7c8451075390ec4e (diff) | |
Merge pull request #11260 from dhwz/dev
fix very slow loading speed of .safetensors files
| -rw-r--r-- | modules/sd_models.py | 7 | ||||
| -rw-r--r-- | modules/shared.py | 1 |
2 files changed, 6 insertions, 2 deletions
diff --git a/modules/sd_models.py b/modules/sd_models.py index 6ff5d17d..0391398a 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -246,8 +246,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)
diff --git a/modules/shared.py b/modules/shared.py index da4918b5..4743a428 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -376,6 +376,7 @@ options_templates.update(options_section(('system', "System"), { "multiple_tqdm": OptionInfo(True, "Add a second progress bar to the console that shows progress for an entire job."),
"print_hypernet_extra": OptionInfo(False, "Print extra hypernetwork information to console."),
"list_hidden_files": OptionInfo(True, "Load models/files in hidden directories").info("directory is hidden if its name starts with \".\""),
+ "disable_mmap_load_safetensors": OptionInfo(False, "Disable memmapping for loading .safetensors files (fixes very slow loading speed in some cases)."),
}))
options_templates.update(options_section(('training', "Training"), {
|