diff --git a/src/diffusers/loaders/ip_adapter.py b/src/diffusers/loaders/ip_adapter.py index 93959b9f0a..a4593ec694 100644 --- a/src/diffusers/loaders/ip_adapter.py +++ b/src/diffusers/loaders/ip_adapter.py @@ -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 diff --git a/src/diffusers/loaders/lora.py b/src/diffusers/loaders/lora.py index c6077f3a8e..973fcbcf12 100644 --- a/src/diffusers/loaders/lora.py +++ b/src/diffusers/loaders/lora.py @@ -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 diff --git a/src/diffusers/loaders/textual_inversion.py b/src/diffusers/loaders/textual_inversion.py index aaaf4b68bb..eb727990af 100644 --- a/src/diffusers/loaders/textual_inversion.py +++ b/src/diffusers/loaders/textual_inversion.py @@ -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 diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py index 0a9544d0db..c3449d53fc 100644 --- a/src/diffusers/loaders/unet.py +++ b/src/diffusers/loaders/unet.py @@ -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 diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index bc5dfaee04..95ae6f4fc4 100644 --- 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: