Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
Совместный зрительно-языковой трансформер Mixture-of-Transformer (MoT), представленный в рамках семейства мировых фундаментальных моделей Cosmos3 от NVIDIA. Модель выполняет два параллельных вычислительных потока по упакованной совместной последовательности:
- a каузальный «понимающий» путь, который выполняет само-внимание по текстовым токенам с каузальной маскировкой, и
- a bi-directional “generation” pathway that cross-attends from generation tokens (vision + optional sound latents) over the full understanding-plus-generation key/value set.
Оба пути имеют одинаковый размер скрытого слоя и количество слоёв, но поддерживают отдельные проекции Q/K/V/O, MLP и параметры RMSNorm, что и делает архитектуру 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
класс 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в объединённой последовательности. - text_indexes — Индексы текстовых токенов в объединённой последовательности.
- position_ids — идентификаторы позиций mRoPE для полной объединённой последовательности,
[3, sequence_length]. - und_len — Длина префикса каузального текста (понимание); за ним следуют токены генерации.
- sequence_length — Общая длина объединённой упакованной последовательности.
- vision_tokens — тензоры визуальных латентов для каждого элемента до применения patchify.
- 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 — Необязательные индексы, используемые для чтения предсказаний действий после backbone.
- action_timesteps — Необязательные временные шаги диффузии для каждого токена действия.
- action_noisy_frame_indexes — Необязательные индексы шумных кадров для каждого элемента действия.
- action_domain_ids — Необязательные идентификаторы доменов для каждого элемента, определяющие веса головок действий.
- return_dict — Возвращать ли [Cosmos3OmniTransformerOutput] вместо кортежа.
Выполните полный прямой проход на шаге денойзинга.