diff --git a/wanvideo/wan_video_vae.py b/wanvideo/wan_video_vae.py index 3fe2e14..a5758f1 100644 --- a/wanvideo/wan_video_vae.py +++ b/wanvideo/wan_video_vae.py @@ -271,7 +271,7 @@ class ResidualBlock(nn.Module): CausalConv3d(out_dim, out_dim, 3, padding=1)) self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ if in_dim != out_dim else nn.Identity() - + def forward(self, x, feat_cache=None, feat_idx=[0]): if self.cpu_cache: return self._forward_cpu_cache(x, feat_cache, feat_idx) @@ -297,7 +297,7 @@ class ResidualBlock(nn.Module): else: x = layer(x) return x + h - + def _forward_cpu_cache(self, x, feat_cache=None, feat_idx=[0]): h = self.shortcut(x) for layer in self.residual: @@ -307,11 +307,11 @@ class ResidualBlock(nn.Module): if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cached_frame = feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device) cache_x = torch.cat([cached_frame, cache_x], dim=2) - + prev_cache = feat_cache[idx].to(x.device) if feat_cache[idx] is not None else None x = layer(x, prev_cache) - + feat_cache[idx] = cache_x.to("cpu", non_blocking=True) feat_idx[0] += 1 else: @@ -1097,9 +1097,9 @@ class VideoVAE_(nn.Module): #modification originally by @raindrop313 https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames def decode_2(self, z): # z: [b,c,t,h,w] - + z = z / self.inv_std.to(z) + self.mean.to(z) - + iter_ = z.shape[2] z_head=z[:,:,:-1,:,:] z_tail=z[:,:,-1,:,:].unsqueeze(2) @@ -1288,7 +1288,7 @@ class WanVideoVAE(nn.Module): def tiled_encode(self, video, device, tile_size, tile_stride, end_=False, pbar=True): _, _, T, H, W = video.shape - + if tile_size is None and tile_stride is None: size_h, size_w = H //2, W // 2 stride_h, stride_w = size_h // 2, size_w // 2 @@ -1492,7 +1492,7 @@ class VideoVAE38_(VideoVAE_): except: pass z = z / self.inv_std.to(z) + self.mean.to(z) - + iter_ = z.shape[2] if pbar: pbar = ProgressBar(iter_)