Diffusers documentation
UNetMotionModel
UNetMotionModel
UNet 模型最初由 Ronneberger 等人提出,用于生物医学图像分割,但由于其输出的图像尺寸与输入相同,因此在 🤗 Diffusers 中也被广泛使用。它是扩散系统中最关键的组件之一,负责实现实际的扩散过程。🤗 Diffusers 中包含多种 UNet 模型变体,具体取决于其维度数量以及是否为条件模型。这里介绍的是一个 2D UNet 模型。
该论文的摘要如下:
人们普遍认为,成功训练深度网络需要数千个带标注的训练样本。在本文中,我们提出了一种网络和训练策略,该策略通过充分利用数据增强,更高效地利用现有的带标注样本。该架构由一条用于捕获上下文的收缩路径和一条对称的扩展路径组成,后者能够实现精确定位。我们证明,这样的网络可以从极少量的图像中进行端到端训练,并在 ISBI 电子显微镜堆叠图像中神经元结构分割挑战赛中,优于此前最佳方法(滑动窗口卷积网络)。使用在透射光显微镜图像(相差和 DIC)上训练的同一网络,我们在 2015 年 ISBI 细胞追踪挑战赛中以较大优势赢得了这些类别的冠军。此外,该网络速度很快。在最新的 GPU 上,对一张 512x512 图像进行分割所需时间不到一秒。完整的实现(基于 Caffe)和训练好的网络可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net 获取。
UNetMotionModel
class diffusers.UNetMotionModel
< source >( sample_size: int | None = Nonein_channels: int = 4out_channels: int = 4down_block_types: tuple = ('CrossAttnDownBlockMotion', 'CrossAttnDownBlockMotion', 'CrossAttnDownBlockMotion', 'DownBlockMotion')up_block_types: tuple = ('UpBlockMotion', 'CrossAttnUpBlockMotion', 'CrossAttnUpBlockMotion', 'CrossAttnUpBlockMotion')block_out_channels: tuple = (320, 640, 1280, 1280)layers_per_block: int | tuple[int] = 2downsample_padding: int = 1mid_block_scale_factor: float = 1act_fn: str = 'silu'norm_num_groups: int = 32norm_eps: float = 1e-05cross_attention_dim: int = 1280transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1reverse_transformer_layers_per_block: int | tuple[int] | tuple[tuple] | None = Nonetemporal_transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1reverse_temporal_transformer_layers_per_block: int | tuple[int] | tuple[tuple] | None = Nonetransformer_layers_per_mid_block: int | tuple[int] | None = Nonetemporal_transformer_layers_per_mid_block: int | tuple[int] | None = 1use_linear_projection: bool = Falsenum_attention_heads: int | tuple[int, ...] = 8motion_max_seq_length: int = 32motion_num_attention_heads: int | tuple[int, ...] = 8reverse_motion_num_attention_heads: int | tuple[int, ...] | tuple[tuple[int, ...], ...] | None = Noneuse_motion_mid_block: bool = Truemid_block_layers: int = 1encoder_hid_dim: int | None = Noneencoder_hid_dim_type: str | None = Noneaddition_embed_type: str | None = Noneaddition_time_embed_dim: int | None = Noneprojection_class_embeddings_input_dim: int | None = Nonetime_cond_proj_dim: int | None = None )
一个经过修改的条件 2D UNet 模型,它接收一个带噪声的样本、条件状态和时间步,并返回一个具有样本形状的输出。
该模型继承自 ModelMixin。请查阅其父类文档,了解为所有模型实现的通用方法(例如下载或保存)。
禁用 FreeU 机制。
enable_forward_chunking
< source >( chunk_size: int | None = Nonedim: int = 0 )
将注意力处理器设置为使用 前馈分块。
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: Tensortimestep_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] = Nonereturn_dict: bool = True ) → UNetMotionOutput or tuple
参数
- sample (
torch.Tensor) — 具有如下形状(batch, num_frames, channel, height, width的含噪输入张量。 - timestep(
torch.Tensor或float或int)— 用于对输入进行去噪的时间步数。 - encoder_hidden_states (
torch.Tensor) — 形状为(batch, sequence_length, feature_dim)的编码器隐藏状态。 - 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, 可选) — 用于条件化模型的额外嵌入(例如文本嵌入和时间嵌入)字典。 - down_block_additional_residuals — (
tupleoftorch.Tensor, optional): 一个张量元组,如果指定,将添加到下采样 unet 块的残差中。 - mid_block_additional_residual — (
torch.Tensor, 可选): 如果指定,该张量将被添加到中间 unet 块的残差中。 - return_dict (
bool,可选,默认为True) — 是否返回UNetMotionOutput而不是普通元组。
返回值
UNetMotionOutput 或 tuple
如果 return_dict 为 True,则返回 UNetMotionOutput,否则返回 tuple,其中第一个元素为样本张量。
UNetMotionModel 的前向方法。
仅冻结 UNet2DConditionModel 的权重,而保持运动模块未冻结,以便进行微调。
启用融合的 QKV 投影。对于自注意力模块,所有投影矩阵(即 query、key、value)均会被融合。对于交叉注意力模块,key 和 value 投影矩阵会被融合。
> 此 API 处于 🧪 实验阶段。
禁用自定义注意力处理器并设置默认的注意力实现。
UNet3DConditionOutput
类 diffusers.models.unets.unet_3d_condition.UNet3DConditionOutput
< source >( sample: Tensor )
UNet3DConditionModel 的输出。