wehub-resource-sync
862 lines
39 KiB
Python
862 lines
39 KiB
Python
# coding=utf-8
|
|
# Copyright 2023 The HuggingFace Inc. team.
|
|
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import inspect
|
|
import os
|
|
import warnings
|
|
from functools import partial
|
|
from typing import Callable, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
from huggingface_hub import hf_hub_download
|
|
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError
|
|
from packaging import version
|
|
from requests import HTTPError
|
|
from torch import Tensor, device
|
|
|
|
from .. import __version__
|
|
from ..utils import (
|
|
CONFIG_NAME,
|
|
DEPRECATED_REVISION_ARGS,
|
|
DIFFUSERS_CACHE,
|
|
FLAX_WEIGHTS_NAME,
|
|
HF_HUB_OFFLINE,
|
|
HUGGINGFACE_CO_RESOLVE_ENDPOINT,
|
|
SAFETENSORS_WEIGHTS_NAME,
|
|
WEIGHTS_NAME,
|
|
is_accelerate_available,
|
|
is_safetensors_available,
|
|
is_torch_version,
|
|
logging,
|
|
)
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
if is_torch_version(">=", "1.9.0"):
|
|
_LOW_CPU_MEM_USAGE_DEFAULT = True
|
|
else:
|
|
_LOW_CPU_MEM_USAGE_DEFAULT = False
|
|
|
|
|
|
if is_accelerate_available():
|
|
import accelerate
|
|
from accelerate.utils import set_module_tensor_to_device
|
|
from accelerate.utils.versions import is_torch_version
|
|
|
|
if is_safetensors_available():
|
|
import safetensors
|
|
|
|
|
|
def get_parameter_device(parameter: torch.nn.Module):
|
|
try:
|
|
return next(parameter.parameters()).device
|
|
except StopIteration:
|
|
# For torch.nn.DataParallel compatibility in PyTorch 1.5
|
|
|
|
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
|
|
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
|
|
return tuples
|
|
|
|
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
|
|
first_tuple = next(gen)
|
|
return first_tuple[1].device
|
|
|
|
|
|
def get_parameter_dtype(parameter: torch.nn.Module):
|
|
try:
|
|
return next(parameter.parameters()).dtype
|
|
except StopIteration:
|
|
# For torch.nn.DataParallel compatibility in PyTorch 1.5
|
|
|
|
def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
|
|
tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
|
|
return tuples
|
|
|
|
gen = parameter._named_members(get_members_fn=find_tensor_attributes)
|
|
first_tuple = next(gen)
|
|
return first_tuple[1].dtype
|
|
|
|
|
|
def load_state_dict(checkpoint_file: Union[str, os.PathLike], variant: Optional[str] = None):
|
|
"""
|
|
Reads a checkpoint file, returning properly formatted errors if they arise.
|
|
"""
|
|
try:
|
|
if os.path.basename(checkpoint_file) == _add_variant(WEIGHTS_NAME, variant):
|
|
return torch.load(checkpoint_file, map_location="cpu")
|
|
else:
|
|
return safetensors.torch.load_file(checkpoint_file, device="cpu")
|
|
except Exception as e:
|
|
try:
|
|
with open(checkpoint_file) as f:
|
|
if f.read().startswith("version"):
|
|
raise OSError(
|
|
"You seem to have cloned a repository without having git-lfs installed. Please install "
|
|
"git-lfs and run `git lfs install` followed by `git lfs pull` in the folder "
|
|
"you cloned."
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained "
|
|
"model. Make sure you have saved the model properly."
|
|
) from e
|
|
except (UnicodeDecodeError, ValueError):
|
|
raise OSError(
|
|
f"Unable to load weights from checkpoint file for '{checkpoint_file}' "
|
|
f"at '{checkpoint_file}'. "
|
|
"If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
|
|
)
|
|
|
|
|
|
def _load_state_dict_into_model(model_to_load, state_dict):
|
|
# Convert old format to new format if needed from a PyTorch state_dict
|
|
# copy state_dict so _load_from_state_dict can modify it
|
|
state_dict = state_dict.copy()
|
|
error_msgs = []
|
|
|
|
# PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
|
|
# so we need to apply the function recursively.
|
|
def load(module: torch.nn.Module, prefix=""):
|
|
args = (state_dict, prefix, {}, True, [], [], error_msgs)
|
|
module._load_from_state_dict(*args)
|
|
|
|
for name, child in module._modules.items():
|
|
if child is not None:
|
|
load(child, prefix + name + ".")
|
|
|
|
load(model_to_load)
|
|
|
|
return error_msgs
|
|
|
|
|
|
def _add_variant(weights_name: str, variant: Optional[str] = None) -> str:
|
|
if variant is not None:
|
|
splits = weights_name.split(".")
|
|
splits = splits[:-1] + [variant] + splits[-1:]
|
|
weights_name = ".".join(splits)
|
|
|
|
return weights_name
|
|
|
|
|
|
class ModelMixin(torch.nn.Module):
|
|
r"""
|
|
Base class for all models.
|
|
|
|
[`ModelMixin`] takes care of storing the configuration of the models and handles methods for loading, downloading
|
|
and saving models.
|
|
|
|
- **config_name** ([`str`]) -- A filename under which the model should be stored when calling
|
|
[`~models.ModelMixin.save_pretrained`].
|
|
"""
|
|
config_name = CONFIG_NAME
|
|
_automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]
|
|
_supports_gradient_checkpointing = False
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
@property
|
|
def is_gradient_checkpointing(self) -> bool:
|
|
"""
|
|
Whether gradient checkpointing is activated for this model or not.
|
|
|
|
Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
|
|
activations".
|
|
"""
|
|
return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())
|
|
|
|
def enable_gradient_checkpointing(self):
|
|
"""
|
|
Activates gradient checkpointing for the current model.
|
|
|
|
Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
|
|
activations".
|
|
"""
|
|
if not self._supports_gradient_checkpointing:
|
|
raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
|
|
self.apply(partial(self._set_gradient_checkpointing, value=True))
|
|
|
|
def disable_gradient_checkpointing(self):
|
|
"""
|
|
Deactivates gradient checkpointing for the current model.
|
|
|
|
Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
|
|
activations".
|
|
"""
|
|
if self._supports_gradient_checkpointing:
|
|
self.apply(partial(self._set_gradient_checkpointing, value=False))
|
|
|
|
def set_use_memory_efficient_attention_xformers(
|
|
self, valid: bool, attention_op: Optional[Callable] = None
|
|
) -> None:
|
|
# Recursively walk through all the children.
|
|
# Any children which exposes the set_use_memory_efficient_attention_xformers method
|
|
# gets the message
|
|
def fn_recursive_set_mem_eff(module: torch.nn.Module):
|
|
if hasattr(module, "set_use_memory_efficient_attention_xformers"):
|
|
module.set_use_memory_efficient_attention_xformers(valid, attention_op)
|
|
|
|
for child in module.children():
|
|
fn_recursive_set_mem_eff(child)
|
|
|
|
for module in self.children():
|
|
if isinstance(module, torch.nn.Module):
|
|
fn_recursive_set_mem_eff(module)
|
|
|
|
def enable_xformers_memory_efficient_attention(self, attention_op: Optional[Callable] = None):
|
|
r"""
|
|
Enable memory efficient attention as implemented in xformers.
|
|
|
|
When this option is enabled, you should observe lower GPU memory usage and a potential speed up at inference
|
|
time. Speed up at training time is not guaranteed.
|
|
|
|
Warning: When Memory Efficient Attention and Sliced attention are both enabled, the Memory Efficient Attention
|
|
is used.
|
|
|
|
Parameters:
|
|
attention_op (`Callable`, *optional*):
|
|
Override the default `None` operator for use as `op` argument to the
|
|
[`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention)
|
|
function of xFormers.
|
|
|
|
Examples:
|
|
|
|
```py
|
|
>>> import torch
|
|
>>> from diffusers import UNet2DConditionModel
|
|
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
|
|
|
|
>>> model = UNet2DConditionModel.from_pretrained(
|
|
... "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16
|
|
... )
|
|
>>> model = model.to("cuda")
|
|
>>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
|
|
```
|
|
"""
|
|
self.set_use_memory_efficient_attention_xformers(True, attention_op)
|
|
|
|
def disable_xformers_memory_efficient_attention(self):
|
|
r"""
|
|
Disable memory efficient attention as implemented in xformers.
|
|
"""
|
|
self.set_use_memory_efficient_attention_xformers(False)
|
|
|
|
def save_pretrained(
|
|
self,
|
|
save_directory: Union[str, os.PathLike],
|
|
is_main_process: bool = True,
|
|
save_function: Callable = None,
|
|
safe_serialization: bool = False,
|
|
variant: Optional[str] = None,
|
|
):
|
|
"""
|
|
Save a model and its configuration file to a directory, so that it can be re-loaded using the
|
|
`[`~models.ModelMixin.from_pretrained`]` class method.
|
|
|
|
Arguments:
|
|
save_directory (`str` or `os.PathLike`):
|
|
Directory to which to save. Will be created if it doesn't exist.
|
|
is_main_process (`bool`, *optional*, defaults to `True`):
|
|
Whether the process calling this is the main process or not. Useful when in distributed training like
|
|
TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
|
|
the main process to avoid race conditions.
|
|
save_function (`Callable`):
|
|
The function to use to save the state dictionary. Useful on distributed training like TPUs when one
|
|
need to replace `torch.save` by another method. Can be configured with the environment variable
|
|
`DIFFUSERS_SAVE_MODE`.
|
|
safe_serialization (`bool`, *optional*, defaults to `False`):
|
|
Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
|
|
variant (`str`, *optional*):
|
|
If specified, weights are saved in the format pytorch_model.<variant>.bin.
|
|
"""
|
|
if safe_serialization and not is_safetensors_available():
|
|
raise ImportError("`safe_serialization` requires the `safetensors library: `pip install safetensors`.")
|
|
|
|
if os.path.isfile(save_directory):
|
|
logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
|
|
return
|
|
|
|
os.makedirs(save_directory, exist_ok=True)
|
|
|
|
model_to_save = self
|
|
|
|
# Attach architecture to the config
|
|
# Save the config
|
|
if is_main_process:
|
|
model_to_save.save_config(save_directory)
|
|
|
|
# Save the model
|
|
state_dict = model_to_save.state_dict()
|
|
|
|
weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME
|
|
weights_name = _add_variant(weights_name, variant)
|
|
|
|
# Save the model
|
|
if safe_serialization:
|
|
safetensors.torch.save_file(
|
|
state_dict, os.path.join(save_directory, weights_name), metadata={"format": "pt"}
|
|
)
|
|
else:
|
|
torch.save(state_dict, os.path.join(save_directory, weights_name))
|
|
|
|
logger.info(f"Model weights saved in {os.path.join(save_directory, weights_name)}")
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
|
|
r"""
|
|
Instantiate a pretrained pytorch model from a pre-trained model configuration.
|
|
|
|
The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
|
|
the model, you should first set it back in training mode with `model.train()`.
|
|
|
|
The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
|
|
pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
|
|
task.
|
|
|
|
The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
|
|
weights are discarded.
|
|
|
|
Parameters:
|
|
pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
|
|
Can be either:
|
|
|
|
- A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
|
|
Valid model ids should have an organization name, like `google/ddpm-celebahq-256`.
|
|
- A path to a *directory* containing model weights saved using [`~ModelMixin.save_config`], e.g.,
|
|
`./my_model_directory/`.
|
|
|
|
cache_dir (`Union[str, os.PathLike]`, *optional*):
|
|
Path to a directory in which a downloaded pretrained model configuration should be cached if the
|
|
standard cache should not be used.
|
|
torch_dtype (`str` or `torch.dtype`, *optional*):
|
|
Override the default `torch.dtype` and load the model under this dtype. If `"auto"` is passed the dtype
|
|
will be automatically derived from the model's weights.
|
|
force_download (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
|
|
cached versions if they exist.
|
|
resume_download (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to delete incompletely received files. Will attempt to resume the download if such a
|
|
file exists.
|
|
proxies (`Dict[str, str]`, *optional*):
|
|
A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
|
|
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
|
|
output_loading_info(`bool`, *optional*, defaults to `False`):
|
|
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
|
|
local_files_only(`bool`, *optional*, defaults to `False`):
|
|
Whether or not to only look at local files (i.e., do not try to download the model).
|
|
use_auth_token (`str` or *bool*, *optional*):
|
|
The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
|
|
when running `diffusers-cli login` (stored in `~/.huggingface`).
|
|
revision (`str`, *optional*, defaults to `"main"`):
|
|
The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
|
|
git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
|
|
identifier allowed by git.
|
|
from_flax (`bool`, *optional*, defaults to `False`):
|
|
Load the model weights from a Flax checkpoint save file.
|
|
subfolder (`str`, *optional*, defaults to `""`):
|
|
In case the relevant files are located inside a subfolder of the model repo (either remote in
|
|
huggingface.co or downloaded locally), you can specify the folder name here.
|
|
|
|
mirror (`str`, *optional*):
|
|
Mirror source to accelerate downloads in China. If you are from China and have an accessibility
|
|
problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
|
|
Please refer to the mirror site for more information.
|
|
device_map (`str` or `Dict[str, Union[int, str, torch.device]]`, *optional*):
|
|
A map that specifies where each submodule should go. It doesn't need to be refined to each
|
|
parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
|
|
same device.
|
|
|
|
To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For
|
|
more information about each option see [designing a device
|
|
map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
|
|
low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
|
|
Speed up model loading by not initializing the weights and only loading the pre-trained weights. This
|
|
also tries to not use more than 1x model size in CPU memory (including peak memory) while loading the
|
|
model. This is only supported when torch version >= 1.9.0. If you are using an older version of torch,
|
|
setting this argument to `True` will raise an error.
|
|
variant (`str`, *optional*):
|
|
If specified load weights from `variant` filename, *e.g.* pytorch_model.<variant>.bin. `variant` is
|
|
ignored when using `from_flax`.
|
|
use_safetensors (`bool`, *optional* ):
|
|
If set to `True`, the pipeline will forcibly load the models from `safetensors` weights. If set to
|
|
`None` (the default). The pipeline will load using `safetensors` if safetensors weights are available
|
|
*and* if `safetensors` is installed. If the to `False` the pipeline will *not* use `safetensors`.
|
|
|
|
<Tip>
|
|
|
|
It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
|
|
models](https://huggingface.co/docs/hub/models-gated#gated-models).
|
|
|
|
</Tip>
|
|
|
|
<Tip>
|
|
|
|
Activate the special ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use
|
|
this method in a firewalled environment.
|
|
|
|
</Tip>
|
|
|
|
"""
|
|
cache_dir = kwargs.pop("cache_dir", DIFFUSERS_CACHE)
|
|
ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
|
|
force_download = kwargs.pop("force_download", False)
|
|
from_flax = kwargs.pop("from_flax", False)
|
|
resume_download = kwargs.pop("resume_download", False)
|
|
proxies = kwargs.pop("proxies", None)
|
|
output_loading_info = kwargs.pop("output_loading_info", False)
|
|
local_files_only = kwargs.pop("local_files_only", HF_HUB_OFFLINE)
|
|
use_auth_token = kwargs.pop("use_auth_token", None)
|
|
revision = kwargs.pop("revision", None)
|
|
torch_dtype = kwargs.pop("torch_dtype", None)
|
|
subfolder = kwargs.pop("subfolder", None)
|
|
device_map = kwargs.pop("device_map", None)
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
|
|
variant = kwargs.pop("variant", None)
|
|
use_safetensors = kwargs.pop("use_safetensors", None)
|
|
|
|
if use_safetensors and not is_safetensors_available():
|
|
raise ValueError(
|
|
"`use_safetensors`=True but safetensors is not installed. Please install safetensors with `pip install safetenstors"
|
|
)
|
|
|
|
allow_pickle = False
|
|
if use_safetensors is None:
|
|
use_safetensors = is_safetensors_available()
|
|
allow_pickle = True
|
|
|
|
if low_cpu_mem_usage and not is_accelerate_available():
|
|
low_cpu_mem_usage = False
|
|
logger.warning(
|
|
"Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
|
|
" environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
|
|
" `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
|
|
" install accelerate\n```\n."
|
|
)
|
|
|
|
if device_map is not None and not is_accelerate_available():
|
|
raise NotImplementedError(
|
|
"Loading and dispatching requires `accelerate`. Please make sure to install accelerate or set"
|
|
" `device_map=None`. You can install accelerate with `pip install accelerate`."
|
|
)
|
|
|
|
# Check if we can handle device_map and dispatching the weights
|
|
if device_map is not None and not is_torch_version(">=", "1.9.0"):
|
|
raise NotImplementedError(
|
|
"Loading and dispatching requires torch >= 1.9.0. Please either update your PyTorch version or set"
|
|
" `device_map=None`."
|
|
)
|
|
|
|
if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
|
|
raise NotImplementedError(
|
|
"Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
|
|
" `low_cpu_mem_usage=False`."
|
|
)
|
|
|
|
if low_cpu_mem_usage is False and device_map is not None:
|
|
raise ValueError(
|
|
f"You cannot set `low_cpu_mem_usage` to `False` while using device_map={device_map} for loading and"
|
|
" dispatching. Please make sure to set `low_cpu_mem_usage=True`."
|
|
)
|
|
|
|
# Load config if we don't provide a configuration
|
|
config_path = pretrained_model_name_or_path
|
|
|
|
user_agent = {
|
|
"diffusers": __version__,
|
|
"file_type": "model",
|
|
"framework": "pytorch",
|
|
}
|
|
|
|
# load config
|
|
config, unused_kwargs, commit_hash = cls.load_config(
|
|
config_path,
|
|
cache_dir=cache_dir,
|
|
return_unused_kwargs=True,
|
|
return_commit_hash=True,
|
|
force_download=force_download,
|
|
resume_download=resume_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
device_map=device_map,
|
|
user_agent=user_agent,
|
|
**kwargs,
|
|
)
|
|
|
|
# load model
|
|
model_file = None
|
|
if from_flax:
|
|
model_file = _get_model_file(
|
|
pretrained_model_name_or_path,
|
|
weights_name=FLAX_WEIGHTS_NAME,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
resume_download=resume_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
commit_hash=commit_hash,
|
|
)
|
|
model = cls.from_config(config, **unused_kwargs)
|
|
|
|
# Convert the weights
|
|
from .modeling_pytorch_flax_utils import load_flax_checkpoint_in_pytorch_model
|
|
|
|
model = load_flax_checkpoint_in_pytorch_model(model, model_file)
|
|
else:
|
|
if use_safetensors:
|
|
try:
|
|
model_file = _get_model_file(
|
|
pretrained_model_name_or_path,
|
|
weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
resume_download=resume_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
commit_hash=commit_hash,
|
|
)
|
|
except IOError as e:
|
|
if not allow_pickle:
|
|
raise e
|
|
pass
|
|
if model_file is None: # deactivate low_cpu_mem_usage mode
|
|
model_file = _get_model_file(
|
|
pretrained_model_name_or_path,
|
|
weights_name=_add_variant(WEIGHTS_NAME, variant),
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
resume_download=resume_download,
|
|
proxies=proxies,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
revision=revision,
|
|
subfolder=subfolder,
|
|
user_agent=user_agent,
|
|
commit_hash=commit_hash,
|
|
)
|
|
|
|
model = cls.from_config(config, **unused_kwargs)
|
|
|
|
state_dict = load_state_dict(model_file, variant=variant)
|
|
|
|
model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
|
|
model,
|
|
state_dict,
|
|
model_file,
|
|
pretrained_model_name_or_path,
|
|
ignore_mismatched_sizes=True,
|
|
)
|
|
|
|
loading_info = {
|
|
"missing_keys": missing_keys,
|
|
"unexpected_keys": unexpected_keys,
|
|
"mismatched_keys": mismatched_keys,
|
|
"error_msgs": error_msgs,
|
|
}
|
|
|
|
if torch_dtype is not None and not isinstance(torch_dtype, torch.dtype):
|
|
raise ValueError(
|
|
f"{torch_dtype} needs to be of type `torch.dtype`, e.g. `torch.float16`, but is {type(torch_dtype)}."
|
|
)
|
|
elif torch_dtype is not None:
|
|
model = model.to(torch_dtype)
|
|
|
|
model.register_to_config(_name_or_path=pretrained_model_name_or_path)
|
|
|
|
# Set model in evaluation mode to deactivate DropOut modules by default
|
|
model.eval()
|
|
if output_loading_info:
|
|
return model, loading_info
|
|
|
|
return model
|
|
|
|
@classmethod
|
|
def _load_pretrained_model(
|
|
cls,
|
|
model,
|
|
state_dict,
|
|
resolved_archive_file,
|
|
pretrained_model_name_or_path,
|
|
ignore_mismatched_sizes=False,
|
|
):
|
|
# Retrieve missing & unexpected_keys
|
|
model_state_dict = model.state_dict()
|
|
loaded_keys = list(state_dict.keys())
|
|
|
|
expected_keys = list(model_state_dict.keys())
|
|
|
|
original_loaded_keys = loaded_keys
|
|
|
|
missing_keys = list(set(expected_keys) - set(loaded_keys))
|
|
unexpected_keys = list(set(loaded_keys) - set(expected_keys))
|
|
|
|
# Make sure we are able to load base models as well as derived models (with heads)
|
|
model_to_load = model
|
|
|
|
def _find_mismatched_keys(
|
|
state_dict,
|
|
model_state_dict,
|
|
loaded_keys,
|
|
ignore_mismatched_sizes,
|
|
):
|
|
mismatched_keys = []
|
|
if ignore_mismatched_sizes:
|
|
for checkpoint_key in loaded_keys:
|
|
model_key = checkpoint_key
|
|
|
|
if (
|
|
model_key in model_state_dict
|
|
and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
|
|
):
|
|
mismatched_keys.append(
|
|
(checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
|
|
)
|
|
del state_dict[checkpoint_key]
|
|
return mismatched_keys
|
|
|
|
if state_dict is not None:
|
|
# Whole checkpoint
|
|
mismatched_keys = _find_mismatched_keys(
|
|
state_dict,
|
|
model_state_dict,
|
|
original_loaded_keys,
|
|
ignore_mismatched_sizes,
|
|
)
|
|
error_msgs = _load_state_dict_into_model(model_to_load, state_dict)
|
|
|
|
if len(error_msgs) > 0:
|
|
error_msg = "\n\t".join(error_msgs)
|
|
if "size mismatch" in error_msg:
|
|
error_msg += (
|
|
"\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
|
|
)
|
|
raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")
|
|
|
|
if len(unexpected_keys) > 0:
|
|
logger.warning(
|
|
f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
|
|
f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
|
|
f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task"
|
|
" or with another architecture (e.g. initializing a BertForSequenceClassification model from a"
|
|
" BertForPreTraining model).\n- This IS NOT expected if you are initializing"
|
|
f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly"
|
|
" identical (initializing a BertForSequenceClassification model from a"
|
|
" BertForSequenceClassification model)."
|
|
)
|
|
else:
|
|
logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
|
|
if len(missing_keys) > 0:
|
|
logger.warning(
|
|
f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
|
|
f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
|
|
" TRAIN this model on a down-stream task to be able to use it for predictions and inference."
|
|
)
|
|
elif len(mismatched_keys) == 0:
|
|
logger.info(
|
|
f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
|
|
f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the"
|
|
f" checkpoint was trained on, you can already use {model.__class__.__name__} for predictions"
|
|
" without further training."
|
|
)
|
|
if len(mismatched_keys) > 0:
|
|
mismatched_warning = "\n".join(
|
|
[
|
|
f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
|
|
for key, shape1, shape2 in mismatched_keys
|
|
]
|
|
)
|
|
logger.warning(
|
|
f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
|
|
f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
|
|
f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be"
|
|
" able to use it for predictions and inference."
|
|
)
|
|
|
|
return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
|
|
|
|
@property
|
|
def device(self) -> device:
|
|
"""
|
|
`torch.device`: The device on which the module is (assuming that all the module parameters are on the same
|
|
device).
|
|
"""
|
|
return get_parameter_device(self)
|
|
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
"""
|
|
`torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
|
|
"""
|
|
return get_parameter_dtype(self)
|
|
|
|
def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
|
|
"""
|
|
Get number of (optionally, trainable or non-embeddings) parameters in the module.
|
|
|
|
Args:
|
|
only_trainable (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to return only the number of trainable parameters
|
|
|
|
exclude_embeddings (`bool`, *optional*, defaults to `False`):
|
|
Whether or not to return only the number of non-embeddings parameters
|
|
|
|
Returns:
|
|
`int`: The number of parameters.
|
|
"""
|
|
|
|
if exclude_embeddings:
|
|
embedding_param_names = [
|
|
f"{name}.weight"
|
|
for name, module_type in self.named_modules()
|
|
if isinstance(module_type, torch.nn.Embedding)
|
|
]
|
|
non_embedding_parameters = [
|
|
parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
|
|
]
|
|
return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
|
|
else:
|
|
return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
|
|
|
|
|
|
def _get_model_file(
|
|
pretrained_model_name_or_path,
|
|
*,
|
|
weights_name,
|
|
subfolder,
|
|
cache_dir,
|
|
force_download,
|
|
proxies,
|
|
resume_download,
|
|
local_files_only,
|
|
use_auth_token,
|
|
user_agent,
|
|
revision,
|
|
commit_hash=None,
|
|
):
|
|
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
|
|
if os.path.isfile(pretrained_model_name_or_path):
|
|
return pretrained_model_name_or_path
|
|
elif os.path.isdir(pretrained_model_name_or_path):
|
|
if os.path.isfile(os.path.join(pretrained_model_name_or_path, weights_name)):
|
|
# Load from a PyTorch checkpoint
|
|
model_file = os.path.join(pretrained_model_name_or_path, weights_name)
|
|
return model_file
|
|
elif subfolder is not None and os.path.isfile(
|
|
os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
|
|
):
|
|
model_file = os.path.join(pretrained_model_name_or_path, subfolder, weights_name)
|
|
return model_file
|
|
else:
|
|
raise EnvironmentError(
|
|
f"Error no file named {weights_name} found in directory {pretrained_model_name_or_path}."
|
|
)
|
|
else:
|
|
# 1. First check if deprecated way of loading from branches is used
|
|
if (
|
|
revision in DEPRECATED_REVISION_ARGS
|
|
and (weights_name == WEIGHTS_NAME or weights_name == SAFETENSORS_WEIGHTS_NAME)
|
|
and version.parse(version.parse(__version__).base_version) >= version.parse("0.17.0")
|
|
):
|
|
try:
|
|
model_file = hf_hub_download(
|
|
pretrained_model_name_or_path,
|
|
filename=_add_variant(weights_name, revision),
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
resume_download=resume_download,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
user_agent=user_agent,
|
|
subfolder=subfolder,
|
|
revision=revision or commit_hash,
|
|
)
|
|
warnings.warn(
|
|
f"Loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'` is deprecated. Loading instead from `revision='main'` with `variant={revision}`. Loading model variants via `revision='{revision}'` will be removed in diffusers v1. Please use `variant='{revision}'` instead.",
|
|
FutureWarning,
|
|
)
|
|
return model_file
|
|
except: # noqa: E722
|
|
warnings.warn(
|
|
f"You are loading the variant {revision} from {pretrained_model_name_or_path} via `revision='{revision}'`. This behavior is deprecated and will be removed in diffusers v1. One should use `variant='{revision}'` instead. However, it appears that {pretrained_model_name_or_path} currently does not have a {_add_variant(weights_name, revision)} file in the 'main' branch of {pretrained_model_name_or_path}. \n The Diffusers team and community would be very grateful if you could open an issue: https://github.com/huggingface/diffusers/issues/new with the title '{pretrained_model_name_or_path} is missing {_add_variant(weights_name, revision)}' so that the correct variant file can be added.",
|
|
FutureWarning,
|
|
)
|
|
try:
|
|
# 2. Load model file as usual
|
|
model_file = hf_hub_download(
|
|
pretrained_model_name_or_path,
|
|
filename=weights_name,
|
|
cache_dir=cache_dir,
|
|
force_download=force_download,
|
|
proxies=proxies,
|
|
resume_download=resume_download,
|
|
local_files_only=local_files_only,
|
|
use_auth_token=use_auth_token,
|
|
user_agent=user_agent,
|
|
subfolder=subfolder,
|
|
revision=revision or commit_hash,
|
|
)
|
|
return model_file
|
|
|
|
except RepositoryNotFoundError:
|
|
raise EnvironmentError(
|
|
f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier "
|
|
"listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a "
|
|
"token having permission to this repo with `use_auth_token` or log in with `huggingface-cli "
|
|
"login`."
|
|
)
|
|
except RevisionNotFoundError:
|
|
raise EnvironmentError(
|
|
f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for "
|
|
"this model name. Check the model page at "
|
|
f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
|
|
)
|
|
except EntryNotFoundError:
|
|
raise EnvironmentError(
|
|
f"{pretrained_model_name_or_path} does not appear to have a file named {weights_name}."
|
|
)
|
|
except HTTPError as err:
|
|
raise EnvironmentError(
|
|
f"There was a specific connection error when trying to load {pretrained_model_name_or_path}:\n{err}"
|
|
)
|
|
except ValueError:
|
|
raise EnvironmentError(
|
|
f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
|
|
f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
|
|
f" directory containing a file named {weights_name} or"
|
|
" \nCheckout your internet connection or see how to run the library in"
|
|
" offline mode at 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
|
|
)
|
|
except EnvironmentError:
|
|
raise EnvironmentError(
|
|
f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from "
|
|
"'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
|
|
f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
|
|
f"containing a file named {weights_name}"
|
|
)
|