LLM/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2025-12-06 12:34:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove insecure torch.load calls (#7393)

Browse Source 
update
This commit is contained in:
Dhruv Nair
2024-03-20 04:11:50 +05:30
committed by GitHub
parent 161c6e14b6
commit 4da810b943
5 changed files with 14 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/diffusers/loaders/ip_adapter.py
View File

@@ -19,7 +19,7 @@ import torch
from huggingface_hub.utils import validate_hf_hub_args
from safetensors import safe_open
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_state_dict
from ..utils import (
_get_model_file,
is_accelerate_available,
@@ -182,7 +182,7 @@ class IPAdapterMixin:
elif key.startswith("ip_adapter."):
state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
else:
state_dict = torch.load(model_file, map_location="cpu")
state_dict = load_state_dict(model_file)
else:
state_dict = pretrained_model_name_or_path_or_dict

4
src/diffusers/loaders/lora.py
View File

@@ -25,7 +25,7 @@ from packaging import version
from torch import nn
from .. import __version__
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_state_dict
from ..utils import (
USE_PEFT_BACKEND,
_get_model_file,
@@ -281,7 +281,7 @@ class LoraLoaderMixin:
subfolder=subfolder,
user_agent=user_agent,
)
state_dict = torch.load(model_file, map_location="cpu")
state_dict = load_state_dict(model_file)
else:
state_dict = pretrained_model_name_or_path_or_dict

3
src/diffusers/loaders/textual_inversion.py
View File

@@ -18,6 +18,7 @@ import torch
from huggingface_hub.utils import validate_hf_hub_args
from torch import nn
from ..models.modeling_utils import load_state_dict
from ..utils import _get_model_file, is_accelerate_available, is_transformers_available, logging
@@ -100,7 +101,7 @@ def load_textual_inversion_state_dicts(pretrained_model_name_or_paths, **kwargs)
subfolder=subfolder,
user_agent=user_agent,
)
state_dict = torch.load(model_file, map_location="cpu")
state_dict = load_state_dict(model_file)
else:
state_dict = pretrained_model_name_or_path

4
src/diffusers/loaders/unet.py
View File

@@ -31,7 +31,7 @@ from ..models.embeddings import (
IPAdapterPlusImageProjection,
MultiIPAdapterImageProjection,
)
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_model_dict_into_meta, load_state_dict
from ..utils import (
USE_PEFT_BACKEND,
_get_model_file,
@@ -214,7 +214,7 @@ class UNet2DConditionLoadersMixin:
subfolder=subfolder,
user_agent=user_agent,
)
state_dict = torch.load(model_file, map_location="cpu")
state_dict = load_state_dict(model_file)
else:
state_dict = pretrained_model_name_or_path_or_dict

7
src/diffusers/models/modeling_utils.py
View File

@@ -108,7 +108,12 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[
if file_extension == SAFETENSORS_FILE_EXTENSION:
return safetensors.torch.load_file(checkpoint_file, device="cpu")
else:
return torch.load(checkpoint_file, map_location="cpu")
weights_only_kwarg = {"weights_only": True} if is_torch_version(">=", "1.13") else {}
return torch.load(
checkpoint_file,
map_location="cpu",
**weights_only_kwarg,
)
except Exception as e:
try:
with open(checkpoint_file) as f: