Hugging Face's logo

Diffusers documentation

CosmosTransformer3DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

CosmosTransformer3DModel

Một mô hình Diffusion Transformer cho dữ liệu dạng video 3D đã được giới thiệu trong Cosmos World Foundation Model Platform for Physical AI của NVIDIA.

Mô hình có thể được tải bằng đoạn mã sau.

from diffusers import CosmosTransformer3DModel

transformer = CosmosTransformer3DModel.from_pretrained("nvidia/Cosmos-1.0-Diffusion-7B-Text2World", subfolder="transformer", dtype=torch.bfloat16)

CosmosTransformer3DModel

lớp diffusers.CosmosTransformer3DModel

< >

( in_channels: int = 16out_channels: int = 16num_attention_heads: int = 32attention_head_dim: int = 128num_layers: int = 28mlp_ratio: float = 4.0text_embed_dim: int = 1024adaln_lora_dim: int = 256max_size: tuple = (128, 240, 240)patch_size: tuple = (1, 2, 2)rope_scale: tuple = (2.0, 1.0, 1.0)concat_padding_mask: bool = Trueextra_pos_embed_type: str | None = 'learnable'use_crossattn_projection: bool = Falsecrossattn_proj_in_channels: int = 1024encoder_hidden_states_channels: int = 1024controlnet_block_every_n: int | None = Noneimg_context_dim_in: int | None = Noneimg_context_num_tokens: int = 256img_context_dim_out: int = 2048 )

Tham số

  • in_channels (int, mặc định là 16) — Số lượng kênh trong đầu vào.
  • out_channels (int, mặc định là 16) — Số lượng kênh trong đầu ra.
  • num_attention_heads (int, mặc định là 32) — Số lượng đầu (heads) được sử dụng cho cơ chế chú ý đa đầu (multi-head attention).
  • attention_head_dim (int, mặc định là 128) — Số lượng kênh trong mỗi đầu chú ý.
  • num_layers (int, mặc định là 28) — Số lượng lớp của các khối transformer cần sử dụng.
  • mlp_ratio (float, mặc định là 4.0) — Tỷ lệ giữa kích thước lớp ẩn và kích thước đầu vào trong mạng feedforward.
  • text_embed_dim (int, mặc định là 4096) — Kích thước đầu vào của các vector nhúng văn bản từ bộ mã hóa văn bản.
  • adaln_lora_dim (int, mặc định là 256) — Chiều ẩn của lớp Adaptive LayerNorm LoRA.
  • max_size (tuple[int, int, int], mặc định là (128, 240, 240)) — Kích thước tối đa của các tensor latent đầu vào theo các chiều thời gian, chiều cao và chiều rộng.
  • patch_size (tuple[int, int, int], mặc định là (1, 2, 2)) — Kích thước patch được sử dụng để patchify các tensor latent đầu vào theo các chiều thời gian, chiều cao và chiều rộng.
  • rope_scale (tuple[float, float, float], mặc định là (2.0, 1.0, 1.0)) — Hệ số tỷ lệ để sử dụng cho RoPE trong các chiều thời gian, chiều cao và chiều rộng.
  • concat_padding_mask (bool, mặc định là True) — Có nên nối padding mask vào các tensor latent đầu vào hay không.
  • extra_pos_embed_type (str, tùy chọn, mặc định là learnable) — Loại nhúng vị trí bổ sung cần sử dụng. Có thể là None hoặc learnable.
  • controlnet_block_every_n (int, tùy chọn) — Khoảng cách giữa các khối transformer cần nhận phần dư điều khiển (ví dụ, 7 để tiêm sau mỗi khối thứ bảy). Bắt buộc đối với Cosmos Transfer2.5.
  • img_context_dim_in (int, tùy chọn) — Kích thước của vector đặc trưng ngữ cảnh ảnh đầu vào, tức là D trong [B, N, D].
  • img_context_num_tokens (int) — Số lượng token trong vector đặc trưng ngữ cảnh của hình ảnh, tức là N trong [B, N, D]. Nếu img_context_dim_in không được cung cấp, tham số này sẽ bị bỏ qua.
  • img_context_dim_out (int) — Kích thước đầu ra của lớp chiếu ngữ cảnh hình ảnh. Nếu img_context_dim_in không được cung cấp, thì tham số này sẽ bị bỏ qua.

Một mô hình Transformer cho dữ liệu dạng video được sử dụng trong Cosmos.

forward

< >

( hidden_states: Tensortimestep: Tensorencoder_hidden_states: Tensorblock_controlnet_hidden_states: list[torch.Tensor] | None = Noneattention_mask: typing.Optional[torch.Tensor] = Nonefps: int | None = Nonecondition_mask: typing.Optional[torch.Tensor] = Nonepadding_mask: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True )

Tham số

  • hidden_states (torch.Tensor có hình dạng (batch_size, num_channels, num_frames, height, width)) — hidden_states đầu vào.
  • timestep (torch.LongTensor) — Dùng để chỉ bước khử nhiễu.
  • encoder_hidden_states (torch.Tensor có hình dạng (batch_size, sequence_len, embed_dims)) — Các embedding điều kiện (embedding được tính từ các điều kiện đầu vào như prompt) cần sử dụng.
  • block_controlnet_hidden_states (list of torch.Tensor, tùy chọn) — Một danh sách các tensor, nếu được chỉ định, sẽ được cộng vào phần dư (residuals) của các khối transformer.
  • attention_mask (torch.Tensor, tùy chọn) — Mặt nạ được áp dụng cho encoder_hidden_states trong quá trình chú ý.
  • fps (int, tùy chọn) — Số khung hình mỗi giây của video đầu vào được sử dụng để tính các nhúng vị trí quay (rotary positional embeddings).
  • condition_mask (torch.Tensor, tùy chọn) — Kênh mask được nối vào hidden_states để chỉ định vùng điều kiện.
  • padding_mask (torch.Tensor, tùy chọn) — Mặt nạ đệm được nối vào hidden_states khi concat_padding_mask được bật.
  • return_dict (bool, tùy chọn, mặc định là True) — Có nên trả về một ~models.transformer_2d.Transformer2DModelOutput thay vì một tuple thông thường hay không.

Phương thức forward của CosmosTransformer3DModel.

Transformer2DModelOutput

lớp diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Tham số

  • sample (torch.Tensor có hình dạng (batch_size, num_channels, height, width) hoặc (batch size, num_vector_embeds - 1, num_latent_pixels) nếu Transformer2DModel là rời rạc) — Trạng thái ẩn đầu ra được điều kiện bởi đầu vào encoder_hidden_states. Nếu là rời rạc, trả về các phân phối xác suất cho các pixel tiềm ẩn không nhiễu.

Đầu ra của Transformer2DModel.

Update on GitHub

←Cosmos3OmniTransformer DiTTransformer2DModel→