Diffusers

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

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[int, int, int] = (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.Tensor de forma (batch_size, image_seq_len, in_channels)) — Latentes de imagem ruidosos empacotados (patchificados).
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)) — (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.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.

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

Update on GitHub

←KandinskyS2VTransformer3DModel LatteTransformer3DModel→