Diffusers 文档
CosmosTransformer3DModel
加入 Hugging Face 社区
并获得增强的文档体验
开始使用
CosmosTransformer3DModel
NVIDIA 在 Cosmos 世界基础模型平台物理 AI 中引入了用于 3D 视频状数据的扩散 Transformer 模型。
该模型可以通过以下代码片段加载。
from diffusers import CosmosTransformer3DModel
transformer = CosmosTransformer3DModel.from_pretrained("nvidia/Cosmos-1.0-Diffusion-7B-Text2World", subfolder="transformer", torch_dtype=torch.bfloat16)CosmosTransformer3DModel
class 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: typing.Tuple[int, int, int] = (128, 240, 240) patch_size: typing.Tuple[int, int, int] = (1, 2, 2) rope_scale: typing.Tuple[float, float, float] = (2.0, 1.0, 1.0) concat_padding_mask: bool = True extra_pos_embed_type: typing.Optional[str] = 'learnable' )
参数
- 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)) — 用于在时间、高度和宽度维度上对输入潜在张量进行分块的补丁大小。 - 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。
在 Cosmos 中使用的视频状数据的 Transformer 模型。
Transformer2DModelOutput
class diffusers.models.modeling_outputs.Transformer2DModelOutput
< 源 >( 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 的输出。