Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
NVIDIAのCosmos3ワールドファウンデーションモデルファミリーの一部として導入された、Mixture-of-Transformer (MoT) によるビジョン・言語統合トランスフォーマーです。このモデルは、パッキングされた統合シーケンスに対して2つの並列計算パスを実行します:
- 因果的な「理解」パスである causal “understanding” pathway は、因果的マスキングを適用してテキストトークンに対して自己注意を適用し、
- 双方向の「生成」パスであり、生成トークン（ビジョン＋任意のサウンド潜在変数）から、理解と生成のキー/バリューセット全体に対してクロスアテンションを行います。
2つのパスウェイは同じ隠れ層サイズと層数を共有しますが、個別の Q/K/V/O 投影、MLP、および RMSNorm パラメータを保持しています。これが、このアーキテクチャを標準的な Mixture-of-Experts ではなく Mixture-of-Transformer とする理由です。位置情報は、3D マルチモーダル RoPE (mRoPE) によって供給されます。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: 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 )
パラメータ
- input_ids — 結合シーケンス内の
text_indexesに配置されるテキストのトークン ID。 - text_indexes — 結合シーケンスにおけるテキストトークンのインデックス。
- position_ids — 完全な結合シーケンスに対する mRoPE の位置 ID。
[3, sequence_length] - und_len — 因果テキスト（理解）プレフィックスの長さ。生成トークンはこれに続きます。
- sequence_length — 連結パックシーケンスの総長。
- vision_tokens — パッチ化前のアイテムごとのビジョン潜在テンソル。
- vision_token_shapes — 各ビジョンアイテムごとのパッチグリッドの形状
(T, H, W)。 - vision_sequence_indexes — 結合シーケンス内のビジョントークンのインデックス。
- vision_mse_loss_indexes — バックボーン後にビジュアル予測を読み取るために使用されるインデックスです。
- vision_timesteps — ビジョントークン用のパッチごとの拡散タイムステップ。
- vision_noisy_frame_indexes — 各ビジョンアイテムごとのノイズフレームのインデックス。
- sound_tokens — パッキング前の任意のサウンド潜在テンソル。
- sound_token_shapes — サウンドアイテム用の任意のパッチグリッド形状。
- sound_sequence_indexes — 結合シーケンス内のサウンドトークンの任意のインデックス。
- sound_mse_loss_indexes — 音声の予測を読み取るために使用する任意のインデックス。
- sound_timesteps — 音声用の、トークンごとの任意の拡散タイムステップ。
- sound_noisy_frame_indexes — サウンドアイテムごとの任意のノイズフレームインデックス。
- action_tokens — パッキング前の任意のアクション潜在テンソル。
- action_token_shapes — 各アクション項目ごとの任意のパッチグリッド形状
(T, H, W)。 - action_sequence_indexes — 結合シーケンス内のアクショントークンの任意のインデックス。
- action_mse_loss_indexes — バックボーン後にアクション予測を読み取るために使用する任意のインデックス。
- action_timesteps — アクショントークン用の、トークンごとの任意の拡散タイムステップ。
- action_noisy_frame_indexes — アクション項目ごとの任意のノイズ付きフレームのインデックス。
- action_domain_ids — 各アイテムのアクションヘッドの重みを決定するための、任意のドメイン ID。
- return_dict — タプルではなく [Cosmos3OmniTransformerOutput] を返すかどうか。
完全なデノイジングステップのフォワードパスを実行します。