Diffusers documentation
Flux3Transformer2DModel
Flux3Transformer2DModel
Un modèle Transformer pour des données de type image issu de FLUX.3, destiné à la génération texte-vers-image et image-vers-image. Trois modes de jetons partagent une séquence conjointe : les jetons d'image bruités, les jetons d'image de référence encodés par le VAE et le contexte texte. Chaque mode passe d'abord par ses propres blocs Transformer spécifiques au mode, puis les modes sont concaténés pour les blocs uniques conjoints, et les jetons d'image passent par les couches de sortie. Les pas de temps sont spécifiques à chaque jeton : les jetons d'image portent le pas de temps du flux, tandis que les jetons de référence et le contexte texte sont positionnés à 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
< source >( 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 = (32, 32, 32, 32)rope_theta: int = 10000eps: float = 1e-06flow_map: bool = False )
Paramètres
- in_channels (
int, par défaut128) — Largeur des canaux des jetons d'image (les latents VAE empaquetés,latent_channels * 4). - cond_in_channels (
int, par défaut128) — Largeur des canaux des jetons de l'image de référence. - num_layers (
int, par défaut5) — Le nombre de blocs de transformateur par mode. - num_single_layers (
int, par défaut28) — Le nombre de blocs de transformateur simples conjoints. - attention_head_dim (
int, par défaut128) — Le nombre de dimensions à utiliser pour chaque tête d'attention. - num_attention_heads (
int, par défaut24) — Le nombre de têtes d'attention à utiliser. - joint_attention_dim (
int, par défaut20480) — La dimension d’embedding/canal deencoder_hidden_states(états cachés empilés de Qwen3-VL). - timestep_channels (
int, par défaut256) — Le nombre de canaux de l'encodage sinusoïdal du pas de temps. - mlp_ratio (
float, par défaut3.0) — Le ratio d'expansion des sous-blocs feedforward. - axes_dims_rope (
tuple[int], par défaut(32, 32, 32, 32)) — Les dimensions à utiliser pour les embeddings positionnels rotatifs, par axe de position(t, h, w, l). - rope_theta (
int, par défaut10000) — La fréquence de base des embeddings positionnels rotatifs. - eps (
float, par défaut1e-6) — L'épsilon des normalisations de couche et des normalisations QK. - flow_map (
bool, par défautFalse) — Indique si le modèle est une flow map (distillée par étape) : ajoute letarget_timestep_embedderet exigetarget_timestepdans le passage en avant.
Le modèle Transformer introduit dans FLUX.3 pour la génération d'images à partir de texte et la génération d'images à partir d'images.
Trois modes de jetons partagent une séquence conjointe : les jetons d’image bruités (hidden_states), les jetons d’image de référence encodés par VAE, facultatifs (image_cond_hidden_states) et le contexte texte (encoder_hidden_states). Chaque mode passe d’abord par ses propres blocs de transformateur spécifiques au mode (mode_blocks pour les modes image, txt_mode_blocks pour le contexte texte), puis le contexte texte et les modes image actifs sont concaténés pour les blocs uniques conjoints, et seuls les jetons d’image passent par les couches de sortie. La modulation est spécifique à chaque mode : les jetons d’image sont modulés par l’instant de flux, les jetons d’image de référence et le contexte texte par l’instant 0. Un point de contrôle de carte de flux (flow_map=True, les modèles distillés par étape) encode en outre l’instant vers lequel chaque étape de débruitage saute et l’ajoute à la modulation de chaque mode.
forward
< source >( 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 )
Paramètres
- hidden_states (
torch.Tensorde forme(batch_size, image_seq_len, in_channels)) — Tokens latents d'image bruités. - encoder_hidden_states (
torch.Tensorde forme(batch_size, text_seq_len, joint_attention_dim)) — Embeddings de contexte textuel (états cachés empilés de Qwen3-VL). - timestep (
torch.Tensorde forme(batch_size,)) — Pas de temps du flux dans[0, 1]des jetons d'image. Les jetons de l'image de référence et le contexte texte sont situés au pas de temps0. - img_ids (
torch.Tensorde forme(image_seq_len, 4)) — Identifiants de position(t, h, w, l)des jetons d’image, utilisés pour les embeddings positionnels rotatifs et partagés par tous les échantillons du lot. - txt_ids (
torch.Tensorde forme(text_seq_len, 4)) — Identifiants de position des jetons du contexte texte. - image_cond_hidden_states (
torch.Tensorde forme(batch_size, cond_seq_len, cond_in_channels), facultatif) — Jetons de l'image de référence encodés par le VAE (contenu de conditionnement propre, pour la génération d'image à partir d'image). - image_cond_ids (
torch.Tensorde forme(cond_seq_len, 4), facultatif) — Identifiants de position des jetons de l'image de référence. Requis avecimage_cond_hidden_states. - target_timestep (
torch.Tensorde forme(batch_size,), facultatif) — Pas de flux dans[0, 1]vers lequel saute l'étape de débruitage actuelle. Requis uniquement par un modèle de carte de flux (flow_map=True). - attention_kwargs (
dict, facultatif) — Dictionnaire de kwargs dont l'entréescale, si spécifiée, met à l'échelle les couches LoRA pour ce passage en avant. - return_dict (
bool, facultatif, par défautTrue) — Indique s'il faut retourner un~models.transformer_2d.Transformer2DModelOutputplutôt qu'un simple tuple.
La méthode forward du Flux3Transformer2DModel.
Transformer2DModelOutput
classe diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Paramètres
- sample (
torch.Tensorde forme(batch_size, num_channels, height, width)ou(batch size, num_vector_embeds - 1, num_latent_pixels)si Transformer2DModel est discret) — Les états cachés en sortie, conditionnés par l'entréeencoder_hidden_states. Si discret, renvoie les distributions de probabilité pour les pixels latents non bruités.
La sortie de Transformer2DModel.