实战指南:10× 推理加速与 5.2× 训练吞吐提升)
Open-Sora 高压缩自编码器Video DC-AE实战指南10× 推理加速与 5.2× 训练吞吐提升【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-SoraOpen-Sora v2 通过引入高压缩视频自编码器 Video DC-AE4×32×32下采样大幅削减视频扩散模型的 token 数量与注意力计算开销实现训练吞吐约 5.2×、推理速度约 10× 的提升。本文基于仓库 docs/hcae.md 展开结合 configs/diffusion/inference/high_compression.py、configs/diffusion/train/high_compression.py 与 opensora/models/dc_ae/models/dc_ae.py 等源码完整讲解模型下载、推理与训练的具体命令、配置参数含义、以及 DC-AE 的核心实现原理与权衡帮助你基于该模型自行复现与继续探索。一、背景为什么需要高压缩自编码器视频生成模型训练成本高昂其根源在于两大因素token 数量庞大与注意力计算占绝对主导地位。扩散模型如 DiT / FLUX 类架构的复杂度随序列长度token 数呈二次增长帧数越多、分辨率越高注意力矩阵规模增长越快。传统 VAE 方案如 HunyuanVideo 的 VAE下采样比通常为8×8×8或更低会在 latent 空间保留较多 token导致训练成本居高不下。为了进一步压缩训练与推理开销Open-Sora v2 探索了高压缩自编码器 Video DC-AE将下采样比提高到4×32×32时间维度 4×、空间 H/W 各 32×即dc-ae-f32t4c128模型。下采样比提升的直接收益是 latent token 数量急剧减少。即便把 patch size 进一步降到1原先需要较大的 patch 来压低 token 数整体的 token 数量仍然远低于传统方案。根据 docs/hcae.md 的对比训练吞吐提升约 5.2×推理速度提升约 10×。二、权衡与挑战大通道数下的收敛问题高压缩并非没有代价。更极端的压缩意味着 latent 通道数显著增大本模型为128 通道而大通道数会拖慢收敛速度。文档中给出的观察数据采用 128 通道 Video DC-AE 适配的生成模型训练 25K iterations 后 loss 仅降到0.5的水平相比之下初始化模型原始低压缩方案可以达到0.1的 loss 水平。因此快速视频生成模型在重建/生成质量上低于原始模型但它仍然能够捕捉时空关系spatial-temporal relationships具备可用的生成能力。作者将其以开源形式发布给研究社区供进一步探索与改进。这意味着本模型更适合对推理/训练成本敏感、或用于快速原型与二次研究的场景而不是无条件替换原有高质量模型。三、模型下载模型权重由hpcai-tech/Open-Sora-v2-Video-DC-AE发布可通过huggingface-cli下载到本地ckpts目录pip install huggingface_hub[cli] huggingface-cli download hpcai-tech/Open-Sora-v2-Video-DC-AE --local-dir ./ckpts下载完成后ckpts目录下应至少包含两个权重文件权重文件用途Open_Sora_v2_Video_DC_AE.safetensors快速视频生成模型DiT权重F32T4C128_AE.safetensorsDC-AE 自编码器权重这两个路径正是 configs/diffusion/inference/high_compression.py 与 configs/diffusion/train/high_compression.py 中model.from_pretrained与ae.from_pretrained所引用的位置需与配置保持一致。四、推理单卡运行快速视频生成模型使用 scripts/diffusion/inference.py 配合推理配置 configs/diffusion/inference/high_compression.py 即可生成视频torchrun --nproc_per_node 1 --standalone scripts/diffusion/inference.py configs/diffusion/inference/high_compression.py --prompt The story of a robots life in a cyberpunk setting.命令要点--nproc_per_node 1 --standalone单机单卡即可完成推理因为推理配置中显式关闭了并行插件见下文--prompt直接以命令行方式传入文本提示词推理脚本会将其写入临时 CSV 作为数据集见 scripts/diffusion/inference.py 中create_tmp_csv的逻辑输出视频默认写入配置的save_dir目录。4.1 推理配置逐项解析推理配置 configs/diffusion/inference/high_compression.py 结构如下_base_ [t2i2v_768px.py] # no need for parallelism plugin None plugin_config None plugin_ae None plugin_config_ae None # model settings patch_size 1 model dict( from_pretrained./ckpts/Open_Sora_v2_Video_DC_AE.safetensors, in_channels128, cond_embedTrue, patch_size1, ) # AE settings ae dict( _delete_True, typedc_ae, from_scratchTrue, model_namedc-ae-f32t4c128, from_pretrained./ckpts/F32T4C128_AE.safetensors, use_spatial_tilingTrue, use_temporal_tilingTrue, spatial_tile_size256, temporal_tile_size32, tile_overlap_factor0.25, ) ae_spatial_compression 32 sampling_option dict( num_frames128, )关键配置含义如下配置项值说明_base_[t2i2v_768px.py]继承 768px 文生图/图生视频基础配置见 configs/diffusion/inference/plugins/t2i2v.py本文件在其上覆盖模型与 AE 定义plugin/plugin_configNone关闭生成模型侧的并行插件。因为 token 数已大幅减少单卡即可推理无需张量/序列并行plugin_ae/plugin_config_aeNone同时关闭 AE 侧的并行patch_size1latent patch 大小设为 1充分利用高压缩带来的 token 红利model.in_channels128DiT 的输入通道数与 DC-AE 的 latent 通道数 128 严格对应model.from_pretrained./ckpts/Open_Sora_v2_Video_DC_AE.safetensors快速生成模型权重model.cond_embedTrue启用条件嵌入层用于传入 i2v/v2v 信息t2v 时该层置零见 plugins/t2i2v.py 注释ae.typedc_ae注册于 opensora/models/dc_ae/ae_model_zoo.py 的模块构建函数ae.model_namedc-ae-f32t4c128模型 zoo 中注册的型号对应f32空间 32×、t4时间 4×、c128128 通道见 ae_model_zoo.pyae.from_scratchTrue从零构建网络结构不通过 HuggingFace Hub 加载预训练结构权重ae.from_pretrained./ckpts/F32T4C128_AE.safetensors再从本地 safetensors 恢复 DC-AE 权重load_checkpoint见 ae_model_zoo.pyae.use_spatial_tiling/use_temporal_tilingTrue空间/时间维度分块tiling推理避免显存爆炸ae.spatial_tile_size256空间分块尺寸像素ae.temporal_tile_size32时间分块尺寸帧数ae.tile_overlap_factor0.25相邻 tile 重叠比例用于消除拼接接缝ae_spatial_compression32空间压缩比供训练/推理流程计算 latent 尺寸sampling_option.num_frames128生成 128 帧视频4.2 分块Tiling机制的源码级说明高分辨率/长时序下DC-AE 的编解码显存开销很大因此配置默认开启use_spatial_tiling与use_temporal_tiling。在 opensora/models/dc_ae/models/dc_ae.py 中encode入口按如下优先级分派def encode(self, x: torch.Tensor) - torch.Tensor: if self.use_temporal_tiling and x.shape[2] self.temporal_tile_size: return self.temporal_tiled_encode(x) elif self.use_spatial_tiling and (x.shape[-1] self.spatial_tile_size or x.shape[-2] self.spatial_tile_size): return self.spatial_tiled_encode(x) else: return self._encode(x)时间分块沿帧轴以temporal_tile_size * (1 - tile_overlap_factor)为步长切块逐块编码后在 latent 空间用blend_t做线性渐变融合见 dc_ae.py空间分块在 H/W 维度以spatial_tile_size * (1 - tile_overlap_factor)为步长切块重叠区域按blend_v/blend_h融合见 dc_ae.pytile_overlap_factor 0.25意味着每块边缘 25% 的区域参与加权过渡从而抑制 tile 边界处的接缝伪影。解码侧decode见 dc_ae.py遵循完全对称的分块与融合逻辑。另外注意tile 尺寸必须能被压缩比整除spatial_tile_size // spatial_compression_ratio与temporal_tile_size // time_compression_ratio在构造DCAE时会被断言见 dc_ae.py。五、训练8 卡训练你自己的快速生成模型5.1 数据准备训练前需要准备符合格式的视频-文本数据集。文档要求先阅读 docs/train.md 中的Prepare dataset章节数据集须为csv或parquet文件至少包含以下列path,text,num_frames,height,width,aspect_ratio,resolution,fps文档以 45k Pexels 数据集为例下载后置于datasets/pexels_45k其中pexels_45k_necessary.csv已包含训练所需的全部信息若只有原始pexels_45k.csv可先用 scripts/cnv/meta.py 处理生成# 并行处理 python scripts/cnv/meta.py --input datasets/pexels_45k.csv --output datasets/pexels_45k_nec.csv --num_workers 645.2 启动训练torchrun --nproc_per_node 8 scripts/diffusion/train.py configs/diffusion/train/high_compression.py --dataset.data-path datasets/pexels_45k_necessary.csv要点与推理单卡不同训练默认采用8 卡--nproc_per_node 8--dataset.data-path用于命令行覆盖配置中的数据集路径mmengine 配置体系支持命令行覆盖字典字段详见 docs/train.md训练脚本入口为 scripts/diffusion/train.py。5.3 训练配置逐项解析训练配置 configs/diffusion/train/high_compression.py 的关键字段_base_ [image.py] bucket_config { _delete_: True, 768px: { 1: (1.0, 20), 16: (1.0, 8), 20: (1.0, 8), ... 128: (1.0, 2), # 30s }, } condition_config dict( t2v1, i2v_head7, ) grad_ckpt_settings (100, 100) patch_size 1 model dict( from_pretrainedNone, grad_ckpt_settingsgrad_ckpt_settings, in_channels128, cond_embedTrue, patch_sizepatch_size, ) ae dict( _delete_True, typedc_ae, model_namedc-ae-f32t4c128, from_pretrained./ckpts/F32T4C128_AE.safetensors, from_scratchTrue, scaling_factor0.493, use_spatial_tilingTrue, use_temporal_tilingTrue, spatial_tile_size256, temporal_tile_size32, tile_overlap_factor0.25, ) is_causal_vae False ae_spatial_compression 32 ckpt_every 250 lr 3e-5 optim dict(lrlr)与推理配置相比训练配置的关键差异配置项值说明_base_[image.py]继承 configs/diffusion/train/image.py 的 FLUX DiT 基础结构hidden_size3072、depth19、38 个 single block 等bucket_config_delete_True后重定义抛弃基础配置中的多分辨率 bucket只保留 768px 一档键为帧数值为(采样概率, batch size)元组例如128: (1.0, 2)表示 128 帧视频以 100% 概率采样、batch size 2condition_configt2v1, i2v_head7同时训练文生视频t2v与图生视频i2v条件grad_ckpt_settings(100, 100)梯度检查点开关阈值在满足条件的位置启用 activation checkpointing 以省显存model.in_channels128与 DC-AE latent 通道数一致model.from_pretrainedNone从零训练生成模型finetune 时改为权重路径ae.scaling_factor0.493latent 的缩放系数编码时除以该值、解码时乘回见 dc_ae.py 与 decode_single用于稳定扩散训练is_causal_vaeFalse本方案非因果 VAE区别于 image.py 中的 HunyuanVideo 因果 VAE时序依赖由 DC-AE 的 3D 卷积结构自身处理ckpt_every250每 250 步保存一次 checkpoint高压缩方案训练更快因此保存更频繁lr3e-5学习率通过optim dict(lrlr)同步到优化器训练过程中可观察到的现象由于 128 通道 latent 带来的收敛困难loss 曲线会比低压缩模型下降得更慢文档给出的参考25K 步后约 0.5 对 0.1需要更多迭代或调整学习率/损失权重来逼近目标质量。六、DC-AE 模型结构与注册机制6.1 模型注册与构建DC-AE 通过注册表接入 Open-Sora 的模块体系。构建入口为 opensora/models/dc_ae/ae_model_zoo.py 中的DC_AE工厂函数它按from_scratch决定是从结构构建还是从 HuggingFace 拉取预训练结构若提供from_pretrained本地路径则调用load_checkpoint恢复权重将 tiling 相关参数use_spatial_tiling、use_temporal_tiling、spatial_tile_size、temporal_tile_size、tile_overlap_factor挂到模型实例上供编解码分派使用若提供scaling_factor则覆盖默认值。dc-ae-f32t4c128是 ae_model_zoo.py 中唯一注册的 DC-AE 型号。6.2 网络结构细节源码级dc_ae_f32工厂函数见 opensora/models/dc_ae/models/dc_ae.py定义了该型号的完整结构参数结构参数编码器 Encoder解码器 Decoder压缩比时间 4×time_compression_ratio4、空间 32×spatial_compression_ratio32对称上采样width_list[128, 256, 512, 512, 1024, 1024][128, 256, 512, 512, 1024, 1024]depth_list[2, 2, 2, 3, 3, 3][3, 3, 3, 3, 3, 3]block_type前 3 阶段ResBlock后 3 阶段EViTS5_GLU同左下/上采样方式Conv带temporal_downsample[False,False,False,True,True,False]InterpolateConv带temporal_upsample[False,False,False,True,True,False]norm / actrms3d/ 默认rms3d/silu输出rms3drelulatent_channels128128解读时间压缩在第 4、5 阶段完成temporal_downsample在前三个 stage 为False第 4、5 个 stage 为True共 2 次时间 2× 下采样合计 4×空间维度则由 5 个 stage 的 2× 下采样累计为 32×混合 block 设计浅层使用ResBlock保持细节深层使用EViTS5_GLUEfficientViT 风格的 GLU MBConv带 5×5 局部卷积提升效率与感受野见 dc_ae.py 的EViTS5_GLU分支编解码器均为纯 3D 结构is_videoTrue、rms3d归一化空间采样使用 Conv/InterpolateConv 的 3D 变体从而直接在时空联合维度上压缩。6.3 数据流与 latent 尺寸计算forward见 dc_ae.py完成encode → decode的完整重建链路get_latent_size见 dc_ae.py给出 latent 尺寸换算公式latent_T (T - 1) // 4 1 latent_H (H - 1) // 32 1 latent_W (W - 1) // 32 1例如 768px、128 帧的输入约 64×64 空间 latent、33 帧时间 latent对比 8× 压缩方案token 数仅为后者的约 1/64这正是训练与推理加速的直接来源。七、进阶从零训练 Video DC-AE 自编码器仓库同时提供了 DC-AE 自身的训练配置可用于从零训练/微调自编码器或作为理解损失函数与对抗训练的参考。基础重建训练配置 configs/vae/train/video_dc_ae.py 关键项datasetvideo_text类型resize_crop预处理fps_max24data_pathdatasets/pexels_45k_necessary.csvbucket_config256px_ar1:1: {32: (1.0, 1)}即 256px、32 帧、batch size 1优化器HybridAdamlr5e-5、betas(0.9, 0.98)、adamw_modeTruemixed_strategy mixed_video_image且mixed_image_ratio 0.2训练 batch 中约 1/5 为图像、4/5 为视频稳定空间维重建dtype bf16、plugin zero2ZeRO-2 并行见plugin_config损失配置vae_loss_configperceptual_loss_weight0.5、kl_loss_weight0KL 权重为 0说明该模型为确定性 AE不使用 KL 正则项。对抗训练变体 configs/vae/train/video_dc_ae_disc.py 增加 3D 判别器discriminator dict( typeN_Layer_discriminator_3D, from_pretrainedNone, input_nc3, n_layers5, conv_clsconv3d, ) gen_loss_config dict( gen_start0, disc_weight0.05, ) disc_loss_config dict( disc_start0, disc_loss_typehinge, ) optim_discriminator dict( clsHybridAdam, lr1e-4, eps1e-8, weight_decay0.0, adamw_modeTrue, betas(0.9, 0.98), ) grad_checkpoint True model dict( disc_off_grad_ckptTrue, # 开启 grad_checkpoint 时必须为 True )判别器为 5 层 3D 卷积N_Layer_discriminator_3D在视频帧的时空维度上做对抗判别生成器损失在训练开始即启用对抗项gen_start0对抗权重disc_weight0.05判别器使用 Hinge 损失disc_loss_typehinge生成器与判别器使用独立优化器optimvsoptim_discriminator判别器学习率1e-4高于生成器5e-5开启grad_checkpoint True时必须同步设置model.disc_off_grad_ckpt True将判别器相关计算排除在 activation checkpointing 之外。八、FAQ 与注意事项为什么推理配置要关闭并行插件因为高压缩方案将 token 数压缩至原来的约 1/64单卡即可承载pluginNone避免不必要的并行通信开销见 configs/diffusion/inference/high_compression.py。生成质量比原始模型差吗文档明确说明快速模型质量低于原始模型25K 步 loss 0.5 vs 0.1但能保留时空关系适合成本敏感场景与研究用途。tile 参数需要调整吗若推理更大分辨率/更长视频需保证spatial_tile_size能被 32 整除、temporal_tile_size能被 4 整除见 dc_ae.py 的断言并可按需增大tile_overlap_factor以进一步减轻接缝。权重放哪里推理/训练配置默认从./ckpts/读取Open_Sora_v2_Video_DC_AE.safetensors与F32T4C128_AE.safetensors务必保持目录结构与配置一致。九、参考资料技术报告Open-Sora v2 论文详见 docs/hcae.md 中引用的 arXiv 链接中文报告docs/zh_CN/report_v4.mdv2 相关技术细节训练完整流程docs/train.md模型结构实现opensora/models/dc_ae/models/dc_ae.py模型注册与加载opensora/models/dc_ae/ae_model_zoo.py推理脚本scripts/diffusion/inference.py、训练脚本scripts/diffusion/train.py【免费下载链接】Open-SoraOpen-Sora: Democratizing Efficient Video Production for All项目地址: https://gitcode.com/GitHub_Trending/op/Open-Sora创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考