Diffusers documentation
Krea2Transformer2DModel
Krea2Transformer2DModel
El transformer de flow-matching MMDiT de un solo flujo utilizado por Krea 2.
Krea2Transformer2DModel
clase diffusers.Krea2Transformer2DModel
< source >( 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 esattention_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
< source >( 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.Tensorde forma(batch_size, image_seq_len, in_channels)) — Latentes de imagen ruidosos empaquetados (patchificados). - encoder_hidden_states (
torch.Tensorde forma(batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — Pila de estados ocultos del codificador de texto por token. - timestep (
torch.Tensorde forma(batch_size,)) — Tiempo de flow-matching en[0, 1](1 es ruido puro, 0 es datos limpios). - position_ids (
torch.Tensorde 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.Tensorde forma(batch_size, text_seq_len), opcional) — Máscara booleana que marca los tokens de texto válidos. PasaNonecuando todos los tokens de texto son válidos. - attention_kwargs (
dict, opcional) — Un diccionario de kwargs que, si contiene una entradascale, 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 predeterminadoTrue) — Si se debe devolver un Transformer2DModelOutput en lugar de una tupla simple.
Predice la velocidad de flow-matching para los tokens de la imagen.