UNet3DConditionModel

UNet 模型最初由 Ronneberger 等人提出，用于生物医学图像分割，但由于其输出的图像尺寸与输入相同，因此在 🤗 Diffusers 中也被广泛使用。它是扩散系统中最关键的组件之一，负责实现实际的扩散过程。🤗 Diffusers 中包含多种 UNet 模型变体，具体取决于其维度数量以及是否为条件模型。这里介绍的是一个 3D UNet 条件模型。

该论文的摘要如下：

人们普遍认为，成功训练深度网络需要数千个带标注的训练样本。在本文中，我们提出了一种网络和训练策略，该策略通过充分利用数据增强，更高效地利用现有的带标注样本。该架构由一条用于捕获上下文的收缩路径和一条对称的扩展路径组成，后者能够实现精确定位。我们证明，这样的网络可以从极少量的图像中进行端到端训练，并在 ISBI 电子显微镜堆叠图像中神经元结构分割挑战赛中，优于此前最佳方法（滑动窗口卷积网络）。使用在透射光显微镜图像（相差和 DIC）上训练的同一网络，我们在 2015 年 ISBI 细胞追踪挑战赛中以较大优势赢得了这些类别的冠军。此外，该网络速度很快。在最新的 GPU 上，对一张 512x512 图像进行分割所需时间不到一秒。完整的实现（基于 Caffe）和训练好的网络可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 获取。

UNet3DConditionModel

类 diffusers.UNet3DConditionModel

< source >

( sample_size: int | None = Nonein_channels: int = 4out_channels: int = 4down_block_types: tuple = ('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D')up_block_types: tuple = ('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D')block_out_channels: tuple = (320, 640, 1280, 1280)layers_per_block: int = 2downsample_padding: int = 1mid_block_scale_factor: float = 1act_fn: str = 'silu'norm_num_groups: int | None = 32norm_eps: float = 1e-05cross_attention_dim: int = 1024attention_head_dim: int | tuple[int] = 64num_attention_heads: int | tuple[int] | None = Nonetime_cond_proj_dim: int | None = None )

参数

sample_size (int 或 tuple[int, int]，可选，默认为 None) — 输入/输出样本的高度和宽度。
in_channels (int，可选，默认为 4) — 输入样本的通道数。
out_channels (int，可选，默认为 4) — 输出中的通道数。
down_block_types (tuple[str]，可选，默认为 ("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")) — 要使用的下采样块元组。
up_block_types (tuple[str]，可选，默认为 ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")) — 要使用的上采样块元组。
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）—— 用于中间块的缩放因子。
act_fn (str，可选，默认为 "silu") — 要使用的激活函数。
norm_num_groups (int，可选，默认为 32) — 归一化所使用的组数。如果 None，后处理中将跳过归一化和激活层。
norm_eps (float，可选，默认值为 1e-5) — 归一化所使用的 epsilon 值。
cross_attention_dim (int，可选，默认值为 1024) — 交叉注意力特征的维度。
attention_head_dim (int，可选，默认值为 64) — 注意力头的维度。
num_attention_heads (int，可选) — 注意力头的数量。
time_cond_proj_dim (int，可选，默认为 None) — 时间步嵌入中 cond_proj 层的维度。

一个条件 3D UNet 模型，接收带噪样本、条件状态和时间步，并返回具有样本形状的。

该模型继承自 ModelMixin。请查阅其父类文档，了解为所有模型实现的通用方法（例如下载或保存）。

disable_freeu

< source >

( )

禁用 FreeU 机制。

enable_forward_chunking

< source >

( chunk_size: int | None = Nonedim: int = 0 )

参数

chunk_size (int, 可选) — 前馈层的块大小。如果未指定，将对每个 dim=dim 的张量单独运行前馈层。
dim (int，可选，默认为 0) — 前馈计算应进行分块的维度。可选择 dim=0（批次）或 dim=1（序列长度）。

将注意力处理器设置为使用前馈分块。

enable_freeu

< source >

( s1s2b1b2 )

参数

s1 (float) — 用于第一阶段的比例因子，用于衰减跳跃特征的贡献。这样做是为了缓解增强去噪过程中的“过度平滑效应”。
s2 (float) — 用于衰减 skip features 贡献的第二阶段缩放因子。这样做是为了缓解增强去噪过程中的“过度平滑效应”。
b1 (float) — 用于放大骨干网络特征贡献的第一阶段缩放因子。
b2 (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 = Nonedown_block_additional_residuals: tuple[torch.Tensor] | None = Nonemid_block_additional_residual: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → UNet3DConditionOutput or tuple

参数

sample (torch.Tensor) — 具有如下形状 (batch, num_channels, num_frames, 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。
down_block_additional_residuals — (tuple of torch.Tensor, optional)：一个张量元组，如果指定，将添加到下采样 unet 块的残差中。
mid_block_additional_residual — (torch.Tensor, 可选)：如果指定，该张量将被添加到中间 unet 块的残差中。
return_dict (bool，可选，默认为 True) — 是否返回 UNet3DConditionOutput 而不是普通元组。
cross_attention_kwargs (dict，可选) — 一个 kwargs 字典，如果指定，将传递给 AttnProcessor。

返回值

UNet3DConditionOutput 或 tuple

如果 return_dict 为 True，则返回 UNet3DConditionOutput，否则返回一个 tuple，其中第一个元素是样本张量。

UNet3DConditionModel 的前向方法。

fuse_qkv_projections

< source >

( )

启用融合的 QKV 投影。对于自注意力模块，所有投影矩阵（即 query、key、value）均会被融合。对于交叉注意力模块，key 和 value 投影矩阵会被融合。

> 此 API 处于 🧪 实验阶段。

set_attention_slice

< source >

( slice_size: str | int | list[int] )

参数

slice_size (str 或 int 或 list(int)，可选，默认为 "auto") — 当值为 "auto" 时，输入到注意力头的张量会被减半，因此注意力计算分两步进行。如果值为 "max"，则每次只运行一个切片，以节省最大量的内存。如果提供了数值，则使用 attention_head_dim // slice_size 个切片。在这种情况下，attention_head_dim 必须是 slice_size 的倍数。

启用切片注意力计算。

启用此选项后，注意力模块会将输入张量切分为多个切片，分步计算注意力。这有助于节省部分内存，但会略微降低速度。

set_default_attn_processor

< source >

( )

禁用自定义注意力处理器并设置默认的注意力实现。

unfuse_qkv_projections

< source >

( )

如果已启用，则禁用融合的 QKV 投影。

> 此 API 处于 🧪 实验阶段。

UNet3DConditionOutput

类 diffusers.models.unets.unet_3d_condition.UNet3DConditionOutput

< source >

( sample: Tensor )

参数

sample (torch.Tensor，形状为 (batch_size, num_channels, num_frames, height, width)) — 基于 encoder_hidden_states 输入条件生成的隐藏状态输出。即模型最后一层的输出。

UNet3DConditionModel 的输出。

Update on GitHub