Hugging Face's logo

Diffusers documentation

Krea2Transformer2DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Krea2Transformer2DModel

Transformer flow-matching MMDiT single-stream yang digunakan oleh Krea 2.

Krea2Transformer2DModel

kelas diffusers.Krea2Transformer2DModel

< >

( 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 )

Parameter

  • in_channels (int, defaultnya 64) — Jumlah channel laten setelah patchification (vae_channels * patch_size ** 2).
  • num_layers (int, defaultnya 28) — Jumlah blok transformer.
  • attention_head_dim (int, default 128) — Dimensi setiap kepala perhatian; ukuran tersembunyi total adalah attention_head_dim * num_attention_heads.
  • num_attention_heads (int, defaultnya 48) — Jumlah kepala query.
  • num_key_value_heads (int, defaultnya 12) — Jumlah kepala key/value untuk grouped-query attention.
  • intermediate_size (int, defaultnya 16384) — Ukuran tersembunyi feed-forward dari MLP SwiGLU di dalam setiap blok.
  • timestep_embed_dim (int, defaultnya 256) — Lebar dari embedding timestep sinusoidal sebelum MLP-nya.
  • text_hidden_dim (int, defaultnya 2560) — Ukuran tersembunyi dari encoder teks yang state tersembunyinya dikonsumsi.
  • num_text_layers (int, defaultnya 12) — Jumlah hidden states text-encoder yang ditap dan ditumpuk per token.
  • text_num_attention_heads (int, defaultnya 20) — Jumlah kepala query dalam blok fusi teks.
  • text_num_key_value_heads (int, defaultnya 20) — Jumlah kepala kunci/nilai dalam blok fusi teks.
  • text_intermediate_size (int, defaultnya 6912) — Ukuran tersembunyi feed-forward dari MLP SwiGLU di dalam blok fusi teks.
  • num_layerwise_text_blocks (int, defaultnya 2) — Jumlah blok fusi teks yang diterapkan di sepanjang sumbu lapisan yang ditap (per token).
  • num_refiner_text_blocks (int, defaultnya 2) — Jumlah blok fusi teks yang diterapkan di sepanjang urutan token.
  • axes_dims_rope (tuple[int, int, int], defaultnya (32, 48, 48)) — Pembagian head-dim di sepanjang sumbu posisi rotasi (t, h, w).
  • rope_theta (float, defaultnya 1000.0) — Basis yang digunakan oleh rotary position embedding.
  • norm_eps (float, defaultnya 1e-5) — Epsilon yang digunakan oleh semua modul RMSNorm.

Backbone flow-matching MMDiT single-stream yang digunakan oleh pipeline Krea 2.

Kondisi teks masuk sebagai tumpukan hidden state yang diambil dari beberapa lapisan encoder teks multimodal. Transformer fusi teks yang kecil meruntuhkan sumbu lapisan dan menyempurnakan urutan token; hasilnya digabungkan dengan latens gambar yang telah dipatch menjadi satu urutan [text, image] tunggal yang diproses oleh blok-blok transformer. Timestep mengondisikan setiap blok melalui satu vektor modulasi bersama ditambah tabel yang dipelajari per blok.

forward

< >

( 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 )

Parameter

  • hidden_states (torch.Tensor dengan bentuk (batch_size, image_seq_len, in_channels)) — Latens gambar berisik yang sudah dikemas (patchified).
  • encoder_hidden_states (torch.Tensor dengan bentuk (batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — Tumpukan hidden states text-encoder yang diambil per token.
  • timestep (torch.Tensor dengan bentuk (batch_size,)) — Waktu flow-matching dalam [0, 1] (1 adalah murni noise, 0 adalah data bersih).
  • position_ids (torch.Tensor dengan bentuk (text_seq_len + image_seq_len, 3)) — koordinat rotasi (t, h, w) untuk urutan gabungan. Baris teks semuanya nol; baris gambar menyimpan koordinat kisi laten.
  • encoder_attention_mask (torch.Tensor dengan bentuk (batch_size, text_seq_len), opsional) — Mask boolean yang menandai token teks yang valid. Berikan None ketika semua token teks valid.
  • attention_kwargs (dict, opsional) — Kamus kwargs yang, jika berisi entri scale, mengatur skala LoRA yang diterapkan pada adapter transformer ini selama forward pass.
  • return_dict (bool, opsional, default True) — Apakah akan mengembalikan Transformer2DModelOutput alih-alih tuple biasa.

Prediksi kecepatan flow-matching untuk token gambar.

Update on GitHub

←KandinskyS2VTransformer3DModel LatteTransformer3DModel→