# Krea2Transformer2DModel

Однопоточный трансформер flow-matching MMDiT, используемый в [Krea 2](https://github.com/krea-ai/krea-2).

## Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

#### diffusers.Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

```python
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 = (32, 48, 48), rope_theta: float = 1000.0, norm_eps: float = 1e-05)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_krea2.py#L339)

**Параметры:**

in_channels (`int`, по умолчанию 64) : Количество латентных каналов после патчификации (`vae_channels * patch_size ** 2`).

num_layers (`int`, по умолчанию 28) : Количество блоков трансформера.

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) : Размер скрытого слоя feed-forward в 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) : Размер скрытого слоя feed-forward в SwiGLU MLP внутри блоков текстового слияния.

num_layerwise_text_blocks (`int`, по умолчанию 2) : Количество блоков текстовой фузии, применяемых вдоль оси тэпнутых слоёв (на токен).

num_refiner_text_blocks (`int`, по умолчанию 2) : Количество блоков текстовой фузии, применяемых к последовательности токенов.

axes_dims_rope (`tuple[int, int, int]`, по умолчанию `(32, 48, 48)`) : Разделение head-dim по осям ротационных позиций (t, h, w).

rope_theta (`float`, по умолчанию 1000.0) : Основание, используемое в вращательном позиционном встраивании.

norm_eps (`float`, по умолчанию 1e-5) : Эпсилон, используемый всеми модулями RMSNorm.

Однопоточный backbone flow-matching MMDiT, используемый в пайплайне Krea 2.

Условие по тексту подаётся в виде стека скрытых состояний, извлечённых из нескольких слоёв мультимодального текстового энкодера. Небольшой трансформер для слияния текста сворачивает ось слоёв и уточняет последовательность токенов; результат конкатенируется с патчифицированными латентами изображения в единую последовательность `[text, image]`, которая обрабатывается блоками трансформера. Шаг времени (timestep) влияет на каждый блок через один общий вектор модуляции и обученные таблицы для каждого блока.

#### forward[[diffusers.Krea2Transformer2DModel.forward]]

```python
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)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_krea2.py#L456)

**Параметры:**

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` формы `(batch_size,)`) : Время flow-matching в диапазоне `[0, 1]` (1 — чистый шум, 0 — чистые данные).

position_ids (`torch.Tensor` формы `(text_seq_len + image_seq_len, 3)`) : `(t, h, w)` координаты вращения для объединённой последовательности. Строки текста заполнены нулями; строки изображений содержат координаты латентной сетки.

encoder_attention_mask (`torch.Tensor` формы `(batch_size, text_seq_len)`, *необязательный*) : Булева маска, отмечающая валидные токены текста. Передайте `None`, если все токены текста валидны.

attention_kwargs (`dict`, *необязательный*) : Словарь kwargs, который при наличии в нём записи `scale` задаёт масштаб LoRA, применяемый к адаптерам данного трансформера на протяжении прямого прохода.

return_dict (`bool`, *необязательный*, по умолчанию `True`) : следует ли возвращать [Transformer2DModelOutput](/docs/diffusers/main/ru/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) вместо обычного кортежа.

**Возвращает:**

[Transformer2DModelOutput](/docs/diffusers/main/ru/api/models/cogview3plus_transformer2d#diffusers.models.modeling_outputs.Transformer2DModelOutput) или `tuple`, первым элементом которого является тензор скорости формы `(batch_size, image_seq_len, in_channels)`.

Предскажите скорость flow-matching для токенов изображения.