Hugging Face's logo

Diffusers documentation

AuraFlowTransformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

AuraFlowTransformer2DModel

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

AuraFlowTransformer2DModel

clase diffusers.AuraFlowTransformer2DModel

< >

( sample_size: int = 64patch_size: int = 2in_channels: int = 4num_mmdit_layers: int = 4num_single_dit_layers: int = 32attention_head_dim: int = 256num_attention_heads: int = 12joint_attention_dim: int = 2048caption_projection_dim: int = 3072out_channels: int = 4pos_embed_max_size: int = 1024 )

Parámetros

  • sample_size (int) — El ancho de las imágenes latentes. Esto se fija durante el entrenamiento, ya que se utiliza para aprender un número de embeddings de posición.
  • patch_size (int) — Tamaño del parche para dividir los datos de entrada en pequeños parches.
  • in_channels (int, opcional, valor predeterminado 4) — El número de canales en la entrada.
  • num_mmdit_layers (int, opcional, valor predeterminado 4) — El número de capas de bloques Transformer MMDiT a utilizar.
  • num_single_dit_layers (int, opcional, valor predeterminado 32) — El número de capas de bloques Transformer a utilizar. Estos bloques usan representaciones de imagen y texto concatenadas.
  • attention_head_dim (int, opcional, valor predeterminado 256) — El número de canales en cada cabeza.
  • num_attention_heads (int, opcional, valor predeterminado 12) — El número de cabezas a utilizar para la atención multi-cabeza.
  • joint_attention_dim (int, opcional) — El número de dimensiones de encoder_hidden_states a utilizar.
  • caption_projection_dim (int) — Número de dimensiones a utilizar al proyectar encoder_hidden_states.
  • out_channels (int, por defecto 4) — Número de canales de salida.
  • pos_embed_max_size (int, por defecto 1024) — Posiciones máximas para incrustar desde los latentes de la imagen.

Un modelo Transformer 2D tal como se introdujo en AuraFlow (https://blog.fal.ai/auraflow/).

forward

< >

( hidden_states: FloatTensorencoder_hidden_states: FloatTensor = Nonetimestep: LongTensor = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )

Parámetros

  • hidden_states (torch.FloatTensor de forma (batch size, channel, height, width)) — Estados ocultos de entrada hidden_states.
  • encoder_hidden_states (torch.FloatTensor de forma (batch size, sequence_len, embed_dims)) — 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.
  • 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.

El método forward de AuraFlowTransformer2DModel.

fuse_qkv_projections

< >

( )

Habilita las proyecciones QKV fusionadas. Para los módulos de autoatención, todas las matrices de proyección (es decir, consulta, clave y valor) se fusionan. Para los módulos de atención cruzada, las matrices de proyección de clave y valor se fusionan.

> Esta API es 🧪 experimental.

unfuse_qkv_projections

< >

( )

Desactiva la proyección QKV fusionada si está habilitada.

> Esta API es 🧪 experimental.

Update on GitHub

←AnyFlowTransformer3DModel BriaFibo2Transformer2DModel→