参数

• 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) — 是否在时间嵌入中将正弦翻转为余弦。
• 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) — 是否在基本转换器块中包含自注意力，请参阅 BasicTransformerBlock。
• 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) — 中间模块的缩放因子。
• 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 类型转换块的数目。仅适用于 CrossAttnDownBlock2D、CrossAttnUpBlock2D 和 UNetMidBlock2DCrossAttn 。
• reverse_transformer_layers_per_block — (Tuple[Tuple]，可选，默认为None）：U-Net上采样块中 BasicTransformerBlock 类型的转换块数目。仅在 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_hid_dim_type 将 encoder_hidden_states 和潜在的其他嵌入下投影到维度为 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。

参数

• sample (torch.Tensor形状为(batch_size, num_channels, height, width)) — 在encoder_hidden_states输入下条件化的隐藏状态输出。模型的最后一层输出。

参数

• sample_size (int，可选) — 输入样本的大小。
• in_channels (int, 可选，默认为4) — 输入样本中的通道数。
• out_channels (int, 可选，默认为4) — 输出中的通道数。
• down_block_types (Tuple[str], 可选，默认为 ("FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxCrossAttnDownBlock2D", "FlaxDownBlock2D")) — 要使用的下采样块的元组。
• up_block_types (Tuple[str], 可选，默认为 ("FlaxUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D", "FlaxCrossAttnUpBlock2D")) — 要使用的上采样块的元组。
• mid_block_type (str, 可选，默认为 "UNetMidBlock2DCrossAttn") — UNet中间部分使用的块类型，可以是 UNetMidBlock2DCrossAttn 之一。如果为 None，则跳过中间块层。
• block_out_channels (Tuple[int], 可选，默认为 (320, 640, 1280, 1280)) — 每个块的输出通道的元组。
• layers_per_block (int, 可选，默认为 2) — 每个块中的层数量。
• attention_head_dim (int 或 Tuple[int]，可选，默认值为 8) — 注意力头部的维度。
• num_attention_heads (int 或 Tuple[int]，可选) — 注意力头部的数量。
• cross_attention_dim (int，可选，默认值为 768) — 跨注意力特征的维度。
• dropout (float，可选，默认值为 0) — 下、上和瓶颈块的dropout概率。
• flip_sin_to_cos (bool, 可选，默认为 True) — 是否在时间嵌入中将正弦函数转换为余弦函数。
• freq_shift (int, 可选，默认为 0) — 对时间嵌入应用频率偏移。
• use_memory_efficient_attention (bool, 可选，默认为 False) — 启用内存高效注意力，见此处描述。
• split_head_dim （bool，可选，默认值为 False） — 是否将头维度拆分为一个新的轴进行自我注意力计算。在大多数情况下，启用此标志应会加快稳定扩散 2.x 和稳定扩散 XL 的计算速度。

参数

• sample (jnp.ndarray of shape (batch_size, num_channels, height, width)) — 根据encoder_hidden_states输入条件输出的隐藏状态。模型最后一层的输出。