Diffusers documentation
Flux3Transformer2DModel
Flux3Transformer2DModel
Un modelo Transformer para datos tipo imagen de FLUX.3, para la generación de texto a imagen y de imagen a imagen. Tres modos de token comparten una secuencia conjunta: los tokens de imagen con ruido, los tokens de la imagen de referencia codificada por VAE y el contexto de texto. Cada modo primero pasa por sus propios bloques de transformer por modo, luego los modos se concatenan para los bloques únicos conjuntos, y los tokens de imagen pasan por las capas de salida. Los pasos de tiempo son por token: los tokens de imagen llevan el paso de tiempo de flujo, mientras que los tokens de referencia y el contexto de texto se sitúan en 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
clase 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 )
Parámetros
- in_channels (
int, por defecto128) — Ancho de canal de los tokens de imagen (los latentes VAE empaquetados,latent_channels * 4). - cond_in_channels (
int, por defecto128) — Ancho de canal de los tokens de la imagen de referencia. - num_layers (
int, por defecto5) — El número de bloques de transformador por modo. - num_single_layers (
int, por defecto28) — El número de bloques de transformador individuales conjuntos. - attention_head_dim (
int, por defecto128) — El número de dimensiones a utilizar para cada cabeza de atención. - num_attention_heads (
int, por defecto24) — El número de cabezas de atención a utilizar. - joint_attention_dim (
int, por defecto20480) — La dimensión de incrustación/canal deencoder_hidden_states(estados ocultos apilados de Qwen3-VL). - timestep_channels (
int, por defecto256) — El número de canales de la incrustación sinusoidal del paso de tiempo. - mlp_ratio (
float, por defecto3.0) — La proporción de expansión de los subbloques de retroalimentación. - axes_dims_rope (
tuple[int], por defecto(32, 32, 32, 32)) — Las dimensiones a utilizar para los embeddings posicionales rotatorios, por eje de posición(t, h, w, l). - rope_theta (
int, por defecto10000) — La frecuencia base de los embeddings posicionales rotatorios. - eps (
float, por defecto1e-6) — La epsilon de las normalizaciones de capa y de las normas QK. - flow_map (
bool, por defectoFalse) — Indica si el modelo es un flow map (destilado por pasos): añade eltarget_timestep_embeddery requieretarget_timestepen la pasada hacia adelante.
El modelo Transformer introducido en FLUX.3 para la generación de texto a imagen y de imagen a imagen.
Tres modos de tokens comparten una secuencia conjunta: los tokens de imagen con ruido (hidden_states), los tokens opcionales de imagen de referencia codificados por VAE (image_cond_hidden_states) y el contexto de texto (encoder_hidden_states). Cada modo primero pasa por sus propios bloques de transformador por modo (mode_blocks para los modos de imagen, txt_mode_blocks para el contexto de texto), luego el contexto de texto y los modos de imagen activos se concatenan para los bloques únicos conjuntos, y solo los tokens de imagen pasan por las capas de salida. La modulación es por modo: los tokens de imagen se modulan por el paso de tiempo del flujo, mientras que los tokens de imagen de referencia y el contexto de texto se modulan por el paso de tiempo 0. Un punto de control de mapa de flujo (flow_map=True, los modelos destilados por pasos) incrusta adicionalmente el paso de tiempo al que salta cada paso de desruido y lo suma a la modulación de cada modo.
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 )
Parámetros
- hidden_states (
torch.Tensorde forma(batch_size, image_seq_len, in_channels)) — Tokens latentes de imagen con ruido. - encoder_hidden_states (
torch.Tensorde forma(batch_size, text_seq_len, joint_attention_dim)) — Incrustaciones de contexto de texto (estados ocultos apilados de Qwen3-VL). - timestep (
torch.Tensorde forma(batch_size,)) — Paso de flujo en[0, 1]de los tokens de la imagen. Los tokens de la imagen de referencia y el contexto de texto se sitúan en el paso0. - img_ids (
torch.Tensorde forma(image_seq_len, 4)) — ids de posición(t, h, w, l)de los tokens de imagen, utilizados para los embeddings posicionales rotatorios y compartidos por todas las muestras del lote. - txt_ids (
torch.Tensorde forma(text_seq_len, 4)) — Ids de posición de los tokens del contexto de texto. - image_cond_hidden_states (
torch.Tensorde forma(batch_size, cond_seq_len, cond_in_channels), opcional) — Tokens de imagen de referencia codificados por VAE (contenido de condicionamiento limpio, para generación de imagen a imagen). - image_cond_ids (
torch.Tensorde forma(cond_seq_len, 4), opcional) — Ids de posición de los tokens de la imagen de referencia. Obligatorio conimage_cond_hidden_states. - target_timestep (
torch.Tensorde forma(batch_size,), opcional) — Paso de flujo en[0, 1]al que salta el paso actual de desruido. Obligatorio y solo para un modelo de mapa de flujo (flow_map=True). - attention_kwargs (
dict, opcional) — Un diccionario de kwargs cuya entradascale, si se especifica, escala las capas LoRA para esta pasada de adelante. - return_dict (
bool, opcional, valor predeterminadoTrue) — Indica si se debe devolver un~models.transformer_2d.Transformer2DModelOutputen lugar de una tupla simple.
El método forward de Flux3Transformer2DModel.
Transformer2DModelOutput
clase diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Parámetros
- sample (
torch.Tensorde forma(batch_size, num_channels, height, width)o(batch size, num_vector_embeds - 1, num_latent_pixels)si Transformer2DModel es discreto) — La salida de estados ocultos condicionada a la entradaencoder_hidden_states. Si es discreto, devuelve distribuciones de probabilidad para los píxeles latentes sin ruido.
La salida de Transformer2DModel.