Diffusers documentation
AuraFlowTransformer2DModel
AuraFlowTransformer2DModel
Un modelo Transformer para datos similares a imágenes de AuraFlow.
AuraFlowTransformer2DModel
clase diffusers.AuraFlowTransformer2DModel
< source >( 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 deencoder_hidden_statesa utilizar. - caption_projection_dim (
int) — Número de dimensiones a utilizar al proyectarencoder_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
< source >( hidden_states: FloatTensorencoder_hidden_states: FloatTensor = Nonetimestep: LongTensor = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )
Parámetros
- hidden_states (
torch.FloatTensorde forma(batch size, channel, height, width)) — Estados ocultos de entradahidden_states. - encoder_hidden_states (
torch.FloatTensorde 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 alAttentionProcessordefinido enself.processordentro de diffusers.models.attention_processor. - 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 AuraFlowTransformer2DModel.
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.
Desactiva la proyección QKV fusionada si está habilitada.
> Esta API es 🧪 experimental.