Diffusers documentation

Flux3Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Flux3Transformer2DModel

Um modelo Transformer para dados semelhantes a imagens do FLUX.3, para geração de texto para imagem e imagem para imagem. Três modos de token compartilham uma única sequência conjunta: os tokens de imagem com ruído, os tokens de imagem de referência codificados pelo VAE e o contexto de texto. Cada modo primeiro passa pelos seus próprios blocos de transformer específicos do modo, em seguida, os modos são concatenados para os blocos únicos conjuntos, e os tokens de imagem passam pelas camadas de saída. Os timesteps são por token: os tokens de imagem carregam o timestep de fluxo, enquanto os tokens de referência e o contexto de texto ficam em 0.

import torch
from diffusers import Flux3Transformer2DModel

transformer = Flux3Transformer2DModel.from_pretrained(
    "black-forest-labs/flux-3-image-klein-base", subfolder="transformer", dtype=torch.bfloat16
)

Flux3Transformer2DModel

classe diffusers.Flux3Transformer2DModel

< >

( in_channels: int = 128cond_in_channels: int = 128num_layers: int = 5num_single_layers: int = 28attention_head_dim: int = 128num_attention_heads: int = 24joint_attention_dim: int = 20480timestep_channels: int = 256mlp_ratio: float = 3.0axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)rope_theta: int = 10000eps: float = 1e-06flow_map: bool = False )

Parâmetros

  • in_channels (int, padrão 128) — Largura dos canais dos tokens de imagem (os latentes do VAE empacotados, latent_channels * 4).
  • cond_in_channels (int, padrão 128) — Largura dos canais dos tokens da imagem de referência.
  • num_layers (int, padrão 5) — O número de blocos de transformer por modo.
  • num_single_layers (int, padrão 28) — O número de blocos de transformador únicos conjuntos.
  • attention_head_dim (int, padrão 128) — O número de dimensões a serem usadas para cada cabeça de atenção.
  • num_attention_heads (int, padrão 24) — O número de cabeças de atenção a serem usadas.
  • joint_attention_dim (int, padrão 20480) — A dimensão de incorporação/canal de encoder_hidden_states (estados ocultos empilhados do Qwen3-VL).
  • timestep_channels (int, padrão 256) — O número de canais da incorporação sinusoidal do passo de tempo.
  • mlp_ratio (float, padrão 3.0) — A razão de expansão dos sub-blocos feedforward.
  • axes_dims_rope (tuple[int], padrão (32, 32, 32, 32)) — As dimensões a serem usadas para os embeddings posicionais rotativos, por eixo de posição (t, h, w, l).
  • rope_theta (int, padrão 10000) — A frequência base dos embeddings posicionais rotativos.
  • eps (float, padrão 1e-6) — O epsilon das normalizações de camada e das normalizações QK.
  • flow_map (bool, padrão False) — Indica se o modelo é um flow map (step-distilled): adiciona o target_timestep_embedder e exige target_timestep na forward pass.

O modelo Transformer introduzido no FLUX.3 para geração de texto para imagem e imagem para imagem.

Três modos de tokens compartilham uma única sequência conjunta: os tokens de imagem com ruído (hidden_states), os tokens opcionais de imagem de referência codificados pelo VAE (image_cond_hidden_states) e o contexto de texto (encoder_hidden_states). Cada modo primeiro passa por seus próprios blocos de transformer específicos do modo (mode_blocks para os modos de imagem, txt_mode_blocks para o contexto de texto), em seguida, o contexto de texto e os modos de imagem ativos são concatenados para os blocos únicos conjuntos, e apenas os tokens de imagem passam pelas camadas de saída. A modulação é por modo: os tokens de imagem são modulados pelo timestep do fluxo, enquanto os tokens de imagem de referência e o contexto de texto são modulados pelo timestep 0. Um checkpoint de flow-map (flow_map=True, os modelos destilados por passo) embute adicionalmente o timestep para o qual cada passo de desruidização salta e o adiciona à modulação de cada modo.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorimg_ids: Tensortxt_ids: Tensorimage_cond_hidden_states: typing.Optional[torch.Tensor] = Noneimage_cond_ids: typing.Optional[torch.Tensor] = Nonetarget_timestep: 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)) — Tokens latentes de imagem com ruído.
  • encoder_hidden_states (torch.Tensor de forma (batch_size, text_seq_len, joint_attention_dim)) — Incorporações de contexto de texto (estados ocultos empilhados do Qwen3-VL).
  • timestep (torch.Tensor de forma (batch_size,)) — Passo de fluxo em [0, 1] dos tokens da imagem. Os tokens da imagem de referência e o contexto de texto estão no passo 0.
  • img_ids (torch.Tensor de forma (image_seq_len, 4)) — Ids de posição (t, h, w, l) dos tokens de imagem, usados para os embeddings posicionais rotativos e compartilhados por todas as amostras no lote.
  • txt_ids (torch.Tensor de forma (text_seq_len, 4)) — Ids de posição dos tokens de contexto de texto.
  • image_cond_hidden_states (torch.Tensor de forma (batch_size, cond_seq_len, cond_in_channels), opcional) — Tokens da imagem de referência codificados pelo VAE (conteúdo de condicionamento limpo, para geração de imagem a partir de imagem).
  • image_cond_ids (torch.Tensor de forma (cond_seq_len, 4), opcional) — Ids de posição dos tokens da imagem de referência. Obrigatório com image_cond_hidden_states.
  • target_timestep (torch.Tensor de forma (batch_size,), opcional) — Passo de fluxo em [0, 1] para o qual o passo atual de desruidamento salta. Obrigatório apenas para um modelo de mapa de fluxo (flow_map=True).
  • attention_kwargs (dict, opcional) — Um dicionário de kwargs cuja entrada scale, se especificada, escala as camadas LoRA para esta passagem de frente.
  • return_dict (bool, opcional, padrão True) — Se deve ou não retornar um ~models.transformer_2d.Transformer2DModelOutput em vez de uma tupla simples.

O método forward do Flux3Transformer2DModel.

Transformer2DModelOutput

classe diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Parâmetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width) ou (batch size, num_vector_embeds - 1, num_latent_pixels) se Transformer2DModel for discreto) — A saída dos estados ocultos condicionada à entrada encoder_hidden_states. Se discreto, retorna distribuições de probabilidade para os pixels latentes sem ruído.

A saída de Transformer2DModel.

Update on GitHub