Diffusers 文档
UNet3DConditionModel
并获取增强的文档体验
开始使用
UNet3DConditionModel
The 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
class diffusers.UNet3DConditionModel
< 源代码 >( sample_size: typing.Optional[int] = None in_channels: int = 4 out_channels: int = 4 down_block_types: typing.Tuple[str, ...] = ('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D') up_block_types: typing.Tuple[str, ...] = ('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D') block_out_channels: typing.Tuple[int, ...] = (320, 640, 1280, 1280) layers_per_block: int = 2 downsample_padding: int = 1 mid_block_scale_factor: float = 1 act_fn: str = 'silu' norm_num_groups: typing.Optional[int] = 32 norm_eps: float = 1e-05 cross_attention_dim: int = 1024 attention_head_dim: typing.Union[int, typing.Tuple[int]] = 64 num_attention_heads: typing.Union[int, typing.Tuple[int], NoneType] = None time_cond_proj_dim: typing.Optional[int] = 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)) — 每个块的输出通道元组。 - 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。查看超类文档，了解为其所有模型实现的通用方法（例如下载或保存）。
禁用 FreeU 机制。
enable_forward_chunking
< source >( chunk_size: typing.Optional[int] = None dim: int = 0 )
设置注意力处理器以使用 前馈分块。
enable_freeu
< source >( s1 s2 b1 b2 )
启用来自 https://arxiv.org/abs/2309.11497 的 FreeU 机制。
缩放因子后的后缀表示应用它们的阶段块。
请参考 官方存储库，了解已知适用于不同管道（例如 Stable Diffusion v1、v2 和 Stable Diffusion XL）的值组合。
forward
< source >( sample: Tensor timestep: typing.Union[torch.Tensor, float, int] encoder_hidden_states: Tensor class_labels: typing.Optional[torch.Tensor] = None timestep_cond: typing.Optional[torch.Tensor] = None attention_mask: typing.Optional[torch.Tensor] = None cross_attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None down_block_additional_residuals: typing.Optional[typing.Tuple[torch.Tensor]] = None mid_block_additional_residual: typing.Optional[torch.Tensor] = None return_dict: bool = True ) → UNet3DConditionOutput 或 tuple
参数
- sample (
torch.Tensor) — 噪声输入张量，形状如下(batch, num_channels, num_frames, height, width。 - timestep (
torch.Tensor或float或int) — 用于去噪输入的 timestep 数量。 - 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 字典，如果指定，则作为self.processor下定义的AttentionProcessor的参数传递，详见 diffusers.models.attention_processor。 - down_block_additional_residuals — (
tupleoftorch.Tensor, 可选): 张量元组，如果指定，则添加到下采样 unet 块的残差中。 - mid_block_additional_residual — (
torch.Tensor, 可选): 张量，如果指定，则添加到中间 unet 块的残差中。 - return_dict (
bool, 可选, 默认为True) — 是否返回 UNet3DConditionOutput 而不是普通元组。 - cross_attention_kwargs (
dict, 可选) — 关键字参数字典，如果指定，则传递给AttnProcessor。
返回
UNet3DConditionOutput 或 tuple
如果 return_dict 为 True，则返回 UNet3DConditionOutput，否则返回 tuple，其中第一个元素是样本张量。
UNet3DConditionModel 前向方法。
启用融合 QKV 投影。对于自注意力模块，所有投影矩阵（即，查询、键、值）都融合在一起。对于交叉注意力模块，键和值投影矩阵融合在一起。
此 API 处于 🧪 实验阶段。
set_attention_slice
< source >( slice_size: typing.Union[str, int, typing.List[int]] )
启用切片注意力计算。
启用此选项后，注意力模块会将输入张量分割成切片，以分几个步骤计算注意力。 这对于节省一些内存以换取速度的少量降低非常有用。
set_attn_processor
< source >( processor: typing.Union[diffusers.models.attention_processor.AttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor2_0, diffusers.models.attention_processor.JointAttnProcessor2_0, diffusers.models.attention_processor.PAGJointAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGJointAttnProcessor2_0, diffusers.models.attention_processor.FusedJointAttnProcessor2_0, diffusers.models.attention_processor.AllegroAttnProcessor2_0, diffusers.models.attention_processor.AuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FusedAuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.CogVideoXAttnProcessor2_0, diffusers.models.attention_processor.FusedCogVideoXAttnProcessor2_0, diffusers.models.attention_processor.XFormersAttnAddedKVProcessor, diffusers.models.attention_processor.XFormersAttnProcessor, diffusers.models.attention_processor.XLAFlashAttnProcessor2_0, diffusers.models.attention_processor.AttnProcessorNPU, diffusers.models.attention_processor.AttnProcessor2_0, diffusers.models.attention_processor.MochiVaeAttnProcessor2_0, diffusers.models.attention_processor.MochiAttnProcessor2_0, diffusers.models.attention_processor.StableAudioAttnProcessor2_0, diffusers.models.attention_processor.HunyuanAttnProcessor2_0, diffusers.models.attention_processor.FusedHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.LuminaAttnProcessor2_0, diffusers.models.attention_processor.FusedAttnProcessor2_0, diffusers.models.attention_processor.CustomDiffusionXFormersAttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor2_0, diffusers.models.attention_processor.SlicedAttnProcessor, diffusers.models.attention_processor.SlicedAttnAddedKVProcessor, diffusers.models.attention_processor.SanaLinearAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGSanaLinearAttnProcessor2_0, diffusers.models.attention_processor.PAGIdentitySanaLinearAttnProcessor2_0, diffusers.models.attention_processor.SanaMultiscaleLinearAttention, diffusers.models.attention_processor.SanaMultiscaleAttnProcessor2_0, diffusers.models.attention_processor.SanaMultiscaleAttentionProjection, diffusers.models.attention_processor.IPAdapterAttnProcessor, diffusers.models.attention_processor.IPAdapterAttnProcessor2_0, diffusers.models.attention_processor.IPAdapterXFormersAttnProcessor, diffusers.models.attention_processor.SD3IPAdapterJointAttnProcessor2_0, diffusers.models.attention_processor.PAGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.LoRAAttnProcessor, diffusers.models.attention_processor.LoRAAttnProcessor2_0, diffusers.models.attention_processor.LoRAXFormersAttnProcessor, diffusers.models.attention_processor.LoRAAttnAddedKVProcessor, typing.Dict[str, typing.Union[diffusers.models.attention_processor.AttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor2_0, diffusers.models.attention_processor.JointAttnProcessor2_0, diffusers.models.attention_processor.PAGJointAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGJointAttnProcessor2_0, diffusers.models.attention_processor.FusedJointAttnProcessor2_0, diffusers.models.attention_processor.AllegroAttnProcessor2_0, diffusers.models.attention_processor.AuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FusedAuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.CogVideoXAttnProcessor2_0, diffusers.models.attention_processor.FusedCogVideoXAttnProcessor2_0, diffusers.models.attention_processor.XFormersAttnAddedKVProcessor, diffusers.models.attention_processor.XFormersAttnProcessor, diffusers.models.attention_processor.XLAFlashAttnProcessor2_0, diffusers.models.attention_processor.AttnProcessorNPU, diffusers.models.attention_processor.AttnProcessor2_0, diffusers.models.attention_processor.MochiVaeAttnProcessor2_0, diffusers.models.attention_processor.MochiAttnProcessor2_0, diffusers.models.attention_processor.StableAudioAttnProcessor2_0, diffusers.models.attention_processor.HunyuanAttnProcessor2_0, diffusers.models.attention_processor.FusedHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.LuminaAttnProcessor2_0, diffusers.models.attention_processor.FusedAttnProcessor2_0, diffusers.models.attention_processor.CustomDiffusionXFormersAttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor2_0, diffusers.models.attention_processor.SlicedAttnProcessor, diffusers.models.attention_processor.SlicedAttnAddedKVProcessor, diffusers.models.attention_processor.SanaLinearAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGSanaLinearAttnProcessor2_0, diffusers.models.attention_processor.PAGIdentitySanaLinearAttnProcessor2_0, diffusers.models.attention_processor.SanaMultiscaleLinearAttention, diffusers.models.attention_processor.SanaMultiscaleAttnProcessor2_0, diffusers.models.attention_processor.SanaMultiscaleAttentionProjection, diffusers.models.attention_processor.IPAdapterAttnProcessor, diffusers.models.attention_processor.IPAdapterAttnProcessor2_0, diffusers.models.attention_processor.IPAdapterXFormersAttnProcessor, diffusers.models.attention_processor.SD3IPAdapterJointAttnProcessor2_0, diffusers.models.attention_processor.PAGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.LoRAAttnProcessor, diffusers.models.attention_processor.LoRAAttnProcessor2_0, diffusers.models.attention_processor.LoRAXFormersAttnProcessor, diffusers.models.attention_processor.LoRAAttnAddedKVProcessor]]] )
设置用于计算注意力的注意力处理器。
禁用自定义注意力处理器并设置默认注意力实现。
UNet3DConditionOutput
class diffusers.models.unets.unet_3d_condition.UNet3DConditionOutput
< source >( sample: Tensor )
UNet3DConditionModel 的输出。