Hugging Face's logo

Diffusers documentation

Flux2Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Flux2Transformer2DModel

Un modelo Transformer para datos similares a imágenes de Flux2.

Flux2Transformer2DModel

clase diffusers.Flux2Transformer2DModel

< >

( patch_size: int = 1in_channels: int = 128out_channels: int | None = Nonenum_layers: int = 8num_single_layers: int = 48attention_head_dim: int = 128num_attention_heads: int = 48joint_attention_dim: int = 15360timestep_guidance_channels: int = 256mlp_ratio: float = 3.0axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)rope_theta: int = 2000eps: float = 1e-06guidance_embeds: bool = True )

Parámetros

  • patch_size (int, por defecto 1) — Tamaño del parche para dividir los datos de entrada en parches pequeños.
  • in_channels (int, por defecto 128) — El número de canales en la entrada.
  • out_channels (int, opcional, valor predeterminado None) — El número de canales en la salida. Si no se especifica, el valor predeterminado es in_channels.
  • num_layers (int, por defecto 8) — El número de capas de bloques DiT de doble flujo a utilizar.
  • num_single_layers (int, por defecto 48) — El número de capas de bloques DiT de un solo flujo que se van a usar.
  • attention_head_dim (int, por defecto 128) — El número de dimensiones a utilizar para cada cabeza de atención.
  • num_attention_heads (int, por defecto 48) — El número de cabezas de atención a utilizar.
  • joint_attention_dim (int, por defecto 15360) — El número de dimensiones a utilizar para la atención conjunta (dimensión de incrustación/canal de encoder_hidden_states).
  • pooled_projection_dim (int, por defecto 768) — El número de dimensiones a utilizar para la proyección agrupada.
  • guidance_embeds (bool, por defecto True) — Si se deben usar embeddings de guía para la variante destilada con guía del modelo.
  • axes_dims_rope (tuple[int], por defecto (32, 32, 32, 32)) — Las dimensiones a utilizar para los embeddings posicionales rotatorios.

El modelo Transformer introducido en Flux 2.

Referencia: https://blackforestlabs.ai/announcing-black-forest-labs/

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensor = Nonetimestep: LongTensor = Noneimg_ids: Tensor = Nonetxt_ids: Tensor = Noneguidance: Tensor = Nonejoint_attention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = Truekv_cache: Flux2KVCache | None = Nonekv_cache_mode: str | None = Nonenum_ref_tokens: int = 0ref_fixed_timestep: float = 0.0 )

Parámetros

  • hidden_states (torch.Tensor de forma (batch_size, image_sequence_length, in_channels)) — Estados ocultos de entrada hidden_states.
  • encoder_hidden_states (torch.Tensor de forma (batch_size, text_sequence_length, joint_attention_dim)) — Incrustaciones condicionales (incrustaciones calculadas a partir de las condiciones de entrada, como los prompts) que se deben utilizar.
  • timestep (torch.LongTensor) — Se utiliza para indicar el paso de desruido.
  • img_ids (torch.Tensor) — ids de posición de imagen utilizados para calcular los embeddings posicionales rotatorios.
  • txt_ids (torch.Tensor) — ids de posición de texto utilizados para calcular los embeddings posicionales rotatorios.
  • guidance (torch.Tensor, opcional) — Incrustación de escala de orientación utilizada para las variantes destiladas por orientación del modelo.
  • joint_attention_kwargs (dict, opcional) — Un diccionario de kwargs que, si se especifica, se pasa al AttentionProcessor definido en self.processor dentro de diffusers.models.attention_processor.
  • return_dict (bool, opcional, valor predeterminado True) — Indica si se debe devolver un ~models.transformer_2d.Transformer2DModelOutput en lugar de una tupla simple.
  • kv_cache (Flux2KVCache, opcional) — Cache KV para los tokens de la imagen de referencia. Cuando kv_cache_mode es “extract”, se crea y devuelve una nueva cache. Cuando es “cached”, se utiliza la cache proporcionada para inyectar los K/V de referencia durante la atención.
  • kv_cache_mode (str, opcional) — Uno de “extract” (primer paso con tokens de referencia) o “cached” (pasos posteriores que usan K/V de referencia en caché). Cuando None, paso de propagación hacia adelante estándar sin caché de KV.
  • num_ref_tokens (int, por defecto 0) — Número de tokens de imagen de referencia añadidos al inicio de hidden_states (solo se usa cuando kv_cache_mode="extract").
  • ref_fixed_timestep (float, por defecto 0.0) — Paso de tiempo fijo para la modulación de tokens de referencia (solo se usa cuando kv_cache_mode="extract").

El método forward de Flux2Transformer2DModel.

Flux2Transformer2DModelOutput

clase diffusers.models.transformers.transformer_flux2.Flux2Transformer2DModelOutput

< >

( sample: torch.Tensorkv_cache: Flux2KVCache | None = None )

Parámetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width)) — Los estados ocultos de salida condicionados a la entrada encoder_hidden_states.
  • kv_cache (Flux2KVCache, opcional) — La caché KV poblada para los tokens de la imagen de referencia. Solo se devuelve cuando kv_cache_mode="extract".

La salida de Flux2Transformer2DModel.

Update on GitHub

←ErnieImageTransformer2DModel Flux3Transformer2DModel→