# Krea2Transformer2DModel

[Krea 2](https://github.com/krea-ai/krea-2) 所使用的单流 MMDiT 流匹配 Transformer。

## Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

#### diffusers.Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

```python
diffusers.Krea2Transformer2DModel(in_channels: int = 64, num_layers: int = 28, attention_head_dim: int = 128, num_attention_heads: int = 48, num_key_value_heads: int = 12, intermediate_size: int = 16384, timestep_embed_dim: int = 256, text_hidden_dim: int = 2560, num_text_layers: int = 12, text_num_attention_heads: int = 20, text_num_key_value_heads: int = 20, text_intermediate_size: int = 6912, num_layerwise_text_blocks: int = 2, num_refiner_text_blocks: int = 2, axes_dims_rope: tuple[int, int, int] = (32, 48, 48), rope_theta: float = 1000.0, norm_eps: float = 1e-05)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_krea2.py#L339)

**参数：**

in_channels (`int`，默认为 64)：分块化后的潜在通道数（`vae_channels * patch_size ** 2`）。

num_layers (`int`，默认为 28)：Transformer 块的数量。

attention_head_dim (`int`，默认为 128)：每个注意力头的维度；总隐藏大小为 `attention_head_dim * num_attention_heads`。

num_attention_heads (`int`，默认为 48)：查询头的数量。

num_key_value_heads (`int`，默认为 12)：分组查询注意力机制中键/值头的数量。

intermediate_size (`int`，默认值为 16384)：每个块内部 SwiGLU MLP 的前馈隐藏层大小。

timestep_embed_dim (`int`，默认为 256)：正弦时间步嵌入在经过其 MLP 之前的宽度。

text_hidden_dim (`int`，默认为 2560)：被消费的隐藏状态所属文本编码器的隐藏层大小。

num_text_layers (`int`，默认为 12)：每个 token 堆叠的 tapped 文本编码器隐藏状态数量。

text_num_attention_heads (`int`，默认为 20)：文本融合块中的查询头数量。

text_num_key_value_heads (`int`，默认为 20)：文本融合块中键/值头的数量。

text_intermediate_size (`int`，默认值为 6912)：文本融合块中 SwiGLU MLP 的前馈隐藏层大小。

num_layerwise_text_blocks (`int`，默认为 2)：在 tapped-layer 轴上（每个 token）应用的文本融合块数量。

num_refiner_text_blocks (`int`，默认为 2)：应用于 token 序列的文本融合块数量。

axes_dims_rope (`tuple[int, int, int]`，默认为 `(32, 48, 48)`)：在 (t, h, w) 旋转位置轴上对头维度进行划分。

rope_theta (`float`，默认值为 1000.0)：旋转位置嵌入所使用的基数。

norm_eps (`float`，默认值为 1e-5)：所有 RMSNorm 模块使用的 Epsilon。

Krea 2 流水线所使用的单流 MMDiT 流匹配骨干网络。

文本条件以从多模态文本编码器多个层中抽取的隐藏状态堆栈形式输入。一个小型的文本融合 Transformer 会压缩层维度并优化 token 序列；其结果与分块后的图像潜变量拼接，形成单一的 `[text, image]` 序列，由 Transformer 块进行处理。时间步通过一个共享的调制向量以及每个块各自的学习表，对每个块进行条件调制。

#### forward[[diffusers.Krea2Transformer2DModel.forward]]

```python
forward(hidden_states: Tensor, encoder_hidden_states: Tensor, timestep: Tensor, position_ids: Tensor, encoder_attention_mask: typing.Optional[torch.Tensor] = None, attention_kwargs: dict[str, typing.Any] | None = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_krea2.py#L456)

**参数：**

hidden_states (`torch.Tensor`，形状为 `(batch_size, image_seq_len, in_channels)`)：打包（分块化）的含噪图像潜在表示。

encoder_hidden_states (`torch.Tensor`，形状为 `(batch_size, text_seq_len, num_text_layers, text_hidden_dim)`)：每个 token 对应的文本编码器各层隐藏状态的堆叠。

timestep (`torch.Tensor`，形状为 `(batch_size,)`) : `[0, 1]` 范围内的流匹配时间（1 表示纯噪声，0 表示干净数据）。

position_ids (`torch.Tensor`，形状为 `(text_seq_len + image_seq_len, 3)`) : 组合序列的 `(t, h, w)` 旋转坐标。文本行全为零；图像行包含潜在网格坐标。

encoder_attention_mask (`torch.Tensor`，形状为 `(batch_size, text_seq_len)`，*可选*)：布尔掩码，用于标记有效的文本 token。当所有文本 token 均有效时，请传入 `None`。

attention_kwargs (`dict`, *可选*) : 一个 kwargs 字典，当其中包含 `scale` 条目时，会在前向传播期间设置应用于此 transformer 适配器的 LoRA 缩放。

return_dict (`bool`，*可选*，默认为 `True`)：是否返回 [Transformer2DModelOutput](/docs/diffusers/v0.43.0/zh/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) 而非普通元组。

**返回值：**

[Transformer2DModelOutput](/docs/diffusers/v0.43.0/zh/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) 或一个 `tuple`，其第一个元素是形状为 `(batch_size, image_seq_len, in_channels)` 的速度张量。

预测图像 token 的 flow-matching 速度。