Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
NVIDIA Cosmos3 世界基础模型系列中引入的一种 Mixture-of-Transformer (MoT) 联合视觉-语言 Transformer。该模型在打包的联合序列上运行两条并行的计算路径：
- 一条因果“理解”路径，它对文本 token 执行带有因果掩码的自注意力，并且
- 一个双向“生成”路径，它从生成 token（视觉 + 可选的音频潜变量）出发，对完整的理解加生成键/值集合执行交叉注意力。
这两条路径共享相同的隐藏层大小和层数，但各自维护独立的 Q/K/V/O 投影、MLP 以及 RMSNorm 参数，这正是该架构被称为 Transformer 混合体（Mixture-of-Transformer）而非标准专家混合体（Mixture-of-Experts）的原因。位置信息通过 3D 多模态 RoPE（mRoPE）提供，该机制为视频潜变量交错排列时间 / 高度 / 宽度频率，并为文本和音频复用时间轴。
可以按以下方式加载模型。
import torch
from diffusers import Cosmos3OmniTransformer
transformer = Cosmos3OmniTransformer.from_pretrained(
"nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)Cosmos3OmniTransformer
class 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 )
参数
- input_ids — 放置在联合序列中
text_indexes位置的文本 token ID。 - text_indexes — 联合序列中文本令牌的索引。
- position_ids —
[3, sequence_length]完整联合序列的 mRoPE 位置 ID。 - und_len — 因果文本（理解）前缀的长度；生成 token 紧随其后。
- sequence_length — 联合打包序列的总长度。
- vision_tokens — 在 patchify 之前，每个样本的视觉潜在张量。
- vision_token_shapes — 每个视觉项的 Patch 网格形状
(T, H, W)。 - vision_sequence_indexes — 联合序列中视觉令牌的索引。
- vision_mse_loss_indexes — 用于在主干网络之后读取视觉预测的索引。
- vision_timesteps — 视觉 token 的逐 patch 扩散时间步。
- vision_noisy_frame_indexes — 每个视觉项的含噪帧索引。
- sound_tokens — 打包前的可选声音潜在张量。
- sound_token_shapes — 声音项的可选 patch 网格形状。
- sound_sequence_indexes — 联合序列中声音令牌的可选索引。
- sound_mse_loss_indexes — 用于读取声音预测的可选索引。
- sound_timesteps — 可选的，用于声音的逐 token 扩散时间步。
- sound_noisy_frame_indexes — 每个声音项的可选噪声帧索引。
- action_tokens — 打包前的可选动作潜在张量。
- action_token_shapes — 每个 action item 的可选 patch 网格形状
(T, H, W)。 - action_sequence_indexes — 联合序列中动作令牌的可选索引。
- action_mse_loss_indexes — 可选的索引，用于在主干网络之后读取动作预测。
- action_timesteps — 可选的，用于 action token 的逐 token 扩散时间步。
- action_noisy_frame_indexes — 每个 action 项的可选噪声帧索引。
- action_domain_ids — 可选的逐项域 ID，用于选择动作头权重。
- return_dict — 是否返回 [Cosmos3OmniTransformerOutput] 而不是元组。
执行完整的去噪步骤前向传播。