Update utils.py: fix very slow loading speed of .safetensors files

This commit is contained in:
Yaruze66 2023-07-08 00:07:47 +05:00 committed by GitHub
parent c2d407b0f7
commit 37f93134d0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -6,7 +6,7 @@ import safetensors.torch
def load_torch_file(ckpt, safe_load=False):
if ckpt.lower().endswith(".safetensors"):
sd = safetensors.torch.load_file(ckpt, device="cpu")
sd = safetensors.torch.load(open(ckpt, 'rb').read())
else:
if safe_load:
if not 'weights_only' in torch.load.__code__.co_varnames: