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:

return next(parameter.parameters()).dtype

except StopIteration:

try:

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

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

_no_split_modules = None

_keep_in_fp32_modules = 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) -> None:

"""

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

"""

if self._supports_gradient_checkpointing:

self.apply(partial(self._set_gradient_checkpointing, value=False))

def set_use_npu_flash_attention(self, valid: bool) -> None:

r"""

def fn_recursive_set_npu_flash_attention(module: torch.nn.Module):

if hasattr(module, "set_use_npu_flash_attention"):

module.set_use_npu_flash_attention(valid)

for child in module.children():

fn_recursive_set_npu_flash_attention(child)

for module in self.children():

if isinstance(module, torch.nn.Module):

fn_recursive_set_npu_flash_attention(module)

def enable_npu_flash_attention(self) -> None:

r"""

self.set_use_npu_flash_attention(True)

def disable_npu_flash_attention(self) -> None:

r"""

self.set_use_npu_flash_attention(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) -> None:

r"""

self.set_use_memory_efficient_attention_xformers(True, attention_op)

def disable_xformers_memory_efficient_attention(self) -> None:

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: Optional[Callable] = None,

safe_serialization: bool = True,

variant: Optional[str] = None,

max_shard_size: Union[int, str] = "10GB",

push_to_hub: bool = False,

**kwargs,

):

"""

if os.path.isfile(save_directory):

logger.error(f"Provided path ({save_directory}) should be a directory, not a file")

return

hf_quantizer = getattr(self, "hf_quantizer", None)

if hf_quantizer is not None:

quantization_serializable = (

hf_quantizer is not None

and isinstance(hf_quantizer, DiffusersQuantizer)

and hf_quantizer.is_serializable

)

if not quantization_serializable:

raise ValueError(

f"The model is quantized with {hf_quantizer.quantization_config.quant_method} and is not serializable - check out the warnings from"

" the logger on the traceback to understand the reason why the quantized model is not serializable."

)

weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME

weights_name = _add_variant(weights_name, variant)

weights_name_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(

".safetensors", "{suffix}.safetensors"

)

os.makedirs(save_directory, exist_ok=True)

if push_to_hub:

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

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

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

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

repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])

repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id

# Only save the model itself if we are using distributed training

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

# Save the model

state_dict_split = split_torch_state_dict_into_shards(

state_dict, max_shard_size=max_shard_size, filename_pattern=weights_name_pattern

)

# Clean the folder from a previous save

if is_main_process:

for filename in os.listdir(save_directory):

if filename in state_dict_split.filename_to_tensors.keys():

continue

full_filename = os.path.join(save_directory, filename)

if not os.path.isfile(full_filename):

continue

weights_without_ext = weights_name_pattern.replace(".bin", "").replace(".safetensors", "")

weights_without_ext = weights_without_ext.replace("{suffix}", "")

filename_without_ext = filename.replace(".bin", "").replace(".safetensors", "")

# make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005

if (

filename.startswith(weights_without_ext)

and _REGEX_SHARD.fullmatch(filename_without_ext) is not None

):

os.remove(full_filename)

for filename, tensors in state_dict_split.filename_to_tensors.items():

shard = {tensor: state_dict[tensor] for tensor in tensors}

filepath = os.path.join(save_directory, filename)

if safe_serialization:

# At some point we will need to deal better with save_function (used for TPU and other distributed

# joyfulness), but for now this enough.

safetensors.torch.save_file(shard, filepath, metadata={"format": "pt"})

else:

torch.save(shard, filepath)

if state_dict_split.is_sharded:

index = {

"metadata": state_dict_split.metadata,

"weight_map": state_dict_split.tensor_to_filename,

}

save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME

save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))

# Save the index as well

with open(save_index_file, "w", encoding="utf-8") as f:

content = json.dumps(index, indent=2, sort_keys=True) + "\n"

f.write(content)

logger.info(

f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "

f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the "

f"index located at {save_index_file}."

)

else:

path_to_weights = os.path.join(save_directory, weights_name)

logger.info(f"Model weights saved in {path_to_weights}")

if push_to_hub:

# Create a new empty model card and eventually tag it

model_card = load_or_create_model_card(repo_id, token=token)

model_card = populate_model_card(model_card)

model_card.save(Path(save_directory, "README.md").as_posix())

self._upload_folder(

save_directory,

repo_id,

token=token,

commit_message=commit_message,

create_pr=create_pr,

)

def dequantize(self):

"""

hf_quantizer = getattr(self, "hf_quantizer", None)

if hf_quantizer is None:

raise ValueError("You need to first quantize your model in order to dequantize it")

return hf_quantizer.dequantize(self)

@classmethod

@validate_hf_hub_args

def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):

r"""

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

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

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

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

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

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

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

token = kwargs.pop("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)

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

allow_pickle = False

if use_safetensors is None:

use_safetensors = True

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

)

# change device_map into a map if we passed an int, a str or a torch.device

if isinstance(device_map, torch.device):

device_map = {"": device_map}

elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:

try:

device_map = {"": torch.device(device_map)}

except RuntimeError:

raise ValueError(

"When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or "

f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}."

)

elif isinstance(device_map, int):

if device_map < 0:

raise ValueError(

"You can't pass device_map as a negative int. If you want to put the model on the cpu, pass device_map = 'cpu' "

)

else:

device_map = {"": device_map}

if device_map is not None:

if low_cpu_mem_usage is None:

low_cpu_mem_usage = True

elif not low_cpu_mem_usage:

raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`")

if low_cpu_mem_usage:

if device_map is not None and not is_torch_version(">=", "1.10"):

# The max memory utils require PyTorch >= 1.10 to have torch.cuda.mem_get_info.

raise ValueError("`low_cpu_mem_usage` and `device_map` require PyTorch >= 1.10.")

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

proxies=proxies,

local_files_only=local_files_only,

token=token,

revision=revision,

subfolder=subfolder,

user_agent=user_agent,

**kwargs,

)

# no in-place modification of the original config.

config = copy.deepcopy(config)

# determine initial quantization config.

#######################################

pre_quantized = "quantization_config" in config and config["quantization_config"] is not None

if pre_quantized or quantization_config is not None:

if pre_quantized:

config["quantization_config"] = DiffusersAutoQuantizer.merge_quantization_configs(

config["quantization_config"], quantization_config

)

else:

config["quantization_config"] = quantization_config

hf_quantizer = DiffusersAutoQuantizer.from_config(

config["quantization_config"], pre_quantized=pre_quantized

)

else:

hf_quantizer = None

if hf_quantizer is not None:

if device_map is not None:

raise NotImplementedError(

"Currently, `device_map` is automatically inferred for quantized models. Support for providing `device_map` as an input will be added in the future."

)

hf_quantizer.validate_environment(torch_dtype=torch_dtype, from_flax=from_flax, device_map=device_map)

torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype)

# In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry`

user_agent["quant"] = hf_quantizer.quantization_config.quant_method.value

# Force-set to `True` for more mem efficiency

if low_cpu_mem_usage is None:

low_cpu_mem_usage = True

logger.info("Set `low_cpu_mem_usage` to True as `hf_quantizer` is not None.")

elif not low_cpu_mem_usage:

raise ValueError("`low_cpu_mem_usage` cannot be False or None when using quantization.")

# Check if `_keep_in_fp32_modules` is not None

use_keep_in_fp32_modules = (cls._keep_in_fp32_modules is not None) and (

(torch_dtype == torch.float16) or hasattr(hf_quantizer, "use_keep_in_fp32_modules")

)

if use_keep_in_fp32_modules:

keep_in_fp32_modules = cls._keep_in_fp32_modules

if not isinstance(keep_in_fp32_modules, list):

keep_in_fp32_modules = [keep_in_fp32_modules]

if low_cpu_mem_usage is None:

low_cpu_mem_usage = True

logger.info("Set `low_cpu_mem_usage` to True as `_keep_in_fp32_modules` is not None.")

elif not low_cpu_mem_usage:

raise ValueError("`low_cpu_mem_usage` cannot be False when `keep_in_fp32_modules` is True.")

else:

keep_in_fp32_modules = []

#######################################

# Determine if we're loading from a directory of sharded checkpoints.

is_sharded = False

index_file = None

is_local = os.path.isdir(pretrained_model_name_or_path)

index_file_kwargs = {

"is_local": is_local,

"pretrained_model_name_or_path": pretrained_model_name_or_path,

"subfolder": subfolder or "",

"use_safetensors": use_safetensors,

"cache_dir": cache_dir,

"variant": variant,

"force_download": force_download,

"proxies": proxies,

"local_files_only": local_files_only,

"token": token,

"revision": revision,

"user_agent": user_agent,

"commit_hash": commit_hash,

}

index_file = _fetch_index_file(**index_file_kwargs)

# In case the index file was not found we still have to consider the legacy format.

# this becomes applicable when the variant is not None.

if variant is not None and (index_file is None or not os.path.exists(index_file)):

index_file = _fetch_index_file_legacy(**index_file_kwargs)

if index_file is not None and index_file.is_file():

is_sharded = True

if is_sharded and from_flax:

raise ValueError("Loading of sharded checkpoints is not supported when `from_flax=True`.")

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

proxies=proxies,

local_files_only=local_files_only,

token=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 is_sharded:

sharded_ckpt_cached_folder, sharded_metadata = _get_checkpoint_shard_files(

pretrained_model_name_or_path,

index_file,

cache_dir=cache_dir,

proxies=proxies,

local_files_only=local_files_only,

token=token,

user_agent=user_agent,

revision=revision,

subfolder=subfolder or "",

)

if hf_quantizer is not None:

model_file = _merge_sharded_checkpoints(sharded_ckpt_cached_folder, sharded_metadata)

logger.info("Merged sharded checkpoints as `hf_quantizer` is not None.")

is_sharded = False

elif use_safetensors and not is_sharded:

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,

proxies=proxies,

local_files_only=local_files_only,

token=token,

revision=revision,

subfolder=subfolder,

user_agent=user_agent,

commit_hash=commit_hash,

)

except IOError as e:

logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")

if not allow_pickle:

raise

logger.warning(

"Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."

)

if model_file is None and not is_sharded:

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,

proxies=proxies,

local_files_only=local_files_only,

token=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 hf_quantizer is not None:

hf_quantizer.preprocess_model(

model=model, device_map=device_map, keep_in_fp32_modules=keep_in_fp32_modules

)

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

if device_map is None and not is_sharded:

# `torch.cuda.current_device()` is fine here when `hf_quantizer` is not None.

# It would error out during the `validate_environment()` call above in the absence of cuda.

is_quant_method_bnb = (

getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES

)

if hf_quantizer is None:

param_device = "cpu"

# TODO (sayakpaul, SunMarc): remove this after model loading refactor

elif is_quant_method_bnb:

param_device = torch.cuda.current_device()

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 hf_quantizer is not None:

missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix="")

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 = load_model_dict_into_meta(

model,

state_dict,

device=param_device,

dtype=torch_dtype,

model_name_or_path=pretrained_model_name_or_path,

hf_quantizer=hf_quantizer,

keep_in_fp32_modules=keep_in_fp32_modules,

)

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.warning(

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

force_hook = True

device_map = _determine_device_map(

model, device_map, max_memory, torch_dtype, keep_in_fp32_modules, hf_quantizer

)

if device_map is None and is_sharded:

# we load the parameters on the cpu

device_map = {"": "cpu"}

force_hook = False

try:

accelerate.load_checkpoint_and_dispatch(

model,

model_file if not is_sharded else index_file,

device_map,

max_memory=max_memory,

offload_folder=offload_folder,

offload_state_dict=offload_state_dict,

dtype=torch_dtype,

force_hooks=force_hook,

strict=True,

)

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.warning(

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 if not is_sharded else index_file,

device_map,

max_memory=max_memory,

offload_folder=offload_folder,

offload_state_dict=offload_state_dict,

dtype=torch_dtype,

force_hooks=force_hook,

strict=True,

)

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 hf_quantizer is not None:

hf_quantizer.postprocess_model(model)

model.hf_quantizer = hf_quantizer

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

)

# When using `use_keep_in_fp32_modules` if we do a global `to()` here, then we will

# completely lose the effectivity of `use_keep_in_fp32_modules`.

elif torch_dtype is not None and hf_quantizer is None and not use_keep_in_fp32_modules:

model = model.to(torch_dtype)

if hf_quantizer is not None:

# We also make sure to purge `_pre_quantization_dtype` when we serialize

# the model config because `_pre_quantization_dtype` is `torch.dtype`, not JSON serializable.

model.register_to_config(_name_or_path=pretrained_model_name_or_path, _pre_quantization_dtype=torch_dtype)

else:

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

# Adapted from `transformers`.

@wraps(torch.nn.Module.cuda)

def cuda(self, *args, **kwargs):

# Checks if the model has been loaded in 4-bit or 8-bit with BNB

if getattr(self, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES:

if getattr(self, "is_loaded_in_8bit", False):

raise ValueError(

"Calling `cuda()` is not supported for `8-bit` quantized models. "

" Please use the model as it is, since the model has already been set to the correct devices."

)

elif is_bitsandbytes_version("<", "0.43.2"):

raise ValueError(

"Calling `cuda()` is not supported for `4-bit` quantized models with the installed version of bitsandbytes. "

f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.43.2."

)

return super().cuda(*args, **kwargs)