Diffusers documentation
AuraFlowTransformer2DModel
AuraFlowTransformer2DModel
用于类图像数据的 Transformer 模型，来自 AuraFlow。
AuraFlowTransformer2DModel
class diffusers.AuraFlowTransformer2DModel
< source >( sample_size: int = 64patch_size: int = 2in_channels: int = 4num_mmdit_layers: int = 4num_single_dit_layers: int = 32attention_head_dim: int = 256num_attention_heads: int = 12joint_attention_dim: int = 2048caption_projection_dim: int = 3072out_channels: int = 4pos_embed_max_size: int = 1024 )
参数
- sample_size (
int) — 潜在图像的宽度。由于该参数用于学习位置嵌入的数量，因此在训练过程中是固定的。 - patch_size (
int) — 用于将输入数据划分为小块的块大小。 - in_channels (
int，可选，默认为 4) — 输入中的通道数。 - num_mmdit_layers (
int，可选，默认为 4) — 要使用的 MMDiT Transformer 块层数。 - num_single_dit_layers (
int, 可选，默认为 32) — 要使用的 Transformer 块层数。这些块使用拼接后的图像和文本表示。 - attention_head_dim (
int，可选，默认值为 256) — 每个注意力头中的通道数。 - num_attention_heads (
int，可选，默认为 12) — 多头注意力机制中使用的注意力头数量。 - joint_attention_dim (
int，可选) — 要使用的encoder_hidden_states维度数量。 - caption_projection_dim (
int) — 用于投影encoder_hidden_states的维度数。 - out_channels (
int，默认为 4) — 输出通道数。 - pos_embed_max_size (
int，默认为 1024) — 从图像潜在变量中嵌入的最大位置数。
AuraFlow（https://blog.fal.ai/auraflow/）中介绍的 2D Transformer 模型。
forward
< source >( hidden_states: FloatTensorencoder_hidden_states: FloatTensor = Nonetimestep: LongTensor = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )
参数
- hidden_states（形状为
(batch size, channel, height, width)的torch.FloatTensor）— 输入的hidden_states。 - encoder_hidden_states (
torch.FloatTensor，形状为(batch size, sequence_len, embed_dims)) — 要使用的条件嵌入（由输入条件（例如提示词）计算得出的嵌入）。 - timestep (
torch.LongTensor) — 用于指示去噪步骤。 - attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定，将传递给 diffusers.models.attention_processor 中定义的AttentionProcessor，该处理器在self.processor下定义。 - return_dict (
bool，可选，默认为True) — 是否返回~models.transformer_2d.Transformer2DModelOutput而不是普通的元组。
AuraFlowTransformer2DModel 的前向方法。
启用融合的 QKV 投影。对于自注意力模块，所有投影矩阵（即 query、key、value）均会被融合。对于交叉注意力模块，key 和 value 投影矩阵会被融合。
> 此 API 处于 🧪 实验阶段。