Seed-vr2在ncnn_llm上的前向推理实现 #6968
Chisato623
started this conversation in
Show and tell
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
在ncnn_llm上实现seedvr2-3B模型的前向推理
注意:本文档用于详细描述seedvr2在ncnn的前向推理实现过程,如果想要快速跑起项目,请直接进入仓库https://github.com/Chisato623/seedvr2_ncnn
关联和参考仓库:
https://github.com/futz12/ncnn_llm
https://github.com/Tencent/ncnn
https://github.com/nihui/zimage-ncnn-vulkan
https://github.com/ByteDance-Seed/SeedVR
1. 模型导出
SeedVR2 是字节跳动开源的通用图像/视频恢复模型,用于超分、去噪、去压缩伪影、去模糊等恢复任务。根据官方 README 说明,单张 H100 80GB 可处理约
100×720×1280的视频;使用 4 张 H100 80GB(sp_size=4)可进一步支持 1080p 和 2K 视频。作为一个图像处理模型,它的结构如图:

这个模型的主体部分由32个DiT Block和前后VAE编解码器组成,我们需要基于ncnn和ncnn_llm的框架实现它的前向推理,同时自定义一些原有的算子并实现,不过在这之前,我们先把它的权重导出:
模型导出主要分为 DiT 和 VAE 两部分。DiT 负责在 latent 空间中完成恢复推理,VAE 则负责视频与 latent 之间的编解码。下面介绍 DiT 的自定义算子和 VAE 的导出。该模型的 VAE 在传统图像 VAE 的基础上增加了时间维度,用于处理视频输入。它的输入为
[B,3,T,H,W],输出为 16 通道的视频 latent;空间尺寸压缩约 8 倍,时间尺寸压缩约 4 倍。1.1 DiT 的自定义算子导出
DiT 导出过程中不会把所有操作都展开为普通 NCNN 层。以下两个模块会保留为自定义层:
AdaptiveWindowAttention:根据运行时的video_shape=[T,H,W]和文本长度动态划分窗口,执行 Q/K RMSNorm、三轴 MM-RoPE、窗口内 SDPA,并将 video 和 text 的结果写回原始顺序。SharedLinear:对共享的 video/text 投影权重只保存一份,并分别执行投影,避免在生产分辨率下创建巨大的 Concat/Crop 中间张量。导出器会将自定义层的输入顺序和参数写入 NCNN
.param。AdaptiveWindowAttention的输入包括:每个 Block 的奇偶编号决定窗口是否 shifted;目标窗口数量为
(4,3,3),并记录在manifest.json中。导出完成后,运行时必须注册AdaptiveWindowAttention和SharedLinear,否则 NCNN 无法加载 blocks 图。1.2 VAE 的部分算子替换
官方 VAE 使用 Conv3D、3D 上采样和 3D 下采样,而 NCNN/Vulkan 运行时使用项目自定义算子(因为本项目需要实现GPU加速下的三维卷积)。导出脚本会递归替换:
Conv3D→DecomposedConv3d:沿时间维度拆成多个 Conv2D 切片计算,再累加结果,同时保留 causal padding、stride、dilation、groups 和 bias。Upsample3D→ExportUpsample3D:处理时空 pixel shuffle 和后置卷积。Downsample3D→ExportDownsample3D:处理官方使用的非对称空间 padding。随后脚本通过 wrapper 分别导出两张图:
Encoder 的 32 通道 moments 中,前 16 个通道是
mean,后 16 个通道是logvar。后验采样、mean/logvar拆分以及官方 scaling factor0.9152由 C++ 运行时完成,不在导出的 encoder/decoder 图中完成。更多的导出细节请参考https://github.com/Chisato623/seedvr2_export_ncnn
2. 部分特殊算子的自定义实现
查阅
ncnn_llm和ncnn的算子后,本项目只对无法直接由现有 NCNN 图表达、或需要在运行时根据视频形状调度的部分进行封装。当前自定义层分为 VAE 和 DiT 两组:DecomposedConv3dDecomposedConv3DExportUpsample3DVAEUpsample3DExportDownsample3DVAEDownsample3DSeedVR2ChunkedMHAVAEChunkedMHAAdaptiveWindowAttentionAdaptiveWindowAttentionSharedLinearSharedLinear2.1 VAE 自定义算子
2.1.1
DecomposedConv3dDecomposedConv3d是 VAE 中最主要的兼容层。NCNN 的常规卷积接口不能直接表达导出脚本所需的带 causal temporal memory 的 5D Conv3D 图,因此导出脚本中的nn.Conv3d会被DecomposedConv3d替换。NCNN 将 VAE 的
[B,C,T,H,W]以单 batch 的 4DMat表示,物理布局为[C,T,H,W],其中T位于Mat::d。一个kernel_t为k的 3D 卷积被保存为k组二维空间卷积,每组权重对应 3D 卷积在一个时间位置上的切片。前向时,对每一个输出时间位置:stride_t和dilation_t选择时间窗口中的输入帧;Convolution;该层从参数和权重中读取输出/输入通道数、时间和空间 kernel、stride、dilation、padding、groups、bias 以及首段
head_frames。首段 causal padding 采用首帧复制;当 VAE 使用时间切片时,层会根据MemoryState保存末尾的kernel_t - stride_t帧,在下一段输入中复用这些帧。完整序列推理时不需要跨段缓存,但 causal 首帧处理仍然由该层保留。Vulkan 路径使用多个 Vulkan
Convolution子层,并通过三个内部 compute pipeline 完成:gather_frame:从对应时间位置收集所有输入通道;zero_output:将输出缓冲区清零;accumulate_frame:把当前时间切片的空间卷积结果累加到输出。当输入通道数可按 4 打包时,Vulkan 路径会在二维卷积前使用 pack4,并在累加前恢复为 pack1。大分辨率下还会按帧提交命令并清理临时 allocator,以限制单次 Vulkan descriptor 和显存占用。
2.1.2
ExportUpsample3DExportUpsample3D是导出脚本为 SeedVR2 的时空上采样模块建立的稳定 PNNX 边界。它只接受当前模型使用的use_conv=True、slicing=False变体,参数中保存基础通道数、时间放大倍率、空间放大倍率和是否为 temporal upsample。其执行顺序为:
DecomposedConv3d执行上采样卷积;spatial_ratio * spatial_ratio * temporal_ratio * channels;T/H/W;DecomposedConv3d执行后置卷积。2.1.3
ExportDownsample3DExportDownsample3D对应官方 VAE 中使用空间下采样的分支。当前模型要求在卷积前仅沿右侧和底部各补一个零,不对左侧和顶部补零,也不额外复制时间帧。层参数保存通道数和空间补零尺寸,卷积权重作为内部DecomposedConv3d读取。2.1.4
SeedVR2ChunkedMHASeedVR2ChunkedMHA用于 VAE 中间层的高分辨率多头注意力。它不重新实现注意力数学,而是封装 NCNN 的MultiHeadAttention,加载 Q/K/V/output 的权重和 bias,并在 Vulkan 路径改变 query 的执行方式。输入是展平后的二维 token 矩阵。Vulkan 前向将 query 按
chunk_size切分,默认每块 128 行;每次只复制当前 query 块,K/V 仍来自完整输入,调用内部MultiHeadAttention后再把结果写回输出对应的行。每个块完成后提交并重置命令缓冲,释放本块临时空间。这样可以避免一次性生成完整 query-attention 中间结果,主要作用是降低峰值显存,而不是改变注意力结果。2.2 DiT 自定义算子
2.2.1
AdaptiveWindowAttentionAdaptiveWindowAttention对应官方 DiT 的自适应窗口注意力。它的窗口参数(4,3,3)表示目标窗口数量,而不是固定的窗口 token 尺寸。运行时根据每个样本的video_shape=[T,H,W]计算实际窗口边界,因此同一张导出的图可以处理不同的时间长度和空间尺寸。该层有两种输入 ABI:旧的融合 QKV 形式有 4 个输入,新的 split-QKV 形式有 8 个输入。split 形式依次为:
输出为 video 和 text 两个 token 矩阵。前向流程如下:
T/H/W和文本长度;SDPA/SDPA_vulkan完成窗口内注意力;Vulkan 路径将 QKV 准备、归一化、RoPE、窗口索引 gather、SDPA、video scatter 和 text pooling 保持在 GPU 上。该层的模型权重包括 video/text 的 Q/K norm 和 MM-RoPE 频率;参数包括 head 数、head dimension、目标窗口数量、shifted 标志、epsilon、RoPE 宽度、是否共享 norm、shape 标量类型和官方 BF16 attention 选项。它是动态 shape 的关键实现,也是当前 DiT 中不能简单替换为固定窗口
MultiHeadAttention的部分。关键代码如下。首先根据输入视频的
T/H/W和目标窗口数量计算实际窗口尺寸。(4,3,3)是目标窗口数量,wt、wh、ww会随输入 shape 变化:每个窗口会记录原始范围和重排后的 token 索引。这样窗口内可以连续读取 token,而输出时再通过同一组索引写回原始顺序:
Vulkan 路径使用相同的窗口计划,将 Q/K/V 准备、RMSNorm、MM-RoPE、SDPA 和 scatter 记录到同一条
VkCompute命令流中:cmd.record_pipeline( split_qkv ? pipeline_prepare_qkv_split : pipeline_prepare_qkv, prepare_bindings, prepare_constants, prepare_dispatcher); std::vector<ncnn::VkMat> sdpa_inputs(3); sdpa_inputs[0] = query; sdpa_inputs[1] = key; sdpa_inputs[2] = value; std::vector<ncnn::VkMat> sdpa_outputs(1); ret = sdpa_vulkan->forward(sdpa_inputs, sdpa_outputs, cmd, opt); if (ret != 0) return ret; scatter_bindings[0] = sdpa_outputs[0]; scatter_bindings[1] = plan_indices_gpu; scatter_bindings[2] = video_output; scatter_bindings[3] = text_output; cmd.record_pipeline(pipeline_scatter, scatter_bindings, scatter_constants, scatter_dispatcher);2.2.2
SharedLinearSharedLinear用于 DiT 中 video/text 共用投影权重的线性层。普通导出方式可能先把两路 tokenConcat,执行一次线性投影,再通过Crop拆回两路;如果输出分辨率很大,可能产生很大的临时 activation。该层改为接收两个二维 token 矩阵,内部只加载一份权重,并分别调用两次 NCNNGemm,输出两路投影结果。其参数包括输入维度、输出维度和 bias 标志,权重布局转换为 NCNN
Gemm所需的二维矩阵。Vulkan 路径复用 NCNNGemmpipeline,并支持输入输出的 NCNN packing。3. 推理主程序
自定义算子实现后,推理主程序负责组织模型加载、输入预处理、VAE 编解码、DiT 去噪和媒体输出。模型计算统一通过 NCNN Vulkan 路径执行,主程序只负责数据准备、阶段调度和结果落盘。
3.1 参数解析与模型加载
主程序入口为
examples/seedvr2_main.cpp。一次推理至少需要指定 DiT 模型目录、VAE 模型目录、输入媒体、输出路径以及正向文本 embedding:其中
--text-pos和可选的--text-neg是预先生成的文本 embedding 文件,不属于 DiT 的.bin权重。程序按照每行 5120 个float读取它们,得到形状为[text_length, 5120]的文本 token。只有在cfg_scale不等于1时,才需要提供负向 embedding。其余常用参数如下:
--gpu0--steps1--cfg-scale1--cfg-rescale0--seed666--res-w、--res-h1280、720--vae-mean--no-low-memory加载模型时,程序首先创建 Vulkan GPU 实例并验证 GPU 编号。完整加载模式通过
SeedVR2Pipeline::load()同时加载 DiT 和 VAE;默认的低显存模式通过load_staged()只保存模型目录和 GPU 信息,具体网络在对应阶段再加载。DiT 目录包含 frontend、32 个 blocks 和 tail 三部分图。加载 blocks 图时注册
AdaptiveWindowAttention和SharedLinear。VAE 目录包含 encoder 和 decoder 两张图,加载时注册 VAE 使用的自定义层。各 NCNN 网络均启用 Vulkan compute,并将对应 Vulkan device 绑定到网络。3.2 输入预处理
输入媒体由
seedvr2_video_io.cpp中的read_video()统一读取,图片和视频使用相同的数据接口。程序通过 FFmpeg 解码每一帧,并根据目标面积保持原始宽高比缩放:缩放后,宽度和高度分别向下裁剪到 16 的倍数,再进行中心裁剪。这样可以同时满足 VAE 的空间下采样和 DiT 的 2×2 patchify 要求。每个 RGB 像素会从
[0, 255]归一化到[-1, 1]。NCNN 中的媒体张量使用单 batch 的四维布局
[C,T,H,W]:图片输入只有一帧,即
T=1。视频输入会保留解码出的全部帧,并记录原始帧率供输出阶段使用。输入的空间尺寸必须满足H % 16 == 0和W % 16 == 0。3.3 VAE 编码
SeedVR2Pipeline::restore()首先将视频帧数补齐到4n+1。当输入不是单张图片时,补齐帧通过复制最后一帧完成;图片输入本身保持一帧。补齐的原因是 VAE 的时间下采样倍率为 4,只有4n+1的帧数可以在编码后得到完整的时间 latent。VAE encoder 输出 32 通道 moments,空间尺寸缩小为输入的
1/8,时间尺寸缩小为:moments 的通道按前后两段排列:前 16 个通道是 posterior mean,后 16 个通道是 log variance。运行时通过
moments_to_latent()构造 16 通道 latent:当传入
--vae-mean时省略随机项,直接使用 mean,再乘以官方 scaling factor0.9152。未启用该选项时,后验采样由--seed控制。当前 VAE Vulkan 路径支持低显存时间切片。encoder 按
5 + 4n的帧段处理,decoder 按2 + 1n的 latent 段处理;跨段的 causal temporal 状态由 VAE 自定义层维护。该时间切片只改变 VAE 的执行调度,不改变 latent 的形状和模型计算结果。3.4 DiT 输入构造与 patchify
VAE 编码得到的 latent 是待去噪的当前 latent,VAE 编码得到的输入 latent 则作为固定 condition。主程序在每个 DiT 时间步构造 33 通道的视频 token:
对于每个
T,H,W位置,33 通道数据会按空间 2×2 邻域重新排列为一个 132 维 token:同时,主程序根据当前 timestep 生成 256 维正弦余弦时间 embedding,并将视频形状和文本长度分别构造成运行时输入:
这些 shape 信息不会固化在
.param中,而是随每次 forward 传入,使同一份 DiT 权重可以处理不同的时间长度和空间尺寸。文本 embedding 保持[text_length, 5120]的 token 矩阵,与 patchify 后的视频 token 一起送入 DiT frontend。3.5 32 个 DiT Block 与 Euler 推理
DiT forward 由三段图组成:frontend、完整的 32 个 DiT blocks 和 tail。
AdaptiveWindowAttention,根据当前video_shape动态生成窗口,并完成 shifted window、MM-RoPE、SDPA 以及 video/text 结果回写。每个去噪步的时间区间从
1000线性走到0。第step步的时间值为:得到 DiT prediction 后使用 Euler 规则更新 latent:
当
cfg_scale == 1时,每个时间步只执行一次正向 prediction。当cfg_scale != 1时,还会使用负向文本执行一次 prediction,并按下式进行 guidance:如果启用了
cfg_rescale,程序会再根据正向 prediction 和 guided prediction 的标准差进行 rescale。完成所有 Euler 步后,得到供 VAE decoder 使用的去噪 latent。3.6 VAE 解码与媒体输出
DiT 输出 latent 仍然带有 VAE 的 scaling factor,送入 decoder 前先执行:
VAE decoder 将 16 通道 latent 恢复为
[3,T,H,W]的 RGB 视频,空间尺寸放大 8 倍,时间尺寸恢复为:如果前面为满足
4n+1约束补充了帧,解码完成后只保留原始输入的帧数。对于图片输入,输出只包含第一帧;对于视频输入,输出保留原始帧率,除非通过--out-fps显式指定新的帧率。输出张量中的 RGB 值从
[-1, 1]截断到有效范围,再转换为 8 位图像数据。单张图片通过write_image()写入,输出格式由输出文件的扩展名选择;视频通过write_video()写入,并在 MP4、MOV、WebM 等支持的容器中尽可能保留源媒体的非视频流和元数据。--output-f32可额外保存未量化的 RGBfloat32结果,用于数值分析。3.7 低显存与动态 shape 调度
默认启用低显存 staged 模式时,Pipeline 按以下顺序调度模型:
这样可以避免 VAE 和完整 DiT 网络同时占用显存。每个阶段使用独立的 Vulkan blob/staging allocator,并在阶段结束时回收网络和临时资源。关闭 staged 模式后,
load()会同时保留 DiT 和 VAE,适合显存充足且需要重复处理多个输入的场景。动态 shape 由三层共同保证:
video_shape、text_length和 token 数量,不依赖固定的T/H/W。AdaptiveWindowAttention根据每个样本的运行时 shape 重新规划窗口,DiT blocks 图本身不需要重新导出。因此,VAE 的空间尺寸、VAE 的时间切片长度和 DiT 的注意力窗口是三个独立的调度层次:VAE 负责显存受限下的时序编解码,DiT 负责动态时空 token 的窗口注意力,Pipeline 负责在两者之间传递正确的 latent 和 shape 信息。
4. 项目构建和运行
本项目用cmake构建,目前只在Ubuntu22.04系统中运行过,对于一台机器,可以快速构建的命令如下:
sudo apt update sudo apt install -y build-essential cmake pkg-config libvulkan-dev \ libavformat-dev libavcodec-dev libavutil-dev libswscale-dev git clone --recurse-submodules https://github.com/Chisato623/seedvr2_ncnn.git seedvr2-ncnn cd seedvr2-ncnn git submodule update --init --recursive git submodule status cmake -S . -B build -DCMAKE_BUILD_TYPE=Release cmake --build build -j构建成功后,我们开始运行:
运行命令中,vae,dit,text-pos等文件可以在https://huggingface.co/Chisato623/seedvr2_ncnn 下载
5.实测推理效果
分别在NVIDIA A100,NVIDIA RTX 5090 32G,NVIDIA 4090 24G中运行了该项目。对于两块消费级显卡5090和4090,执行了同一项720P超分到2k分辨率的任务,5090耗时约50秒,4090耗时约100秒








我们比较一下超分的效果:
下图是一张1280x720分辨率的图片,可以看到富冈义勇的下巴处有很多锯齿,很适合我们测试超分效果
下图为超分到2k后的效果,可以看到锯齿减少,清晰度明显提升:
然而,这份推理框架和官方的推理结果有一定误差,经过排查发现可能是官方权重转化为ncnn的FP16/BF16有一定误差,但转为FP32又有显存压力,下面是官方推理的结果和本推理框架推理的结果对比。
720P原图:
官方推理结果:
ncnn推理结果:
720P原图:
官方推理结果:
ncnn推理结果:
发现有一定色差,经过排查发现,主要的误差来源来自DiT误差会逐层累计,但排查许久也没能发现原因,有可能来自FP16和官方权重的转化部分,或是导出脚本的精度问题。
固定输入:frames=1, height=4, width=8,完整执行 frontend、32 层 DiT、tail。
更新:NCNN的最新版本降低了FP32->BF16的误差,经过导出脚本的优化和NCNN框架的更新,我们的实测误差变为:
输出示例图:


All reactions