diffusers/main/pt/api/models/krea2_transformer2d.md

Krea2Transformer2DModel

O transformer de flow-matching MMDiT de fluxo único usado pelo Krea 2.

Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

diffusers.Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

diffusers.Krea2Transformer2DModel(in_channels: int = 64, num_layers: int = 28, attention_head_dim: int = 128, num_attention_heads: int = 48, num_key_value_heads: int = 12, intermediate_size: int = 16384, timestep_embed_dim: int = 256, 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, num_layerwise_text_blocks: int = 2, num_refiner_text_blocks: int = 2, axes_dims_rope: tuple[int, int, int] = (32, 48, 48), rope_theta: float = 1000.0, norm_eps: float = 1e-05)

Fonte

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 amostrados 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 ao longo da sequência de tokens.

axes_dims_rope (tuple[int, int, int], padrão (32, 48, 48)) : Divisão do head-dim entre os eixos de posição rotativa (t, h, w).

rope_theta (float, padrão 1000.0) : Base utilizada 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[[diffusers.Krea2Transformer2DModel.forward]]

forward(hidden_states: Tensor, encoder_hidden_states: Tensor, timestep: Tensor, position_ids: Tensor, encoder_attention_mask: typing.Optional[torch.Tensor] = None, attention_kwargs: dict[str, typing.Any] | None = None, return_dict: bool = True)

Fonte

Parâmetros:

hidden_states (torch.Tensor de forma (batch_size, image_seq_len, in_channels)) : Latentes de imagem ruidosa empacotados (em patches).

encoder_hidden_states (torch.Tensor de 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.Tensor de forma (batch_size,)) : Tempo de flow-matching em [0, 1] (1 é ruído puro, 0 é dados limpos).

position_ids (torch.Tensor de forma (text_seq_len + image_seq_len, 3)) : coordenadas rotativas (t, h, w) 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.Tensor de forma (batch_size, text_seq_len), opcional) : Máscara booleana que marca os tokens de texto válidos. Passe None quando todos os tokens de texto forem válidos.

attention_kwargs (dict, opcional) : Um dicionário de kwargs que, quando contém uma entrada scale, define a escala LoRA aplicada aos adaptadores deste transformer durante a passagem de frente.

return_dict (bool, opcional, padrão True) : Se deve retornar um Transformer2DModelOutput em vez de uma tupla simples.

Retorna:

Transformer2DModelOutput ou uma tuple cujo primeiro elemento é o tensor de velocidade com a forma (batch_size, image_seq_len, in_channels).

Preveja a velocidade de flow-matching para os tokens da imagem.