diff options
| author | AUTOMATIC <16777216c@gmail.com> | 2023-06-27 06:19:04 +0000 |
|---|---|---|
| committer | AUTOMATIC <16777216c@gmail.com> | 2023-06-27 06:19:04 +0000 |
| commit | 24129368f1b732be25ef486edb2cf5a6ace66737 (patch) | |
| tree | e50fa0b140810eeb2698df5bdd49f2acdcc7d6d7 /modules/sd_models.py | |
| parent | 14196548c55dfe4775c96bdb939ce1a150933393 (diff) | |
| download | stable-diffusion-webui-gfx803-24129368f1b732be25ef486edb2cf5a6ace66737.tar.gz stable-diffusion-webui-gfx803-24129368f1b732be25ef486edb2cf5a6ace66737.tar.bz2 stable-diffusion-webui-gfx803-24129368f1b732be25ef486edb2cf5a6ace66737.zip | |
send tensors to the correct device when loading from safetensors file with memmap disabled for #11260
Diffstat (limited to 'modules/sd_models.py')
| -rw-r--r-- | modules/sd_models.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/modules/sd_models.py b/modules/sd_models.py index 0391398a..f65f4e36 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -246,11 +246,13 @@ 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()
+
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())
+ pl_sd = {k: v.to(device) for k, v in pl_sd.items()}
else:
pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location)
|