模型
🤗 Diffusers 提供了针对流行算法的预训练模型,以及用于构建自定义扩散系统的模块。模型的主要功能是根据分布 .
所有模型都基于 ModelMixin 基类构建,该类是一个 torch.nn.Module,提供了在本地和 Hugging Face Hub 上保存和加载模型的基本功能。
ModelMixin
所有模型的基类。
ModelMixin 负责存储模型配置,并提供用于加载、下载和保存模型的方法。
- config_name (
str) — 调用 save_pretrained() 时保存模型所使用的文件名。
仅编译模型中频繁重复的子模块(例如 Transformer 层),而不是编译整个模型。这种技术通常被称为区域编译(参见 PyTorch 教程 https://docs.pytorch.org/tutorials/recipes/regional_compilation.html),可以大幅缩短端到端的编译时间,同时保留与完整 torch.compile 相当的运行时加速效果。
要编译的子模块集合是通过模型定义中是否存在 _repeated_blocks 属性来确定的。请在您的模型子类中将此属性定义为一个类名(字符串)列表/元组。所有类名与之匹配的模块都将被编译。
发现后,每个匹配的子模块都会通过调用 submodule.compile(*args, **kwargs) 进行编译。您提供给 compile_repeated_blocks 的任何位置参数或关键字参数都会原样转发给 torch.compile。
点分路径(例如 torch.float32)的模拟值:属性访问链和调用行为如同透传装饰器,repr 为点分路径,将其用作基类时会替换为普通的 type 基类(PEP 560 __mro_entries__),因此真实的子类保留正常的元类,且 inspect.signature 读取的是它们真实的 __init__,而非模拟对象的。
如果模型已通过支持反量化的量化方法进行了量化,则可能需要对其进行反量化。
停用当前模型的梯度检查点(在其他框架中可能被称为 activation checkpointing 或 checkpoint activations)。
从 torch_npu 中禁用 npu flash attention
禁用来自 xFormers 的内存高效注意力。
禁用 torch_xla 的 flash attention pallals 内核。
enable_gradient_checkpointing
< source >( gradient_checkpointing_func: typing.Optional[typing.Callable] = None )
为当前模型启用梯度检查点(在其他框架中可能被称为 activation checkpointing 或 checkpoint activations)。
enable_group_offload
< source >( onload_device: deviceoffload_device: device = torch.device()offload_type: str = 'block_level'num_blocks_per_group: int | None = Nonenon_blocking: bool = Falseuse_stream: bool = Falserecord_stream: bool = Falselow_cpu_mem_usage = Falseoffload_to_disk_path: str | None = Noneblock_modules: str | None = Noneexclude_kwargs: str | None = None )
为当前模型启用组卸载。
有关更多信息,请参阅 apply_group_offloading()。
示例:
>>> from diffusers import CogVideoXTransformer3DModel
>>> transformer = CogVideoXTransformer3DModel.from_pretrained(
... "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16
... )
>>> transformer.enable_group_offload(
... onload_device=torch.device("cuda"),
... offload_device=torch.device("cpu"),
... offload_type="leaf_level",
... use_stream=True,
... )enable_layerwise_casting
< source >( storage_dtype: dtype = torch.float8_e4m3fncompute_dtype: typing.Optional[torch.dtype] = Noneskip_modules_pattern: tuple[str, ...] | None = Noneskip_modules_classes: tuple[typing.Type[torch.nn.Module], ...] | None = Nonenon_blocking: bool = False )
参数
- storage_dtype (
torch.dtype) — 用于存储时模型应转换到的 dtype。 - compute_dtype (
torch.dtype) — 在前向传播过程中,模型权重应转换到的数据类型。 - skip_modules_pattern (
tuple[str, ...], 可选) — 用于匹配在逐层转换过程中需要跳过的模块名称的模式列表。如果设置为None,则使用默认的跳过模式来忽略模块的某些内部层和 PEFT 层。 - skip_modules_classes (
tuple[Type[torch.nn.Module], ...],可选) — 在逐层转换过程中需要跳过的模块类列表。 - non_blocking (
bool,可选,默认为False) — 如果为True,则权重转换操作是非阻塞的。
为当前模型启用逐层类型转换。
逐层转换是一种技术,它将模型权重转换为较低精度的 dtype 进行存储,但在计算时即时将其上转为较高精度的 dtype。此过程可以显著降低模型权重带来的内存占用,但可能会导致输出质量出现一定程度的下降。大多数质量下降微乎其微,主要源于归一化层和调制层中的权重转换。
默认情况下,diffusers 中的大多数模型会将 _skip_layerwise_casting_patterns 属性设置为忽略 patch 嵌入、位置嵌入和归一化层。这是因为这些层对质量而言最可能是精度敏感的。如果您希望更改此行为,可以将 _skip_layerwise_casting_patterns 属性设置为 None,或者使用自定义参数调用 apply_layerwise_casting()。
示例:
使用 enable_layerwise_casting():
>>> from diffusers import CogVideoXTransformer3DModel
>>> transformer = CogVideoXTransformer3DModel.from_pretrained(
... "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16
... )
>>> # Enable layerwise casting via the model, which ignores certain modules by default
>>> transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)从 torch_npu 启用 npu flash attention
enable_xformers_memory_efficient_attention
< source >( attention_op: typing.Optional[typing.Callable] = None )
参数
- attention_op (
Callable, 可选) — 覆盖默认的None算子,用作 xFormers 的memory_efficient_attention()函数的op参数。
启用来自 xFormers 的内存高效注意力机制。
启用此选项后,您应能观察到更低的 GPU 内存占用,并可能在推理阶段获得加速。训练阶段的加速并不保证。
> ⚠️ 当同时启用 memory efficient attention 和 sliced attention 时,memory efficient attention 具有 > 优先权。
示例:
>>> import torch
>>> from diffusers import UNet2DConditionModel
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
>>> model = UNet2DConditionModel.from_pretrained(
... "stabilityai/stable-diffusion-2-1", subfolder="unet", torch_dtype=torch.float16
... )
>>> model = model.to("cuda")
>>> model.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)enable_xla_flash_attention
< source >( partition_spec: typing.Optional[typing.Callable] = None**kwargs )
为 torch_xla 启用 flash attention pallals 内核。
from_pretrained
< source >( pretrained_model_name_or_path: str | os.PathLike | None**kwargs )
参数
- pretrained_model_name_or_path (
str或os.PathLike, 可选) — 可以是以下之一:- 一个字符串,即托管在 Hub 上的预训练模型的 model id(例如
google/ddpm-celebahq-256)。 - 指向一个 目录(例如
./my_model_directory)的路径,该目录包含使用 save_pretrained() 保存的模型权重。
- 一个字符串,即托管在 Hub 上的预训练模型的 model id(例如
- cache_dir (
str | os.PathLike, 可选) — 如果未使用标准缓存,则下载的预训练模型配置将被缓存到的目录路径。 - dtype (
torch.dtype,可选) — 覆盖默认的torch.dtype,并使用其他 dtype 加载模型。 - force_download (
bool,可选,默认为False) — 是否强制(重新)下载模型权重和配置文件,如果存在缓存版本则覆盖它们。 - proxies (
dict[str, str],可选) — 按协议或端点使用的代理服务器字典,例如{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}。每次请求都会使用这些代理。 - output_loading_info (
bool,可选,默认为False) — 是否同时返回一个包含缺失键、意外键和错误消息的字典。 - local_files_only(
bool, 可选,默认为False) — 是否仅加载本地的模型权重和配置文件。如果设置为True,则不会从 Hub 下载模型。 - token (
str或 bool,可选) — 用于远程文件的 HTTP bearer 授权令牌。如果为True,则使用由diffusers-cli login生成的令牌(存储在~/.huggingface中)。 - revision (
str, 可选, 默认为"main") — 要使用的特定模型版本。它可以是分支名称、标签名称、提交 id,或 Git 允许的任何标识符。 - subfolder (
str,可选,默认为"") — 模型文件在 Hub 或本地更大的模型仓库中的子文件夹位置。 - mirror (
str, 可选) — 镜像源,用于解决在中国下载模型时可能遇到的访问问题。我们不保证该源的时效性或安全性,详情请参阅镜像站点。 - device_map (
int | str | torch.device或dict[str, int | str | torch.device], 可选) — 一个映射,用于指定每个子模块应放置的位置。无需为每个参数/缓冲区名称单独定义;一旦某个模块名称被包含在内,其所有子模块都将被发送到同一设备。默认为None, ,表示模型将在 CPU 上加载。示例:
从预训练模型配置实例化一个预训练的 PyTorch 模型。
模型默认处于评估模式 - model.eval() -,且 dropout 模块已停用。若要训练模型,请使用 model.train() 将其切换回训练模式。
> 若要使用私有或 受控模型,请先通过
hf > auth login登录。您还可以启用特殊的 > “离线模式”,以便在防火墙环境中使用此方法。
示例:
from diffusers import UNet2DConditionModel
unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet")如果出现以下错误消息,您需要针对下游任务微调权重:
Some weights of UNet2DConditionModel were not initialized from the model checkpoint at stable-diffusion-v1-5/stable-diffusion-v1-5 and are newly initialized because the shapes did not match:
- conv_in.weight: found shape torch.Size([320, 4, 3, 3]) in the checkpoint and torch.Size([320, 9, 3, 3]) in the model instantiated
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.get_memory_footprint
< source >( return_buffers = True )
参数
- return_buffers (
bool,可选,默认为True) — 是否在计算内存占用时返回缓冲区张量的大小。缓冲区是指不需要梯度且未注册为参数的张量。例如,批归一化层中的均值和标准差。请参阅:https://discuss.pytorch.org/t/what-pytorch-means-by-buffers/120266/2
获取模型的内存占用。这将返回当前模型的内存占用(以字节为单位)。 可用于对当前模型的内存占用进行基准测试并设计一些测试。该解决方案灵感来自 PyTorch 讨论区:https://discuss.pytorch.org/t/gpu-memory-that-model-uses/56822/2
num_parameters
< source >( only_trainable: bool = Falseexclude_embeddings: bool = False ) → int
获取模块中(可训练或非嵌入)参数的数量。
重置模型的注意力后端。后续对 forward 的调用将使用环境变量中设置的默认值(如果已设置),否则使用 torch 原生的缩放点积注意力。
save_pretrained
< source >( save_directory: str | os.PathLikeis_main_process: bool = Truesave_function: typing.Optional[typing.Callable] = Nonesafe_serialization: bool = Truevariant: str | None = Nonemax_shard_size: int | str = '10GB'push_to_hub: bool = Falseuse_flashpack: bool = False**kwargs )
参数
- save_directory (
str或os.PathLike) — 用于保存模型及其配置文件的目录。如果该目录不存在,将会自动创建。 - is_main_process (
bool,可选,默认为True) — 调用此函数的进程是否为主进程。在分布式训练期间,当需要在所有进程上调用此函数时非常有用。在这种情况下,仅在主进程上设置is_main_process=True,以避免竞态条件。 - save_function (
Callable) — 用于保存状态字典的函数。在分布式训练期间,当您需要用其他方法替换torch.save时非常有用。可以通过环境变量DIFFUSERS_SAVE_MODE进行配置。 - safe_serialization(
bool,可选,默认为True)—— 是否使用safetensors保存模型,还是使用传统的基于pickle的 PyTorch 方式。 - variant (
str, 可选) — 如果指定,权重将以pytorch_model.<variant>.bin格式保存。 - max_shard_size (
int或str,默认为"10GB") — 分片前检查点(checkpoint)的最大大小。分片后的每个检查点大小都将小于该值。如果以字符串形式表示,需要是数字后跟单位(例如"5GB")。如果以整数形式表示,单位为字节。请注意,从 2024 年 10 月开始,经过一段时间后,该限制将会降低,以便用户升级到diffusers的最新版本。这是为了在 Hugging Face 生态系统的不同库(例如transformers和accelerate)之间为该参数建立统一的默认大小。 - push_to_hub (
bool,可选,默认为False) — 是否在保存模型后将其推送到 Hugging Face Hub。您可以使用repo_id指定要推送到的仓库(默认使用您命名空间中save_directory的名称)。 - kwargs (
dict[str, Any],可选) — 传递给 push_to_hub() 方法的额外关键字参数。
将模型及其配置文件保存到目录中,以便使用 from_pretrained() 类方法重新加载。
set_attention_backend
< source >( backend: str )
设置模型的注意力后端。
设置 npu 快速注意力的开关。
PushToHubMixin
一个用于将模型、调度器或流水线推送到 Hugging Face Hub 的 Mixin。
push_to_hub
< source >( repo_id: strcommit_message: str | None = Noneprivate: bool | None = Nonetoken: str | None = Nonecreate_pr: bool = Falsesafe_serialization: bool = Truevariant: str | None = Nonesubfolder: str | None = None )
参数
- repo_id (
str) — 要推送模型、调度器或流水线文件的目标仓库名称。推送到组织时,应包含组织名称。repo_id也可以是本地目录的路径。 - commit_message (
str, 可选) — 推送时提交的提交信息。默认为"Upload {object}"。 - private (
bool,可选) — 是否将仓库设为私有。如果为None(默认值),除非组织的默认设置为私有,否则仓库将是公开的。如果仓库已存在,则忽略此值。 - token (
str, 可选) — 用于远程文件的 HTTP bearer 授权令牌。即运行hf auth login时生成的令牌(存储在~/.huggingface中)。 - create_pr (
bool,可选,默认为False) — 是否使用上传的文件创建 PR,还是直接提交。 - safe_serialization(
bool,可选,默认为True)—— 是否将模型权重转换为safetensors格式。 - variant (
str, 可选) — 如果指定,权重将以pytorch_model.<variant>.bin格式保存。
将模型、调度器或流水线文件上传到 🤗 Hugging Face Hub。
示例:
from diffusers import UNet2DConditionModel
unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-2", subfolder="unet")
# Push the `unet` to your namespace with the name "my-finetuned-unet".
unet.push_to_hub("my-finetuned-unet")
# Push the `unet` to an organization with the name "my-finetuned-unet".
unet.push_to_hub("your-org/my-finetuned-unet")