ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Diffusers 中的 CosmosTransformer3DModel:NVIDIA Cosmos 3D 视频扩散 Transformer 架构与实战

Diffusers 中的 CosmosTransformer3DModel:NVIDIA Cosmos 3D 视频扩散 Transformer 架构与实战 Diffusers 中的 CosmosTransformer3DModelNVIDIA Cosmos 3D 视频扩散 Transformer 架构与实战【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本文围绕 Hugging Face Diffusers 仓库中CosmosTransformer3DModel这一核心组件展开系统讲解其在 NVIDIA Cosmos「世界基础模型平台」World Foundation Model Platform for Physical AI论文 Cosmos World Foundation Model Platform for Physical AI中的应用如何加载预训练权重、理解其 3D Patch Embedding / 3D 旋转位置编码 / AdaIN 自适应归一化 / 图像上下文双流注意力等关键设计并给出结合CosmosTextToWorldPipeline的文本生成世界Text2World完整调用链路。读完本文你将掌握CosmosTransformer3DModel的配置项语义、前向接口、在仓库中的配套管线与测试验证方式能够独立完成模型加载、参数调优与视频生成任务。模型概述面向 3D 视频类数据的扩散 TransformerCosmosTransformer3DModel是一个用于处理3D 视频类数据张量形状为(batch_size, num_channels, num_frames, height, width)即多一维时间轴的 Diffusion TransformerDiT模型。它由 NVIDIA 在 Cosmos 世界基础模型平台中提出是 Cosmos 系列文生视频 / 视频生世界Video2World管线中的去噪骨干网络。在仓库中模型定义位于 src/diffusers/models/transformers/transformer_cosmos.py并从diffusers顶层命名空间导出见 src/diffusers/init.py 中的CosmosTransformer3DModel注册。该模型继承自ModelMixin、ConfigMixin、FromOriginalModelMixin与PeftAdapterMixin因此天然支持from_pretrained/save_pretrained标准权重与配置管理从原始 Cosmos 检查点直接加载FromOriginalModelMixinPEFT LoRA 适配器挂载PeftAdapterMixin。快速加载示例官方 API 文档给出的最小加载方式如下docs/source/en/api/models/cosmos_transformer3d.mdfrom diffusers import CosmosTransformer3DModel import torch transformer CosmosTransformer3DModel.from_pretrained( nvidia/Cosmos-1.0-Diffusion-7B-Text2World, subfoldertransformer, dtypetorch.bfloat16, )要点说明subfoldertransformerCosmos 仓库卡repo是多组件复合仓库transformer 权重单独存放在transformer子目录下dtypetorch.bfloat167B 规模模型建议以 bf16 精度加载以降低显存占用加载后可继续调用transformer.to(cuda)、transformer.enable_gradient_checkpointing()源码中_supports_gradient_checkpointing True等标准方法。构造参数详解从默认值看 Cosmos 7B 的模型骨架从 transformer_cosmos.py 的__init__签名可以看到全部可配置参数。默认配置对应 Cosmos-1.0-Diffusion-7B-Text2World 的实际结构参数默认值含义in_channels16输入 latent 的通道数Cosmos 的 VAE 压缩后为 16 通道out_channels16输出 latent 通道数与in_channels一致num_attention_heads32多头注意力的头数attention_head_dim128每个注意力头的通道数hidden_size 32 × 128 4096num_layers28Transformer 块堆叠层数mlp_ratio4.0FFN 隐藏层维度与输入维度的比值即 FFN 维度为4 × 4096text_embed_dim1024文本编码器T5输出向量的维度用于交叉注意力adaln_lora_dim256Adaptive LayerNorm 中间隐藏维度对时间步条件做降维投影max_size(128, 240, 240)latent 在 (时间, 高, 宽) 三个维度的最大支持尺寸patch_size(1, 2, 2)三个维度的 patch 化尺寸时间不压缩、空间 2×2rope_scale(2.0, 1.0, 1.0)时间/高/宽三个维度的 RoPE 缩放系数时间维度 NTK 外推 2 倍concat_padding_maskTrue是否将 padding mask 作为额外通道拼接到输入 latentextra_pos_embed_typelearnable额外位置编码类型可为None或learnableuse_crossattn_projectionFalse是否对文本特征做交叉注意力前投影Linear GELUcontrolnet_block_every_nNone每隔多少层接收 ControlNet 残差如7表示每 7 层注入一次用于 Cosmos Transfer2.5img_context_dim_inNone图像上下文特征向量输入维度[B, N, D]中的 D启用图像上下文双流注意力img_context_num_tokens256图像上下文 token 数量[B, N, D]中的 N未提供img_context_dim_in时忽略img_context_dim_out2048图像上下文投影层输出维度源码中还有几个派生自上述参数的内部模块属性_no_split_modules [CosmosTransformerBlock]、_keep_in_fp32_modules [learnable_pos_embed]分别服务于设备分片/FSDP 的模块拆分以及数值稳定性可学习位置编码保持 fp32。这也解释了为什么测试套件中learnable_pos_embed会被单独保持在 fp32。架构拆解六大组成模块的前向数据流CosmosTransformer3DModel.forwardtransformer_cosmos.py的输入输出签名如下def forward( self, hidden_states: torch.Tensor, # (B, C, T, H, W) 视频 latent timestep: torch.Tensor, # 去噪步支持 [B] 或 [B, 1, T, 1, 1] encoder_hidden_states: torch.Tensor, # (B, S, D) 文本/条件 embedding block_controlnet_hidden_states: list[torch.Tensor] | None None, # ControlNet 逐块残差 attention_mask: torch.Tensor | None None, # 文本注意力掩码 fps: int | None None, # 视频帧率用于 RoPE 时间轴缩放 condition_mask: torch.Tensor | None None, # 条件区域掩码通道 padding_mask: torch.Tensor | None None, # padding 掩码 return_dict: bool True, # 返回 Transformer2DModelOutput 或 tuple ) - tuple[torch.Tensor] | Transformer2DModelOutput:前向过程分 8 步拼接掩码 → 生成位置编码 → patchify → 时间步嵌入 → 处理文本/图像上下文 → 构建 ControlNet 块索引映射 → 28 层 Transformer 块 → 输出归一化与反 patchify。下面按模块逐一说明。1. CosmosPatchEmbed3D 非重叠 patch 化CosmosPatchEmbedtransformer_cosmos.py用单个线性层完成时间×高×宽的 3D patch 化将(B, C, T, H, W)张量按(p_t, p_h, p_w)切块后把每个 patch 内所有元素拼成向量做Linear投影。默认patch_size(1, 2, 2)意味着时间维度不压缩仅对空间做 2×2 切块因此 latent 序列长度为T × (H/2) × (W/2)。注意源码中patch_embed_in_channels in_channels 1 if concat_padding_mask else in_channels——当concat_padding_maskTrue时padding mask 会先被缩放到 latent 空间尺寸作为额外通道与 latent 拼接后再进入 patch 化见 forward 第 1 步。2. 双轨位置编码3D RoPE 可学习位置嵌入Cosmos 采用两种互补的位置编码CosmosRotaryPosEmbed3D 旋转位置编码将隐藏维度按dim_t : dim_h : dim_w 2 : 2 : 2的比例即各占 1/3切分到时间、高、宽三个轴分别计算 3D 旋转频率。两个值得注意的细节NTK 缩放h_ntk_factor rope_scale[1] ** (dim_h / (dim_h - 2))等通过rope_scale放大基础频率theta实现超出训练分辨率的外推帧率自适应当传入fps时视频场景时间轴频率会乘以base_fps / fpsbase_fps24即对时间序列做缩放实现跨帧率推理见 transformer_cosmos.py。CosmosLearnablePositionalEmbed可学习位置嵌入三个轴各维护一组可学习参数pos_emb_t/h/w相加后做归一化emb / (eps norm)为每个 patch 提供补充的绝对位置信息。该模块保持 fp32 运算。3. CosmosEmbedding 与 AdaIN 自适应归一化CosmosEmbeddingtransformer_cosmos.py将时间步通过Timesteps正弦/余弦投影 双线性层SiLU 激活映射为条件向量temb同时保留归一化后的时间步嵌入embedded_timestep。CosmosAdaLayerNormZero对时间步嵌入做SiLU → Linear → Linear输出(shift, scale, gate)三元组其中gate用于门控残差hidden_states hidden_states gate * attn_output。源码中linear_1为adaln_lora_dim维度的低秩瓶颈即adaln_lora_dim参数名的由来。CosmosAdaLayerNorm输出(shift, scale)二元组用于最终输出归一化。4. CosmosTransformerBlock自注意力 交叉注意力 FFN 的三段式门控残差每个CosmosTransformerBlocktransformer_cosmos.py依次执行AdaIN 门控的自注意力attn1使用CosmosAttnProcessor2_0AdaIN 门控的交叉注意力attn2对文本条件做交叉注意力AdaIN 门控的FeedForwardGELU 激活mlp_ratio4.0。块级设计要点QK 归一化自注意力与交叉注意力均启用qk_normrms_norm对 Q、K 做 RMSNorm稳定大模型训练与推理ControlNet 残差注入forward支持controlnet_residual由外部 ControlNet 缩放后直接加到隐藏状态上配合controlnet_block_every_n参数决定注入间隔。测试 tests/models/controlnets/test_models_controlnet_cosmos.py 对该机制做了覆盖before/after 投影before_proj/after_proj用于 Cosmos Transfer 类任务的 latent 注入与输出例如为视频到视频转换提供参考帧 latent 的跳跃连接。5. 双上下文交叉注意力图像 文本的分离式 QKV当配置了img_context_dim_in时CosmosTransformerBlock使用CosmosAttentiontransformer_cosmos.py——在标准文本交叉注意力之上额外增加一组针对图像上下文的q_img / k_img / v_img投影及其 RMSNorm。配套的CosmosAttnProcessor2_5transformer_cosmos.py要求encoder_hidden_states为(text_context, img_context)元组分别对文本和图像上下文做 GQAGrouped Query Attention注意力后将两条路径的输出相加。这一设计支撑了 Cosmos 的 Video2World / Image2World 能力——把首帧图像编码成图像 token 序列作为条件。所有注意力处理器统一走dispatch_attention_fn见 src/diffusers/models/attention_dispatch.py可通过AttnProcessor机制替换为 Flash Attention、SageAttention 等后端实现。6. 输出层AdaIN 归一化 反 patchify最后经norm_outCosmosAdaLayerNorm与proj_out输出维度为p_t×p_h×p_w×out_channels投影后通过unflattenpermute(0, 7, 1, 6, 2, 4, 3, 5)重新排列并展平将 token 序列还原为(B, C, T, H, W)的视频 latent。若return_dictTrue返回Transformer2DModelOutput定义于 src/diffusers/models/modeling_outputs.py其sample字段即为去噪后的 latent否则返回裸 tuple。模型输出Transformer2DModelOutputCosmosTransformer3DModel的输出类型复用的是通用Transformer2DModelOutput[[autodoc]] models.modeling_outputs.Transformer2DModelOutput。该输出类的核心字段为sample形状(batch_size, num_channels, num_frames, height, width)的去噪后视频 latent后续交给 Cosmos 专用 VAEAutoencoderKLCosmos见 src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py解码为像素级视频帧当return_dictFalse时返回(sample,)形式的 tuple两种方式均可与DiffusionPipeline的标准调用约定无缝对接。实战在 Text2World 管线中使用 CosmosTransformer3DModelCosmosTransformer3DModel并非孤立组件它作为去噪骨干被 Cosmos 系列管线组合使用。以 src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py 中的CosmosTextToWorldPipeline为例其组件构成清晰地展示了模型在整体架构中的位置text_encoderT5EncoderModelt5-11b 变体T5TokenizerFast将提示词编码为(B, S, 1024)文本特征transformer即CosmosTransformer3DModel负责在EDMEulerScheduler的调度下逐步去噪vaeAutoencoderKLCosmos负责视频与 latent 之间的编解码时间压缩 8 倍、空间压缩 8 倍safety_checker可选安全检测器需pip install cosmos_guardrail启用。管线的完整调用示例摘自该文件 docstringimport torch from diffusers import CosmosTextToWorldPipeline from diffusers.utils import export_to_video model_id nvidia/Cosmos-1.0-Diffusion-7B-Text2World pipe CosmosTextToWorldPipeline.from_pretrained(model_id, torch_dtypetorch.bfloat16) pipe.to(cuda) prompt A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes... output pipe(promptprompt).frames[0] export_to_video(output, output.mp4, fps30)结合源码可进一步理解内部机制latent 形状推导prepare_latentspipeline_cosmos_text2world.py中num_latent_frames (num_frames - 1) // 8 1、latent_height height // 8默认生成704×1280×121帧的视频对应 latent 为16 × 16 × 88 × 160初始噪声乘以scheduler.config.sigma_max完成 EDM 风格缩放CFG 默认负提示词管线内置了DEFAULT_NEGATIVE_PROMPT描述抖动、模糊、低分辨率等视频缺陷未显式提供negative_prompt时自动使用设备卸载策略model_cpu_offload_seq text_encoder-transformer-vae配合enable_model_cpu_offload()可按序卸载 T5 → Transformer → VAE显著降低峰值显存。测试与验证如何确认模型行为仓库为CosmosTransformer3DModel提供了完整的单元测试tests/models/transformers/test_models_transformer_cosmos.py可用于本地验证CosmosTransformerTesterConfig以小规模配置num_attention_heads2、attention_head_dim12、num_layers2、patch_size(1,2,2)等构造微型模型前向输入形状为(1, 4, 1, 16, 16)的 latent验证输出形状与输入一致TestCosmosTransformer继承ModelTesterMixin覆盖前向输出、配置序列化等核心行为TestCosmosTransformerMemory继承MemoryTesterMixin验证显存/内存相关优化TestCosmosTransformerTraining继承TrainingTesterMixin并额外断言gradient_checkpointing正确应用于CosmosTransformer3DModel。此外tests/pipelines/cosmos/test_cosmos.py 等管线级测试验证了 Transformer 与 VAE、T5、调度器组合后的端到端生成行为。如果你在本地基于该模型做二次开发如接入 ControlNet、训练 LoRA可参照这些测试快速搭建回归验证。小结CosmosTransformer3DModel是 Diffusers 中 Cosmos 系列视频生成能力的核心骨干它以 3D patch 化 双轨位置编码3D RoPE 与可学习位置嵌入 AdaIN 门控残差 图像/文本双上下文 GQA 注意力为技术骨架配合EDMEulerScheduler、AutoencoderKLCosmos与 T5 文本编码器组成完整的 Text2World / Video2World 管线。理解其构造参数与前向语义是上手 Cosmos 系列模型微调、部署与二次开发ControlNet、图像上下文、LoRA的起点。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表