from einops import rearrange, repeat import torch import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm from comfy.utils import ProgressBar from ..utils import print_memory, log from comfy import model_management as mm device = mm.get_torch_device() import comfy.ops ops = comfy.ops.disable_weight_init CACHE_T = 2 def check_is_instance(model, module_class): if isinstance(model, module_class): return True if hasattr(model, "module") and isinstance(model.module, module_class): return True return False class CausalConv3d(ops.Conv3d): """ Causal 3d convolusion. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 2 * self.padding[0], 0) self.padding = (0, 0, 0) def forward(self, x, cache_x=None): padding = list(self._padding) if cache_x is not None and self._padding[4] > 0: cache_x = cache_x.to(x.device) x = torch.cat([cache_x, x], dim=2) padding[4] -= cache_x.shape[2] x = F.pad(x, padding) return super().forward(x) class RMS_norm(nn.Module): def __init__(self, dim, channel_first=True, images=True, bias=False): super().__init__() broadcastable_dims = (1, 1, 1) if not images else (1, 1) shape = (dim, *broadcastable_dims) if channel_first else (dim,) self.channel_first = channel_first self.scale = dim**0.5 self.gamma = nn.Parameter(torch.ones(shape)) self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. def forward(self, x): return F.normalize( x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias class Upsample(nn.Upsample): def forward(self, x): """ Fix bfloat16 support for nearest neighbor interpolation. """ return super().forward(x.float()).type_as(x) class Resample(nn.Module): def __init__(self, dim, mode): assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', 'downsample3d') super().__init__() self.dim = dim self.mode = mode # layers if mode == 'upsample2d': self.resample = nn.Sequential( Upsample(scale_factor=(2., 2.), mode='nearest-exact'), nn.Conv2d(dim, dim // 2, 3, padding=1)) elif mode == 'upsample3d': self.resample = nn.Sequential( Upsample(scale_factor=(2., 2.), mode='nearest-exact'), nn.Conv2d(dim, dim // 2, 3, padding=1)) self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == 'downsample2d': self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) elif mode == 'downsample3d': self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) else: self.resample = nn.Identity() def forward(self, x, feat_cache=None, feat_idx=[0]): b, c, t, h, w = x.size() if self.mode == 'upsample3d': if feat_cache is not None: idx = feat_idx[0] if feat_cache[idx] is None: feat_cache[idx] = 'Rep' feat_idx[0] += 1 else: cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[ idx] is not None and feat_cache[idx] != 'Rep': # cache last frame of last two chunk cache_x = torch.cat([ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( cache_x.device), cache_x ], dim=2) if cache_x.shape[2] < 2 and feat_cache[ idx] is not None and feat_cache[idx] == 'Rep': cache_x = torch.cat([ torch.zeros_like(cache_x).to(cache_x.device), cache_x ], dim=2) if feat_cache[idx] == 'Rep': x = self.time_conv(x) else: x = self.time_conv(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 x = x.reshape(b, 2, c, t, h, w) x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) t = x.shape[2] x = rearrange(x, 'b c t h w -> (b t) c h w') x = self.resample(x) x = rearrange(x, '(b t) c h w -> b c t h w', t=t) if self.mode == 'downsample3d': if feat_cache is not None: idx = feat_idx[0] if feat_cache[idx] is None: feat_cache[idx] = x.clone() feat_idx[0] += 1 else: cache_x = x[:, :, -1:, :, :].clone() x = self.time_conv( torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) feat_cache[idx] = cache_x feat_idx[0] += 1 return x def init_weight(self, conv): conv_weight = conv.weight nn.init.zeros_(conv_weight) c1, c2, t, h, w = conv_weight.size() one_matrix = torch.eye(c1, c2) init_matrix = one_matrix nn.init.zeros_(conv_weight) conv_weight.data[:, :, 1, 0, 0] = init_matrix conv.weight.data.copy_(conv_weight) nn.init.zeros_(conv.bias.data) def init_weight2(self, conv): conv_weight = conv.weight.data nn.init.zeros_(conv_weight) c1, c2, t, h, w = conv_weight.size() init_matrix = torch.eye(c1 // 2, c2) conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix conv.weight.data.copy_(conv_weight) nn.init.zeros_(conv.bias.data) def patchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) else: raise ValueError(f"Invalid input shape: {x.shape}") return x def unpatchify(x, patch_size): if patch_size == 1: return x if x.dim() == 4: x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) elif x.dim() == 5: x = rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size) return x class Resample38(Resample): def __init__(self, dim, mode): assert mode in ( "none", "upsample2d", "upsample3d", "downsample2d", "downsample3d", ) super(Resample, self).__init__() self.dim = dim self.mode = mode # layers if mode == "upsample2d": self.resample = nn.Sequential( Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim, 3, padding=1), ) elif mode == "upsample3d": self.resample = nn.Sequential( Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim, 3, padding=1), ) self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) elif mode == "downsample2d": self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) ) elif mode == "downsample3d": self.resample = nn.Sequential( nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) ) self.time_conv = CausalConv3d( dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) ) else: self.resample = nn.Identity() class ResidualBlock(nn.Module): def __init__(self, in_dim, out_dim, dropout=0.0, cpu_cache=False): super().__init__() self.in_dim = in_dim self.out_dim = out_dim self.cpu_cache = cpu_cache # layers self.residual = nn.Sequential( RMS_norm(in_dim, images=False), nn.SiLU(), CausalConv3d(in_dim, out_dim, 3, padding=1), RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), 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) else: return self._forward(x, feat_cache, feat_idx) def _forward(self, x, feat_cache=None, feat_idx=[0]): h = self.shortcut(x) for layer in self.residual: if check_is_instance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( cache_x.device), cache_x ], dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 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: if check_is_instance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() 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: x = layer(x) return x + h class AttentionBlock(nn.Module): """ Causal self-attention with a single head. """ def __init__(self, dim): super().__init__() self.dim = dim # layers self.norm = RMS_norm(dim) self.to_qkv = nn.Conv2d(dim, dim * 3, 1) self.proj = nn.Conv2d(dim, dim, 1) # zero out the last layer params nn.init.zeros_(self.proj.weight) def forward(self, x): identity = x b, c, t, h, w = x.size() x = rearrange(x, 'b c t h w -> (b t) c h w') x = self.norm(x) # compute query, key, value q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( 0, 1, 3, 2).contiguous().chunk(3, dim=-1) # apply attention x = F.scaled_dot_product_attention(q, k, v) x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) # output x = self.proj(x) x = rearrange(x, '(b t) c h w-> b c t h w', t=t) return x + identity class AvgDown3D(nn.Module): def __init__( self, in_channels, out_channels, factor_t, factor_s=1, ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s self.factor = self.factor_t * self.factor_s * self.factor_s assert in_channels * self.factor % out_channels == 0 self.group_size = in_channels * self.factor // out_channels def forward(self, x: torch.Tensor) -> torch.Tensor: pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t pad = (0, 0, 0, 0, pad_t, 0) x = F.pad(x, pad) B, C, T, H, W = x.shape x = x.view( B, C, T // self.factor_t, self.factor_t, H // self.factor_s, self.factor_s, W // self.factor_s, self.factor_s, ) x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() x = x.view( B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s, ) x = x.view( B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s, ) x = x.mean(dim=2) return x class DupUp3D(nn.Module): def __init__( self, in_channels: int, out_channels: int, factor_t, factor_s=1, ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.factor_t = factor_t self.factor_s = factor_s self.factor = self.factor_t * self.factor_s * self.factor_s assert out_channels * self.factor % in_channels == 0 self.repeats = out_channels * self.factor // in_channels def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: x = x.repeat_interleave(self.repeats, dim=1) x = x.view( x.size(0), self.out_channels, self.factor_t, self.factor_s, self.factor_s, x.size(2), x.size(3), x.size(4), ) x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() x = x.view( x.size(0), self.out_channels, x.size(2) * self.factor_t, x.size(4) * self.factor_s, x.size(6) * self.factor_s, ) if first_chunk: x = x[:, :, self.factor_t - 1 :, :, :] return x class Down_ResidualBlock(nn.Module): def __init__( self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False ): super().__init__() # Shortcut path with downsample self.avg_shortcut = AvgDown3D( in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1, ) # Main path with residual blocks and downsample downsamples = [] for _ in range(mult): downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) in_dim = out_dim # Add the final downsample block if down_flag: mode = "downsample3d" if temperal_downsample else "downsample2d" downsamples.append(Resample38(out_dim, mode=mode)) self.downsamples = nn.Sequential(*downsamples) def forward(self, x, feat_cache=None, feat_idx=[0]): x_copy = x.clone() for module in self.downsamples: x = module(x, feat_cache, feat_idx) return x + self.avg_shortcut(x_copy) class Up_ResidualBlock(nn.Module): def __init__( self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False ): super().__init__() # Shortcut path with upsample if up_flag: self.avg_shortcut = DupUp3D( in_dim, out_dim, factor_t=2 if temperal_upsample else 1, factor_s=2 if up_flag else 1, ) else: self.avg_shortcut = None # Main path with residual blocks and upsample upsamples = [] for _ in range(mult): upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) in_dim = out_dim # Add the final upsample block if up_flag: mode = "upsample3d" if temperal_upsample else "upsample2d" upsamples.append(Resample38(out_dim, mode=mode)) self.upsamples = nn.Sequential(*upsamples) def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): x_main = x.clone() for module in self.upsamples: x_main = module(x_main, feat_cache, feat_idx) if self.avg_shortcut is not None: x_shortcut = self.avg_shortcut(x, first_chunk) return x_main + x_shortcut else: return x_main class Encoder3d(nn.Module): def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_downsample=[True, True, False], dropout=0.0, pruning_rate=0.0, cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_downsample = temperal_downsample self.cpu_cache = cpu_cache # dimensions dims = [dim * u for u in [1] + dim_mult] dims = [int(d * (1 - pruning_rate)) for d in dims] scale = 1.0 # init block self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) # downsample blocks downsamples = [] for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): # residual (+attention) blocks for _ in range(num_res_blocks): downsamples.append(ResidualBlock(in_dim, out_dim, dropout, cpu_cache=cpu_cache)) if scale in attn_scales: downsamples.append(AttentionBlock(out_dim)) in_dim = out_dim # downsample block if i != len(dim_mult) - 1: mode = 'downsample3d' if temperal_downsample[ i] else 'downsample2d' downsamples.append(Resample(out_dim, mode=mode)) scale /= 2.0 self.downsamples = nn.Sequential(*downsamples) # middle blocks self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache), AttentionBlock(out_dim), ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache)) # output blocks self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, z_dim, 3, padding=1)) def forward(self, x, feat_cache=None, feat_idx=[0]): if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( cache_x.device), cache_x ], dim=2) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv1(x) ## downsamples for layer in self.downsamples: if feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## middle for layer in self.middle: if check_is_instance(layer, ResidualBlock) and feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## head for layer in self.head: if check_is_instance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( cache_x.device), cache_x ], dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = layer(x) return x class Encoder3d_38(nn.Module): def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_downsample=[False, True, True], dropout=0.0, pruning_rate=0.0, cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_downsample = temperal_downsample self.cpu_cache = cpu_cache # dimensions dims = [dim * u for u in [1] + dim_mult] dims = [int(d * (1 - pruning_rate)) for d in dims] scale = 1.0 # init block self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) # 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) # middle blocks self.middle = nn.Sequential( ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache), AttentionBlock(out_dim), ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache), ) # # output blocks self.head = nn.Sequential( RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, z_dim, 3, padding=1), ) def forward(self, x, feat_cache=None, feat_idx=[0]): if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv1(x) ## downsamples for layer in self.downsamples: if feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## middle for layer in self.middle: if isinstance(layer, ResidualBlock) and feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## head for layer in self.head: if isinstance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = layer(x) return x class Decoder3d(nn.Module): def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_upsample=[False, True, True], dropout=0.0, pruning_rate=0.0, cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_upsample = temperal_upsample self.cpu_cache = cpu_cache # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] dims = [int(d * (1 - pruning_rate)) for d in dims] scale = 1.0 / 2**(len(dim_mult) - 2) # init block self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache), AttentionBlock(dims[0]), ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache)) # upsample blocks upsamples = [] for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): # residual (+attention) blocks if i == 1 or i == 2 or i == 3: in_dim = in_dim // 2 for _ in range(num_res_blocks + 1): upsamples.append(ResidualBlock(in_dim, out_dim, dropout, cpu_cache=cpu_cache)) if scale in attn_scales: upsamples.append(AttentionBlock(out_dim)) in_dim = out_dim # upsample block if i != len(dim_mult) - 1: mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' upsamples.append(Resample(out_dim, mode=mode)) scale *= 2.0 self.upsamples = nn.Sequential(*upsamples) # output blocks self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, 3, 3, padding=1)) def forward(self, x, feat_cache=None, feat_idx=[0]): ## conv1 if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( cache_x.device), cache_x ], dim=2) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv1(x) ## middle for layer in self.middle: if check_is_instance(layer, ResidualBlock) and feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## upsamples for layer in self.upsamples: if feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## head for layer in self.head: if check_is_instance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: # cache last frame of last two chunk cache_x = torch.cat([ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( cache_x.device), cache_x ], dim=2) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = layer(x) return x class Decoder3d_38(nn.Module): def __init__(self, dim=128, z_dim=4, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_upsample=[False, True, True], dropout=0.0, pruning_rate=0.0, cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_upsample = temperal_upsample # dimensions dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] dims = [int(d * (1 - pruning_rate)) for d in dims] # init block self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache), AttentionBlock(dims[0]), ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache)) # upsample blocks upsamples = [] for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False upsamples.append( Up_ResidualBlock(in_dim=in_dim, out_dim=out_dim, dropout=dropout, mult=num_res_blocks + 1, temperal_upsample=t_up_flag, up_flag=i != len(dim_mult) - 1)) self.upsamples = nn.Sequential(*upsamples) # output blocks self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), CausalConv3d(out_dim, 12, 3, padding=1)) def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): if feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x, ], dim=2, ) x = self.conv1(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = self.conv1(x) for layer in self.middle: if check_is_instance(layer, ResidualBlock) and feat_cache is not None: x = layer(x, feat_cache, feat_idx) else: x = layer(x) ## upsamples for layer in self.upsamples: if feat_cache is not None: x = layer(x, feat_cache, feat_idx, first_chunk) else: x = layer(x) ## head for layer in self.head: if check_is_instance(layer, CausalConv3d) and feat_cache is not None: idx = feat_idx[0] cache_x = x[:, :, -CACHE_T:, :, :].clone() if cache_x.shape[2] < 2 and feat_cache[idx] is not None: cache_x = torch.cat( [ feat_cache[idx][:, :, -1, :, :] .unsqueeze(2) .to(cache_x.device), cache_x, ], dim=2, ) x = layer(x, feat_cache[idx]) feat_cache[idx] = cache_x feat_idx[0] += 1 else: x = layer(x) return x def count_conv3d(model): count = 0 for m in model.modules(): if check_is_instance(m, CausalConv3d): count += 1 return count class VideoVAE_(nn.Module): def __init__(self, dim=96, z_dim=16, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_downsample=[False, True, True], dropout=0.0, mean=None, inv_std=None, pruning_rate=0.0, cpu_cache=False, verbose=False): super().__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_downsample = temperal_downsample self.temperal_upsample = temperal_downsample[::-1] self.mean = mean self.inv_std = inv_std self.verbose = verbose # modules self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout, pruning_rate, cpu_cache=cpu_cache) self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) self.conv2 = CausalConv3d(z_dim, z_dim, 1) self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout, pruning_rate, cpu_cache=cpu_cache) def forward(self, x): mu, log_var = self.encode(x) z = self.reparameterize(mu, log_var) x_recon = self.decode(z) return x_recon, mu, log_var #modification originally by @raindrop313 https://github.com/raindrop313/ComfyUI-WanVideoStartEndFrames def encode_2(self, x, pbar=True, sample=False): t = x.shape[2] iter_ = 2 + (t - 2) // 4 if pbar: pbar = ProgressBar(iter_) for i in range(iter_): self._enc_conv_idx = [0] if i == 0: out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) elif i== iter_-1: out_ = self.encoder(x[:, :, -1:, :, :], feat_cache=[None] * self._enc_conv_num, feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) else: out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) self.clear_cache() out_head = out[:, :, :iter_ - 1, :, :] out_tail = out[:, :, -1, :, :].unsqueeze(2) mu = torch.cat([self.conv1(out_head), self.conv1(out_tail)], dim=2).chunk(2, dim=1)[0] mu = (mu - self.mean.to(mu)) * self.inv_std.to(mu) if pbar: pbar.update_absolute(0) return mu def encode(self, x, pbar=True, sample=False): t = x.shape[2] iter_ = 1 + (t - 1) // 4 input_shape = x.shape if pbar: pbar = ProgressBar(iter_) try: torch.cuda.reset_peak_memory_stats(device) except Exception: pass for i in tqdm(range(iter_), desc="WanVAE encoding frames", disable=not pbar): self._enc_conv_idx = [0] if i == 0: out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) else: out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) if pbar: pbar.update(1) self.clear_cache() mu, log_var = self.conv1(out).chunk(2, dim=1) mu = (mu - self.mean.to(mu)) * self.inv_std.to(mu) if pbar: pbar.update_absolute(0) if sample: std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) eps = torch.randn_like(std) return mu + std * eps if self.verbose: try: log.info(f"WanVAE encoded input:{input_shape} to {out.shape}") print_memory(device, process="WanVAE encode") torch.cuda.reset_peak_memory_stats(device) except Exception: pass return mu #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) x = torch.cat([self.conv2(z_head), self.conv2(z_tail)], dim=2) for i in range(iter_): self._conv_idx = [0] if i == 0: out = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) elif i==iter_-1: out_ = self.decoder(x[:, :, -1, :, :].unsqueeze(2), feat_cache=None, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) else: out_ = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) self.clear_cache() return out def decode(self, z, pbar=True): # z: [b,c,t,h,w] z = z / self.inv_std.to(z) + self.mean.to(z) input_shape = z.shape iter_ = z.shape[2] if pbar: pbar = ProgressBar(iter_) try: torch.cuda.reset_peak_memory_stats(device) except Exception: pass x = self.conv2(z) for i in tqdm(range(iter_), desc="WanVAE decoding frames", disable=not pbar): self._conv_idx = [0] if i == 0: out = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) else: out_ = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) if pbar: pbar.update(1) if pbar: pbar.update_absolute(0) self.clear_cache() if self.verbose: try: log.info(f"WanVAE decoded input:{input_shape} to {out.shape}") print_memory(device, process="WanVAE decode") torch.cuda.reset_peak_memory_stats(device) except Exception: pass return out def reparameterize(self, mu, log_var): std = torch.exp(0.5 * log_var) eps = torch.randn_like(std) return eps * std + mu def sample(self, imgs, deterministic=False): mu, log_var = self.encode(imgs) if deterministic: return mu std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) return mu + std * torch.randn_like(std) def clear_cache(self): self._conv_num = count_conv3d(self.decoder) self._conv_idx = [0] self._feat_map = [None] * self._conv_num # cache encode self._enc_conv_num = count_conv3d(self.encoder) self._enc_conv_idx = [0] self._enc_feat_map = [None] * self._enc_conv_num class WanVideoVAE(nn.Module): def __init__(self, z_dim=16, dtype=torch.float32, pruning_rate=0.0, cpu_cache=False, verbose=False): super().__init__() self.dtype = dtype self.cpu_cache = cpu_cache self.verbose = verbose mean = [ -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 ] std = [ 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 ] self.mean = torch.tensor(mean).view(1, z_dim, 1, 1, 1) self.inv_std = (1.0 / torch.tensor(std)).view(1, z_dim, 1, 1, 1) self.z_dim = z_dim # init model self.model = VideoVAE_(z_dim=z_dim, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=self.cpu_cache, verbose=self.verbose).eval().requires_grad_(False) self.upsampling_factor = 8 def build_1d_mask(self, length, left_bound, right_bound, border_width): x = torch.ones((length,)) if not left_bound: x[:border_width] = (torch.arange(border_width) + 1) / border_width if not right_bound: x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) return x def build_mask(self, data, is_bound, border_width): _, _, _, H, W = data.shape h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) h = repeat(h, "H -> H W", H=H, W=W) w = repeat(w, "W -> H W", H=H, W=W) mask = torch.stack([h, w]).min(dim=0).values mask = rearrange(mask, "H W -> 1 1 1 H W") return mask def tiled_decode(self, hidden_states, device, tile_size, tile_stride, end_=False, pbar=True): _, _, T, H, W = hidden_states.shape size_h, size_w = tile_size stride_h, stride_w = tile_stride # Split tasks tasks = [] for h in range(0, H, stride_h): if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue for w in range(0, W, stride_w): if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue h_, w_ = h + size_h, w + size_w tasks.append((h, h_, w, w_)) data_device = "cpu" computation_device = device weight, values = None, None if pbar: pbar = ProgressBar(len(tasks)) for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device) if end_: hidden_states_batch = self.model.decode_2(hidden_states_batch).to(data_device) else: hidden_states_batch = self.model.decode(hidden_states_batch).to(data_device) if weight is None: weight = torch.zeros((1, 1, hidden_states_batch.shape[2], H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) if values is None: values = torch.zeros((1, 3, hidden_states_batch.shape[2], H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) mask = self.build_mask( hidden_states_batch, is_bound=(h==0, h_>=H, w==0, w_>=W), border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) ).to(dtype=hidden_states.dtype, device=data_device) target_h = h * self.upsampling_factor target_w = w * self.upsampling_factor values[ :, :, :, target_h:target_h + hidden_states_batch.shape[3], target_w:target_w + hidden_states_batch.shape[4], ] += hidden_states_batch * mask weight[ :, :, :, target_h: target_h + hidden_states_batch.shape[3], target_w: target_w + hidden_states_batch.shape[4], ] += mask if pbar: pbar.update(1) values = values / weight values = values.float().clamp_(-1, 1) return values 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 else: size_h, size_w = tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor stride_h, stride_w = tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor # Split tasks tasks = [] for h in range(0, H, stride_h): if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue for w in range(0, W, stride_w): if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue h_, w_ = h + size_h, w + size_w tasks.append((h, h_, w, w_)) data_device = device computation_device = device out_T = (T + 3) // 4 if end_: out_T += 1 weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) if pbar: pbar = ProgressBar(len(tasks)) for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device) if end_: hidden_states_batch = self.model.encode_2(hidden_states_batch).to(data_device) else: hidden_states_batch = self.model.encode(hidden_states_batch).to(data_device) mask = self.build_mask( hidden_states_batch, is_bound=(h==0, h_>=H, w==0, w_>=W), border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) ).to(dtype=video.dtype, device=data_device) target_h = h // self.upsampling_factor target_w = w // self.upsampling_factor values[ :, :, :, target_h:target_h + hidden_states_batch.shape[3], target_w:target_w + hidden_states_batch.shape[4], ] += hidden_states_batch * mask weight[ :, :, :, target_h: target_h + hidden_states_batch.shape[3], target_w: target_w + hidden_states_batch.shape[4], ] += mask if pbar: pbar.update(1) values = values / weight values = values.float() return values def single_encode(self, video, device, pbar=True, sample=False): video = video.to(device) x = self.model.encode(video, pbar=pbar, sample=sample) return x.float() def single_decode(self, hidden_state, device, pbar=True): hidden_state = hidden_state.to(device) video = self.model.decode(hidden_state, pbar=pbar) return video def double_encode(self, video, device, pbar=True, sample=False): print('double_encode') video = video.to(device) x = self.model.encode_2(video, pbar=pbar, sample=sample) return x.float() def double_decode(self, hidden_state, device): print('double_decode') hidden_state = hidden_state.to(device) video = self.model.decode_2(hidden_state) return video def encode(self, videos, device, tiled=False,end_=False, tile_size=None, tile_stride=None, pbar=True, sample=False): self.model.clear_cache() videos = [video.to("cpu") for video in videos] hidden_states = [] for video in videos: video = video.unsqueeze(0) if tiled: hidden_state = self.tiled_encode(video, device, tile_size, tile_stride, end_=end_, pbar=pbar) else: if end_: hidden_state = self.double_encode(video, device, pbar=pbar, sample=sample) else: hidden_state = self.single_encode(video, device, pbar=pbar, sample=sample) hidden_state = hidden_state.squeeze(0) hidden_states.append(hidden_state) hidden_states = torch.stack(hidden_states) return hidden_states def decode(self, hidden_states, device, tiled=False, end_=False, tile_size=(34, 34), tile_stride=(18, 16), pbar=True): self.model.clear_cache() hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states] videos = [] for hidden_state in hidden_states: hidden_state = hidden_state.unsqueeze(0) if tiled: video = self.tiled_decode(hidden_state, device, tile_size, tile_stride, end_=end_, pbar=pbar) else: if end_: video = self.double_decode(hidden_state, device) else: video = self.single_decode(hidden_state, device, pbar=pbar) video = video.squeeze(0) videos.append(video) return videos class VideoVAE38_(VideoVAE_): def __init__(self, dim=160, z_dim=48, dec_dim=256, dim_mult=[1, 2, 4, 4], num_res_blocks=2, attn_scales=[], temperal_downsample=[False, True, True], dropout=0.0, dtype=torch.bfloat16, mean=None, inv_std=None, pruning_rate=0.0, cpu_cache=False, verbose=False): super(VideoVAE_, self).__init__() self.dim = dim self.z_dim = z_dim self.dim_mult = dim_mult self.num_res_blocks = num_res_blocks self.attn_scales = attn_scales self.temperal_downsample = temperal_downsample self.temperal_upsample = temperal_downsample[::-1] self.dtype = dtype self.mean = mean self.inv_std = inv_std self.cpu_cache = cpu_cache self.verbose = verbose # modules self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout, pruning_rate, cpu_cache=cpu_cache) self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) self.conv2 = CausalConv3d(z_dim, z_dim, 1) self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout, pruning_rate, cpu_cache=cpu_cache) def encode(self, x, pbar=True, sample=False): input_shape = x.shape self.clear_cache() try: torch.cuda.reset_peak_memory_stats(device) except Exception: pass x = patchify(x, patch_size=2) t = x.shape[2] iter_ = 1 + (t - 1) // 4 if pbar: pbar = ProgressBar(iter_) for i in tqdm(range(iter_), desc="WanVAE encoding frames", disable=not pbar): self._enc_conv_idx = [0] if i == 0: out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) else: out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) if pbar: pbar.update(1) mu = self.conv1(out).chunk(2, dim=1)[0] mu = (mu - self.mean.to(mu)) * self.inv_std.to(mu) self.clear_cache() if self.verbose: try: log.info(f"WanVAE decoded input:{input_shape} to {out.shape}") print_memory(device, process="WanVAE decode") torch.cuda.reset_peak_memory_stats(device) except Exception: pass return mu def decode(self, z, pbar=True): self.clear_cache() input_shape = z.shape try: torch.cuda.reset_peak_memory_stats(device) except Exception: pass z = z / self.inv_std.to(z) + self.mean.to(z) iter_ = z.shape[2] if pbar: pbar = ProgressBar(iter_) x = self.conv2(z) for i in tqdm(range(iter_), desc="WanVAE decoding frames", disable=not pbar): self._conv_idx = [0] if i == 0: out = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True) else: out_ = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) if pbar: pbar.update(1) out = unpatchify(out, patch_size=2) self.clear_cache() if self.verbose: try: log.info(f"WanVAE decoded input:{input_shape} to {out.shape}") print_memory(device, process="WanVAE decode") torch.cuda.reset_peak_memory_stats(device) except Exception: pass return out class WanVideoVAE38(WanVideoVAE): def __init__(self, z_dim=48, dim=160, dtype=torch.bfloat16, pruning_rate=0.0, cpu_cache=False, verbose=False): super(WanVideoVAE, self).__init__() mean = [ -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 ] std = [ 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 ] self.mean = torch.tensor(mean).view(1, z_dim, 1, 1, 1) self.inv_std = (1.0 / torch.tensor(std)).view(1, z_dim, 1, 1, 1) self.dtype = dtype self.z_dim = z_dim self.cpu_cache = cpu_cache self.verbose = verbose # init model self.model = VideoVAE38_(z_dim=z_dim, dim=dim, dtype=dtype, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=cpu_cache, verbose=verbose).eval().requires_grad_(False) self.upsampling_factor = 16