MochiTransformer3DModel 深度解析Diffusers 中的 Mochi-1 视频生成 DiT 模型【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本文以 diffusers 仓库中的 MochiTransformer3DModel API 文档 为核心骨架结合 transformer_mochi.py 源码实现与 MochiPipeline、官方测试 等仓库证据系统讲解 Mochi-1 Preview 视频生成模型中 Diffusion TransformerDiT组件的加载方式、全部可配置参数、内部架构与注意力机制并给出从独立加载到完整文生视频管线的实战用法。读完本文你将能独立完成 MochiTransformer3DModel 的加载、参数调优、内存优化低精度 / 量化 / 多 GPU 切分以及单文件single-file权重导入。MochiTransformer3DModel 是什么MochiTransformer3DModel 是 Genmo 推出的 Mochi-1 Preview 视频生成模型中的核心去噪网络一个面向 3D 视频类数据帧 × 高 × 宽的 Diffusion Transformer。从仓库中 pipeline_mochi.py 的注释可以看到整个文生视频管线由三部分组成transformerMochiTransformer3DModel条件 Transformer负责在潜空间latent space中逐步去噪编码后的视频潜变量schedulerFlowMatchEulerDiscreteScheduler调度去噪步进vaeAutoencoderKLMochi在视频与潜变量之间编码 / 解码。其中MochiTransformer3DModel承担了绝大部分计算量属于约 10B 参数规模的模型该模型信息来自 Mochi 1 管线文档其架构属于非对称 Diffusion Transformer并在注意力层采用了非方形的 QKV 与输出投影用于降低推理时的显存占用。一行代码加载模型API 文档给出了最直接的加载方式mochi_transformer3d.mdfrom diffusers import MochiTransformer3DModel transformer MochiTransformer3DModel.from_pretrained( genmo/mochi-1-preview, subfoldertransformer, dtypetorch.float16, ).to(cuda) # 或 mps、xpu、cpu几个要点subfoldertransformer表示从仓库的transformer子目录加载权重dtypetorch.float16可将权重降到半精度实测中也可使用torch.bfloat16仓库官方管线示例即默认 bf16设备端支持cuda、mpsApple Silicon、xpuIntel 独立显卡与cpu该模型继承自ModelMixin、ConfigMixin、PeftAdapterMixin、FromOriginalModelMixin、CacheMixin见 transformer_mochi.py因此天然支持from_pretrained、from_single_file、LoRA 适配器注入与缓存管理等能力。核心配置参数全解通过MochiTransformer3DModel.__init__的注册配置transformer_mochi.py可完整列出该模型的构造参数及默认值参数默认值说明patch_size2Patch 嵌入层的 patch 边长视频潜变量按2×2空间块切分num_attention_heads24多头注意力头数attention_head_dim128每个注意力头的通道数num_layers48Transformer 块MochiTransformerBlock层数pooled_projection_dim1536文本池化投影维度caption 条件分支in_channels12输入通道数与 Mochi VAE 的 12 通道潜变量对应out_channelsNone输出通道数缺省时取in_channels源码out_channels out_channels or in_channelsqk_normrms_normQ/K 投影后的归一化方式text_embed_dim4096文本编码器T5-XXL输出的嵌入维度time_embed_dim256时间步嵌入的输出维度activation_fnswigluFeed-Forward 网络激活函数max_sequence_length256支持的最大文本序列长度其中inner_dim num_attention_heads * attention_head_dim 24 * 128 3072是整个 Transformer 的隐层宽度。测试用例 test_models_transformer_mochi.py 提供了同构的迷你配置num_layers2、num_attention_heads2、in_channels4等可用于验证这些参数之间的约束关系attention_head_dim与num_attention_heads的乘积构成块内维度pooled_projection_dim需与文本分支维度一致。模型内部还声明了三个元属性transformer_mochi.py_supports_gradient_checkpointing True支持梯度检查点训练_no_split_modules [MochiTransformerBlock]设备映射device_map切分时以 Transformer 块为最小粒度_skip_layerwise_casting_patterns [patch_embed, norm]层间精度转换时跳过 patch 嵌入与归一化层。内部架构五段式数据流从 forward 实现 可以还原完整的计算流程。输入hidden_states的形状为(batch_size, num_channels, num_frames, height, width)输出形状为(batch_size, out_channels, num_frames, height, width)中间经历五个阶段1. 时间步与文本条件嵌入MochiCombinedTimestepCaptionEmbeddingtimestep与encoder_hidden_states一起进入该模块定义于 embeddings.py 的MochiCombinedTimestepCaptionEmbedding生成时间步嵌入temb并更新文本嵌入同时把encoder_attention_mask一并传入用于后续注意力阶段过滤无效文本 token。2. Patch 嵌入PatchEmbed源码将输入先permute(0, 2, 1, 3, 4).flatten(0, 1)把帧维提到 batch 维之上再切 patchpost_patch_height height // p、post_patch_width width // p其中p patch_size。即每帧独立按p×p空间块嵌入到inner_dim维随后unflatten(0, (batch_size, -1)).flatten(1, 2)恢复为(batch, seq, dim)的 token 序列。3. 3D 旋转位置编码MochiRoPEMochiRoPEtransformer_mochi.py以base_height192、base_width192为基准面积按当前height * width计算插值缩放scale (target_area / (height * width)) ** 0.5再对帧维t、高维h、宽维w生成网格坐标并施加旋转编码。频率参数pos_frequencies为可学习参数形状为(3, num_attention_heads, attention_head_dim // 2)。值得注意的实现细节_create_rope强制在torch.float32下计算 cos/sin 频率避免低精度带来的数值误差。4. 48 层MochiTransformerBlock核心计算全部发生在堆叠的 Transformer 块中每块同时更新视频 token 与文本 token 两条流详见下一节。最后一个块以context_pre_onlyTrue构建transformer_mochi.py仅处理视频流不再对文本流做完整的前馈更新这是非对称架构的关键体现。5. 输出投影AdaLayerNormContinuousproj_out最后对视频 token 做 AdaIN 式 LayerNormnorm_typelayer_norm条件来自temb经nn.Linear(inner_dim, patch_size * patch_size * out_channels)反投影再通过reshape与permute把 patch 还原为(batch, out_channels, num_frames, height, width)的视频潜变量张量最终以Transformer2DModelOutput(sampleoutput)返回return_dictFalse时返回裸元组。注意力机制非对称设计与 SDPA 实现MochiTransformerBlocktransformer_mochi.py是理解整个模型的关键它采用 DiT 中常见的“条件调制 门控残差”模式调制归一化norm1使用MochiRMSNormZero——先对时间步嵌入做 SiLU 线性投影再chunk(4)拆出scale_msa / gate_msa / scale_mlp / gate_mlp四路调制信号分别控制注意力与前馈分支的缩放与门控门控信号在残差相加时经过torch.tanh归一化到(-1, 1)。非对称注意力MochiAttention见 attention_processor.py视频 token 的 Q/K/V 由to_q/to_k/to_v从query_dim投影到inner_dim文本 token 额外通过add_q_proj/add_k_proj/add_v_proj从pooled_projection_dim1536投影到同一inner_dim。所有 Q/K 在投影后都施加MochiRMSNorm即qk_normrms_norm的实现。与常规 DiT 的“方形”投影不同这里的 Q/K/V 输出维与输入维并不相等而是统一对齐到inner_dim这正是 Mochi 1 管线文档 中所说的非方形 QKV 与输出投影层——它让文本侧不必保持 4096 维的大投影矩阵从而显著降低推理显存。SDPA 处理器MochiAttnProcessor2_0见 attention_processor.py将视频 Q 与文本 Q 拼接后调用F.scaled_dot_product_attention。实现上先按attention_mask逐 batch 过滤掉被 padding 的文本 tokentorch.nonzero(mask.flatten())取有效下标注意力计算完成后再F.pad补回原长度从而避免无效 token 干扰注意力分布。前馈分支FeedForward使用swiglu激活隐藏维按(4 * dim * 2) // 3计算文本分支非最后一层使用独立的ff_context其隐藏维按(4 * pooled_projection_dim * 2) // 3计算。双流残差块内视频流与文本流各自带 4 个归一化层norm1..norm4与对应的*_context最后返回(hidden_states, encoder_hidden_states)二元组继续下一层。由于MochiAttnProcessor2_0在构造时校验F.scaled_dot_product_attention存在性attention_processor.py该模型要求PyTorch 2.0 及以上。与 MochiPipeline 组合的实战用法在真实推理中MochiTransformer3DModel 通常不单独使用而是作为MochiPipeline的一个组件注入。仓库在 pipeline_mochi.py 中给出了官方示例import torch from diffusers import MochiPipeline from diffusers.utils import export_to_video pipe MochiPipeline.from_pretrained(genmo/mochi-1-preview, torch_dtypetorch.bfloat16) pipe.enable_model_cpu_offload() pipe.vae.enable_tiling() prompt Close-up of a chameleons eye, with its scaly skin changing color. Ultra high resolution 4k. frames pipe(prompt, num_inference_steps28, guidance_scale3.5).frames[0] export_to_video(frames, mochi.mp4)对应的推荐推理参数源自 Mochi 1 管线文档num_inference_steps28默认档或64复现原版质量档guidance_scale3.5默认档或4.5复现原版num_frames85为通用档163帧档需要约 70GB 显存用于全精度解码输出经VideoProcessor后处理fps30导出 mp4。若想手动把自加载的 Transformer 注入管线可参考多 GPU 示例的组件替换写法mochi.mdfrom diffusers import MochiPipeline, MochiTransformer3DModel model_id genmo/mochi-1-preview transformer MochiTransformer3DModel.from_pretrained( model_id, subfoldertransformer, device_mapauto, max_memory{0: 24GB, 1: 24GB}, ) pipe MochiPipeline.from_pretrained(model_id, transformertransformer) pipe.enable_model_cpu_offload() pipe.vae.enable_tiling() with torch.autocast(device_typecuda, dtypetorch.bfloat16, cache_enabledFalse): frames pipe( promptClose-up of a chameleons eye, with its scaly skin changing color. Ultra high resolution 4k., negative_prompt, height480, width848, num_frames85, num_inference_steps50, guidance_scale4.5, max_sequence_length256, output_typepil, ).frames[0] export_to_video(frames, output.mp4, fps30)由于_no_split_modules [MochiTransformerBlock]device_mapauto可以把 48 个 Transformer 块按显存预算自动切分到多张 GPU 上配合max_memory精确控制每卡上限。内存优化与进阶用法低精度变体加载variantbf16的权重并配合torch.autocast(cuda, torch.bfloat16)显存需求可从全精度的约 42GB 降至约 22GB数据来自 Mochi 1 管线文档代价是轻微的画质下降。bitsandbytes 量化该模型继承FromOriginalModelMixin并可通过 diffusers 的BitsAndBytesConfig做 8-bit 量化加载from diffusers import BitsAndBytesConfig, MochiTransformer3DModel quant_config BitsAndBytesConfig(load_in_8bitTrue) transformer_8bit MochiTransformer3DModel.from_pretrained( genmo/mochi-1-preview, subfoldertransformer, quantization_configquant_config, dtypetorch.float16, )单文件single-file加载支持直接加载原始格式权重如 Comfy-Org 打包的mochi_preview_bf16.safetensors加载逻辑注册于 single_file_model.py 的MochiTransformer3DModel条目。注意官方提示目前不支持 FP8 缩放版本的单文件权重。仓库测试 TestMochiTransformer3DSingleFile 即验证了从genmo/mochi-1-preview加载配置 从单文件加载权重的组合路径。LoRA 微调与推理模型继承PeftAdapterMixinMochiPipeline也组合了Mochi1LoraLoaderMixin见 pipeline_mochi.pytest_lora_layers_mochi.py 覆盖了 LoRA 注入测试说明可以对 Mochi 视频生成模型做 LoRA 定制。此外_supports_gradient_checkpointing True意味着训练时可通过梯度检查点大幅降低激活显存。测试与可靠性验证仓库提供了完整的测试矩阵来验证模型行为test_models_transformer_mochi.pyTestMochiTransformer基于ModelTesterMixin验证前向输出形状测试配置下输出(4, 2, 16, 16)、attention processor 兼容性等TestMochiTransformerTraining验证梯度检查点确实作用于MochiTransformer3DModelTestMochiTransformerLoRA验证 LoRA 适配器可注入TestMochiTransformer3DSingleFile验证单文件权重加载。这些测试使用的迷你配置num_attention_heads2、attention_head_dim8、num_layers2是学习模型前向流程、做单元级调试的良好起点。若需将自定义权重转为 diffusers 格式可参考 convert_mochi_to_diffusers.py 转换脚本。小结MochiTransformer3DModel 是 Mochi-1 Preview 文生视频系统的引擎它以 PatchEmbed 切分 12 通道视频潜变量通过MochiCombinedTimestepCaptionEmbedding融合时间步与 T5 文本条件在 48 层非对称 Transformer 块中以QK RMSNorm 非方形投影 SDPA 3D RoPE完成去噪最终经 AdaIN 输出层还原视频潜变量。无论是直接from_pretrained加载、from_single_file导入原始权重还是通过量化、多 GPU 切分与 LoRA 进行扩展仓库都提供了开箱即用的支持。建议读者结合 API 文档、管线文档 与 transformer_mochi.py 源码交叉阅读以获得对模型调参最完整的把握。【免费下载链接】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），仅供参考