ALSv/self-forcing
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import logging3 4import torch5import torch.cuda.amp as amp6import torch.nn as nn7import torch.nn.functional as F8from einops import rearrange9 10__all__ = [11 'WanVAE',12]13 14CACHE_T = 215 16 17class CausalConv3d(nn.Conv3d):18 """19 Causal 3d convolusion.20 """21 22 def __init__(self, *args, **kwargs):23 super().__init__(*args, **kwargs)24 self._padding = (self.padding[2], self.padding[2], self.padding[1],25 self.padding[1], 2 * self.padding[0], 0)26 self.padding = (0, 0, 0)27 28 def forward(self, x, cache_x=None):29 padding = list(self._padding)30 if cache_x is not None and self._padding[4] > 0:31 cache_x = cache_x.to(x.device)32 x = torch.cat([cache_x, x], dim=2)33 padding[4] -= cache_x.shape[2]34 x = F.pad(x, padding)35 36 return super().forward(x)37 38 39class RMS_norm(nn.Module):40 41 def __init__(self, dim, channel_first=True, images=True, bias=False):42 super().__init__()43 broadcastable_dims = (1, 1, 1) if not images else (1, 1)44 shape = (dim, *broadcastable_dims) if channel_first else (dim,)45 46 self.channel_first = channel_first47 self.scale = dim**0.548 self.gamma = nn.Parameter(torch.ones(shape))49 self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.50 51 def forward(self, x):52 return F.normalize(53 x, dim=(1 if self.channel_first else54 -1)) * self.scale * self.gamma + self.bias55 56 57class Upsample(nn.Upsample):58 59 def forward(self, x):60 """61 Fix bfloat16 support for nearest neighbor interpolation.62 """63 return super().forward(x.float()).type_as(x)64 65 66class Resample(nn.Module):67 68 def __init__(self, dim, mode):69 assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',70 'downsample3d')71 super().__init__()72 self.dim = dim73 self.mode = mode74 75 # layers76 if mode == 'upsample2d':77 self.resample = nn.Sequential(78 Upsample(scale_factor=(2., 2.), mode='nearest'),79 nn.Conv2d(dim, dim // 2, 3, padding=1))80 elif mode == 'upsample3d':81 self.resample = nn.Sequential(82 Upsample(scale_factor=(2., 2.), mode='nearest'),83 nn.Conv2d(dim, dim // 2, 3, padding=1))84 self.time_conv = CausalConv3d(85 dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))86 87 elif mode == 'downsample2d':88 self.resample = nn.Sequential(89 nn.ZeroPad2d((0, 1, 0, 1)),90 nn.Conv2d(dim, dim, 3, stride=(2, 2)))91 elif mode == 'downsample3d':92 self.resample = nn.Sequential(93 nn.ZeroPad2d((0, 1, 0, 1)),94 nn.Conv2d(dim, dim, 3, stride=(2, 2)))95 self.time_conv = CausalConv3d(96 dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))97 98 else:99 self.resample = nn.Identity()100 101 def forward(self, x, feat_cache=None, feat_idx=[0]):102 b, c, t, h, w = x.size()103 if self.mode == 'upsample3d':104 if feat_cache is not None:105 idx = feat_idx[0]106 if feat_cache[idx] is None:107 feat_cache[idx] = 'Rep'108 feat_idx[0] += 1109 else:110 111 cache_x = x[:, :, -CACHE_T:, :, :].clone()112 if cache_x.shape[2] < 2 and feat_cache[113 idx] is not None and feat_cache[idx] != 'Rep':114 # cache last frame of last two chunk115 cache_x = torch.cat([116 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(117 cache_x.device), cache_x118 ],119 dim=2)120 if cache_x.shape[2] < 2 and feat_cache[121 idx] is not None and feat_cache[idx] == 'Rep':122 cache_x = torch.cat([123 torch.zeros_like(cache_x).to(cache_x.device),124 cache_x125 ],126 dim=2)127 if feat_cache[idx] == 'Rep':128 x = self.time_conv(x)129 else:130 x = self.time_conv(x, feat_cache[idx])131 feat_cache[idx] = cache_x132 feat_idx[0] += 1133 134 x = x.reshape(b, 2, c, t, h, w)135 x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),136 3)137 x = x.reshape(b, c, t * 2, h, w)138 t = x.shape[2]139 x = rearrange(x, 'b c t h w -> (b t) c h w')140 x = self.resample(x)141 x = rearrange(x, '(b t) c h w -> b c t h w', t=t)142 143 if self.mode == 'downsample3d':144 if feat_cache is not None:145 idx = feat_idx[0]146 if feat_cache[idx] is None:147 feat_cache[idx] = x.clone()148 feat_idx[0] += 1149 else:150 151 cache_x = x[:, :, -1:, :, :].clone()152 # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':153 # # cache last frame of last two chunk154 # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)155 156 x = self.time_conv(157 torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))158 feat_cache[idx] = cache_x159 feat_idx[0] += 1160 return x161 162 def init_weight(self, conv):163 conv_weight = conv.weight164 nn.init.zeros_(conv_weight)165 c1, c2, t, h, w = conv_weight.size()166 one_matrix = torch.eye(c1, c2)167 init_matrix = one_matrix168 nn.init.zeros_(conv_weight)169 # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5170 conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5171 conv.weight.data.copy_(conv_weight)172 nn.init.zeros_(conv.bias.data)173 174 def init_weight2(self, conv):175 conv_weight = conv.weight.data176 nn.init.zeros_(conv_weight)177 c1, c2, t, h, w = conv_weight.size()178 init_matrix = torch.eye(c1 // 2, c2)179 # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)180 conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix181 conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix182 conv.weight.data.copy_(conv_weight)183 nn.init.zeros_(conv.bias.data)184 185 186class ResidualBlock(nn.Module):187 188 def __init__(self, in_dim, out_dim, dropout=0.0):189 super().__init__()190 self.in_dim = in_dim191 self.out_dim = out_dim192 193 # layers194 self.residual = nn.Sequential(195 RMS_norm(in_dim, images=False), nn.SiLU(),196 CausalConv3d(in_dim, out_dim, 3, padding=1),197 RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),198 CausalConv3d(out_dim, out_dim, 3, padding=1))199 self.shortcut = CausalConv3d(in_dim, out_dim, 1) \200 if in_dim != out_dim else nn.Identity()201 202 def forward(self, x, feat_cache=None, feat_idx=[0]):203 h = self.shortcut(x)204 for layer in self.residual:205 if isinstance(layer, CausalConv3d) and feat_cache is not None:206 idx = feat_idx[0]207 cache_x = x[:, :, -CACHE_T:, :, :].clone()208 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:209 # cache last frame of last two chunk210 cache_x = torch.cat([211 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(212 cache_x.device), cache_x213 ],214 dim=2)215 x = layer(x, feat_cache[idx])216 feat_cache[idx] = cache_x217 feat_idx[0] += 1218 else:219 x = layer(x)220 return x + h221 222 223class AttentionBlock(nn.Module):224 """225 Causal self-attention with a single head.226 """227 228 def __init__(self, dim):229 super().__init__()230 self.dim = dim231 232 # layers233 self.norm = RMS_norm(dim)234 self.to_qkv = nn.Conv2d(dim, dim * 3, 1)235 self.proj = nn.Conv2d(dim, dim, 1)236 237 # zero out the last layer params238 nn.init.zeros_(self.proj.weight)239 240 def forward(self, x):241 identity = x242 b, c, t, h, w = x.size()243 x = rearrange(x, 'b c t h w -> (b t) c h w')244 x = self.norm(x)245 # compute query, key, value246 q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,247 -1).permute(0, 1, 3,248 2).contiguous().chunk(249 3, dim=-1)250 251 # apply attention252 x = F.scaled_dot_product_attention(253 q,254 k,255 v,256 )257 x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)258 259 # output260 x = self.proj(x)261 x = rearrange(x, '(b t) c h w-> b c t h w', t=t)262 return x + identity263 264 265class Encoder3d(nn.Module):266 267 def __init__(self,268 dim=128,269 z_dim=4,270 dim_mult=[1, 2, 4, 4],271 num_res_blocks=2,272 attn_scales=[],273 temperal_downsample=[True, True, False],274 dropout=0.0):275 super().__init__()276 self.dim = dim277 self.z_dim = z_dim278 self.dim_mult = dim_mult279 self.num_res_blocks = num_res_blocks280 self.attn_scales = attn_scales281 self.temperal_downsample = temperal_downsample282 283 # dimensions284 dims = [dim * u for u in [1] + dim_mult]285 scale = 1.0286 287 # init block288 self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)289 290 # downsample blocks291 downsamples = []292 for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):293 # residual (+attention) blocks294 for _ in range(num_res_blocks):295 downsamples.append(ResidualBlock(in_dim, out_dim, dropout))296 if scale in attn_scales:297 downsamples.append(AttentionBlock(out_dim))298 in_dim = out_dim299 300 # downsample block301 if i != len(dim_mult) - 1:302 mode = 'downsample3d' if temperal_downsample[303 i] else 'downsample2d'304 downsamples.append(Resample(out_dim, mode=mode))305 scale /= 2.0306 self.downsamples = nn.Sequential(*downsamples)307 308 # middle blocks309 self.middle = nn.Sequential(310 ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),311 ResidualBlock(out_dim, out_dim, dropout))312 313 # output blocks314 self.head = nn.Sequential(315 RMS_norm(out_dim, images=False), nn.SiLU(),316 CausalConv3d(out_dim, z_dim, 3, padding=1))317 318 def forward(self, x, feat_cache=None, feat_idx=[0]):319 if feat_cache is not None:320 idx = feat_idx[0]321 cache_x = x[:, :, -CACHE_T:, :, :].clone()322 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:323 # cache last frame of last two chunk324 cache_x = torch.cat([325 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(326 cache_x.device), cache_x327 ],328 dim=2)329 x = self.conv1(x, feat_cache[idx])330 feat_cache[idx] = cache_x331 feat_idx[0] += 1332 else:333 x = self.conv1(x)334 335 # downsamples336 for layer in self.downsamples:337 if feat_cache is not None:338 x = layer(x, feat_cache, feat_idx)339 else:340 x = layer(x)341 342 # middle343 for layer in self.middle:344 if isinstance(layer, ResidualBlock) and feat_cache is not None:345 x = layer(x, feat_cache, feat_idx)346 else:347 x = layer(x)348 349 # head350 for layer in self.head:351 if isinstance(layer, CausalConv3d) and feat_cache is not None:352 idx = feat_idx[0]353 cache_x = x[:, :, -CACHE_T:, :, :].clone()354 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:355 # cache last frame of last two chunk356 cache_x = torch.cat([357 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(358 cache_x.device), cache_x359 ],360 dim=2)361 x = layer(x, feat_cache[idx])362 feat_cache[idx] = cache_x363 feat_idx[0] += 1364 else:365 x = layer(x)366 return x367 368 369class Decoder3d(nn.Module):370 371 def __init__(self,372 dim=128,373 z_dim=4,374 dim_mult=[1, 2, 4, 4],375 num_res_blocks=2,376 attn_scales=[],377 temperal_upsample=[False, True, True],378 dropout=0.0):379 super().__init__()380 self.dim = dim381 self.z_dim = z_dim382 self.dim_mult = dim_mult383 self.num_res_blocks = num_res_blocks384 self.attn_scales = attn_scales385 self.temperal_upsample = temperal_upsample386 387 # dimensions388 dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]389 scale = 1.0 / 2**(len(dim_mult) - 2)390 391 # init block392 self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)393 394 # middle blocks395 self.middle = nn.Sequential(396 ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),397 ResidualBlock(dims[0], dims[0], dropout))398 399 # upsample blocks400 upsamples = []401 for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):402 # residual (+attention) blocks403 if i == 1 or i == 2 or i == 3:404 in_dim = in_dim // 2405 for _ in range(num_res_blocks + 1):406 upsamples.append(ResidualBlock(in_dim, out_dim, dropout))407 if scale in attn_scales:408 upsamples.append(AttentionBlock(out_dim))409 in_dim = out_dim410 411 # upsample block412 if i != len(dim_mult) - 1:413 mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'414 upsamples.append(Resample(out_dim, mode=mode))415 scale *= 2.0416 self.upsamples = nn.Sequential(*upsamples)417 418 # output blocks419 self.head = nn.Sequential(420 RMS_norm(out_dim, images=False), nn.SiLU(),421 CausalConv3d(out_dim, 3, 3, padding=1))422 423 def forward(self, x, feat_cache=None, feat_idx=[0]):424 # conv1425 if feat_cache is not None:426 idx = feat_idx[0]427 cache_x = x[:, :, -CACHE_T:, :, :].clone()428 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:429 # cache last frame of last two chunk430 cache_x = torch.cat([431 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(432 cache_x.device), cache_x433 ],434 dim=2)435 x = self.conv1(x, feat_cache[idx])436 feat_cache[idx] = cache_x437 feat_idx[0] += 1438 else:439 x = self.conv1(x)440 441 # middle442 for layer in self.middle:443 if isinstance(layer, ResidualBlock) and feat_cache is not None:444 x = layer(x, feat_cache, feat_idx)445 else:446 x = layer(x)447 448 # upsamples449 for layer in self.upsamples:450 if feat_cache is not None:451 x = layer(x, feat_cache, feat_idx)452 else:453 x = layer(x)454 455 # head456 for layer in self.head:457 if isinstance(layer, CausalConv3d) and feat_cache is not None:458 idx = feat_idx[0]459 cache_x = x[:, :, -CACHE_T:, :, :].clone()460 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:461 # cache last frame of last two chunk462 cache_x = torch.cat([463 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(464 cache_x.device), cache_x465 ],466 dim=2)467 x = layer(x, feat_cache[idx])468 feat_cache[idx] = cache_x469 feat_idx[0] += 1470 else:471 x = layer(x)472 return x473 474 475def count_conv3d(model):476 count = 0477 for m in model.modules():478 if isinstance(m, CausalConv3d):479 count += 1480 return count481 482 483class WanVAE_(nn.Module):484 485 def __init__(self,486 dim=128,487 z_dim=4,488 dim_mult=[1, 2, 4, 4],489 num_res_blocks=2,490 attn_scales=[],491 temperal_downsample=[True, True, False],492 dropout=0.0):493 super().__init__()494 self.dim = dim495 self.z_dim = z_dim496 self.dim_mult = dim_mult497 self.num_res_blocks = num_res_blocks498 self.attn_scales = attn_scales499 self.temperal_downsample = temperal_downsample500 self.temperal_upsample = temperal_downsample[::-1]501 502 # modules503 self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,504 attn_scales, self.temperal_downsample, dropout)505 self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)506 self.conv2 = CausalConv3d(z_dim, z_dim, 1)507 self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,508 attn_scales, self.temperal_upsample, dropout)509 self.clear_cache()510 511 def forward(self, x):512 mu, log_var = self.encode(x)513 z = self.reparameterize(mu, log_var)514 x_recon = self.decode(z)515 return x_recon, mu, log_var516 517 def encode(self, x, scale):518 self.clear_cache()519 # cache520 t = x.shape[2]521 iter_ = 1 + (t - 1) // 4522 # 对encode输入的x，按时间拆分为1、4、4、4....523 for i in range(iter_):524 self._enc_conv_idx = [0]525 if i == 0:526 out = self.encoder(527 x[:, :, :1, :, :],528 feat_cache=self._enc_feat_map,529 feat_idx=self._enc_conv_idx)530 else:531 out_ = self.encoder(532 x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],533 feat_cache=self._enc_feat_map,534 feat_idx=self._enc_conv_idx)535 out = torch.cat([out, out_], 2)536 mu, log_var = self.conv1(out).chunk(2, dim=1)537 if isinstance(scale[0], torch.Tensor):538 mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(539 1, self.z_dim, 1, 1, 1)540 else:541 mu = (mu - scale[0]) * scale[1]542 self.clear_cache()543 return mu544 545 def decode(self, z, scale):546 self.clear_cache()547 # z: [b,c,t,h,w]548 if isinstance(scale[0], torch.Tensor):549 z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(550 1, self.z_dim, 1, 1, 1)551 else:552 z = z / scale[1] + scale[0]553 iter_ = z.shape[2]554 x = self.conv2(z)555 for i in range(iter_):556 self._conv_idx = [0]557 if i == 0:558 out = self.decoder(559 x[:, :, i:i + 1, :, :],560 feat_cache=self._feat_map,561 feat_idx=self._conv_idx)562 else:563 out_ = self.decoder(564 x[:, :, i:i + 1, :, :],565 feat_cache=self._feat_map,566 feat_idx=self._conv_idx)567 out = torch.cat([out, out_], 2)568 self.clear_cache()569 return out570 571 def cached_decode(self, z, scale):572 # z: [b,c,t,h,w]573 if isinstance(scale[0], torch.Tensor):574 z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(575 1, self.z_dim, 1, 1, 1)576 else:577 z = z / scale[1] + scale[0]578 iter_ = z.shape[2]579 x = self.conv2(z)580 for i in range(iter_):581 self._conv_idx = [0]582 if i == 0:583 out = self.decoder(584 x[:, :, i:i + 1, :, :],585 feat_cache=self._feat_map,586 feat_idx=self._conv_idx)587 else:588 out_ = self.decoder(589 x[:, :, i:i + 1, :, :],590 feat_cache=self._feat_map,591 feat_idx=self._conv_idx)592 out = torch.cat([out, out_], 2)593 return out594 595 def sample(self, imgs, deterministic=False):596 mu, log_var = self.encode(imgs)597 if deterministic:598 return mu599 std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))600 return mu + std * torch.randn_like(std)601 602 def clear_cache(self):603 self._conv_num = count_conv3d(self.decoder)604 self._conv_idx = [0]605 self._feat_map = [None] * self._conv_num606 # cache encode607 self._enc_conv_num = count_conv3d(self.encoder)608 self._enc_conv_idx = [0]609 self._enc_feat_map = [None] * self._enc_conv_num610 611 612def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):613 """614 Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.615 """616 # params617 cfg = dict(618 dim=96,619 z_dim=z_dim,620 dim_mult=[1, 2, 4, 4],621 num_res_blocks=2,622 attn_scales=[],623 temperal_downsample=[False, True, True],624 dropout=0.0)625 cfg.update(**kwargs)626 627 # init model628 with torch.device('meta'):629 model = WanVAE_(**cfg)630 631 # load checkpoint632 logging.info(f'loading {pretrained_path}')633 model.load_state_dict(634 torch.load(pretrained_path, map_location=device), assign=True)635 636 return model637 638 639class WanVAE:640 641 def __init__(self,642 z_dim=16,643 vae_pth='cache/vae_step_411000.pth',644 dtype=torch.float,645 device="cuda"):646 self.dtype = dtype647 self.device = device648 649 mean = [650 -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,651 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921652 ]653 std = [654 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,655 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160656 ]657 self.mean = torch.tensor(mean, dtype=dtype, device=device)658 self.std = torch.tensor(std, dtype=dtype, device=device)659 self.scale = [self.mean, 1.0 / self.std]660 661 # init model662 self.model = _video_vae(663 pretrained_path=vae_pth,664 z_dim=z_dim,665 ).eval().requires_grad_(False).to(device)666 667 def encode(self, videos):668 """669 videos: A list of videos each with shape [C, T, H, W].670 """671 with amp.autocast(dtype=self.dtype):672 return [673 self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)674 for u in videos675 ]676 677 def decode(self, zs):678 with amp.autocast(dtype=self.dtype):679 return [680 self.model.decode(u.unsqueeze(0),681 self.scale).float().clamp_(-1, 1).squeeze(0)682 for u in zs683 ]684