Wan2.2 代码仓库深度解读 — 框架、架构与关键代码分析
·
仓库地址：https://github.com/Wan-Video/Wan2.2
团队：阿里巴巴通义实验室
开源协议：Apache 2.0
核心语言：Python (PyTorch)
一、代码仓库总览
1.1 目录结构
Wan2.2/
├── generate.py # 统一推理入口（CLI）
├── pyproject.toml # 项目配置
├── requirements.txt # 依赖（torch>=2.4.0, flash_attn等）
├── requirements_s2v.txt # S2V任务额外依赖（CosyVoice等）
├── INSTALL.md # 安装指南
├── LICENSE.txt # Apache 2.0
├── assets/ # 资源文件
├── examples/ # 示例文件（图片、音频、视频）
│ ├── wan_animate/ # Animate任务示例
│ ├── i2v_input.JPG # I2V参考图
│ ├── pose.mp4/.png # 姿态驱动示例
│ ├── talk.wav / sing.MP3 # S2V音频示例
│ └── ...
├── tests/ # 单元测试
└── wan/ # 核心代码包 ★
├── __init__.py # 包初始化，导出4个Pipeline类
├── configs/ # 模型配置（分辨率/参数/路径）
│ └── __init__.py # WAN_CONFIGS, SIZE_CONFIGS, MAX_AREA_CONFIGS
├── distributed/ # 分布式推理工具
│ └── util.py # init_distributed_group, FSDP, Ulysses
├── modules/ # 核心模块 ★★
│ ├── vae2_1.py # Wan2.1 VAE（128×压缩，16通道）
│ ├── vae2_2.py # Wan2.2 VAE（1024×压缩，48通道）★★
│ ├── dit/ # DiT模型定义
│ │ ├── wan2_1.py # Wan2.1 DiT（单一14B Dense模型）
│ │ └── wan2_2.py # Wan2.2 DiT（MoE双专家）★★
│ └── animate/ # Animate模块
│ └── preprocess/ # 姿态预处理
├── utils/ # 工具函数
│ ├── prompt_extend.py # Prompt扩展（DashScope/Qwen）
│ └── utils.py # save_video, merge_video_audio等
├── text2video.py # WanT2V Pipeline ★
├── image2video.py # WanI2V Pipeline ★
├── textimage2video.py # WanTI2V Pipeline ★
└── speech2video.py # WanS2V Pipeline ★
1.2 包入口 wan/__init__.py
from . import configs, distributed, modules
from .image2video import WanI2V
from .speech2video import WanS2V
from .text2video import WanT2V
from .textimage2video import WanTI2V
四个核心Pipeline类全部暴露在外，用户可通过 wan.WanT2V 等直接使用。
二、推理入口 generate.py 详解
2.1 整体流程
# generate.py 核心逻辑（简化）
import wan
from wan.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
def main():
# 1. 解析参数
args = parse_args() # task, size, ckpt_dir, prompt, image, audio...
# 2. 初始化分布式
init_distributed_group(args)
# 3. 根据task选择Pipeline
if args.task == "t2v-A14B":
pipeline = wan.WanT2V(
config=WAN_CONFIGS[args.task],
checkpoint_dir=args.ckpt_dir,
device_id=dist.get_rank(),
rank=dist.get_rank(),
t5_cpu=args.t5_cpu, # T5编码器放CPU节省显存
offload_model=args.offload_model, # 模型卸载到CPU
)
elif args.task == "i2v-A14B":
pipeline = wan.WanI2V(...)
elif args.task == "ti2v-5B":
pipeline = wan.WanTI2V(...)
elif args.task == "s2v-14B":
pipeline = wan.WanS2V(...)
# 4. Prompt扩展（可选）
if args.prompt_extend:
prompt = extend_prompt(prompt, args)
# 5. 生成视频
video = pipeline.generate(
prompt=args.prompt,
image=args.image, # I2V/TI2V/S2V
audio=args.audio, # S2V
size=args.size, # "1280*720"
frame_num=args.frame_num,
shift=args.shift, # Flow Matching的shift参数
sampling_steps=args.sampling_steps,
guidance_scale=args.guidance_scale, # CFG权重
seed=args.seed,
offload_model=args.offload_model,
)
# 6. 保存视频
save_video(video, output_path, fps=16, quality=5)
2.2 关键CLI参数
| 参数 | 说明 | 默认值 |
|---|---|---|
--task |
任务类型：t2v-A14B/i2v-A14B/ti2v-5B/s2v-14B | 必填 |
--size |
分辨率：“1280*720”/“832*480” | 依赖模型 |
--ckpt_dir |
模型权重路径 | 必填 |
--prompt |
文本提示词 | 必填 |
--image |
参考图片（I2V/TI2V/S2V） | 可选 |
--audio |
音频（S2V） | 可选 |
--offload_model |
非活跃模型卸载到CPU | True |
--convert_model_dtype |
BF16精度 | True |
--t5_cpu |
T5编码器放CPU | False |
--dit_fsdp |
DiT使用FSDP分片 | False |
--t5_fsdp |
T5使用FSDP分片 | False |
--ulysses_size |
Ulysses上下文并行度 | 1 |
--shift |
Flow Matching的shift调度 | 模型默认 |
--sampling_steps |
采样步数 | 50 |
--guidance_scale |
CFG引导权重 | 5.0 |
--num_clip |
生成片段数 | 1 |
2.3 多GPU推理命令
# 8卡FSDP + Ulysses并行推理
torchrun --nproc_per_node=8 generate.py \
--task t2v-A14B \
--size 1280*720 \
--ckpt_dir ./Wan2.2-T2V-A14B \
--dit_fsdp \ # DiT参数分片
--t5_fsdp \ # T5参数分片
--ulysses_size 8 \ # Ulysses序列并行
--prompt "..."
设计原理：--offload_model True + --convert_model_dtype 组合是单卡推理的标配——MoE两个专家不同时使用，非活跃的卸载到CPU可节省~50%显存；BF16精度进一步减半。
三、核心配置 wan/configs/
3.1 配置结构
# wan/configs/__init__.py（简化）
# 模型配置字典
WAN_CONFIGS = {
"t2v-A14B": {
"dit_model": "Wan2.2-T2V-A14B", # MoE高/低噪声双专家
"vae_type": "wan2.1", # 使用Wan2.1 VAE
"vae_checkpoint": "vae.cp",
"text_encoder": "umt5-xxl",
"param_dtype": torch.bfloat16,
"num_frames": 41, # 5秒@8fps → 41帧(1+T/4编码后)
"image_size": (1280, 720),
"shift": 5.0, # Flow Matching shift
"guidance_scale": 5.0,
"sampling_steps": 50,
"moe": True, # ★ MoE模式
},
"i2v-A14B": { ... },
"ti2v-5B": {
"dit_model": "Wan2.2-TI2V-5B",
"vae_type": "wan2.2", # ★ 使用高压缩VAE
"vae_checkpoint": "vae2_2.cp",
"param_dtype": torch.bfloat16,
"moe": False, # Dense模型
...
},
"s2v-14B": { ... },
}
# 分辨率配置
SIZE_CONFIGS = {
"t2v-A14B": {
"1280*720": {"num_frames": 41, "height": 720, "width": 1280},
"832*480": {"num_frames": 41, "height": 480, "width": 832},
},
...
}
# 最大面积限制
MAX_AREA_CONFIGS = { ... }
关键配置差异：
| 模型 | VAE类型 | MoE | 压缩比 | 帧数 |
|---|---|---|---|---|
| T2V-A14B | wan2.1 | ✅ | 128× | 41 (5s) |
| I2V-A14B | wan2.1 | ✅ | 128× | 41 |
| TI2V-5B | wan2.2 | ❌ | 1024× | 81 (5s@24fps) |
| S2V-14B | wan2.1 | ❌ | 128× | 变长 |
四、VAE模块详解
4.1 wan/modules/vae2_1.py — Wan2.1 VAE（128×压缩）
class Wan2_1_VAE(nn.Module):
"""3D因果VAE — 时间4×压缩, 空间8×压缩, 16通道, 总128倍"""
def __init__(self, vae_pth, device='cuda'):
super().__init__()
self.encoder = Encoder3D(
dim=128, # 基础通道数
z_dim=4, # 单组件潜变量通道(实际×2=8→最终16)
dim_mult=[1, 2, 4, 4], # 4级下采样通道变化: 128→256→512→512
num_res_blocks=2,
attn_scales=[], # 注意力注入尺度(空=不注入)
temperal_downsample=[True, True, False], # ★ 第1,2级做时间下采样
dropout=0.0,
)
self.decoder = Decoder3D(...)
# GroupNorm → RMSNorm (关键：保证因果性)
# 空间上采样层通道减半：推理内存再降33%
编码流程：
输入: x ∈ R^{B×3×(1+T)×H×W} (像素空间)
↓ patchify(1,2,2) — 无patchify，直接3D卷积
↓ Encoder3D:
│─ Level 1: dim 128→256, 时间下采样2×, 空间下采样2×
│─ Level 2: dim 256→512, 时间下采样2×, 空间下采样2×
│─ Level 3: dim 512→512, 无时间下采样, 空间下采样2×
│─ Level 4: dim 512→512, 无下采样(残差块)
↓ Conv3D(z_dim*2) → mu, log_var
↓ 重参数化: z = mu + sigma * noise
↓ 归一化: z = (z - scale[0]) * scale[1]
输出: z ∈ R^{B×16×(1+T/4)×H/8×W/8} (潜变量空间)
分块因果编码：
def encode(self, x):
"""分块编码，支持无限长视频"""
t = x.shape[2]
iter_ = 1 + (t - 1) // 4 # 第0帧单独，之后每4帧一块
for i in range(iter_):
if i == 0:
# ★ 首帧单独编码（兼容图像生成）
out = self.encoder(x[:, :, :1, ...])
else:
# ★ 后续块：4帧一块，带缓存特征
chunk = x[:, :, 1+4*(i-1):1+4*i, ...]
out_ = self.encoder(chunk) # 内部复用feat_cache
out = torch.cat([out, out_], dim=2)
return out
设计原理：
- 第一帧仅空间压缩（遵循MagViT-v2设计），确保图像/视频统一编码
- 因果分块 = RNN式的隐藏状态传递，确保时间维度上的因果性
- 显存占用 O(1)：只与块大小有关，与视频总长度解耦
4.2 wan/modules/vae2_2.py — Wan2.2 VAE（1024×压缩）★★
class Wan2_2_VAE(nn.Module):
"""高压缩VAE — 时间4×, 空间16×(含patchify 2×), 48通道, 总1024倍"""
def __init__(self, vae_pth, device='cuda'):
super().__init__()
self.patch_size = (1, 2, 2) # ★ 时间1, 空间2×2
self.encoder = Encoder3D(
dim=128,
z_dim=4,
dim_mult=[1, 2, 4, 4],
num_res_blocks=2,
temperal_downsample=[True, True, False], # 同Wan2.1
# 但多了patchify预压缩！
)
# 通道数: 16 → 48（补偿高压缩信息损失）
Patchify额外压缩：
def patchify(x, patch_size=(1, 2, 2)):
"""空间2×2打patch → 空间额外4倍压缩"""
if x.dim() == 5: # [B, C, F, H, W]
x = rearrange(x, "b c f (hq wq) w -> b (c hq wq) f h w",
hq=patch_size[1], wq=patch_size[2])
return x
压缩倍率对比：
| Wan2.1 VAE | Wan2.2 VAE | |
|---|---|---|
| 时间下采样 | 4× | 4× |
| 空间卷积下采样 | 8× | 8× |
| Patchify | 无 | 2×2=4× |
| 总空间 | 8× | 16× |
| 总压缩 | 4×8=32× (×通道) | 4×16=64× (×通道) |
| 通道数 | 16 | 48 |
| 最终压缩比 | 128× | 1024× |
| 720P PSNR | ~30.1 dB | ~32.5 dB |
设计原理：更多通道(48 vs 16)补偿高压缩信息损失。Wan2.2 VAE仅用于TI2V-5B（消费级），A14B MoE模型仍用Wan2.1 VAE——大模型有充足容量处理128×潜变量，而5B小模型必须更激进压缩。
五、DiT模块详解
5.1 wan/modules/dit/wan2_1.py — Wan2.1 DiT（单一Dense模型）
class Wan2_1_DiT(nn.Module):
"""Wan2.1 DiT: 单一14B参数的Dense DiT模型"""
def __init__(self, config):
super().__init__()
self.patch_size = (1, 2, 2)
self.in_channels = 16 # VAE潜变量通道
self.out_channels = 16
self.hidden_size = 5120 # 隐藏维度
self.num_heads = 40 # 注意力头数
self.num_layers = 40 # Transformer层数
# Patchify：3D Conv (1,2,2)
self.patch_embedding = nn.Conv3d(
in_channels, hidden_size, kernel_size=(1,2,2), stride=(1,2,2))
# Transformer Blocks
self.blocks = nn.ModuleList([
WanTransformerBlock(hidden_size, num_heads, ...)
for _ in range(num_layers)
])
# Un-patchify
self.unpatch = nn.Linear(hidden_size, patch_size[1]*patch_size[2]*out_channels)
# 共享时间步MLP（★关键设计：参数减少25%）
self.shared_mlp = nn.Sequential(
nn.Linear(hidden_size, hidden_size * 6),
nn.SiLU(),
nn.Linear(hidden_size * 6, hidden_size * 6),
)
Transformer Block结构：
class WanTransformerBlock(nn.Module):
def __init__(self, hidden_size, num_heads, ...):
# 1. Cross-Attention：注入文本条件
self.cross_attn = Attention(hidden_size, num_heads)
# 2. Full Self-Attention：全局时空建模（非因果！）
self.self_attn = Attention(hidden_size, num_heads)
# 3. FFN/MLP
self.ffn = FeedForward(hidden_size)
# 4. 时间步调制（共享MLP + 块专属偏置）
self.shift_bias = nn.Parameter(torch.zeros(hidden_size * 6))
def forward(self, x, text_emb, timestep_emb):
# 时间步调制：6个参数（scale, shift, gate各×2）
mod_params = self.shared_mlp(timestep_emb) + self.shift_bias
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
mod_params.chunk(6, dim=-1)
# Cross-Attention ← 文本
x = x + gate_msa * self.cross_attn(
norm(x) * (1 + scale_msa) + shift_msa, text_emb)
# Self-Attention ← 全局
x = x + gate_msa * self.self_attn(norm(x))
# FFN
x = x + gate_mlp * self.ffn(norm(x) * (1 + scale_mlp) + shift_mlp)
return x
设计原理：
- Cross-Attention注入文本（非拼接）：确保长序列下指令跟随能力
- 全局注意力（无因果Mask）：视频需全局感知，不同于LLM
- 共享时间步MLP：40层共享1个MLP，各层仅学偏置→参数减少~25%
5.2 wan/modules/dit/wan2_2.py — Wan2.2 DiT（MoE双专家）★★
class Wan2_2_DiT(nn.Module):
"""Wan2.2 MoE DiT: 高噪声专家 + 低噪声专家"""
def __init__(self, config):
super().__init__()
# ★ 两个独立的Dense DiT！
self.high_noise_model = Wan2_1_DiT(config) # 高噪声专家 ~14B
self.low_noise_model = Wan2_1_DiT(config) # 低噪声专家 ~14B
# 总参数 ~27B，但推理时仅14B活跃
self.t_moe = None # ★ MoE切换阈值时间步（加载权重时设置）
def select_expert(self, t):
"""根据当前时间步选择专家"""
if t > self.t_moe:
return self.high_noise_model # 高噪声阶段 → 布局专家
else:
return self.low_noise_model # 低噪声阶段 → 细节专家
def forward(self, x, timestep, text_emb, **kwargs):
t = timestep.item() if isinstance(timestep, torch.Tensor) else timestep
# ★ 选择并加载对应专家
if t > self.t_moe:
model = self._load_expert('high_noise')
else:
model = self._load_expert('low_noise')
return model(x, timestep, text_emb, **kwargs)
def _load_expert(self, required_name):
"""显存管理：激活所需专家，卸载另一个"""
offload_name = 'low_noise_model' if required_name == 'high_noise_model' \
else 'high_noise_model'
if self.offload_model:
# 将非活跃专家卸载到CPU
if next(getattr(self, offload_name).parameters()).device.type == 'cuda':
getattr(self, offload_name).to('cpu')
# 将活跃专家加载到GPU
if next(getattr(self, required_name).parameters()).device.type == 'cpu':
getattr(self, required_name).to(self.device)
return getattr(self, required_name)
SNR切换逻辑（简化）：
# MoE切换点：SNR(t_moe) = SNR_min / 2
# t_moe 通常设置在 ~800步（总1000步的前80%）
# 去噪初期(t大, SNR低) → 高噪声专家
# 去噪后期(t小, SNR高) → 低噪声专家
设计原理：
- 这是时间步级路由，不是传统LLM的token级MoE
- 整个去噪阶段统一切换模型，而非每个token独立路由
- 推理时GPU上始终只有一个14B模型，显存占用不变
- 非活跃专家在CPU上待命，切换时通过PCIe传输（~100ms量级）
六、Pipeline详解
6.1 wan/text2video.py — WanT2V
class WanT2V:
"""文本到视频生成Pipeline"""
def __init__(self, config, checkpoint_dir, device_id=0, rank=0,
t5_cpu=False, offload_model=False):
# 1. 加载文本编码器（umT5）
self.text_encoder = load_t5(config, device, t5_cpu=t5_cpu)
# 2. 加载VAE
if config['vae_type'] == 'wan2.1':
self.vae = Wan2_1_VAE(vae_pth, device)
elif config['vae_type'] == 'wan2.2':
self.vae = Wan2_2_VAE(vae_pth, device)
# 3. 加载DiT
if config.get('moe', False):
self.model = Wan2_2_DiT(config) # MoE双专家
else:
self.model = Wan2_1_DiT(config) # 单一Dense
load_checkpoint(self.model, checkpoint_dir)
# 4. 加载Flow Matching调度器
self.scheduler = FlowMatchScheduler(
shift=config.get('shift', 5.0),
num_train_timesteps=1000,
)
def generate(self, prompt, size, frame_num, shift, sampling_steps,
guidance_scale, seed, offload_model, **kwargs):
# Step 1: 文本编码
text_emb = self.text_encoder(prompt) # umT5编码
# Step 2: 初始化噪声
z = torch.randn(batch, channels, frames, h, w) * scheduler.init_noise_sigma
# Step 3: CFG → 两次前向
if guidance_scale > 1:
text_emb_cfg = torch.cat([text_emb, text_null]) # 有条件+无条件
else:
text_emb_cfg = text_emb
# Step 4: Flow Matching去噪循环
for t in scheduler.timesteps:
# 预测速度
v_pred = self.model(z, t, text_emb_cfg)
# CFG缩放
if guidance_scale > 1:
v_cond, v_uncond = v_pred.chunk(2)
v_pred = v_uncond + guidance_scale * (v_cond - v_uncond)
# ODE步进
z = scheduler.step(v_pred, t, z)
# Step 5: VAE解码
video = self.vae.decode(z)
return video
6.2 wan/image2video.py — WanI2V
class WanI2V:
"""图像到视频生成Pipeline"""
def generate(self, prompt, image, size, ...):
# 1. 编码参考图像
image_tensor = preprocess_image(image) # [-1, 1]
image_latent = self.vae.encode(image_tensor) # [B, 16, 1, H/8, W/8]
# 2. 将图像潜变量作为首帧条件
# ★ 关键：第一帧替换为图像潜变量
z = torch.randn(batch, channels, frames, h, w)
z[:, :, 0:1, :, :] = image_latent[:, :, :, :, :] # 首帧初始化
# 3-5: 同T2V的Flow Matching去噪循环
...
6.3 wan/textimage2video.py — WanTI2V
class WanTI2V:
"""文图统一视频生成Pipeline — 使用高压缩VAE"""
def __init__(self, config, ...):
# ★ 使用Wan2.2高压缩VAE
self.vae = Wan2_2_VAE(vae_pth, device) # 1024×压缩, 48通道
self.model = Wan2_1_DiT(config) # 5B Dense模型
def generate(self, prompt, image=None, ...):
if image is not None:
# I2V模式：首帧条件
image_latent = self.vae.encode(preprocess(image))
else:
# T2V模式：纯噪声
image_latent = None
# ★ 高压缩VAE使得5B模型也能生成720P@24fps
...
6.4 wan/speech2video.py — WanS2V
class WanS2V:
"""语音驱动视频生成Pipeline"""
def generate(self, prompt, image, audio, pose_video=None, ...):
# 1. 音频特征提取
audio_feat = extract_audio_features(audio) # 从音频提取特征
# 2. 参考图像编码
image_latent = self.vae.encode(preprocess(image))
# 3. 可选：姿态驱动
if pose_video is not None:
pose_feat = extract_pose_features(pose_video)
# 姿态特征注入DiT
# 4. 音频特征注入DiT的Cross-Attention
# ★ 音频特征与文本特征一起作为条件
combined_emb = torch.cat([text_emb, audio_feat], dim=1)
# 5-7: Flow Matching去噪 + VAE解码
...
七、Flow Matching调度器
class FlowMatchScheduler:
"""Rectified Flow调度器 — 线性插值噪声路径"""
def __init__(self, shift=5.0, num_train_timesteps=1000):
self.num_train_timesteps = num_train_timesteps
self.shift = shift # ★ 时间步偏移，控制采样密度
def set_timesteps(self, num_inference_steps):
"""设置推理时间步 — 应用shift采样"""
timesteps = torch.linspace(
1.0, 0.0, num_inference_steps + 1) # [1.0 → 0.0]
# ★ Shift采样：偏向中间时间步
timesteps = self._shift_schedule(timesteps)
self.timesteps = timesteps[:-1]
def _shift_schedule(self, timesteps):
"""Shift Schedule：对数正态偏移
使采样密度偏向中间时间步，同时关注粗布局和细细节"""
return 1.0 / (1.0 + (1.0 - timesteps) / timesteps * self.shift)
def step(self, model_output, timestep, sample):
"""ODE步进：x_{t-1} = x_t + v_pred * dt"""
dt = timestep - self.prev_timestep # 时间步差
prev_sample = sample + model_output * dt
return prev_sample
设计原理：
- Rectified Flow使用线性路径（而非DDPM的弯曲路径），ODE采样更高效
shift参数控制时间步采样密度：偏向中间步既关注全局布局又照顾细节- Wan2.1 T2V-14B shift=5.0，Wan2.2 MoE可能使用不同的shift值
- 推理仅需20-50步即可收敛（对比DDPM的1000步）
八、分布式推理 wan/distributed/
# wan/distributed/util.py
def init_distributed_group(args):
"""初始化分布式进程组"""
if args.ulysses_size > 1:
# Ulysses序列并行：沿注意力头分片
setup_ulysses(args.ulysses_size)
if args.dit_fsdp:
# DiT模块使用FSDP参数分片
wrap_dit_with_fsdp()
if args.t5_fsdp:
# T5编码器使用FSDP参数分片
wrap_t5_with_fsdp()
并行策略映射：
┌──────────────────────────────────────────────────────────┐
│ 多GPU推理并行策略 │
├──────────────┬───────────────────────────────────────────┤
│ VAE │ DP（数据并行）— 显存小，无需分片 │
│ T5 Encoder │ FSDP（全分片数据并行）— >20GB显存需分片 │
│ DiT (单专家) │ FSDP + 2D-CP (Ulysses×RingAttention) + DP │
│ DiT (MoE) │ 同上 + 专家卸载（CPU↔GPU切换） │
└──────────────┴───────────────────────────────────────────┘
九、推理效率优化
9.1 模型卸载（Model Offloading）
# 在MoE DiT中，非活跃专家自动卸载到CPU
if self.offload_model:
# GPU → CPU：非活跃专家
# CPU → GPU：活跃专家
# 切换延迟：~100ms（PCIe传输）
9.2 BF16精度
# --convert_model_dtype 启用
model = model.to(torch.bfloat16)
# 显存减半，精度损失可忽略（A100/H100原生支持BF16）
9.3 T5 CPU卸载
# --t5_cpu 启用
# T5编码器放在CPU上，仅在文本编码时使用
# 节省~20GB GPU显存
self.text_encoder = load_t5(config, device='cpu')
9.4 性能参考
| 模型 | 配置 | GPU | 720P生成时间 | 显存 |
|---|---|---|---|---|
| T2V-A14B | 8×GPU, FSDP+Ulysses | A100 | ~2min | ~40GB/GPU |
| T2V-A14B | 单GPU, offload+BF16 | A100 | ~15min | ~80GB |
| TI2V-5B | 单GPU, offload | 4090 | ~9min | ~24GB |
| I2V-A14B | 8×GPU | A100 | ~2min | ~40GB/GPU |
| S2V-14B | 8×GPU | A100 | ~3min | ~40GB/GPU |
十、Prompt扩展
# wan/utils/prompt_extend.py
class DashScopePromptExpander:
"""使用阿里DashScope API扩展提示词"""
def expand(self, prompt, task_type="t2v"):
# 调用qwen-plus（T2V）或qwen-vl-max（I2V）
# 将简短提示词扩展为详细描述
# 例如: "猫拳击" → "两只拟人化猫穿着拳击装备..."
...
class QwenPromptExpander:
"""本地Qwen模型扩展提示词（无需API）"""
...
设计原理：短视频prompt通常过于简短，扩展为细节丰富的描述可显著提升生成质量。论文中推荐的prompt扩展包含运动描述、镜头语言、光线和氛围等。
十一、视频保存与后处理
# wan/utils/utils.py
def save_video(video_tensor, output_path, fps=16, quality=5):
"""将张量视频保存为MP4"""
# video_tensor: [B, C, T, H, W] → [T, H, W, C]
frames = video_tensor.permute(2, 0, 1, 3).cpu().numpy()
# 归一化到[0, 255]
frames = (frames + 1) / 2 * 255
frames = frames.clip(0, 255).astype(np.uint8)
# 使用imageio或ffmpeg写入
...
def merge_video_audio(video_path, audio_path, output_path):
"""合并视频和音频（S2V任务）"""
# 使用ffmpeg合并
subprocess.run([
'ffmpeg', '-i', video_path, '-i', audio_path,
'-c:v', 'copy', '-c:a', 'aac', output_path
])
十二、Wan2.2 vs Wan2.1 代码层面差异
| 代码模块 | Wan2.1 | Wan2.2 | 差异说明 |
|---|---|---|---|
generate.py |
task: t2v-14B/t2v-1.3B/i2v-14B | task: t2v-A14B/i2v-A14B/ti2v-5B/s2v-14B | 新增TI2V和S2V任务 |
configs/ |
WAN_CONFIGS无MoE | WAN_CONFIGS含moe: True |
新增MoE配置+切换阈值 |
vae2_1.py |
✅ 128×压缩, 16ch | ✅ 复用 | T2V/I2V MoE模型仍用 |
vae2_2.py |
— | ✅ 新增 1024×压缩, 48ch | TI2V-5B专用高压缩VAE |
dit/wan2_1.py |
✅ 单一Dense 14B | ✅ 复用 | 作为MoE的单组件 |
dit/wan2_2.py |
— | ✅ 新增 MoE双专家 | 高/低噪声路由+专家卸载 |
text2video.py |
WanT2V | WanT2V | 支持MoE模型加载 |
image2video.py |
WanI2V | WanI2V | 支持MoE模型加载 |
textimage2video.py |
— | WanTI2V | ★ 新增TI2V统一管线+高压缩VAE |
speech2video.py |
— | WanS2V | ★ 新增语音驱动管线 |
| Prompt扩展 | 无 | DashScope/Qwen | ★ 新增prompt自动扩展 |
| 分布式 | FSDP+Ulysses | 同+MoE专家卸载 | 新增--offload_model CPU卸载 |
十三、核心流程图
用户输入: prompt + (可选) image/audio
│
▼
┌─────────────────────────────────────────────────────────┐
│ generate.py — 推理入口 │
│ ├── 解析CLI参数 │
│ ├── 初始化分布式 (FSDP / Ulysses / DP) │
│ ├── 选择Pipeline (WanT2V / WanI2V / WanTI2V / WanS2V) │
│ └── Prompt扩展 (可选, DashScope/Qwen) │
└─────────────────────┬───────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────┐
│ Pipeline.generate() │
│ │
│ ┌─ Step 1: 文本编码 ─────────────────────────────────┐ │
│ │ umT5-XXL → text_emb [B, L_text, D] │ │
│ │ (可选) T5放CPU节省显存 │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ ┌─ Step 2: VAE编码 (仅I2V/TI2V/S2V) ───────────────┐ │
│ │ Wan2.1 VAE: x → z [B,16,T/4,H/8,W/8] (128×) │ │
│ │ Wan2.2 VAE: x → z [B,48,T/4,H/16,W/16] (1024×) │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ ┌─ Step 3: 初始化噪声 ───────────────────────────────┐ │
│ │ z_T ~ N(0, I) [B, C, F, H', W'] │ │
│ │ (I2V: 首帧替换为image_latent) │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ ┌─ Step 4: Flow Matching去噪循环 ────────────────────┐ │
│ │ for t in timesteps: │ │
│ │ ┌─ MoE专家选择 (Wan2.2 A14B) ──────────────┐ │ │
│ │ │ if t > t_moe → high_noise_model (14B) │ │ │
│ │ │ else → low_noise_model (14B) │ │ │
│ │ │ (非活跃专家卸载到CPU) │ │ │
│ │ └────────────────────────────────────────────┘ │ │
│ │ │ │
│ │ v_pred = DiT(z_t, t, text_emb) │ │
│ │ │ │
│ │ ┌─ CFG (guidance_scale > 1) ────────────────┐ │ │
│ │ │ v = v_uncond + s * (v_cond - v_uncond) │ │ │
│ │ └────────────────────────────────────────────┘ │ │
│ │ │ │
│ │ z_{t-dt} = z_t + v * dt (Rectified Flow ODE) │ │
│ └──────────────────────────────────────────────────────┘ │
│ │
│ ┌─ Step 5: VAE解码 ──────────────────────────────────┐ │
│ │ z_0 → VAE Decoder → video [B, 3, T, H, W] │ │
│ │ (分块解码，O(1)显存，支持无限长) │ │
│ └──────────────────────────────────────────────────────┘ │
└─────────────────────┬───────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────┐
│ save_video() / merge_video_audio() │
│ → MP4输出 │
└─────────────────────────────────────────────────────────┘
十四、数据集与模型地址
| 资源 | 地址 |
|---|---|
| Wan2.2 GitHub | https://github.com/Wan-Video/Wan2.2 |
| Wan2.1 GitHub | https://github.com/Wan-Video/Wan2.1 |
| HuggingFace模型 | https://huggingface.co/Wan-AI |
| ModelScope模型 | https://modelscope.cn/organization/Wan-AI |
| 在线体验 | https://tongyi.aliyun.com/wanxiang/ |
| T2V-A14B权重 | https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B |
| I2V-A14B权重 | https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B |
| TI2V-5B权重 | https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B |
| S2V-14B权重 | https://huggingface.co/Wan-AI/Wan2.2-S2V-14B |
| Animate-14B权重 | https://huggingface.co/Wan-AI/Wan2.2-Animate-14B |
| umT5文本编码器 | https://huggingface.co/google/umt5-xxl |
| DiffSynth-Studio（训练/LoRA） | https://github.com/modelscope/DiffSynth-Studio |
| ComfyUI WanVideoWrapper | https://github.com/kijai/ComfyUI-WanVideoWrapper |
十五、总结
核心代码设计思想
- 模块解耦：VAE/DiT/Pipeline三层分离，VAE和DiT独立可替换（Wan2.1 VAE ↔ Wan2.2 VAE）
- MoE透明化：
wan2_2.py包装了双专家路由，上层Pipeline无需关心MoE逻辑 - 显存友好：
offload_model+t5_cpu+BF16组合拳让80GB单卡也能跑14B模型 - 因果VAE：分块编码+特征缓存实现O(1)显存，是长视频推理的基石
- Flow Matching：线性路径ODE，20-50步即收敛，比传统DDPM快20-50×
- 共享调制MLP：40层共享1个时间步MLP，参数减少25%无性能损失
代码演进方向
| 方向 | 当前状态 | 未来可能 |
|---|---|---|
| MoE架构 | 时间步级两专家切换 | Token级动态路由（DiffMOE风格） |
| VAE | Wan2.1(128×) + Wan2.2(1024×) | 统一高压缩VAE |
| 推理加速 | offload+BF16+FSDP | 量化(INT8/FP8)+蒸馏+Cache |
| 训练 | 仅开源推理代码 | 可能开源训练代码 |
| 下游任务 | T2V/I2V/TI2V/S2V/Animate | 更多可控生成任务 |
免费领 150 小时云算力，进群参与显卡、AI PC 幸运抽奖
更多推荐
4
0- 0
已为社区贡献2条内容
所有评论(0)