参数

• 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)) — 每个块的输出通道的元组。
• 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): BasicTransformerBlock 类型的 Transformer 块的数量，在 U-Net 的上采样块中。 仅当 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。

参数

• sample (torch.Tensor of shape (batch_size, num_channels, height, width)) — The hidden states output conditioned on encoder_hidden_states input. Output of last layer of model.

参数

• sample_size (int, optional) — The size of the input sample.
• in_channels (int, optional, defaults to 4) — The number of channels in the input sample.
• 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) — 是否在时间嵌入中将 sin 翻转为 cos。
• freq_shift (int, 可选, 默认为 0) — 应用于时间嵌入的频率偏移。
• use_memory_efficient_attention (bool, 可选, 默认为 False) — 启用内存高效注意力，如此处所述。
• split_head_dim (bool, 可选, 默认为 False) — 是否将 head 维度拆分为自注意力计算的新轴。在大多数情况下，启用此标志应加速 Stable Diffusion 2.x 和 Stable Diffusion XL 的计算。

参数

• sample (jnp.ndarray，形状为 (batch_size, num_channels, height, width)) — 以 encoder_hidden_states 输入为条件的隐藏状态输出。模型的最后一层的输出。