Remove insecure torch.load calls (huggingface#7393)

update
EMinsight · Mar 19, 2024 · 4da810b · 4da810b
1 parent 161c6e1
commit 4da810b
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 8 deletions.
diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py
@@ -19,7 +19,7 @@
 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 @@ def load_ip_adapter(
                             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
 

diff --git a/src/diffusers/loaders/lora.py b/src/diffusers/loaders/lora.py
@@ -25,7 +25,7 @@
 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 @@ def lora_state_dict(
                     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
 

diff --git a/src/diffusers/loaders/textual_inversion.py b/src/diffusers/loaders/textual_inversion.py
@@ -18,6 +18,7 @@
 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
 

diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py
@@ -31,7 +31,7 @@
     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 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: Union[str, Dict
                     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
 

diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
@@ -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: