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

El transformer de flow-matching MMDiT de un solo flujo utilizado por Krea 2.

Krea2Transformer2DModel

clase 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 )

Parámetros

  • in_channels (int, por defecto 64) — Número de canales latentes después de la patchification (vae_channels * patch_size ** 2).
  • num_layers (int, por defecto 28) — Número de bloques de transformer.
  • attention_head_dim (int, por defecto 128) — Dimensión de cada cabeza de atención; el tamaño oculto total es attention_head_dim * num_attention_heads.
  • num_attention_heads (int, por defecto 48) — Número de cabezas de consulta.
  • num_key_value_heads (int, por defecto 12) — Número de cabezas de clave/valor para la atención de consultas agrupadas.
  • intermediate_size (int, por defecto 16384) — Tamaño oculto de la red feed-forward del MLP SwiGLU dentro de cada bloque.
  • timestep_embed_dim (int, por defecto 256) — Ancho de la incrustación sinusoidal del paso de tiempo antes de su MLP.
  • text_hidden_dim (int, por defecto 2560) — Tamaño oculto del codificador de texto cuyos estados ocultos se consumen.
  • num_text_layers (int, por defecto 12) — Número de estados ocultos del codificador de texto extraídos y apilados por token.
  • text_num_attention_heads (int, por defecto 20) — Número de cabezas de consulta en los bloques de fusión de texto.
  • text_num_key_value_heads (int, valor predeterminado 20) — Número de cabezas de clave/valor en los bloques de fusión de texto.
  • text_intermediate_size (int, por defecto 6912) — Tamaño oculto de la red feed-forward del MLP SwiGLU dentro de los bloques de fusión de texto.
  • num_layerwise_text_blocks (int, por defecto 2) — Número de bloques de fusión de texto aplicados a lo largo del eje de capas muestreadas (por token).
  • num_refiner_text_blocks (int, valor predeterminado 2) — Número de bloques de fusión de texto aplicados en la secuencia de tokens.
  • axes_dims_rope (tuple[int, int, int], por defecto (32, 48, 48)) — División de la dimensión de la cabeza entre los ejes de posición rotacional (t, h, w).
  • rope_theta (float, por defecto 1000.0) — Base utilizada por el embebido de posición rotativa.
  • norm_eps (float, por defecto 1e-5) — Epsilon utilizado por todos los módulos RMSNorm.

El backbone de flow-matching MMDiT de un solo flujo utilizado por el pipeline Krea 2.

La condición de texto se introduce como una pila de estados ocultos extraídos de varias capas de un codificador de texto multimodal. Un pequeño transformer de fusión de texto colapsa el eje de capas y refina la secuencia de tokens; el resultado se concatena con los latentes de imagen parcheados en una única secuencia [text, image] que procesan los bloques transformer. El paso de tiempo condicional afecta a cada bloque a través de un vector de modulación compartido más tablas aprendidas por bloque.

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 )

Parámetros

  • hidden_states (torch.Tensor de forma (batch_size, image_seq_len, in_channels)) — Latentes de imagen ruidosos empaquetados (patchificados).
  • encoder_hidden_states (torch.Tensor de forma (batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — Pila de estados ocultos del codificador de texto por token.
  • timestep (torch.Tensor de forma (batch_size,)) — Tiempo de flow-matching en [0, 1] (1 es ruido puro, 0 es datos limpios).
  • position_ids (torch.Tensor de forma (text_seq_len + image_seq_len, 3)) — coordenadas rotatorias (t, h, w) para la secuencia combinada. Las filas de texto son todas cero; las filas de imagen contienen las coordenadas de la cuadrícula latente.
  • encoder_attention_mask (torch.Tensor de forma (batch_size, text_seq_len), opcional) — Máscara booleana que marca los tokens de texto válidos. Pasa None cuando todos los tokens de texto son válidos.
  • attention_kwargs (dict, opcional) — Un diccionario de kwargs que, si contiene una entrada scale, establece la escala LoRA aplicada a los adaptadores de este transformer durante la duración de la pasada hacia adelante.
  • return_dict (bool, opcional, valor predeterminado True) — Si se debe devolver un Transformer2DModelOutput en lugar de una tupla simple.

Predice la velocidad de flow-matching para los tokens de la imagen.

Update on GitHub