Diffusers documentation
Cosmos3OmniTransformer
Cosmos3OmniTransformer
Một transformer thị giác-ngôn ngữ kết hợp dạng Mixture-of-Transformer (MoT), được giới thiệu như một phần của dòng mô hình nền tảng thế giới Cosmos3 của NVIDIA. Mô hình chạy hai đường dẫn tính toán song song trên một chuỗi kết hợp đã đóng gói:
- một đường dẫn “hiểu” nhân quả tự chú ý trên các token văn bản với mặt nạ nhân quả, và
- một đường dẫn “sinh” hai chiều thực hiện cross-attention từ các token sinh (latent thị giác + latent âm thanh tùy chọn) trên toàn bộ tập key/value gồm cả phần hiểu và phần sinh.
Hai nhánh này chia sẻ cùng kích thước ẩn và số lớp nhưng duy trì các phép chiếu Q/K/V/O, MLP và tham số RMSNorm riêng biệt, đây chính là điều khiến kiến trúc trở thành Mixture-of-Transformer thay vì Mixture-of-Experts thông thường. Thông tin vị trí được cung cấp thông qua RoPE đa phương thức 3D (mRoPE), xen kẽ các tần số thời gian / chiều cao / chiều rộng cho các latent video và tái sử dụng trục thời gian cho văn bản và âm thanh.
Mô hình có thể được nạp theo cách sau.
import torch
from diffusers import Cosmos3OmniTransformer
transformer = Cosmos3OmniTransformer.from_pretrained(
"nvidia/Cosmos3-Nano", subfolder="transformer", dtype=torch.bfloat16
)Cosmos3OmniTransformer
lớp 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 )
Tham số
- input_ids — Các ID token văn bản được đặt tại
text_indexestrong chuỗi kết hợp. - text_indexes — Chỉ số của các token văn bản trong chuỗi kết hợp.
- position_ids —
[3, sequence_length]các ID vị trí mRoPE cho toàn bộ chuỗi kết hợp. - und_len — Độ dài của phần tiền tố văn bản nhân quả (hiểu); các token sinh ra sẽ đi theo sau.
- sequence_length — Tổng độ dài của chuỗi đóng gói kết hợp.
- vision_tokens — Các tensor latent thị giác cho từng mục trước khi patchify.
- vision_token_shapes — Hình dạng lưới patch
(T, H, W)cho từng mục thị giác. - vision_sequence_indexes — Chỉ số của các token thị giác trong chuỗi kết hợp.
- vision_mse_loss_indexes — Các chỉ số được dùng để đọc các dự đoán thị giác sau backbone.
- vision_timesteps — Các bước khuếch tán (diffusion timesteps) cho từng patch của các token thị giác.
- vision_noisy_frame_indexes — Chỉ số các khung hình nhiễu cho mỗi mục thị giác.
- sound_tokens — Tensor âm thanh tiềm ẩn tùy chọn trước khi đóng gói.
- sound_token_shapes — Hình dạng lưới patch tùy chọn cho các mục âm thanh.
- sound_sequence_indexes — Các chỉ số tùy chọn của các token âm thanh trong chuỗi kết hợp.
- sound_mse_loss_indexes — Các chỉ số tùy chọn được sử dụng để đọc các dự đoán âm thanh.
- sound_timesteps — Các bước khuếch tán tùy chọn cho từng token của âm thanh.
- sound_noisy_frame_indexes — Chỉ số các khung hình nhiễu tùy chọn cho từng mục âm thanh.
- action_tokens — Các tensor ẩn hành động tùy chọn trước khi đóng gói.
- action_token_shapes — Hình dạng lưới patch tùy chọn
(T, H, W)cho từng mục hành động. - action_sequence_indexes — Các chỉ số tùy chọn của các token hành động trong chuỗi kết hợp.
- action_mse_loss_indexes — Các chỉ số tùy chọn được dùng để đọc các dự đoán hành động sau backbone.
- action_timesteps — Các bước khuếch tán (diffusion timesteps) tùy chọn cho từng token của các token hành động.
- action_noisy_frame_indexes — Chỉ số các khung hình nhiễu tùy chọn cho từng mục hành động.
- action_domain_ids — ID miền tùy chọn cho từng mục, dùng để chọn trọng số của đầu action.
- return_dict — Xác định việc trả về [Cosmos3OmniTransformerOutput] thay vì một tuple.
Chạy một bước truyền thẳng đầy đủ cho quá trình khử nhiễu.