Hugging Face's logo

Diffusers documentation

Krea2Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Krea2Transformer2DModel

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

Krea2Transformer2DModel

класс 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[int, int, int] = (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) — Количество блоков трансформера.
  • 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.

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

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

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 )

Параметры

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

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

Update on GitHub

←KandinskyS2VTransformer3DModel LatteTransformer3DModel→