Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
Un transformer de visión-lenguaje conjunto Mixture-of-Transformer (MoT) introducido como parte de la familia de modelos de mundo base Cosmos3 de NVIDIA. El modelo ejecuta dos vías de cómputo paralelas sobre una secuencia conjunta empaquetada:
- una vía de “comprensión” causal que realiza autoatención sobre tokens de texto con enmascarado causal, y
- una vía de “generación” bidireccional que realiza una atención cruzada desde los tokens de generación (latentes de visión + sonido opcionales) sobre el conjunto completo de claves/valores de comprensión más generación.
Los dos caminos comparten el mismo tamaño oculto y número de capas, pero mantienen proyecciones Q/K/V/O, MLPs y parámetros RMSNorm separados, lo que hace que la arquitectura sea una Mixture-of-Transformer en lugar de una Mixture-of-Experts estándar. La información de posición se proporciona a través de un RoPE multimodal 3D (mRoPE) que intercala frecuencias temporales / de altura / de ancho para los latentes de video y reutiliza el eje temporal para el texto y el audio.
El modelo se puede cargar de la siguiente manera.
import torch
from diffusers import Cosmos3OmniTransformer
transformer = Cosmos3OmniTransformer.from_pretrained(
"nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)Cosmos3OmniTransformer
clase 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: list[torch.Tensor]vision_token_shapes: list[tuple[int, int, int]]vision_sequence_indexes: Tensorvision_mse_loss_indexes: Tensorvision_timesteps: Tensorvision_noisy_frame_indexes: list[torch.Tensor]sound_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 )
Parámetros
- input_ids — IDs de tokens de texto colocados en
text_indexesen la secuencia conjunta. - text_indexes — Índices de los tokens de texto en la secuencia conjunta.
- position_ids — IDs de posición mRoPE para la secuencia conjunta completa,
[3, sequence_length]. - und_len — Longitud del prefijo de texto causal (comprensión); los tokens de generación siguen a continuación.
- sequence_length — Longitud total de la secuencia empaquetada conjunta.
- vision_tokens — Tensores latentes de visión por elemento antes de patchify.
- vision_token_shapes — Formas de la cuadrícula de parches
(T, H, W)por elemento de visión. - vision_sequence_indexes — Índices de los tokens de visión en la secuencia conjunta.
- vision_mse_loss_indexes — Índices utilizados para leer las predicciones visuales después del backbone.
- vision_timesteps — Pasos de difusión por parche para los tokens de visión.
- vision_noisy_frame_indexes — Índices de cuadros ruidosos por elemento de visión.
- sound_tokens — Tensores latentes de sonido opcionales antes del empaquetado.
- sound_token_shapes — Formas opcionales de la cuadrícula de parches para elementos de sonido.
- sound_sequence_indexes — Índices opcionales de los tokens de sonido en la secuencia conjunta.
- sound_mse_loss_indexes — Índices opcionales utilizados para leer las predicciones de sonido.
- sound_timesteps — Pasos de difusión opcionales por token para sonido.
- sound_noisy_frame_indexes — Índices de fotogramas ruidosos opcionales por elemento de sonido.
- action_tokens — Tensores latentes de acción opcionales antes del empaquetado.
- action_token_shapes — Formas opcionales de la cuadrícula de parches
(T, H, W)por elemento de acción. - action_sequence_indexes — Índices opcionales de los tokens de acción en la secuencia conjunta.
- action_mse_loss_indexes — Índices opcionales usados para leer las predicciones de acción después del backbone.
- action_timesteps — Pasos de difusión opcionales por token para los tokens de acción.
- action_noisy_frame_indexes — Índices de fotogramas ruidosos opcionales por elemento de acción.
- action_domain_ids — IDs de dominio opcionales por elemento para seleccionar los pesos de la cabeza de acción.
- return_dict — Si se debe devolver un [Cosmos3OmniTransformerOutput] en lugar de una tupla.
Ejecuta una pasada de propagación hacia adelante completa de un paso de desruido.