Diffusers documentation
Krea2Transformer2DModel
Krea2Transformer2DModel
O transformer de flow-matching MMDiT de fluxo único usado pelo Krea 2.
Krea2Transformer2DModel
classe 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, padrão 64) — Número de canais latentes após a patchification (vae_channels * patch_size ** 2). - num_layers (
int, padrão 28) — Número de blocos de transformer. - attention_head_dim (
int, padrão 128) — Dimensão de cada cabeça de atenção; o tamanho oculto total éattention_head_dim * num_attention_heads. - num_attention_heads (
int, padrão 48) — Número de cabeças de consulta. - num_key_value_heads (
int, padrão 12) — Número de cabeças de chave/valor para atenção de consulta agrupada. - intermediate_size (
int, padrão 16384) — Tamanho oculto do feed-forward do MLP SwiGLU dentro de cada bloco. - timestep_embed_dim (
int, padrão 256) — Largura do embedding sinusoidal de timestep antes do seu MLP. - text_hidden_dim (
int, padrão 2560) — Tamanho oculto do codificador de texto cujos estados ocultos são consumidos. - num_text_layers (
int, padrão 12) — Número de estados ocultos do codificador de texto capturados e empilhados por token. - text_num_attention_heads (
int, padrão 20) — Número de cabeças de consulta nos blocos de fusão de texto. - text_num_key_value_heads (
int, padrão 20) — Número de cabeças de chave/valor nos blocos de fusão de texto. - text_intermediate_size (
int, padrão 6912) — Tamanho oculto do feed-forward do MLP SwiGLU dentro dos blocos de fusão de texto. - num_layerwise_text_blocks (
int, padrão 2) — Número de blocos de fusão de texto aplicados ao longo do eixo da camada tapada (por token). - num_refiner_text_blocks (
int, padrão 2) — Número de blocos de fusão de texto aplicados em toda a sequência de tokens. - axes_dims_rope (
tuple[int, int, int], padrão(32, 48, 48)) — Divisão da dimensão da cabeça entre os eixos de posição rotacional (t, h, w). - rope_theta (
float, padrão 1000.0) — Base usada pelo incorporamento de posição rotativa. - norm_eps (
float, padrão 1e-5) — Epsilon usado por todos os módulos RMSNorm.
O backbone de flow-matching MMDiT de fluxo único usado pelo pipeline Krea 2.
A condicionamento por texto entra como uma pilha de estados ocultos extraídos de várias camadas de um codificador de texto multimodal. Um pequeno transformer de fusão de texto reduz o eixo das camadas e refina a sequência de tokens; o resultado é concatenado com os latentes de imagem em patches, formando uma única sequência [text, image] processada pelos blocos do transformer. O timestep condiciona cada bloco por meio de um vetor de modulação compartilhado, além de tabelas aprendidas por bloco.
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 imagem ruidosos empacotados (patchificados). - encoder_hidden_states (
torch.Tensorde forma(batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — Pilha de estados ocultos do codificador de texto capturados por token. - timestep (
torch.Tensorde forma(batch_size,)) — Tempo de flow-matching em[0, 1](1 é ruído puro, 0 é dados limpos). - position_ids (
torch.Tensorde forma(text_seq_len + image_seq_len, 3)) —(t, h, w)coordenadas rotativas para a sequência combinada. As linhas de texto são todas zero; as linhas de imagem contêm as coordenadas da grade latente. - encoder_attention_mask (
torch.Tensorde forma(batch_size, text_seq_len), opcional) — Máscara booleana que marca os tokens de texto válidos. PasseNonequando todos os tokens de texto forem válidos. - attention_kwargs (
dict, opcional) — Um dicionário de kwargs que, quando contém uma entradascale, define a escala LoRA aplicada aos adaptadores deste transformer durante a passagem de frente. - return_dict (
bool, opcional, padrãoTrue) — Se deve retornar um Transformer2DModelOutput em vez de uma tupla simples.
Preveja a velocidade de flow-matching para os tokens da imagem.