Diffusers documentation
Flux2Transformer2DModel
Flux2Transformer2DModel
用于 Flux2 图像类数据的 Transformer 模型。
Flux2Transformer2DModel
类 diffusers.Flux2Transformer2DModel
< source >( patch_size: int = 1in_channels: int = 128out_channels: int | None = Nonenum_layers: int = 8num_single_layers: int = 48attention_head_dim: int = 128num_attention_heads: int = 48joint_attention_dim: int = 15360timestep_guidance_channels: int = 256mlp_ratio: float = 3.0axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)rope_theta: int = 2000eps: float = 1e-06guidance_embeds: bool = True )
参数
- patch_size (
int，默认为1) — 用于将输入数据划分为小块的块大小。 - in_channels (
int，默认为128) — 输入中的通道数。 - out_channels (
int，可选，默认为None) — 输出通道数。如果未指定，则默认为in_channels。 - num_layers (
int，默认为8) — 要使用的双流 DiT 块层数。 - num_single_layers (
int，默认为48) — 要使用的单流 DiT 块层数。 - attention_head_dim (
int，默认为128) — 每个注意力头使用的维度数。 - num_attention_heads (
int，默认为48) — 要使用的注意力头数量。 - joint_attention_dim (
int，默认为15360) — 用于联合注意力的维度数（encoder_hidden_states的嵌入/通道维度）。 - pooled_projection_dim (
int，默认为768) — 用于池化投影的维度数量。 - guidance_embeds (
bool，默认为True) — 是否对模型的 guidance 蒸馏变体使用 guidance 嵌入。 - axes_dims_rope (
tuple[int]，默认为(32, 32, 32, 32)) — 用于旋转位置嵌入的维度。
Flux 2 中引入的 Transformer 模型。
参考：https://blackforestlabs.ai/announcing-black-forest-labs/
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensor = Nonetimestep: LongTensor = Noneimg_ids: Tensor = Nonetxt_ids: Tensor = Noneguidance: Tensor = Nonejoint_attention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = Truekv_cache: Flux2KVCache | None = Nonekv_cache_mode: str | None = Nonenum_ref_tokens: int = 0ref_fixed_timestep: float = 0.0 )
参数
- hidden_states（形状为
(batch_size, image_sequence_length, in_channels)的torch.Tensor）— 输入的hidden_states。 - encoder_hidden_states（形状为
(batch_size, text_sequence_length, joint_attention_dim)的torch.Tensor）— 要使用的条件嵌入（由输入条件（例如提示词）计算得出的嵌入）。 - timestep (
torch.LongTensor) — 用于指示去噪步骤。 - img_ids (
torch.Tensor) — 用于计算旋转位置嵌入的图像位置 ID。 - txt_ids (
torch.Tensor) — 用于计算旋转位置嵌入的文本位置 ID。 - guidance (
torch.Tensor, 可选) — 用于模型引导蒸馏变体的引导尺度嵌入。 - joint_attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定，将作为参数传递给 diffusers.models.attention_processor 中self.processor下定义的AttentionProcessor。 - return_dict (
bool，可选，默认为True) — 是否返回~models.transformer_2d.Transformer2DModelOutput而不是普通的元组。 - kv_cache (
Flux2KVCache, 可选) — 用于参考图像 token 的 KV 缓存。当kv_cache_mode为 “extract” 时，会创建并返回一个新的缓存。当为 “cached” 时，将使用提供的缓存，在注意力计算期间注入参考 K/V。 - kv_cache_mode (
str, 可选) — 取值为 “extract”（使用参考 token 的第一步）或 “cached”（使用缓存的参考 K/V 的后续步骤）。当为None时，执行不带 KV 缓存的标准前向传播。 - num_ref_tokens (
int，默认为0) — 添加到hidden_states前面的参考图像 token 数量（仅在kv_cache_mode="extract"时使用）。 - ref_fixed_timestep (
float，默认为0.0) — 用于参考 token 调制的固定时间步（仅在kv_cache_mode="extract"时使用）。
Flux2Transformer2DModel 的前向方法。
Flux2Transformer2DModelOutput
class diffusers.models.transformers.transformer_flux2.Flux2Transformer2DModelOutput
< source >( sample: torch.Tensorkv_cache: Flux2KVCache | None = None )