Diffusers documentation
AuraFlowTransformer2DModel
AuraFlowTransformer2DModel
Un modèle Transformer pour des données de type image issu de AuraFlow.
AuraFlowTransformer2DModel
classe 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 )
Paramètres
- sample_size (
int) — La largeur des images latentes. Cette valeur est fixe pendant l’entraînement, car elle est utilisée pour apprendre un nombre d’embeddings de position. - patch_size (
int) — Taille des patches pour découper les données d'entrée en petits patches. - in_channels (
int, facultatif, par défaut 4) — Le nombre de canaux de l'entrée. - num_mmdit_layers (
int, facultatif, valeur par défaut 4) — Le nombre de blocs Transformer MMDiT à utiliser. - num_single_dit_layers (
int, facultatif, valeur par défaut 32) — Le nombre de blocs Transformer à utiliser. Ces blocs utilisent des représentations d'image et de texte concaténées. - attention_head_dim (
int, facultatif, valeur par défaut 256) — Le nombre de canaux dans chaque tête. - num_attention_heads (
int, facultatif, valeur par défaut 12) — Le nombre de têtes à utiliser pour l'attention multi-têtes. - joint_attention_dim (
int, facultatif) — Le nombre de dimensions deencoder_hidden_statesà utiliser. - caption_projection_dim (
int) — Nombre de dimensions à utiliser lors de la projection desencoder_hidden_states. - out_channels (
int, par défaut 4) — Nombre de canaux de sortie. - pos_embed_max_size (
int, par défaut 1024) — Nombre maximal de positions à intégrer à partir des latents de l’image.
Un modèle Transformer 2D tel qu'introduit dans 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 )
Paramètres
- hidden_states (
torch.FloatTensorde forme(batch size, channel, height, width)) —hidden_statesd'entrée. - encoder_hidden_states (
torch.FloatTensorde forme(batch size, sequence_len, embed_dims)) — Embeddings conditionnels (embeddings calculés à partir des conditions d’entrée telles que les prompts) à utiliser. - timestep (
torch.LongTensor) — Utilisé pour indiquer l'étape de débruitage. - attention_kwargs (
dict, facultatif) — Un dictionnaire de kwargs qui, s'il est spécifié, est transmis à l'AttentionProcessordéfini sousself.processordans diffusers.models.attention_processor. - 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 AuraFlowTransformer2DModel.
Active les projections QKV fusionnées. Pour les modules d'auto-attention, toutes les matrices de projection (c'est-à-dire requête, clé, valeur) sont fusionnées. Pour les modules d'attention croisée, les matrices de projection des clés et des valeurs sont fusionnées.
> Cette API est 🧪 expérimentale.
Désactive la projection QKV fusionnée si elle est activée.
> Cette API est 🧪 expérimentale.