Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
Un transformer vision-langage conjoint de type Mixture-of-Transformer (MoT), introduit dans le cadre de la famille de modèles de fondation du monde Cosmos3 de NVIDIA. Le modèle exécute deux chemins de calcul parallèles sur une séquence conjointe compactée :
- un chemin causal de « compréhension » qui s'auto-attende sur les jetons de texte avec un masquage causal, et
- un chemin de « génération » bidirectionnel qui effectue une attention croisée depuis les jetons de génération (latents visuels + latents sonores optionnels) sur l’ensemble complet des clés/valeurs de compréhension et de génération.
Les deux chemins partagent la même taille de représentation cachée et le même nombre de couches, mais conservent des projections Q/K/V/O, des MLP et des paramètres RMSNorm distincts, ce qui fait de l'architecture un Mixture-of-Transformer plutôt qu'un Mixture-of-Experts standard. Les informations de position sont fournies par un RoPE multimodal 3D (mRoPE) qui entrelace les fréquences temporelles / hauteur / largeur pour les latents vidéo et réutilise l'axe temporel pour le texte et l'audio.
Le modèle peut être chargé comme suit.
import torch
from diffusers import Cosmos3OmniTransformer
transformer = Cosmos3OmniTransformer.from_pretrained(
"nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)Cosmos3OmniTransformer
classe diffusers.Cosmos3OmniTransformer
< source >( attention_bias: bool = Falseattention_dropout: float = 0.0dtype: str = 'bfloat16'head_dim: int = 128hidden_size: int = 4096intermediate_size: int = 12288base_fps: int = 24enable_fps_modulation: bool = Truelatent_channel: int = 48unified_3d_mrope_reset_spatial_ids: bool = Trueunified_3d_mrope_temporal_modality_margin: int = 15000latent_patch_size: int = 2num_attention_heads: int = 32num_hidden_layers: int = 36num_key_value_heads: int = 8patch_latent_dim: int = 192rms_norm_eps: float = 1e-06rope_scaling: dict | None = Nonerope_theta: float = 5000000.0action_dim: int | None = Noneaction_gen: bool = Falsenum_embodiment_domains: int = 32sound_dim: int | None = Nonesound_gen: bool = Falsesound_latent_fps: float = 25.0timestep_scale: float = 0.001vocab_size: int = 151936hidden_act: str = 'silu'qk_norm_for_text: bool = Trueuse_und_k_norm_for_gen: bool = Falserope_axes_dim: tuple[int, int, int] | list[int] | None = None )
forward
< source >( input_ids: Tensortext_indexes: Tensorposition_ids: Tensorund_len: intsequence_length: intvision_tokens: listvision_token_shapes: listvision_sequence_indexes: Tensorvision_mse_loss_indexes: Tensorvision_timesteps: Tensorvision_noisy_frame_indexes: listsound_tokens: list[torch.Tensor] | None = Nonesound_token_shapes: list[tuple[int, int, int]] | None = Nonesound_sequence_indexes: typing.Optional[torch.Tensor] = Nonesound_mse_loss_indexes: typing.Optional[torch.Tensor] = Nonesound_timesteps: typing.Optional[torch.Tensor] = Nonesound_noisy_frame_indexes: list[torch.Tensor] | None = Noneaction_tokens: list[torch.Tensor] | None = Noneaction_token_shapes: list[tuple[int, int, int]] | None = Noneaction_sequence_indexes: typing.Optional[torch.Tensor] = Noneaction_mse_loss_indexes: typing.Optional[torch.Tensor] = Noneaction_timesteps: typing.Optional[torch.Tensor] = Noneaction_noisy_frame_indexes: list[torch.Tensor] | None = Noneaction_domain_ids: list[torch.Tensor] | None = Nonereturn_dict: bool = True )
Paramètres
- input_ids — Identifiants des jetons de texte placés aux positions
text_indexesdans la séquence conjointe. - text_indexes — Indices des jetons de texte dans la séquence conjointe.
- position_ids — identifiants de position mRoPE pour la séquence jointe complète, de forme
[3, sequence_length]. - und_len — Longueur du préfixe de texte causal (compréhension) ; les jetons de génération suivent.
- sequence_length — Longueur totale de la séquence jointe compactée.
- vision_tokens — Tenseurs latents visuels par élément avant patchify.
- vision_token_shapes — Formes de la grille de patches
(T, H, W)pour chaque élément visuel. - vision_sequence_indexes — Indices des jetons visuels dans la séquence conjointe.
- vision_mse_loss_indexes — Indices utilisés pour lire les prédictions visuelles après le backbone.
- vision_timesteps — Étapes de diffusion par patch pour les jetons visuels.
- vision_noisy_frame_indexes — Indices des images bruitées par élément visuel.
- sound_tokens — Tenseurs latents sonores facultatifs avant l’empaquetage.
- sound_token_shapes — Formes de grille de patch facultatives pour les éléments sonores.
- sound_sequence_indexes — Indices facultatifs des jetons sonores dans la séquence conjointe.
- sound_mse_loss_indexes — Indices facultatifs utilisés pour lire les prédictions sonores.
- sound_timesteps — Étapes de diffusion par jeton facultatives pour le son.
- sound_noisy_frame_indexes — Indices de frames bruitées facultatifs par élément sonore.
- action_tokens — Tenseurs latents d'action optionnels avant l'empaquetage.
- action_token_shapes — Formes de grille de patch facultatives
(T, H, W)pour chaque élément d'action. - action_sequence_indexes — Indices facultatifs des jetons d'action dans la séquence conjointe.
- action_mse_loss_indexes — Indices facultatifs utilisés pour lire les prédictions d’action après le backbone.
- action_timesteps — Étapes de diffusion par jeton facultatives pour les jetons d'action.
- action_noisy_frame_indexes — Indices de frames bruitées par élément d'action, facultatif.
- action_domain_ids — Identifiants de domaine par élément, facultatifs, permettant de sélectionner les poids de la tête d’action.
- return_dict — Indique s'il faut retourner un [Cosmos3OmniTransformerOutput] plutôt qu'un tuple.
Exécuter un passage avant complet d'une étape de débruitage.