ALSv/self-forcing
0
1from typing import List2from einops import rearrange3import torch4import torch.nn as nn5 6from wan.modules.vae import AttentionBlock, CausalConv3d, RMS_norm, ResidualBlock, Upsample7 8 9class Resample(nn.Module):10 11 def __init__(self, dim, mode):12 assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',13 'downsample3d')14 super().__init__()15 self.dim = dim16 self.mode = mode17 self.cache_t = 218 19 # layers20 if mode == 'upsample2d':21 self.resample = nn.Sequential(22 Upsample(scale_factor=(2., 2.), mode='nearest'),23 nn.Conv2d(dim, dim // 2, 3, padding=1))24 elif mode == 'upsample3d':25 self.resample = nn.Sequential(26 Upsample(scale_factor=(2., 2.), mode='nearest'),27 nn.Conv2d(dim, dim // 2, 3, padding=1))28 self.time_conv = CausalConv3d(29 dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))30 31 elif mode == 'downsample2d':32 self.resample = nn.Sequential(33 nn.ZeroPad2d((0, 1, 0, 1)),34 nn.Conv2d(dim, dim, 3, stride=(2, 2)))35 elif mode == 'downsample3d':36 self.resample = nn.Sequential(37 nn.ZeroPad2d((0, 1, 0, 1)),38 nn.Conv2d(dim, dim, 3, stride=(2, 2)))39 self.time_conv = CausalConv3d(40 dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))41 42 else:43 self.resample = nn.Identity()44 45 def forward(self, x, feat_cache=None, feat_idx=[0]):46 b, c, t, h, w = x.size()47 if self.mode == 'upsample3d':48 if feat_cache is not None:49 idx = feat_idx[0]50 if feat_cache[idx] is None:51 feat_cache[idx] = 'Rep'52 feat_idx[0] += 153 else:54 55 cache_x = x[:, :, -self.cache_t:, :, :].clone()56 if cache_x.shape[2] < 2 and feat_cache[57 idx] is not None and feat_cache[idx] != 'Rep':58 # cache last frame of last two chunk59 cache_x = torch.cat([60 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(61 cache_x.device), cache_x62 ],63 dim=2)64 if cache_x.shape[2] < 2 and feat_cache[65 idx] is not None and feat_cache[idx] == 'Rep':66 cache_x = torch.cat([67 torch.zeros_like(cache_x).to(cache_x.device),68 cache_x69 ],70 dim=2)71 if feat_cache[idx] == 'Rep':72 x = self.time_conv(x)73 else:74 x = self.time_conv(x, feat_cache[idx])75 feat_cache[idx] = cache_x76 feat_idx[0] += 177 78 x = x.reshape(b, 2, c, t, h, w)79 x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),80 3)81 x = x.reshape(b, c, t * 2, h, w)82 t = x.shape[2]83 x = rearrange(x, 'b c t h w -> (b t) c h w')84 x = self.resample(x)85 x = rearrange(x, '(b t) c h w -> b c t h w', t=t)86 87 if self.mode == 'downsample3d':88 if feat_cache is not None:89 idx = feat_idx[0]90 if feat_cache[idx] is None:91 feat_cache[idx] = x.clone()92 feat_idx[0] += 193 else:94 95 cache_x = x[:, :, -1:, :, :].clone()96 # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':97 # # cache last frame of last two chunk98 # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)99 100 x = self.time_conv(101 torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))102 feat_cache[idx] = cache_x103 feat_idx[0] += 1104 return x105 106 def init_weight(self, conv):107 conv_weight = conv.weight108 nn.init.zeros_(conv_weight)109 c1, c2, t, h, w = conv_weight.size()110 one_matrix = torch.eye(c1, c2)111 init_matrix = one_matrix112 nn.init.zeros_(conv_weight)113 # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5114 conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5115 conv.weight.data.copy_(conv_weight)116 nn.init.zeros_(conv.bias.data)117 118 def init_weight2(self, conv):119 conv_weight = conv.weight.data120 nn.init.zeros_(conv_weight)121 c1, c2, t, h, w = conv_weight.size()122 init_matrix = torch.eye(c1 // 2, c2)123 # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)124 conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix125 conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix126 conv.weight.data.copy_(conv_weight)127 nn.init.zeros_(conv.bias.data)128 129 130class VAEDecoderWrapper(nn.Module):131 def __init__(self):132 super().__init__()133 self.decoder = VAEDecoder3d()134 mean = [135 -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,136 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921137 ]138 std = [139 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,140 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160141 ]142 self.mean = torch.tensor(mean, dtype=torch.float32)143 self.std = torch.tensor(std, dtype=torch.float32)144 self.z_dim = 16145 self.conv2 = CausalConv3d(self.z_dim, self.z_dim, 1)146 147 def forward(148 self,149 z: torch.Tensor,150 *feat_cache: List[torch.Tensor]151 ):152 # from [batch_size, num_frames, num_channels, height, width]153 # to [batch_size, num_channels, num_frames, height, width]154 z = z.permute(0, 2, 1, 3, 4)155 feat_cache = list(feat_cache)156 print("Length of feat_cache: ", len(feat_cache))157 158 device, dtype = z.device, z.dtype159 scale = [self.mean.to(device=device, dtype=dtype),160 1.0 / self.std.to(device=device, dtype=dtype)]161 162 if isinstance(scale[0], torch.Tensor):163 z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(164 1, self.z_dim, 1, 1, 1)165 else:166 z = z / scale[1] + scale[0]167 iter_ = z.shape[2]168 x = self.conv2(z)169 for i in range(iter_):170 if i == 0:171 out, feat_cache = self.decoder(172 x[:, :, i:i + 1, :, :],173 feat_cache=feat_cache)174 else:175 out_, feat_cache = self.decoder(176 x[:, :, i:i + 1, :, :],177 feat_cache=feat_cache)178 out = torch.cat([out, out_], 2)179 180 out = out.float().clamp_(-1, 1)181 # from [batch_size, num_channels, num_frames, height, width]182 # to [batch_size, num_frames, num_channels, height, width]183 out = out.permute(0, 2, 1, 3, 4)184 return out, feat_cache185 186 187class VAEDecoder3d(nn.Module):188 def __init__(self,189 dim=96,190 z_dim=16,191 dim_mult=[1, 2, 4, 4],192 num_res_blocks=2,193 attn_scales=[],194 temperal_upsample=[True, True, False],195 dropout=0.0):196 super().__init__()197 self.dim = dim198 self.z_dim = z_dim199 self.dim_mult = dim_mult200 self.num_res_blocks = num_res_blocks201 self.attn_scales = attn_scales202 self.temperal_upsample = temperal_upsample203 self.cache_t = 2204 self.decoder_conv_num = 32205 206 # dimensions207 dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]208 scale = 1.0 / 2**(len(dim_mult) - 2)209 210 # init block211 self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)212 213 # middle blocks214 self.middle = nn.Sequential(215 ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),216 ResidualBlock(dims[0], dims[0], dropout))217 218 # upsample blocks219 upsamples = []220 for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):221 # residual (+attention) blocks222 if i == 1 or i == 2 or i == 3:223 in_dim = in_dim // 2224 for _ in range(num_res_blocks + 1):225 upsamples.append(ResidualBlock(in_dim, out_dim, dropout))226 if scale in attn_scales:227 upsamples.append(AttentionBlock(out_dim))228 in_dim = out_dim229 230 # upsample block231 if i != len(dim_mult) - 1:232 mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'233 upsamples.append(Resample(out_dim, mode=mode))234 scale *= 2.0235 self.upsamples = nn.Sequential(*upsamples)236 237 # output blocks238 self.head = nn.Sequential(239 RMS_norm(out_dim, images=False), nn.SiLU(),240 CausalConv3d(out_dim, 3, 3, padding=1))241 242 def forward(243 self,244 x: torch.Tensor,245 feat_cache: List[torch.Tensor]246 ):247 feat_idx = [0]248 249 # conv1250 idx = feat_idx[0]251 cache_x = x[:, :, -self.cache_t:, :, :].clone()252 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:253 # cache last frame of last two chunk254 cache_x = torch.cat([255 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(256 cache_x.device), cache_x257 ],258 dim=2)259 x = self.conv1(x, feat_cache[idx])260 feat_cache[idx] = cache_x261 feat_idx[0] += 1262 263 # middle264 for layer in self.middle:265 if isinstance(layer, ResidualBlock) and feat_cache is not None:266 x = layer(x, feat_cache, feat_idx)267 else:268 x = layer(x)269 270 # upsamples271 for layer in self.upsamples:272 x = layer(x, feat_cache, feat_idx)273 274 # head275 for layer in self.head:276 if isinstance(layer, CausalConv3d) and feat_cache is not None:277 idx = feat_idx[0]278 cache_x = x[:, :, -self.cache_t:, :, :].clone()279 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:280 # cache last frame of last two chunk281 cache_x = torch.cat([282 feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(283 cache_x.device), cache_x284 ],285 dim=2)286 x = layer(x, feat_cache[idx])287 feat_cache[idx] = cache_x288 feat_idx[0] += 1289 else:290 x = layer(x)291 return x, feat_cache292