【代码解读】通义万相最新视频生成模型 Wan 2.2 实现解析
GitHub:https://github.com/Wan-Video/Wan2.2
Huggingface:https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B / https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B / https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B
今天通义万相实验室开源了视频生成模型 Wan 2.2,相比于 Wan 2.1 模型来说,该模型的整体架构变为了混合专家架构,整体参数量为 27B,推理时激活参数量为 14B,在推理成本与 2.1 模型相近的条件下得到了更高参数量。
本次开源的模型共有三种,分别为文生视频 27B MoE 模型、图生视频 27B MoE 模型以及文/图生视频 5B Dense 模型。其中最后一种模型能够在单种架构中同时实现文生视频和图生视频两种不同的任务,且能在消费级显卡上运行。
在模型的功能方面,得益于 65.6% 图像数据、83.2% 视频数据的扩增,以及精细筛选的美学数据,Wan 2.2 能够实现电影级美学生成、复杂运动生成,且指令遵循能力也有所增强。
在技术层面,Wan 2.2 的 MoE 模型和 LLM 中带有 router 和 expert 的网络不同,这里的 Expert 指的是网络不同阶段的子网络。Wan 2.2 采用了两阶段的生成范式,整个网络架构分为两个部分,分别为高噪模型和低噪模型。在生成早期信噪比比较高的阶段,模型激活高噪模型;在生成后期信噪比比较低的阶段,模型激活低噪模型。两个模型均为 14B 的模型,通过分阶段激活不同的模型,能够有效地控制推理时的资源消耗。
高噪模型和低噪模型均复用了 Wan 2.1 的网络架构,其中高噪模型为从头训练,而低噪模型是在 Wan 2.1 权重的基础上进行后训练得到的。
从 Wan 2.2 的代码仓库中我们也可以发现一些实现的技术细节。在模型定义时,分别定义了 low_noise_model
和 high_noise_model
两个子模块:
# Wan2.2/wan/text2video.py
logging.info(f"Creating WanModel from {checkpoint_dir}")
self.low_noise_model = WanModel.from_pretrained(checkpoint_dir, subfolder=config.low_noise_checkpoint)
self.low_noise_model = self._configure_model(model=self.low_noise_model,use_sp=use_sp,dit_fsdp=dit_fsdp,shard_fn=shard_fn,convert_model_dtype=convert_model_dtype)self.high_noise_model = WanModel.from_pretrained(checkpoint_dir, subfolder=config.high_noise_checkpoint)
self.high_noise_model = self._configure_model(model=self.high_noise_model,use_sp=use_sp,dit_fsdp=dit_fsdp,shard_fn=shard_fn,convert_model_dtype=convert_model_dtype)
在推理阶段,可以发现相比 Wan 2.1,Wan 2.2 模型多了一个根据时间步来准备模型的步骤:
# Wan2.2/wan/text2video.py
model = self._prepare_model_for_timestep(t, boundary, offload_model)
在这个函数中,当采样时间步大于某个边界值时,会加载高噪模型、卸载低噪模型;反之亦然。最后将加载的模型返回,用这种方式,在每个推理阶段都只加载高噪模型与低噪模型中的一个。
# Wan2.2/wan/text2video.py
def _prepare_model_for_timestep(self, t, boundary, offload_model):if t.item() >= boundary:required_model_name = 'high_noise_model'offload_model_name = 'low_noise_model'else:required_model_name = 'low_noise_model'offload_model_name = 'high_noise_model'if offload_model or self.init_on_cpu:if next(getattr(self,offload_model_name).parameters()).device.type == 'cuda':getattr(self, offload_model_name).to('cpu')if next(getattr(self,required_model_name).parameters()).device.type == 'cpu':getattr(self, required_model_name).to(self.device)return getattr(self, required_model_name)
除此之外,观察仓库的结构可以发现 VAE 有 vae2_1.py
和 vae2_2.py
两种,前者是文生视频和图生视频使用的 VAE 模型,而后者则是 TI2V 模型使用的高压缩率 VAE。分析这个高压缩率 VAE,首先可以看到编码器的下采样模块:
# downsample blocks
downsamples = []
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):t_down_flag = (temperal_downsample[i]if i < len(temperal_downsample) else False)downsamples.append(Down_ResidualBlock(in_dim=in_dim,out_dim=out_dim,dropout=dropout,mult=num_res_blocks,temperal_downsample=t_down_flag,down_flag=i != len(dim_mult) - 1,))scale /= 2.0
self.downsamples = nn.Sequential(*downsamples)
由于 dim_mult=[1,2,4,4]
、temperal_downsample=[True,True,False]
,所以共有四个下采样模块,其中,前两个模块对时间维度进行下采样,前三个模块对空间维度进行下采样。因此在 Encoder 中时间下采样 4 倍,空间下采样 8 倍,算上在 Encoder 之前的 patchify 操作,整体的下采样倍数为 4x16x16。
总体上来看 Wan 2.2 相比 Wan 2.1 主要的变化是从单阶段模型变成了两阶段模型,并且看样子这个模型暂时还没有适配 VACE 编辑模型。至于这个模型在各个基础任务和下游任务上的表现如何,可以拭目以待。