Diffusers documentation
Krea2Transformer2DModel
Krea2Transformer2DModel
Krea 2 所使用的单流 MMDiT 流匹配 Transformer。
Krea2Transformer2DModel
class diffusers.Krea2Transformer2DModel
< source >( in_channels: int = 64num_layers: int = 28attention_head_dim: int = 128num_attention_heads: int = 48num_key_value_heads: int = 12intermediate_size: int = 16384timestep_embed_dim: int = 256text_hidden_dim: int = 2560num_text_layers: int = 12text_num_attention_heads: int = 20text_num_key_value_heads: int = 20text_intermediate_size: int = 6912num_layerwise_text_blocks: int = 2num_refiner_text_blocks: int = 2axes_dims_rope: tuple[int, int, int] = (32, 48, 48)rope_theta: float = 1000.0norm_eps: float = 1e-05 )
参数
- 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 堆叠的文本编码器隐藏状态层数。 - 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) 旋转位置轴上对 head-dim 进行划分。 - rope_theta (
float，默认为 1000.0) — 旋转位置嵌入所使用的基数。 - norm_eps (
float，默认为 1e-5) — 所有 RMSNorm 模块使用的 epsilon。
Krea 2 流水线所使用的单流 MMDiT 流匹配骨干网络。
文本条件以从多模态文本编码器多个层中抽取的隐藏状态堆栈形式输入。一个小型的文本融合 Transformer 会压缩层维度并精炼 token 序列；其结果与分块后的图像潜变量拼接，形成单一的 [text, image] 序列，由 Transformer 块进行处理。时间步通过一个共享的调制向量以及每个块各自的学习表，对每个块进行条件调制。
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorposition_ids: Tensorencoder_attention_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )
参数
- hidden_states（形状为
(batch_size, image_seq_len, in_channels)的torch.Tensor）— 打包（分块化）的含噪图像潜在表示。 - encoder_hidden_states (
torch.Tensor，形状为(batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — 每个 token 对应的文本编码器各层隐藏状态的堆叠。 - timestep（形状为
(batch_size,)的torch.Tensor）—[0, 1]区间内的 Flow-matching 时间（1 表示纯噪声，0 表示干净数据）。 - position_ids (
torch.Tensor，形状为(text_seq_len + image_seq_len, 3)) —(t, h, w)旋转坐标，用于组合序列。文本行的值均为零；图像行包含潜在网格坐标。 - encoder_attention_mask（形状为
(batch_size, text_seq_len)的torch.Tensor，可选） — 用于标记有效文本 token 的布尔掩码。当所有文本 token 均有效时，请传入None。 - attention_kwargs (
dict, 可选) — 一个 kwargs 字典，当其中包含scale条目时，会在前向传播期间设置应用于该 transformer 适配器的 LoRA 缩放系数。 - return_dict (
bool, 可选，默认为True) — 是否返回 Transformer2DModelOutput 而不是普通元组。
预测图像 token 的 flow-matching 速度。