Diffusers

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

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: 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 )

Parameter

input_ids — ID token teks yang ditempatkan di text_indexes dalam 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.

Update on GitHub

←ConsisIDTransformer3DModel CosmosTransformer3DModel→