Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
Sebuah Mixture-of-Transformer (MoT) transformer gabungan visi-bahasa yang diperkenalkan sebagai bagian dari keluarga model fondasi dunia Cosmos3 dari NVIDIA. Model ini menjalankan dua jalur komputasi paralel di atas urutan gabungan yang dikemas:
- a jalur “pemahaman” kausal yang melakukan self-attention pada token teks dengan masking kausal, dan
- a jalur “generasi” bi-arah yang melakukan cross-attention dari token generasi (laten visual + laten suara opsional) ke seluruh himpunan kunci/nilai pemahaman-plus-generasi.
Kedua jalur berbagi ukuran tersembunyi dan jumlah lapisan yang sama, tetapi mempertahankan proyeksi Q/K/V/O, MLP, dan parameter RMSNorm yang terpisah, yang menjadikan arsitektur ini sebagai Mixture-of-Transformer alih-alih Mixture-of-Experts standar. Informasi posisi disuplai melalui 3D multimodal RoPE (mRoPE) yang menginterleaving frekuensi temporal / tinggi / lebar untuk latens video dan menggunakan kembali sumbu temporal untuk teks dan audio.
Model dapat dimuat sebagai berikut.
import torch
from diffusers import Cosmos3OmniTransformer
transformer = Cosmos3OmniTransformer.from_pretrained(
"nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)Cosmos3OmniTransformer
kelas 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 )
Parameter
- input_ids — ID token teks yang ditempatkan di
text_indexesdalam urutan gabungan. - text_indexes — Indeks token teks dalam urutan gabungan.
- position_ids — ID posisi mRoPE untuk urutan gabungan penuh,
[3, sequence_length]. - und_len — Panjang prefix teks kausal (pemahaman); token generasi mengikuti setelahnya.
- sequence_length — Total panjang dari urutan gabungan yang dikemas.
- vision_tokens — Tensor laten visual per item sebelum patchify.
- vision_token_shapes — Bentuk grid patch
(T, H, W)per item visual. - vision_sequence_indexes — Indeks token visual dalam urutan gabungan.
- vision_mse_loss_indexes — Indeks yang digunakan untuk membaca prediksi visual setelah backbone.
- vision_timesteps — Timesteps difusi per-patch untuk token visual.
- vision_noisy_frame_indexes — Indeks frame berisik per item visual.
- sound_tokens — Tensor laten suara opsional sebelum packing.
- sound_token_shapes — Bentuk grid patch opsional untuk item suara.
- sound_sequence_indexes — Indeks opsional dari token suara dalam urutan gabungan.
- sound_mse_loss_indexes — Indeks opsional yang digunakan untuk membaca prediksi suara.
- sound_timesteps — Timestep difusi per token opsional untuk suara.
- sound_noisy_frame_indexes — Indeks frame berisik opsional per item suara.
- action_tokens — Tensor laten aksi opsional sebelum packing.
- action_token_shapes — Bentuk grid patch opsional
(T, H, W)per item aksi. - action_sequence_indexes — Indeks opsional dari token aksi dalam urutan gabungan.
- action_mse_loss_indexes — Indeks opsional yang digunakan untuk membaca prediksi aksi setelah backbone.
- action_timesteps — Timestep difusi per token opsional untuk token aksi.
- action_noisy_frame_indexes — Indeks frame berisik opsional per item aksi.
- action_domain_ids — ID domain per item opsional untuk memilih bobot kepala aksi.
- return_dict — Apakah akan mengembalikan [Cosmos3OmniTransformerOutput] alih-alih tuple.
Jalankan langkah forward penuh untuk denoising-step.