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

# Krea2Transformer2DModel

The single-stream MMDiT flow-matching transformer used by [Krea 2](https://github.com/krea-ai/krea-2).

## Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

- **in_channels** (`int`, defaults to 64) --
  Latent channel count after patchification (`vae_channels * patch_size ** 2`).
- **num_layers** (`int`, defaults to 28) --
  Number of transformer blocks.
- **attention_head_dim** (`int`, defaults to 128) --
  Dimension of each attention head; the total hidden size is `attention_head_dim * num_attention_heads`.
- **num_attention_heads** (`int`, defaults to 48) --
  Number of query heads.
- **num_key_value_heads** (`int`, defaults to 12) --
  Number of key/value heads for grouped-query attention.
- **intermediate_size** (`int`, defaults to 16384) --
  Feed-forward hidden size of the SwiGLU MLP inside each block.
- **timestep_embed_dim** (`int`, defaults to 256) --
  Width of the sinusoidal timestep embedding before its MLP.
- **text_hidden_dim** (`int`, defaults to 2560) --
  Hidden size of the text encoder whose hidden states are consumed.
- **num_text_layers** (`int`, defaults to 12) --
  Number of tapped text-encoder hidden states stacked per token.
- **text_num_attention_heads** (`int`, defaults to 20) --
  Number of query heads in the text fusion blocks.
- **text_num_key_value_heads** (`int`, defaults to 20) --
  Number of key/value heads in the text fusion blocks.
- **text_intermediate_size** (`int`, defaults to 6912) --
  Feed-forward hidden size of the SwiGLU MLP inside the text fusion blocks.
- **num_layerwise_text_blocks** (`int`, defaults to 2) --
  Number of text fusion blocks applied across the tapped-layer axis (per token).
- **num_refiner_text_blocks** (`int`, defaults to 2) --
  Number of text fusion blocks applied across the token sequence.
- **axes_dims_rope** (`tuple[int, int, int]`, defaults to `(32, 48, 48)`) --
  Head-dim split across the (t, h, w) rotary position axes.
- **rope_theta** (`float`, defaults to 1000.0) --
  Base used by the rotary position embedding.
- **norm_eps** (`float`, defaults to 1e-5) --
  Epsilon used by all RMSNorm modules.

The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline.

Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A
small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated
with the patchified image latents into a single `[text, image]` sequence processed by the transformer blocks. The
timestep conditions every block through one shared modulation vector plus per-block learned tables.

- **hidden_states** (`torch.Tensor` of shape `(batch_size, image_seq_len, in_channels)`) --
  Packed (patchified) noisy image latents.
- **encoder_hidden_states** (`torch.Tensor` of shape `(batch_size, text_seq_len, num_text_layers, text_hidden_dim)`) --
  Stack of tapped text-encoder hidden states per token.
- **timestep** (`torch.Tensor` of shape `(batch_size,)`) --
  Flow-matching time in `[0, 1]` (1 is pure noise, 0 is clean data).
- **position_ids** (`torch.Tensor` of shape `(text_seq_len + image_seq_len, 3)`) --
  `(t, h, w)` rotary coordinates for the combined sequence. Text rows are all-zero; image rows hold the
  latent-grid coordinates.
- **encoder_attention_mask** (`torch.Tensor` of shape `(batch_size, text_seq_len)`, *optional*) --
  Boolean mask marking valid text tokens. Pass `None` when every text token is valid.
- **attention_kwargs** (`dict`, *optional*) --
  A kwargs dictionary that, when it contains a `scale` entry, sets the LoRA scale applied to this
  transformer's adapters for the duration of the forward pass.
- **return_dict** (`bool`, *optional*, defaults to `True`) --
  Whether to return a [Transformer2DModelOutput](/docs/diffusers/main/en/api/models/sana_video_transformer3d#diffusers.models.modeling_outputs.Transformer2DModelOutput) instead of a plain tuple.[Transformer2DModelOutput](/docs/diffusers/main/en/api/models/sana_video_transformer3d#diffusers.models.modeling_outputs.Transformer2DModelOutput) or a `tuple` whose first element is the velocity
tensor of shape `(batch_size, image_seq_len, in_channels)`.

Predict the flow-matching velocity for the image tokens.