# CosmosTransformer3DModel

NVIDIA 在 [Cosmos World Foundation Model Platform for Physical AI](https://huggingface.co/papers/2501.03575) 中介绍了一种用于 3D 类视频数据的 Diffusion Transformer 模型。

可以使用以下代码片段加载该模型。

```python
from diffusers import CosmosTransformer3DModel

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

## CosmosTransformer3DModel[[diffusers.CosmosTransformer3DModel]]

#### diffusers.CosmosTransformer3DModel[[diffusers.CosmosTransformer3DModel]]

```python
diffusers.CosmosTransformer3DModel(in_channels: int = 16, out_channels: int = 16, num_attention_heads: int = 32, attention_head_dim: int = 128, num_layers: int = 28, mlp_ratio: float = 4.0, text_embed_dim: int = 1024, adaln_lora_dim: int = 256, max_size: tuple[int, int, int] = (128, 240, 240), patch_size: tuple[int, int, int] = (1, 2, 2), rope_scale: tuple[float, float, float] = (2.0, 1.0, 1.0), concat_padding_mask: bool = True, extra_pos_embed_type: str | None = 'learnable', use_crossattn_projection: bool = False, crossattn_proj_in_channels: int = 1024, encoder_hidden_states_channels: int = 1024, controlnet_block_every_n: int | None = None, img_context_dim_in: int | None = None, img_context_num_tokens: int = 256, img_context_dim_out: int = 2048)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_cosmos.py#L554)

**参数：**

in_channels (`int`，默认为 `16`)：输入中的通道数。

out_channels (`int`，默认为 `16`)：输出的通道数。

num_attention_heads (`int`，默认为 `32`)：用于多头注意力的注意力头数量。

attention_head_dim (`int`，默认为 `128`)：每个注意力头中的通道数。

num_layers (`int`，默认为 `28`)：要使用的 transformer 块层数。

mlp_ratio (`float`，默认为 `4.0`)：前馈网络中隐藏层大小与输入大小的比率。

text_embed_dim (`int`，默认为 `4096`)：文本编码器输出的文本嵌入的输入维度。

adaln_lora_dim (`int`，默认为 `256`)：Adaptive LayerNorm LoRA 层的隐藏维度。

max_size (`tuple[int, int, int]`，默认为 `(128, 240, 240)`)：输入潜在张量在时间、高度和宽度维度上的最大尺寸。

patch_size (`tuple[int, int, int]`，默认为 `(1, 2, 2)`)：用于在时间、高度和宽度维度上对输入潜在张量进行分块（patchify）的块大小。

rope_scale (`tuple[float, float, float]`，默认为 `(2.0, 1.0, 1.0)`)：用于 RoPE 在时间、高度和宽度维度上的缩放因子。

concat_padding_mask (`bool`，默认为 `True`)：是否将填充掩码拼接到输入潜在张量中。

extra_pos_embed_type (`str`，*可选*，默认为 `learnable`)：要使用的额外位置嵌入类型。可以是 `None` 或 `learnable` 之一。

controlnet_block_every_n (`int`, *可选*) : 应接收控制残差的 Transformer 块之间的间隔（例如，`7` 表示在每第七个块之后注入）。Cosmos Transfer2.5 必需。

img_context_dim_in (`int`, *可选*) : 输入图像上下文特征向量的维度，即 [B, N, D] 中的 D。

img_context_num_tokens (`int`) : 图像上下文特征向量中的 token 数量，即 [B, N, D] 中的 N。如果未提供 `img_context_dim_in`，则忽略此参数。

img_context_dim_out (`int`) : 图像上下文投影层的输出维度。如果未提供 `img_context_dim_in`，则忽略此参数。

用于 [Cosmos](https://github.com/NVIDIA/Cosmos) 的类视频数据 Transformer 模型。

#### forward[[diffusers.CosmosTransformer3DModel.forward]]

```python
forward(hidden_states: Tensor, timestep: Tensor, encoder_hidden_states: Tensor, block_controlnet_hidden_states: list[torch.Tensor] | None = None, attention_mask: typing.Optional[torch.Tensor] = None, fps: int | None = None, condition_mask: typing.Optional[torch.Tensor] = None, padding_mask: typing.Optional[torch.Tensor] = None, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_cosmos.py#L688)

**参数：**

hidden_states (`torch.Tensor`，形状为 `(batch_size, num_channels, num_frames, height, width)`) : 输入 `hidden_states`。

timestep (`torch.LongTensor`)：用于指示去噪步骤。

encoder_hidden_states (`torch.Tensor`，形状为 `(batch_size, sequence_len, embed_dims)`)：要使用的条件嵌入（从输入条件（例如提示词）计算得出的嵌入）。

block_controlnet_hidden_states (`list` of `torch.Tensor`, *可选*) : 一个张量列表，如果指定，将被添加到 Transformer 块的残差中。

attention_mask (`torch.Tensor`, *可选*) : 在注意力机制中应用于 `encoder_hidden_states` 的掩码。

fps (`int`, *可选*) : 用于计算旋转位置嵌入的输入视频的每秒帧数。

condition_mask (`torch.Tensor`，*可选*)：与 `hidden_states` 拼接的掩码通道，用于指示条件区域。

padding_mask (`torch.Tensor`, *可选*) : 当启用 `concat_padding_mask` 时，与 `hidden_states` 拼接的填充掩码。

return_dict (`bool`，*可选*，默认为 `True`)：是否返回 `~models.transformer_2d.Transformer2DModelOutput` 而不是普通元组。

**返回值：**

如果 `return_dict` 为 True，则返回 `~models.transformer_2d.Transformer2DModelOutput`，否则返回一个 `tuple`，其中第一个元素是样本张量。

[CosmosTransformer3DModel](/docs/diffusers/v0.43.0/zh/api/models/cosmos_transformer3d#diffusers.CosmosTransformer3DModel) 的前向方法。

## Transformer2DModelOutput[[diffusers.models.modeling_outputs.Transformer2DModelOutput]]

#### diffusers.models.modeling_outputs.Transformer2DModelOutput[[diffusers.models.modeling_outputs.Transformer2DModelOutput]]

```python
diffusers.models.modeling_outputs.Transformer2DModelOutput(sample: torch.Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/modeling_outputs.py#L21)

**参数：**

sample (`torch.Tensor`，形状为 `(batch_size, num_channels, height, width)`；若 [Transformer2DModel](/docs/diffusers/v0.43.0/zh/api/models/transformer2d#diffusers.Transformer2DModel) 为离散模式，则为 `(batch size, num_vector_embeds - 1, num_latent_pixels)`)：以 `encoder_hidden_states` 输入为条件的隐藏状态输出。若为离散模式，则返回未加噪潜在像素的概率分布。

[Transformer2DModel](/docs/diffusers/v0.43.0/zh/api/models/transformer2d#diffusers.Transformer2DModel) 的输出。