"""

def __init__(self, context_managers: list[ContextManager]):

self.context_managers = context_managers

self.stack = ExitStack()

def __enter__(self):

for context_manager in self.context_managers:

self.stack.enter_context(context_manager)

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

"uniform_": nn.init.uniform_,

"normal_": nn.init.normal_,

"trunc_normal_": nn.init.trunc_normal_,

"constant_": nn.init.constant_,

"xavier_uniform_": nn.init.xavier_uniform_,

"xavier_normal_": nn.init.xavier_normal_,

"kaiming_uniform_": nn.init.kaiming_uniform_,

"kaiming_normal_": nn.init.kaiming_normal_,

"uniform": nn.init.uniform,

"normal": nn.init.normal,

"xavier_uniform": nn.init.xavier_uniform,

"xavier_normal": nn.init.xavier_normal,

"kaiming_uniform": nn.init.kaiming_uniform,

"kaiming_normal": nn.init.kaiming_normal,

from ..hooks.group_offloading import _get_group_onload_device

try:

# Try to get the onload device from the group offloading hook

return _get_group_onload_device(parameter)

except ValueError:

pass

try:

# If the onload device is not available due to no group offloading hooks, try to get the device

# from the first parameter or buffer

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)

"""

# 1. Check if we have attached any dtype modifying hooks (eg. layerwise casting)

if isinstance(parameter, nn.Module):

for name, submodule in parameter.named_modules():

if not hasattr(submodule, "_diffusers_hook"):

continue

registry = submodule._diffusers_hook

hook = registry.get_hook("layerwise_casting")

if hook is not None:

return hook.compute_dtype

# 2. If no dtype modifying hooks are attached, return the dtype of the first floating point parameter/buffer

last_dtype = None

for name, param in parameter.named_parameters():

last_dtype = param.dtype

if (

hasattr(parameter, "_keep_in_fp32_modules")

and parameter._keep_in_fp32_modules

and any(m in name for m in parameter._keep_in_fp32_modules)

):

continue

if param.is_floating_point():

return param.dtype

for buffer in parameter.buffers():

last_dtype = buffer.dtype

if buffer.is_floating_point():

return buffer.dtype

if last_dtype is not None:

# if no floating dtype was found return whatever the first dtype is

return last_dtype

# For nn.DataParallel compatibility in PyTorch > 1.5

def find_tensor_attributes(module: 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)

last_tuple = None

for tuple in gen:

last_tuple = tuple

if tuple[1].is_floating_point():

return tuple[1].dtype

if last_tuple is not None:

# fallback to the last dtype

"""

def _skip_init(*args, **kwargs):

pass

for name, init_func in TORCH_INIT_FUNCTIONS.items():

setattr(torch.nn.init, name, _skip_init)

try:

yield

finally:

# Restore the original initialization functions

for name, init_func in TORCH_INIT_FUNCTIONS.items():

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

_skip_layerwise_casting_patterns = None

_supports_group_offloading = True

_repeated_blocks = []

_parallel_config = None

_cp_plan = None

_skip_keys = None

def __init__(self):

super().__init__()

self._gradient_checkpointing_func = None

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, gradient_checkpointing_func: Callable | None = None) -> None:

"""

if not self._supports_gradient_checkpointing:

raise ValueError(

f"{self.__class__.__name__} does not support gradient checkpointing. Please make sure to set the boolean attribute "

f"`_supports_gradient_checkpointing` to `True` in the class definition."

)

if gradient_checkpointing_func is None:

def _gradient_checkpointing_func(module, *args):

ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}

return torch.utils.checkpoint.checkpoint(

module.__call__,

*args,

**ckpt_kwargs,

)

gradient_checkpointing_func = _gradient_checkpointing_func

self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func)

def disable_gradient_checkpointing(self) -> None:

"""

if self._supports_gradient_checkpointing:

self._set_gradient_checkpointing(enable=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_xla_flash_attention(

self, use_xla_flash_attention: bool, partition_spec: Callable | None = None, **kwargs

) -> None:

# Recursively walk through all the children.

# Any children which exposes the set_use_xla_flash_attention method

# gets the message

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

if hasattr(module, "set_use_xla_flash_attention"):

module.set_use_xla_flash_attention(use_xla_flash_attention, partition_spec, **kwargs)

for child in module.children():

fn_recursive_set_flash_attention(child)

for module in self.children():

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

fn_recursive_set_flash_attention(module)

def enable_xla_flash_attention(self, partition_spec: Callable | None = None, **kwargs):

r"""

self.set_use_xla_flash_attention(True, partition_spec, **kwargs)

def disable_xla_flash_attention(self):

r"""

self.set_use_xla_flash_attention(False)

def set_use_memory_efficient_attention_xformers(self, valid: bool, attention_op: Callable | None = 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: Callable | None = 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 enable_layerwise_casting(

self,

storage_dtype: torch.dtype = torch.float8_e4m3fn,

compute_dtype: torch.dtype | None = None,

skip_modules_pattern: tuple[str, ...] | None = None,

skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None,

non_blocking: bool = False,

) -> None:

r"""

from ..hooks import apply_layerwise_casting

user_provided_patterns = True

if skip_modules_pattern is None:

from ..hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN

skip_modules_pattern = DEFAULT_SKIP_MODULES_PATTERN

user_provided_patterns = False

if self._keep_in_fp32_modules is not None:

skip_modules_pattern += tuple(self._keep_in_fp32_modules)

if self._skip_layerwise_casting_patterns is not None:

skip_modules_pattern += tuple(self._skip_layerwise_casting_patterns)

skip_modules_pattern = tuple(set(skip_modules_pattern))

if is_peft_available() and not user_provided_patterns:

# By default, we want to skip all peft layers because they have a very low memory footprint.

# If users want to apply layerwise casting on peft layers as well, they can utilize the

# `~diffusers.hooks.layerwise_casting.apply_layerwise_casting` function which provides

# them with more flexibility and control.

from peft.tuners.loha.layer import LoHaLayer

from peft.tuners.lokr.layer import LoKrLayer

from peft.tuners.lora.layer import LoraLayer

for layer in (LoHaLayer, LoKrLayer, LoraLayer):

skip_modules_pattern += tuple(layer.adapter_layer_names)

if compute_dtype is None:

logger.info("`compute_dtype` not provided when enabling layerwise casting. Using dtype of the model.")

compute_dtype = self.dtype

apply_layerwise_casting(

self, storage_dtype, compute_dtype, skip_modules_pattern, skip_modules_classes, non_blocking

)

def enable_group_offload(

self,

onload_device: torch.device,

offload_device: torch.device = torch.device("cpu"),

offload_type: str = "block_level",

num_blocks_per_group: int | None = None,

non_blocking: bool = False,

use_stream: bool = False,

record_stream: bool = False,

low_cpu_mem_usage=False,

offload_to_disk_path: str | None = None,

block_modules: str | None = None,

exclude_kwargs: str | None = None,

) -> None:

r"""

from ..hooks import apply_group_offloading

if getattr(self, "enable_tiling", None) is not None and getattr(self, "use_tiling", False) and use_stream:

msg = (

"Applying group offloading on autoencoders, with CUDA streams, may not work as expected if the first "

"forward pass is executed with tiling enabled. Please make sure to either:\n"

"1. Run a forward pass with small input shapes.\n"

"2. Or, run a forward pass with tiling disabled (can still use small dummy inputs)."

)

logger.warning(msg)

if not self._supports_group_offloading:

raise ValueError(

f"{self.__class__.__name__} does not support group offloading. Please make sure to set the boolean attribute "

f"`_supports_group_offloading` to `True` in the class definition. If you believe this is a mistake, please "

f"open an issue at https://github.com/huggingface/diffusers/issues."

)

apply_group_offloading(

module=self,

onload_device=onload_device,

offload_device=offload_device,

offload_type=offload_type,

num_blocks_per_group=num_blocks_per_group,

non_blocking=non_blocking,

use_stream=use_stream,

record_stream=record_stream,

low_cpu_mem_usage=low_cpu_mem_usage,

offload_to_disk_path=offload_to_disk_path,

block_modules=block_modules,

exclude_kwargs=exclude_kwargs,

)

def set_attention_backend(self, backend: str) -> None:

"""

from .attention import AttentionModuleMixin

from .attention_dispatch import (

AttentionBackendName,

_AttentionBackendRegistry,

_check_attention_backend_requirements,

_maybe_download_kernel_for_backend,

)

# TODO: the following will not be required when everything is refactored to AttentionModuleMixin

from .attention_processor import Attention, MochiAttention

logger.warning("Attention backends are an experimental feature and the API may be subject to change.")

attention_classes = (Attention, MochiAttention, AttentionModuleMixin)

parallel_config_set = False

for module in self.modules():

if not isinstance(module, attention_classes):

continue

processor = module.processor

if getattr(processor, "_parallel_config", None) is not None:

parallel_config_set = True

break

backend = backend.lower()

available_backends = {x.value for x in AttentionBackendName.__members__.values()}

if backend not in available_backends:

raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends))

backend = AttentionBackendName(backend)

if parallel_config_set and not _AttentionBackendRegistry._is_context_parallel_available(backend):

compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)

raise ValueError(

f"Context parallelism is enabled but current attention backend '{backend.value}' "

f"does not support context parallelism. "

f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()`."

)

_check_attention_backend_requirements(backend)

_maybe_download_kernel_for_backend(backend)

for module in self.modules():

if not isinstance(module, attention_classes):

continue

processor = module.processor

if processor is None or not hasattr(processor, "_attention_backend"):

continue

processor._attention_backend = backend

# Important to set the active backend so that it propagates gracefully throughout.

_AttentionBackendRegistry.set_active_backend(backend)

def reset_attention_backend(self) -> None:

"""

from .attention import AttentionModuleMixin

from .attention_processor import Attention, MochiAttention

logger.warning("Attention backends are an experimental feature and the API may be subject to change.")

attention_classes = (Attention, MochiAttention, AttentionModuleMixin)

for module in self.modules():

if not isinstance(module, attention_classes):

continue

processor = module.processor

if processor is None or not hasattr(processor, "_attention_backend"):

continue

processor._attention_backend = None

def save_pretrained(

self,

save_directory: str | os.PathLike,

is_main_process: bool = True,

save_function: Callable | None = None,

safe_serialization: bool = True,

variant: str | None = None,

max_shard_size: int | str = "10GB",

push_to_hub: bool = False,

use_flashpack: 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 safe_serialization and quantization_serializable:

quantization_serializable = (

quantization_serializable and hf_quantizer.supports_safetensors_serialization

)

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 = WEIGHTS_NAME

if use_flashpack:

weights_name = FLASHPACK_WEIGHTS_NAME

elif safe_serialization:

weights_name = SAFETENSORS_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", None)

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

quantization_metadata = {}

if hf_quantizer is not None:

state_dict, quantization_metadata = hf_quantizer.get_state_dict_and_metadata(

state_dict, safe_serialization=safe_serialization

)

if use_flashpack:

if is_flashpack_available():

import flashpack

else:

logger.error(

"Saving a FlashPack checkpoint in PyTorch, requires both PyTorch and flashpack to be installed. Please see "

"https://pytorch.org/ and https://github.com/fal-ai/flashpack for installation instructions."

)

raise ImportError("Please install torch and flashpack to save a FlashPack checkpoint in PyTorch.")

flashpack.serialization.pack_to_file(

state_dict_or_model=state_dict,

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

target_dtype=self.dtype,

)

else:

# 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].contiguous() for tensor in tensors}

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

if safe_serialization:

metadata = {"format": "pt"}

if quantization_metadata:

metadata.update(quantization_metadata)

metadata = {k: str(v) if not isinstance(v, str) else v for k, v in metadata.items()}

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

else:

torch.save(shard, filepath)

if state_dict_split.is_sharded:

metadata = dict(state_dict_split.metadata)

if quantization_metadata:

metadata.update(quantization_metadata)

index = {

"metadata": 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: str | os.PathLike | None, **kwargs) -> Self:

r"""