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

Krea2Transformer2DModel

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

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) : Количество блоков трансформера.

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]]

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 формы (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 вместо обычного кортежа.

Возвращает:

Transformer2DModelOutput или tuple, первым элементом которого является тензор скорости формы (batch_size, image_seq_len, in_channels).

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