Diffusers documentation

UNet

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

UNet

某些训练方法(例如 LoRA 和 Custom Diffusion)通常针对 UNet 的注意力层,但这些训练方法也可以针对其他非注意力层。与训练模型的所有参数不同,这里只训练参数的一部分,因此速度更快、效率更高。如果您仅将权重加载到 UNet 中,则此类非常有用。如果需要将权重加载到文本编码器,或同时加载到文本编码器和 UNet 中,请尝试使用 load_lora_weights() 函数。

UNet2DConditionLoadersMixin 类提供了用于加载和保存权重、融合和取消融合 LoRA、禁用和启用 LoRA,以及设置和删除适配器的函数。

要了解如何加载 LoRA 权重,请参阅 LoRA 指南。

UNet2DConditionLoadersMixin

类 diffusers.loaders.UNet2DConditionLoadersMixin

< >

( )

将 LoRA 层加载到 UNet2DCondtionModel 中。

load_attn_procs

< >

( pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor]**kwargs )

参数

  • pretrained_model_name_or_path_or_dict (str 或 os.PathLike 或 dict) — 可以是以下之一: - 一个字符串,即托管在 Hub 上的预训练模型的模型 ID(例如 google/ddpm-celebahq-256)。 - 一个目录路径(例如 ./my_model_directory),其中包含使用 ModelMixin.save_pretrained() 保存的模型权重。 - 一个 torch 状态字典。
    • 一个字符串,即托管在 Hub 上的预训练模型的模型 ID(例如 google/ddpm-celebahq-256)。
    • 指向一个目录的路径(例如 ./my_model_directory),其中包含使用 ModelMixin.save_pretrained() 保存的模型权重。
    • 一个 torch 状态字典。
  • cache_dir (str | os.PathLike, 可选) — 如果未使用标准缓存,则下载的预训练模型配置将被缓存到的目录路径。
  • force_download (bool,可选,默认为 False) — 是否强制(重新)下载模型权重和配置文件,如果存在缓存版本则覆盖它们。
  • proxies (dict[str, str],可选) — 按协议或端点使用的代理服务器字典,例如 {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}。每次请求都会使用这些代理。
  • 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 或本地更大的模型仓库中的子文件夹位置。
  • weight_name (str,可选,默认为 None) — 序列化状态字典文件的名称。

将预训练的 Custom Diffusion 注意力处理器层加载到 UNet2DConditionModel 中。注意力处理器层必须在 attention_processor.py 中定义,并且是一个 torch.nn.Module 类。若要加载 LoRA 层,请改用 load_lora_adapter()。

示例:

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float16,
).to("cuda")
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")

save_attn_procs

< >

( save_directory: str | os.PathLikeis_main_process: bool = Trueweight_name: str = Nonesave_function: typing.Callable = Nonesafe_serialization: bool = True**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 来保存模型。

将自定义 Diffusion 注意力处理器层保存到目录中,以便使用 load_attn_procs() 方法重新加载。若要保存 LoRA 层,请改用 save_lora_adapter()。

示例:

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float16,
).to("cuda")
pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
pipeline.unet.save_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin")
Update on GitHub