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)

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

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

weights_name = SAFETENSORS_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME

weights_name = _add_variant(weights_name, variant)

weight_name_split = weights_name.split(".")

if len(weight_name_split) in [2, 3]:

weights_name_pattern = weight_name_split[0] + "{suffix}." + ".".join(weight_name_split[1:])

else:

raise ValueError(f"Invalid {weights_name} provided.")

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,

)

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

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,

)

# 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 = _fetch_index_file(

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,

)

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

)

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

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

model,

state_dict,

device=param_device,

dtype=torch_dtype,

model_name_or_path=pretrained_model_name_or_path,

)

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)

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 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: OrderedDict,

resolved_archive_file,

pretrained_model_name_or_path: Union[str, os.PathLike],

ignore_mismatched_sizes: bool = 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

@classmethod

def _get_signature_keys(cls, obj):

parameters = inspect.signature(obj.__init__).parameters

required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}

optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})

expected_modules = set(required_parameters.keys()) - {"self"}

return expected_modules, optional_parameters

# Adapted from `transformers` modeling_utils.py

def _get_no_split_modules(self, device_map: str):

"""

_no_split_modules = set()

modules_to_check = [self]

while len(modules_to_check) > 0:

module = modules_to_check.pop(-1)

# if the module does not appear in _no_split_modules, we also check the children

if module.__class__.__name__ not in _no_split_modules:

if isinstance(module, ModelMixin):

if module._no_split_modules is None:

raise ValueError(

f"{module.__class__.__name__} does not support `device_map='{device_map}'`. To implement support, the model "

"class needs to implement the `_no_split_modules` attribute."

)

else: