Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
NVIDIA의 Cosmos3 월드 파운데이션 모델 계열의 일부로 소개된 Mixture-of-Transformer (MoT) 기반의 결합형 비전-언어 transformer입니다. 이 모델은 패킹된 결합 시퀀스 위에서 두 개의 병렬 계산 경로를 실행합니다:
- causal masking을 적용하여 텍스트 토큰에 대해 self-attention을 수행하는 causal “understanding” pathway
- 생성 토큰(비전 + 선택적 사운드 잠재 변수)에서 이해 및 생성 전체 키/값 집합에 대해 교차 주의를 수행하는 양방향 “생성” 경로
두 경로는 동일한 hidden size와 레이어 수를 공유하지만 별도의 Q/K/V/O 투영, MLP, RMSNorm 파라미터를 유지하며, 이것이 아키텍처를 표준 Mixture-of-Experts가 아닌 Mixture-of-Transformer로 만드는 이유입니다. 위치 정보는 3D 멀티모달 RoPE(mRoPE)를 통해 제공되며, 이는 비디오 latent에 대해 시간/높이/너비 주파수를 교차 배치하고 텍스트 및 오디오에는 시간 축을 재사용합니다.
모델은 다음과 같이 로드할 수 있습니다.
import torch
from diffusers import Cosmos3OmniTransformer
transformer = Cosmos3OmniTransformer.from_pretrained(
"nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)Cosmos3OmniTransformer
클래스 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위치에 배치된 텍스트 토큰 ID입니다. - text_indexes — 결합된 시퀀스에서 텍스트 토큰의 인덱스.
- position_ids — 전체 결합 시퀀스에 대한 mRoPE 위치 ID,
[3, sequence_length]. - und_len — 인과적 텍스트(이해) 접두사의 길이; 생성 토큰이 뒤에 따라옵니다.
- sequence_length — 결합된 패킹 시퀀스의 총 길이.
- vision_tokens — patchify 이전의 항목별 vision latent 텐서.
- vision_token_shapes — 각 vision 항목에 대한 패치 그리드 모양
(T, H, W). - vision_sequence_indexes — 결합된 시퀀스 내 비전 토큰의 인덱스.
- vision_mse_loss_indexes — 백본 이후 비전 예측을 읽는 데 사용되는 인덱스입니다.
- vision_timesteps — 시각 토큰의 패치별 diffusion timesteps입니다.
- vision_noisy_frame_indexes — 각 vision 항목에 대한 노이즈 프레임 인덱스.
- sound_tokens — 패킹 전 선택적 사운드 잠재 텐서.
- sound_token_shapes — 사운드 항목에 대한 선택적 패치 그리드 모양입니다.
- sound_sequence_indexes — 결합된 시퀀스에서 사운드 토큰의 선택적 인덱스.
- sound_mse_loss_indexes — 사운드 예측을 읽는 데 사용되는 선택적 인덱스입니다.
- sound_timesteps — 사운드에 대한 선택적 토큰별 diffusion timesteps입니다.
- sound_noisy_frame_indexes — 사운드 항목별 선택적 노이즈 프레임 인덱스.
- action_tokens — 패킹 전 선택적 액션 잠재 텐서.
- action_token_shapes — 각 action item에 대한 선택적 patch grid shape
(T, H, W). - action_sequence_indexes — 결합된 시퀀스에서 액션 토큰의 선택적 인덱스입니다.
- action_mse_loss_indexes — 백본 이후에 액션 예측을 읽는 데 사용되는 선택적 인덱스입니다.
- action_timesteps — action 토큰에 대한 선택적 per-token diffusion timesteps입니다.
- action_noisy_frame_indexes — 각 action 항목에 대한 선택적 노이즈 프레임 인덱스.
- action_domain_ids — 각 항목별로 action head 가중치를 선택하는 선택적 도메인 ID입니다.
- return_dict — 튜플 대신 [Cosmos3OmniTransformerOutput]를 반환할지 여부.
전체 노이즈 제거 단계의 순전파를 실행합니다.