Hugging Face's logo

Diffusers documentation

QwenImage21Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

QwenImage21Transformer2DModel

Le transformer à flux unique utilisé par Qwen-Image 2.1. Les latents texte et image partagent une seule séquence, et une unique projection de modulation partagée alimente chaque bloc.

Deux comportements distinguent 2.1 des transformeurs QwenImage antérieurs :

  • Attention causale par blocs — l’attention suit (q_idx >= kv_idx) or same_image_block, de sorte que la séquence conjointe est causale tandis que chaque bloc d’image reste bidirectionnel en interne. QwenImage21AttnProcessor l’implémente sous forme d’un appel d’attention par segment de préfixe et est le choix par défaut. QwenImage21FlexAttnProcessor l’implémente sous forme d’un unique appel flex_attention piloté par un BlockMask, ce qui est plus rapide une fois le modèle compilé. Les deux produisent les mêmes résultats.
  • causal_condition — les jetons de texte et d’image de condition sont modulés à partir de t = 0 plutôt qu’à partir de l’étape d’échantillonnage. Leurs activations sont indépendantes de l’étape de débruitage, de sorte que les clés et les valeurs de ce préfixe peuvent être mises en cache entre les étapes via l’argument kv_cache.

Chargez-le avec :

import torch
from diffusers import QwenImage21Transformer2DModel

transformer = QwenImage21Transformer2DModel.from_pretrained(
    "Qwen/Qwen-Image-2.1", subfolder="transformer", dtype=torch.bfloat16
)

QwenImage21Transformer2DModel

classe diffusers.QwenImage21Transformer2DModel

< >

( patch_size: int = 1in_channels: int = 64out_channels: int | None = 64num_layers: int = 32attention_head_dim: int = 128num_attention_heads: int = 32context_in_dim: int = 4096mlp_ratio: int = 3axes_dims_rope: tuple[int, int, int] = (16, 56, 56)eps: float = 1e-06causal_condition: bool = True )

Paramètres

  • patch_size (int, par défaut 1) — Longueur du côté du patch latent replié dans la dimension des canaux. La version 2.1 consomme les latents sans patchification.
  • in_channels (int, par défaut 64) — Canaux latents de l'entrée.
  • out_channels (int, facultatif, par défaut 64) — Canaux latents de la sortie. Repli sur in_channels.
  • num_layers (int, par défaut 32) — Nombre de blocs à flux unique.
  • attention_head_dim (int, par défaut 128) — Nombre de canaux par tête d'attention.
  • num_attention_heads (int, par défaut 32) — Nombre de têtes d'attention.
  • context_in_dim (int, par défaut 4096) — Dimension des canaux de encoder_hidden_states.
  • mlp_ratio (int, par défaut 3) — Facteur d'expansion du réseau de neurones en avant.
  • axes_dims_rope (tuple[int], par défaut (16, 56, 56)) — Dimensions rotationnelles pour les axes de la trame, de la hauteur et de la largeur.
  • eps (float, par défaut 1e-6) — Epsilon pour les couches de normalisation.
  • causal_condition (bool, par défaut True) — Moduler les jetons de texte et d’image de conditionnement à partir de t = 0. Requis pour la mise en cache des clés et des valeurs.

Le Transformer à flux unique utilisé par Qwen-Image 2.1.

Les latents texte et image partagent une même séquence : les jetons de l’image de condition sont substitués dans le flux texte aux positions réservées par l’encodeur vision-langage, et les jetons de l’image cible sont ajoutés à la fin. Une unique projection modulation partagée alimente tous les blocs, de sorte que les blocs ne possèdent pas de paramètres de modulation propres.

Deux comportements distinguent la version 2.1 :

  • Attention causale par blocs — l'attention suit (q_idx >= kv_idx) or same_image_block, de sorte que la séquence est causale tandis que chaque bloc d'image reste bidirectionnel en interne. Le backend d'attention flex offre une attention efficace en un seul passage ; sans lui, le modèle utilise un préremplissage SDPA exact en plusieurs passages qui traite chaque bloc séparément.
  • causal_condition — les jetons de texte et d’image de condition sont modulés à partir de t = 0 plutôt qu’à partir de l’étape d’échantillonnage, ce qui rend également leurs activations indépendantes de l’étape et donc cacheables entre les étapes de débruitage.

build_token_metadata

< >

( image_pad_mask: Tensorimg_shapes: list[tuple[int, int, int]] ) → tuple[torch.Tensor, torch.Tensor]

Paramètres

  • image_pad_mask (torch.Tensor) — bool de forme (seq_len,), avec True aux positions des jetons d’image.
  • img_shapes (list[tuple[int, int, int]]) — Pour chaque image, (frame, height, width) en jetons latents, les images de conditionnement d'abord et l'image cible en dernier.

Retourne

tuple[torch.Tensor, torch.Tensor]

image_ids (seq_len,) avec -1 aux positions du texte et un identifiant unique par bloc d'image, et target_token_mask (seq_len,) marquant les jetons de l'image cible.

Attribuez à chaque jeton de la séquence conjointe le bloc d'image auquel il appartient.

Les limites des blocs proviennent des nombres de jetons dans img_shapes, et non des séquences de True dans image_pad_mask : deux images de condition qui se trouvent côte à côte sans texte entre elles forment une seule séquence, mais doivent rester des blocs distincts, sinon elles se feraient attention mutuellement de manière bidirectionnelle.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorimg_shapes: list[list[tuple[int, int, int]]]img_mask: Tensorencoder_hidden_states_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonekv_cache: diffusers.models.transformers.transformer_qwenimage21.QwenImage21KVCache | None = Nonekv_cache_mode: str | None = Nonereturn_dict: bool = True ) → Transformer2DModelOutput or tuple

Paramètres

  • hidden_states (torch.Tensor de forme (batch_size, image_sequence_length, in_channels)) — Latents empilés, avec les images de conditionnement en premier et l'image cible en dernier.
  • encoder_hidden_states (torch.Tensor de forme (batch_size, text_sequence_length, context_in_dim)) — Embeddings de texte issus de l’encodeur vision-langage.
  • timestep (torch.Tensor) — Étape de débruitage actuelle, mise à l'échelle dans [0, 1].
  • img_shapes (list[list[tuple[int, int, int]]]) — Liste par échantillon de (frame, height, width) en jetons latents, les images de condition d’abord et l’image cible en dernier. Tous les échantillons doivent partager une même disposition.
  • img_mask (torch.Tensor de forme (batch_size, vlm_sequence_length)) — True aux emplacements d’image de l’encodeur vision-langage, chacun représentant un groupe de 2x2 jetons latents.
  • encoder_hidden_states_mask (torch.Tensor, facultatif) — (batch_size, text_sequence_length) bool indiquant les jetons de texte valides. Les positions rembourrées sont exclues de l'attention.
  • kv_cache (QwenImage21KVCache, facultatif) — Conteneur de cache. À passer avec kv_cache_mode pour activer le cache KV préfixe.
  • kv_cache_mode (str, facultatif) — "extract" pour préremplir le cache (première étape de débruitage), "cached" pour décoder à partir de celui-ci (étapes suivantes). Nécessite causal_condition=True.
  • attention_kwargs (dict, facultatif) — Transmis aux processeurs d'attention et porte scale pour les couches LoRA.
  • return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un Transformer2DModelOutput plutôt qu'un simple tuplet.

Retourne

Transformer2DModelOutput ou tuple

Transformer2DModelOutput si return_dict est True, sinon un simple tuple dont le premier élément est les latents débruités.

Transformer2DModelOutput

classe diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Paramètres

  • sample (torch.Tensor de forme (batch_size, num_channels, height, width) ou (batch size, num_vector_embeds - 1, num_latent_pixels) si Transformer2DModel est discret) — Les états cachés en sortie, conditionnés par l'entrée encoder_hidden_states. Si discret, renvoie les distributions de probabilité pour les pixels latents non bruités.

La sortie de Transformer2DModel.

Update on GitHub

←PriorTransformer QwenImageTransformer2DModel→