Diffusers documentation
CosmosTransformer3DModel
CosmosTransformer3DModel
NVIDIA 在 Cosmos World Foundation Model Platform for Physical AI 中介绍了一种用于 3D 类视频数据的 Diffusion Transformer 模型。
可以使用以下代码片段加载该模型。
from diffusers import CosmosTransformer3DModel
transformer = CosmosTransformer3DModel.from_pretrained("nvidia/Cosmos-1.0-Diffusion-7B-Text2World", subfolder="transformer", dtype=torch.bfloat16)CosmosTransformer3DModel
class diffusers.CosmosTransformer3DModel
< source >( 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[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 = 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 )
参数
- 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) — 自适应 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 的类视频数据 Transformer 模型。
forward
< source >( 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 )
参数
- hidden_states（形状为
(batch_size, num_channels, num_frames, height, width)的torch.Tensor）— 输入的hidden_states。 - timestep (
torch.LongTensor) — 用于指示去噪步骤。 - encoder_hidden_states (
torch.Tensor，形状为(batch_size, sequence_len, embed_dims)) — 要使用的条件嵌入（从输入条件（例如提示词）计算得出的嵌入）。 - block_controlnet_hidden_states (
listoftorch.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而不是普通的元组。
CosmosTransformer3DModel 的前向方法。
Transformer2DModelOutput
类 diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
参数
- sample（形状为
(batch_size, num_channels, height, width)的torch.Tensor，或者如果 Transformer2DModel 是离散的，则为(batch size, num_vector_embeds - 1, num_latent_pixels)）— 以encoder_hidden_states输入为条件的隐藏状态输出。如果是离散的，则返回未加噪潜在像素的概率分布。
Transformer2DModel 的输出。