diffusers/main/ja/api/models/krea2_transformer2d.md

Krea2Transformer2DModel

Krea 2で使用される、単一ストリームのMMDiTフローマッチングトランスフォーマー。

Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

diffusers.Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

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)

ソース

パラメータ:

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) : トークンごとにスタックされるタップされたテキストエンコーダーの隠れ状態の数。

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) : トークンごとに、タップされた層の軸全体に適用されるテキスト融合ブロックの数。

num_refiner_text_blocks (int、デフォルトは 2) : トークンシーケンス全体に適用されるテキスト融合ブロックの数。

axes_dims_rope (tuple[int, int, int]、デフォルトは (32, 48, 48)) : (t, h, w) の回転位置軸にわたる head-dim の分割。

rope_theta (float、デフォルト値は 1000.0) : rotary position embedding で使用される基数です。

norm_eps (float、デフォルト値は 1e-5) : すべての RMSNorm モジュールで使用する Epsilon。

Krea 2 パイプラインで使用される、単一ストリームの MMDiT フローマッチング基盤です。

テキスト条件は、マルチモーダルなテキストエンコーダーの複数の層から取得された隠れ状態のスタックとして入力されます。小型のテキスト融合トランスフォーマーが層軸を圧縮し、トークンシーケンスを洗練します。その結果は、パッチ化された画像潜在変数と連結され、[text, image] という単一のシーケンスとしてトランスフォーマーブロックで処理されます。タイムステップは、共有される1つのモジュレーションベクトルと、ブロックごとに学習されたテーブルを介して、すべてのブロックに条件付けられます。

forward[[diffusers.Krea2Transformer2DModel.forward]]

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)

ソース

パラメータ:

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)) : トークンごとのタップされたテキストエンコーダーの隠れ状態のスタック。

timestep (torch.Tensor、shape (batch_size,)) : [0, 1] における Flow-matching の時間（1 は純粋なノイズ、0 はクリーンなデータ）。

position_ids (torch.Tensor、shape (text_seq_len + image_seq_len, 3)) : 結合シーケンスの (t, h, w) rotary 座標。テキスト行はすべてゼロであり、画像行には潜在グリッドの座標が格納されています。

encoder_attention_mask (torch.Tensor、shape (batch_size, text_seq_len)、任意) : 有効なテキストトークンを示すブールマスク。すべてのテキストトークンが有効な場合は None を渡します。

attention_kwargs (dict, 任意) : scale エントリを含む場合、この transformer のアダプターに適用される LoRA スケールを、forward pass の間だけ設定するための kwargs 辞書です。

return_dict (bool、任意、デフォルトはTrue) : プレーンなタプルではなく、Transformer2DModelOutputを返すかどうか。

戻り値:

Transformer2DModelOutput または、最初の要素が (batch_size, image_seq_len, in_channels) の形状を持つ速度テンソルである tuple。

画像トークンに対するフローマッチングの速度を予測します。