hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-10 10:50:09 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Use GPU for loading safetensors, disable export

Browse Source
This commit is contained in:
Tim Patton
2022-11-21 16:40:18 -05:00
parent e134b74ce9
commit 210cb4c128
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
modules/sd_models.py
View File

@@ -147,8 +147,9 @@ def torch_load(model_filename, model_info, map_override=None):
map_override=shared.weight_load_location if not map_override else map_override
if(checkpoint_types[model_info.exttype] == 'safetensors'):
# safely load weights
# TODO: safetensors supports zero copy fast load to gpu, see issue #684
return load_file(model_filename, device=map_override)
# TODO: safetensors supports zero copy fast load to gpu, see issue #684.
# GPU only for now, see https://github.com/huggingface/safetensors/issues/95
return load_file(model_filename, device='cuda')
else:
return torch.load(model_filename, map_location=map_override)