try:

parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())

return next(parameters_and_buffers).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)

try:

params = tuple(parameter.parameters())

if len(params) > 0:

return params[0].dtype

buffers = tuple(parameter.buffers())

if len(buffers) > 0:

return buffers[0].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)

"""

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."

# 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)

r"""

config_name = CONFIG_NAME

_automatically_saved_args = ["_diffusers_version", "_class_name", "_name_or_path"]

_supports_gradient_checkpointing = False

_keys_to_ignore_on_load_unexpected = None

def __init__(self):

super().__init__()

def __getattr__(self, name: str) -> Any:

"""The only reason we overwrite `getattr` here is to gracefully deprecate accessing

is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)

is_attribute = name in self.__dict__

if is_in_config and not is_attribute:

deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'unet.config.{name}'."

deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False, stacklevel=3)

return self._internal_dict[name]

# call PyTorch's https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module

return super().__getattr__(name)

@property

def is_gradient_checkpointing(self) -> bool:

"""

return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())

def enable_gradient_checkpointing(self):

"""

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):

"""

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"""

self.set_use_memory_efficient_attention_xformers(True, attention_op)

def disable_xformers_memory_efficient_attention(self):

r"""

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,

):

"""

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"""

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)

max_memory = kwargs.pop("max_memory", None)

offload_folder = kwargs.pop("offload_folder", None)

offload_state_dict = kwargs.pop("offload_state_dict", False)

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,

max_memory=max_memory,

offload_folder=offload_folder,

offload_state_dict=offload_state_dict,

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:

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,

)

if low_cpu_mem_usage:

# Instantiate model with empty weights

with accelerate.init_empty_weights():

model = cls.from_config(config, **unused_kwargs)

# if device_map is None, load the state dict and move the params from meta device to the cpu

if device_map is None:

param_device = "cpu"

state_dict = load_state_dict(model_file, variant=variant)

model._convert_deprecated_attention_blocks(state_dict)

# move the params from meta device to cpu

missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())

if len(missing_keys) > 0:

raise ValueError(

f"Cannot load {cls} from {pretrained_model_name_or_path} because the following keys are"

f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"

" `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"

" those weights or else make sure your checkpoint file is correct."

)

unexpected_keys = []

empty_state_dict = model.state_dict()

for param_name, param in state_dict.items():

accepts_dtype = "dtype" in set(

inspect.signature(set_module_tensor_to_device).parameters.keys()

)

if param_name not in empty_state_dict:

unexpected_keys.append(param_name)

continue

if empty_state_dict[param_name].shape != param.shape:

raise ValueError(

f"Cannot load {pretrained_model_name_or_path} because {param_name} expected shape {empty_state_dict[param_name]}, but got {param.shape}. If you want to instead overwrite randomly initialized weights, please make sure to pass both `low_cpu_mem_usage=False` and `ignore_mismatched_sizes=True`. For more information, see also: https://github.com/huggingface/diffusers/issues/1619#issuecomment-1345604389 as an example."

)

if accepts_dtype:

set_module_tensor_to_device(

model, param_name, param_device, value=param, dtype=torch_dtype

)

else:

set_module_tensor_to_device(model, param_name, param_device, value=param)

if cls._keys_to_ignore_on_load_unexpected is not None:

for pat in cls._keys_to_ignore_on_load_unexpected:

unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

if len(unexpected_keys) > 0:

logger.warn(

f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"

)

else: # else let accelerate handle loading and dispatching.

# Load weights and dispatch according to the device_map

# by default the device_map is None and the weights are loaded on the CPU

try:

accelerate.load_checkpoint_and_dispatch(

model,

model_file,

device_map,

max_memory=max_memory,

offload_folder=offload_folder,

offload_state_dict=offload_state_dict,

dtype=torch_dtype,

)

except AttributeError as e:

# When using accelerate loading, we do not have the ability to load the state

# dict and rename the weight names manually. Additionally, accelerate skips

# torch loading conventions and directly writes into `module.{_buffers, _parameters}`

# (which look like they should be private variables?), so we can't use the standard hooks

# to rename parameters on load. We need to mimic the original weight names so the correct

# attributes are available. After we have loaded the weights, we convert the deprecated

# names to the new non-deprecated names. Then we _greatly encourage_ the user to convert

# the weights so we don't have to do this again.

if "'Attention' object has no attribute" in str(e):

logger.warn(

f"Taking `{str(e)}` while using `accelerate.load_checkpoint_and_dispatch` to mean {pretrained_model_name_or_path}"

" was saved with deprecated attention block weight names. We will load it with the deprecated attention block"

" names and convert them on the fly to the new attention block format. Please re-save the model after this conversion,"

" so we don't have to do the on the fly renaming in the future. If the model is from a hub checkpoint,"

" please also re-upload it or open a PR on the original repository."

)

model._temp_convert_self_to_deprecated_attention_blocks()

accelerate.load_checkpoint_and_dispatch(

model,

model_file,

device_map,

max_memory=max_memory,

offload_folder=offload_folder,

offload_state_dict=offload_state_dict,

dtype=torch_dtype,

)

model._undo_temp_convert_self_to_deprecated_attention_blocks()

else:

raise e

loading_info = {

"missing_keys": [],

"unexpected_keys": [],

"mismatched_keys": [],

"error_msgs": [],

}

else:

model = cls.from_config(config, **unused_kwargs)

state_dict = load_state_dict(model_file, variant=variant)

model._convert_deprecated_attention_blocks(state_dict)

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=ignore_mismatched_sizes,

)

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:

"""

return get_parameter_device(self)

@property

def dtype(self) -> torch.dtype:

"""

return get_parameter_dtype(self)

def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:

"""

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 _convert_deprecated_attention_blocks(self, state_dict):

deprecated_attention_block_paths = []

def recursive_find_attn_block(name, module):

if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:

deprecated_attention_block_paths.append(name)

for sub_name, sub_module in module.named_children():

sub_name = sub_name if name == "" else f"{name}.{sub_name}"

recursive_find_attn_block(sub_name, sub_module)

recursive_find_attn_block("", self)

# NOTE: we have to check if the deprecated parameters are in the state dict

# because it is possible we are loading from a state dict that was already

# converted

for path in deprecated_attention_block_paths:

# group_norm path stays the same

# query -> to_q

if f"{path}.query.weight" in state_dict:

state_dict[f"{path}.to_q.weight"] = state_dict.pop(f"{path}.query.weight")

if f"{path}.query.bias" in state_dict:

state_dict[f"{path}.to_q.bias"] = state_dict.pop(f"{path}.query.bias")

# key -> to_k

if f"{path}.key.weight" in state_dict:

state_dict[f"{path}.to_k.weight"] = state_dict.pop(f"{path}.key.weight")

if f"{path}.key.bias" in state_dict:

state_dict[f"{path}.to_k.bias"] = state_dict.pop(f"{path}.key.bias")

# value -> to_v

if f"{path}.value.weight" in state_dict:

state_dict[f"{path}.to_v.weight"] = state_dict.pop(f"{path}.value.weight")

if f"{path}.value.bias" in state_dict:

state_dict[f"{path}.to_v.bias"] = state_dict.pop(f"{path}.value.bias")

# proj_attn -> to_out.0

if f"{path}.proj_attn.weight" in state_dict:

state_dict[f"{path}.to_out.0.weight"] = state_dict.pop(f"{path}.proj_attn.weight")

if f"{path}.proj_attn.bias" in state_dict:

state_dict[f"{path}.to_out.0.bias"] = state_dict.pop(f"{path}.proj_attn.bias")

def _temp_convert_self_to_deprecated_attention_blocks(self):

deprecated_attention_block_modules = []

def recursive_find_attn_block(module):

if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:

deprecated_attention_block_modules.append(module)

for sub_module in module.children():

recursive_find_attn_block(sub_module)

recursive_find_attn_block(self)

for module in deprecated_attention_block_modules:

module.query = module.to_q

module.key = module.to_k

module.value = module.to_v

module.proj_attn = module.to_out[0]

# We don't _have_ to delete the old attributes, but it's helpful to ensure

# that _all_ the weights are loaded into the new attributes and we're not

# making an incorrect assumption that this model should be converted when

# it really shouldn't be.

del module.to_q

del module.to_k

del module.to_v

del module.to_out

def _undo_temp_convert_self_to_deprecated_attention_blocks(self):

deprecated_attention_block_modules = []

def recursive_find_attn_block(module):

if hasattr(module, "_from_deprecated_attn_block") and module._from_deprecated_attn_block:

deprecated_attention_block_modules.append(module)

for sub_module in module.children():

recursive_find_attn_block(sub_module)

recursive_find_attn_block(self)

for module in deprecated_attention_block_modules:

module.to_q = module.query

module.to_k = module.key

module.to_v = module.value

module.to_out = nn.ModuleList([module.proj_attn, nn.Dropout(module.dropout)])

del module.query

del module.key

del module.value