Diffusers documentation
Krea2Transformer2DModel
Krea2Transformer2DModel
Krea 2で使用される、単一ストリームのMMDiTフローマッチングトランスフォーマー。
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 = (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) — トークンごとにスタックする、タップされたテキストエンコーダーの隠れ状態の層数。 - 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) の回転位置軸にわたるヘッド次元の分割です。 - rope_theta (
float、デフォルト値は 1000.0) — 回転位置埋め込みで使用する基数です。 - norm_eps (
float、デフォルト値は 1e-5) — すべての RMSNorm モジュールで使用する Epsilon。
Krea 2 パイプラインで使用される、単一ストリームの MMDiT フローマッチング基盤です。
テキスト条件は、マルチモーダルなテキストエンコーダーの複数の層から取得された隠れ状態のスタックとして入力されます。小型のテキスト融合トランスフォーマーが層軸を圧縮し、トークンシーケンスを洗練します。その結果は、パッチ化された画像潜在変数と連結され、[text, image] という単一のシーケンスとしてトランスフォーマーブロックで処理されます。タイムステップは、共有される1つのモジュレーションベクトルと、ブロックごとに学習されたテーブルを介して、すべてのブロックに条件付けを行います。
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 (
torch.Tensor、shape(batch_size, image_seq_len, in_channels)) — パッチ化されたノイズ付き画像の潜在変数（packed）。 - encoder_hidden_states (
torch.Tensor、shape(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)回転座標です。テキストの行はすべてゼロであり、画像の行には潜在グリッドの座標が格納されます。 - 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) —Trueの場合、通常のタプルではなく Transformer2DModelOutput を返します。
画像トークンに対するフローマッチングの速度を予測します。