# 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）提供，该机制为视频潜变量交错排列时间 / 高度 / 宽度频率，并为文本和音频复用时间轴。

可以按以下方式加载模型。

```python
import torch
from diffusers import Cosmos3OmniTransformer

transformer = Cosmos3OmniTransformer.from_pretrained(
    "nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)
```

## Cosmos3OmniTransformer[[diffusers.Cosmos3OmniTransformer]]

#### diffusers.Cosmos3OmniTransformer[[diffusers.Cosmos3OmniTransformer]]

```python
diffusers.Cosmos3OmniTransformer(attention_bias: bool = False, attention_dropout: float = 0.0, dtype: str = 'bfloat16', head_dim: int = 128, hidden_size: int = 4096, intermediate_size: int = 12288, base_fps: int = 24, enable_fps_modulation: bool = True, latent_channel: int = 48, unified_3d_mrope_reset_spatial_ids: bool = True, unified_3d_mrope_temporal_modality_margin: int = 15000, latent_patch_size: int = 2, num_attention_heads: int = 32, num_hidden_layers: int = 36, num_key_value_heads: int = 8, patch_latent_dim: int = 192, rms_norm_eps: float = 1e-06, rope_scaling: dict | None = None, rope_theta: float = 5000000.0, action_dim: int | None = None, action_gen: bool = False, num_embodiment_domains: int = 32, sound_dim: int | None = None, sound_gen: bool = False, sound_latent_fps: float = 25.0, timestep_scale: float = 0.001, vocab_size: int = 151936, hidden_act: str = 'silu', qk_norm_for_text: bool = True, use_und_k_norm_for_gen: bool = False, rope_axes_dim: tuple[int, int, int] | list[int] | None = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_cosmos3.py#L374)

#### forward[[diffusers.Cosmos3OmniTransformer.forward]]

```python
forward(input_ids: Tensor, text_indexes: Tensor, position_ids: Tensor, und_len: int, sequence_length: int, vision_tokens: list[torch.Tensor], vision_token_shapes: list[tuple[int, int, int]], vision_sequence_indexes: Tensor, vision_mse_loss_indexes: Tensor, vision_timesteps: Tensor, vision_noisy_frame_indexes: list[torch.Tensor], sound_tokens: list[torch.Tensor] | None = None, sound_token_shapes: list[tuple[int, int, int]] | None = None, sound_sequence_indexes: typing.Optional[torch.Tensor] = None, sound_mse_loss_indexes: typing.Optional[torch.Tensor] = None, sound_timesteps: typing.Optional[torch.Tensor] = None, sound_noisy_frame_indexes: list[torch.Tensor] | None = None, action_tokens: list[torch.Tensor] | None = None, action_token_shapes: list[tuple[int, int, int]] | None = None, action_sequence_indexes: typing.Optional[torch.Tensor] = None, action_mse_loss_indexes: typing.Optional[torch.Tensor] = None, action_timesteps: typing.Optional[torch.Tensor] = None, action_noisy_frame_indexes: list[torch.Tensor] | None = None, action_domain_ids: list[torch.Tensor] | None = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_cosmos3.py#L671)

**参数：**

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 : 每个视觉项的 `(T, H, W)` 补丁网格形状。

vision_sequence_indexes : 联合序列中视觉令牌的索引。

vision_mse_loss_indexes : 用于在主干网络之后读取视觉预测的索引。

vision_timesteps : 用于视觉 token 的逐 patch 扩散时间步。

vision_noisy_frame_indexes : 每个视觉项的噪声帧索引。

sound_tokens : 可选的打包前声音潜在张量。

sound_token_shapes : 声音项的可选补丁网格形状。

sound_sequence_indexes : 联合序列中声音令牌的可选索引。

sound_mse_loss_indexes : 用于读取声音预测的可选索引。

sound_timesteps : 用于声音的可选逐 token 扩散时间步。

sound_noisy_frame_indexes : 每个声音项的可选噪声帧索引。

action_tokens : 打包前的可选动作潜在张量。

action_token_shapes : 每个操作项的可选 patch 网格形状 `(T, H, W)`。

action_sequence_indexes : 联合序列中动作令牌的可选索引。

action_mse_loss_indexes : 用于在主干网络之后读取动作预测的可选索引。

action_timesteps : 可选的，针对动作 token 的逐 token 扩散时间步。

action_noisy_frame_indexes : 每个操作项的可选噪声帧索引。

action_domain_ids : 可选的逐项域 ID，用于选择动作头权重。

return_dict : 是否返回 [*Cosmos3OmniTransformerOutput*] 而不是元组。

**返回值：**

一个 [*Cosmos3OmniTransformerOutput*] 或一个包含各模态预测列表的元组。对于可选模态，如果省略其输入，则返回 `None`。

执行完整的去噪步骤前向传播。