Hugging Face's logo

Diffusers documentation

Krea2Transformer2DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Krea2Transformer2DModel

Le transformateur de correspondance de flux MMDiT à flux unique utilisé par Krea 2.

Krea2Transformer2DModel

classe diffusers.Krea2Transformer2DModel

< >

( in_channels: int = 64num_layers: int = 28attention_head_dim: int = 128num_attention_heads: int = 48num_key_value_heads: int = 12intermediate_size: int = 16384timestep_embed_dim: int = 256text_hidden_dim: int = 2560num_text_layers: int = 12text_num_attention_heads: int = 20text_num_key_value_heads: int = 20text_intermediate_size: int = 6912num_layerwise_text_blocks: int = 2num_refiner_text_blocks: int = 2axes_dims_rope: tuple = (32, 48, 48)rope_theta: float = 1000.0norm_eps: float = 1e-05 )

Paramètres

  • in_channels (int, par défaut 64) — Nombre de canaux latents après patchification (vae_channels * patch_size ** 2).
  • num_layers (int, par défaut 28) — Nombre de blocs de transformer.
  • attention_head_dim (int, par défaut 128) — Dimension de chaque tête d'attention ; la taille cachée totale est attention_head_dim * num_attention_heads.
  • num_attention_heads (int, par défaut 48) — Nombre de têtes de requête.
  • num_key_value_heads (int, par défaut 12) — Nombre de têtes clé/valeur pour l'attention à requêtes groupées.
  • intermediate_size (int, par défaut 16384) — Taille cachée du réseau de feed-forward SwiGLU dans chaque bloc.
  • timestep_embed_dim (int, par défaut 256) — Largeur de l'encodage sinusoïdal du pas de temps avant son MLP.
  • text_hidden_dim (int, par défaut 2560) — Taille cachée de l’encodeur de texte dont les états cachés sont consommés.
  • num_text_layers (int, par défaut 12) — Nombre d'états cachés de l'encodeur de texte tapés et empilés par jeton.
  • text_num_attention_heads (int, par défaut 20) — Nombre de têtes de requête dans les blocs de fusion texte.
  • text_num_key_value_heads (int, par défaut 20) — Nombre de têtes clé/valeur dans les blocs de fusion texte.
  • text_intermediate_size (int, par défaut 6912) — Taille cachée du réseau de feed-forward du MLP SwiGLU dans les blocs de fusion texte.
  • num_layerwise_text_blocks (int, par défaut 2) — Nombre de blocs de fusion de texte appliqués sur l'axe des couches tapées (par jeton).
  • num_refiner_text_blocks (int, par défaut 2) — Nombre de blocs de fusion de texte appliqués sur la séquence de jetons.
  • axes_dims_rope (tuple[int, int, int], par défaut (32, 48, 48)) — Répartition de la dimension des têtes sur les axes de position rotative (t, h, w).
  • rope_theta (float, par défaut 1000.0) — Base utilisée par l'incorporation de position rotative.
  • norm_eps (float, par défaut 1e-5) — Epsilon utilisé par tous les modules RMSNorm.

Le backbone de flow-matching MMDiT à flux unique utilisé par le pipeline Krea 2.

Le conditionnement par le texte est introduit sous forme d'une pile d'états cachés prélevés sur plusieurs couches d'un encodeur de texte multimodal. Un petit transformeur de fusion de texte réduit l'axe des couches et affine la séquence de jetons ; le résultat est concaténé avec les latents d'image patchifiés pour former une unique séquence [text, image] traitée par les blocs transformeur. L'horodatage conditionne chaque bloc via un vecteur de modulation partagé ainsi que des tables apprises par bloc.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorposition_ids: Tensorencoder_attention_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )

Paramètres

  • hidden_states (torch.Tensor de forme (batch_size, image_seq_len, in_channels)) — Latents d'image bruités empilés (patchifiés).
  • encoder_hidden_states (torch.Tensor de forme (batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — Pile des états cachés de l'encodeur de texte tapés par jeton.
  • timestep (torch.Tensor de forme (batch_size,)) — Temps de flow-matching dans [0, 1] (1 correspond au bruit pur, 0 aux données propres).
  • position_ids (torch.Tensor de forme (text_seq_len + image_seq_len, 3)) — coordonnées rotatives (t, h, w) pour la séquence combinée. Les lignes de texte sont toutes nulles ; les lignes d’image contiennent les coordonnées de la grille latente.
  • encoder_attention_mask (torch.Tensor de forme (batch_size, text_seq_len), facultatif) — Masque booléen indiquant les jetons de texte valides. Passez None lorsque tous les jetons de texte sont valides.
  • attention_kwargs (dict, facultatif) — Dictionnaire de kwargs qui, s'il contient une entrée scale, définit l'échelle LoRA appliquée aux adaptateurs de ce transformateur pendant la durée du passage avant.
  • return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un Transformer2DModelOutput plutôt qu'un simple tuple.

Prédire la vitesse de flow-matching pour les jetons d'image.

Update on GitHub

←KandinskyS2VTransformer3DModel LatteTransformer3DModel→