Diffusers documentation
UNet2DConditionModel
UNet2DConditionModel
UNet 模型最初由 Ronneberger 等人提出，用于生物医学图像分割，但由于其输出的图像尺寸与输入相同，因此在 🤗 Diffusers 中也被广泛使用。它是扩散系统中最关键的组件之一，因为它实现了实际的扩散过程。🤗 Diffusers 中包含多种 UNet 模型变体，具体取决于其维度数量以及是否为条件模型。这里介绍的是一个 2D UNet 条件模型。
该论文的摘要如下：
人们普遍认为，成功训练深度网络需要数千个带标注的训练样本。在本文中，我们提出了一种网络和训练策略，该策略通过充分利用数据增强，更高效地利用现有的带标注样本。该架构由一条用于捕获上下文的收缩路径和一条对称的扩展路径组成，后者能够实现精确定位。我们证明，这样的网络可以从极少量的图像中进行端到端训练，并在 ISBI 电子显微镜堆叠图像中神经元结构分割挑战赛中，优于此前最佳方法（滑动窗口卷积网络）。使用在透射光显微镜图像（相差和 DIC）上训练的同一网络，我们在 2015 年 ISBI 细胞追踪挑战赛中以较大优势赢得了这些类别的冠军。此外，该网络速度很快。在最新的 GPU 上，对一张 512x512 图像进行分割所需时间不到一秒。完整的实现（基于 Caffe）和训练好的网络可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 获取。
UNet2DConditionModel
类 diffusers.UNet2DConditionModel
< source >( sample_size: int | tuple[int, int] | None = Nonein_channels: int = 4out_channels: int = 4center_input_sample: bool = Falseflip_sin_to_cos: bool = Truefreq_shift: int = 0down_block_types: tuple = ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D')mid_block_type: str | None = 'UNetMidBlock2DCrossAttn'up_block_types: tuple = ('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D')only_cross_attention: bool | tuple[bool] = Falseblock_out_channels: tuple = (320, 640, 1280, 1280)layers_per_block: int | tuple[int] = 2downsample_padding: int = 1mid_block_scale_factor: float = 1dropout: float = 0.0act_fn: str = 'silu'norm_num_groups: int | None = 32norm_eps: float = 1e-05cross_attention_dim: int | tuple[int] = 1280transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1reverse_transformer_layers_per_block: tuple[tuple[int]] | None = Noneencoder_hid_dim: int | None = Noneencoder_hid_dim_type: str | None = Noneattention_head_dim: int | tuple[int] = 8num_attention_heads: int | tuple[int] | None = Nonedual_cross_attention: bool = Falseuse_linear_projection: bool = Falseclass_embed_type: str | None = Noneaddition_embed_type: str | None = Noneaddition_time_embed_dim: int | None = Nonenum_class_embeds: int | None = Noneupcast_attention: bool = Falseresnet_time_scale_shift: str = 'default'resnet_skip_time_act: bool = Falseresnet_out_scale_factor: float = 1.0time_embedding_type: str = 'positional'time_embedding_dim: int | None = Nonetime_embedding_act_fn: str | None = Nonetimestep_post_act: str | None = Nonetime_cond_proj_dim: int | None = Noneconv_in_kernel: int = 3conv_out_kernel: int = 3projection_class_embeddings_input_dim: int | None = Noneattention_type: str = 'default'class_embeddings_concat: bool = Falsemid_block_only_cross_attention: bool | None = Nonecross_attention_norm: str | None = Noneaddition_embed_type_num_heads: int = 64 )
参数
- sample_size (
int或tuple[int, int]，可选，默认为None) — 输入/输出样本的高度和宽度。 - in_channels (
int，可选，默认为 4) — 输入样本的通道数。 - out_channels (
int，可选，默认为 4) — 输出中的通道数。 - center_input_sample (
bool，可选，默认为False) — 是否对输入样本进行居中处理。 - flip_sin_to_cos (
bool，可选，默认为True) — 是否将时间嵌入中的 sin 翻转为 cos。 - freq_shift（
int，可选，默认值为 0）——应用于时间嵌入的频率偏移量。 - down_block_types (
tuple[str]，可选，默认为("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")) — 要使用的下采样块元组。 - mid_block_type (
str，可选，默认为"UNetMidBlock2DCrossAttn") — UNet 中间部分的块类型，可以是UNetMidBlock2DCrossAttn、UNetMidBlock2D或UNetMidBlock2DSimpleCrossAttn之一。如果为None，则跳过中间块层。 - up_block_types (
tuple[str]，可选，默认为("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")) — 要使用的上采样块元组。 - only_cross_attention(
bool或tuple[bool]，可选，默认为False) — 是否在基础 Transformer 块中包含自注意力，参见BasicTransformerBlock。 - block_out_channels (
tuple[int]，可选，默认为(320, 640, 1280, 1280)) — 每个 block 的输出通道数元组。 - layers_per_block (
int，可选，默认为 2) — 每个块中的层数。 - downsample_padding (
int，可选，默认为 1) — 用于下采样卷积的填充值。 - mid_block_scale_factor（
float，可选，默认值为 1.0）—— 用于中间块的缩放因子。 - dropout (
float，可选，默认值为 0.0) — 要使用的 dropout 概率。 - act_fn (
str，可选，默认为"silu") — 要使用的激活函数。 - norm_num_groups (
int，可选，默认为 32) — 归一化所使用的组数。 如果None，后处理中将跳过归一化和激活层。 - norm_eps (
float，可选，默认值为 1e-5) — 归一化所使用的 epsilon 值。 - cross_attention_dim (
int或tuple[int]，可选，默认值为 1280) — 交叉注意力特征的维度。 - transformer_layers_per_block（
int、tuple[int]或tuple[tuple]，可选，默认为 1）——BasicTransformerBlock类型 transformer 块的数量。仅对CrossAttnDownBlock2D、CrossAttnUpBlock2D、UNetMidBlock2DCrossAttn有效。 - reverse_transformer_layers_per_block — (
tuple[tuple]，可选，默认为 None)： U-Net 的上采样块中类型为BasicTransformerBlock的 transformer 块的数量。仅当transformer_layers_per_block的类型为tuple[tuple]且针对CrossAttnDownBlock2D、CrossAttnUpBlock2D、UNetMidBlock2DCrossAttn时相关。 - encoder_hid_dim (
int，可选，默认为 None) — 如果定义了encoder_hid_dim_type，encoder_hidden_states将从encoder_hid_dim维度投影到cross_attention_dim。 - encoder_hid_dim_type (
str，可选，默认为None) — 如果指定，encoder_hidden_states以及其他潜在嵌入将根据encoder_hid_dim_type下投影到维度为cross_attention的文本嵌入。 - attention_head_dim (
int，可选，默认为 8) — 注意力头的维度。 - num_attention_heads (
int, 可选) — 注意力头的数量。如果未定义，则默认为attention_head_dim - resnet_time_scale_shift (
str，可选，默认为"default") — ResNet 块的时间缩放偏移配置（参见ResnetBlock2D）。可选值为default或scale_shift。 - class_embed_type (
str，可选，默认为None) — 要使用的类别嵌入类型，最终会与时间嵌入相加。可选值包括None、"timestep"、"identity"、"projection"或"simple_projection"。 - addition_embed_type (
str，可选，默认为None) — 配置一个可选的嵌入，该嵌入将与时间嵌入相加。可选值为None或 “text”。选择 “text” 时将使用TextTimeEmbedding层。 - addition_time_embed_dim — (
int，可选，默认为None)： 时间步嵌入的维度。 - num_class_embeds (
int，可选，默认为None) — 当使用class_embed_type等于None的类别条件时，可学习嵌入矩阵的输入维度，该矩阵将被投影到time_embed_dim。 - time_embedding_type (
str，可选，默认为positional) — 用于时间步的位置嵌入类型。可选值为positional或fourier。 - time_embedding_dim (
int，可选，默认为None) — 用于覆盖投影后时间嵌入维度的可选参数。 - time_embedding_act_fn (
str，可选，默认为None) — 可选的激活函数，仅在时间嵌入传递到 UNet 其余部分之前使用一次。可从silu、mish、gelu和swish中选择。 - timestep_post_act (
str，可选，默认为None) — 用于时间步嵌入的第二个激活函数。可从silu、mish和gelu中选择。 - time_cond_proj_dim (
int，可选，默认为None) — 时间步嵌入中cond_proj层的维度。 - conv_in_kernel (
int，可选，默认为3) —conv_in层的卷积核大小。 - conv_out_kernel (
int，可选，默认为3) —conv_out层的卷积核大小。 - projection_class_embeddings_input_dim (
int, 可选) — 当class_embed_type="projection"时，class_labels输入的维度。当class_embed_type="projection"时必填。 - class_embeddings_concat (
bool，可选，默认为False) — 是否将时间嵌入与类别嵌入进行拼接。 - mid_block_only_cross_attention (
bool，可选，默认为None) — 在使用UNetMidBlock2DSimpleCrossAttn时，是否对中间块使用交叉注意力。如果only_cross_attention被指定为单个布尔值，且mid_block_only_cross_attention为None，则 使用only_cross_attention的值作为mid_block_only_cross_attention的值。否则默认为False。
一个条件 2D UNet 模型，接收带噪样本、条件状态和时间步，并返回具有样本形状的。
该模型继承自 ModelMixin。请查阅其父类文档，了解为所有模型实现的通用方法（例如下载或保存）。
禁用 FreeU 机制。
enable_freeu
< source >( s1: floats2: floatb1: floatb2: float )
启用来自 https://huggingface.co/papers/2309.11497 的 FreeU 机制。
缩放因子后面的后缀表示它们所应用的阶段块。
有关适用于 Stable Diffusion v1、v2 和 Stable Diffusion XL 等不同流水线的已知良好参数组合，请参阅官方仓库。
forward
< source >( sample: Tensortimestep: typing.Union[torch.Tensor, float, int]encoder_hidden_states: Tensorclass_labels: typing.Optional[torch.Tensor] = Nonetimestep_cond: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Nonecross_attention_kwargs: dict[str, typing.Any] | None = Noneadded_cond_kwargs: dict[str, torch.Tensor] | None = Nonedown_block_additional_residuals: tuple[torch.Tensor] | None = Nonemid_block_additional_residual: typing.Optional[torch.Tensor] = Nonedown_intrablock_additional_residuals: tuple[torch.Tensor] | None = Noneencoder_attention_mask: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → UNet2DConditionOutput or tuple
参数
- sample (
torch.Tensor) — 形状为(batch, channel, height, width)的含噪输入张量。 - timestep（
torch.Tensor或float或int）— 用于对输入进行去噪的时间步数。 - encoder_hidden_states (
torch.Tensor) — 形状为(batch, sequence_length, feature_dim)的编码器隐藏状态。 - class_labels (
torch.Tensor，可选，默认为None) — 用于条件控制的可选类别标签。它们的嵌入将与时间步嵌入相加。 - timestep_cond — (
torch.Tensor，可选，默认为None)： 时间步的条件嵌入。如果提供，这些嵌入将与通过self.time_embedding层传入的样本相加，以得到时间步嵌入。 - attention_mask (
torch.Tensor，可选，默认为None) — 形状为(batch, key_tokens)的注意力掩码将应用于encoder_hidden_states。如果值为1，则保留该掩码；如果值为0，则将其丢弃。掩码会被转换为偏置，向对应于“丢弃”标记的注意力分数添加较大的负值。 - cross_attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定，将作为参数传递给 diffusers.models.attention_processor 中self.processor下定义的AttentionProcessor。 - added_cond_kwargs — (
dict，可选)： 一个 kwargs 字典，包含额外的嵌入。如果指定了这些嵌入，它们将被添加到传递给 UNet 块的嵌入中。 - down_block_additional_residuals — (
tupleoftorch.Tensor, optional)： 一个张量元组，如果指定，将添加到下采样 unet 块的残差中。 - mid_block_additional_residual — (
torch.Tensor, 可选)： 如果指定，该张量将被添加到中间 unet 块的残差中。 - down_intrablock_additional_residuals（
tupleoftorch.Tensor，可选）— 在 UNet 下采样块内部需要额外添加的残差，例如来自 T2I-Adapter 侧模型的残差 - encoder_attention_mask (
torch.Tensor) — 应用于encoder_hidden_states的交叉注意力掩码，形状为(batch, sequence_length)。如果True，则保留该掩码；否则，如果False，则丢弃该掩码。掩码将被转换为偏置， 向对应于“丢弃”标记的注意力分数添加较大的负值。 - return_dict (
bool，可选，默认为True) — 是否返回 UNet2DConditionOutput 而不是普通元组。
返回值
UNet2DConditionOutput 或 tuple
如果 return_dict 为 True，则返回 UNet2DConditionOutput，否则返回一个 tuple，其中第一个元素是样本张量。
UNet2DConditionModel 的前向方法。
启用融合的 QKV 投影。对于自注意力模块，所有投影矩阵（即 query、key、value）均会被融合。对于交叉注意力模块，key 和 value 投影矩阵会被融合。
> 此 API 处于 🧪 实验阶段。
set_attention_slice
< source >( slice_size: str | int | list[int] = 'auto' )
启用切片注意力计算。
启用此选项后，注意力模块会将输入张量切分为多个切片，分步计算注意力。这有助于节省部分内存，但会略微降低速度。
禁用自定义注意力处理器并设置默认的注意力实现。
UNet2DConditionOutput
类 diffusers.models.unets.unet_2d_condition.UNet2DConditionOutput
< source >( sample: Tensor = None )
UNet2DConditionModel 的输出。