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

Однопоточный трансформер 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 = (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→