diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 477bee4..77d7056 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -1694,6 +1694,7 @@ class WanVideoVAELoader: {"default": "bf16"} ), "compile_args": ("WANCOMPILEARGS", ), + "use_cpu_cache": ("BOOLEAN", {"default": False, "tooltip": "Reduces VRAM usage, but slows the VAE down a lot"}), } } @@ -1703,7 +1704,7 @@ class WanVideoVAELoader: CATEGORY = "WanVideoWrapper" DESCRIPTION = "Loads Wan VAE model from 'ComfyUI/models/vae'" - def loadmodel(self, model_name, precision, compile_args=None): + def loadmodel(self, model_name, precision, compile_args=None, use_cpu_cache=False): dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] model_path = folder_paths.get_full_path_or_raise("vae", model_name) vae_sd = load_torch_file(model_path, safe_load=True) @@ -1720,9 +1721,9 @@ class WanVideoVAELoader: pruning_rate = 0.0 if vae_sd["model.conv2.weight"].shape[0] == 16: - vae = WanVideoVAE(dtype=dtype, pruning_rate=pruning_rate) + vae = WanVideoVAE(dtype=dtype, pruning_rate=pruning_rate, cpu_cache=use_cpu_cache) elif vae_sd["model.conv2.weight"].shape[0] == 48: - vae = WanVideoVAE38(dtype=dtype, pruning_rate=pruning_rate) + vae = WanVideoVAE38(dtype=dtype, pruning_rate=pruning_rate, cpu_cache=use_cpu_cache) vae.load_state_dict(vae_sd) del vae_sd diff --git a/utils.py b/utils.py index 5b155c6..ddfb1c1 100644 --- a/utils.py +++ b/utils.py @@ -108,13 +108,13 @@ def check_diffusers_version(): except importlib.metadata.PackageNotFoundError: raise AssertionError("diffusers is not installed.") -def print_memory(device): +def print_memory(device, process="Sampling"): memory = torch.cuda.memory_allocated(device) / 1024**3 max_memory = torch.cuda.max_memory_allocated(device) / 1024**3 max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 - log.info(f"Allocated memory: {memory=:.3f} GB") - log.info(f"Max allocated memory: {max_memory=:.3f} GB") - log.info(f"Max reserved memory: {max_reserved=:.3f} GB") + log.info(f"[{process}] Allocated memory: {memory=:.3f} GB") + log.info(f"[{process}] Max allocated memory: {max_memory=:.3f} GB") + log.info(f"[{process}] Max reserved memory: {max_reserved=:.3f} GB") #memory_summary = torch.cuda.memory_summary(device=device, abbreviated=False) #log.info(f"Memory Summary:\n{memory_summary}") diff --git a/wanvideo/wan_video_vae.py b/wanvideo/wan_video_vae.py index 740822c..3fe2e14 100644 --- a/wanvideo/wan_video_vae.py +++ b/wanvideo/wan_video_vae.py @@ -5,6 +5,9 @@ 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 @@ -254,10 +257,11 @@ class Resample38(Resample): class ResidualBlock(nn.Module): - def __init__(self, in_dim, out_dim, dropout=0.0): + 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( @@ -267,8 +271,14 @@ 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) + 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: @@ -287,6 +297,26 @@ 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: + 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): @@ -507,7 +537,8 @@ class Encoder3d(nn.Module): attn_scales=[], temperal_downsample=[True, True, False], dropout=0.0, - pruning_rate=0.0): + pruning_rate=0.0, + cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim @@ -515,6 +546,7 @@ class Encoder3d(nn.Module): 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] @@ -529,7 +561,7 @@ class Encoder3d(nn.Module): 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)) + 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 @@ -543,9 +575,9 @@ class Encoder3d(nn.Module): self.downsamples = nn.Sequential(*downsamples) # middle blocks - self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache), AttentionBlock(out_dim), - ResidualBlock(out_dim, out_dim, dropout)) + ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache)) # output blocks self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), @@ -612,7 +644,8 @@ class Encoder3d_38(nn.Module): attn_scales=[], temperal_downsample=[False, True, True], dropout=0.0, - pruning_rate=0.0): + pruning_rate=0.0, + cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim @@ -620,6 +653,7 @@ class Encoder3d_38(nn.Module): 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] @@ -650,9 +684,9 @@ class Encoder3d_38(nn.Module): # middle blocks self.middle = nn.Sequential( - ResidualBlock(out_dim, out_dim, dropout), + ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache), AttentionBlock(out_dim), - ResidualBlock(out_dim, out_dim, dropout), + ResidualBlock(out_dim, out_dim, dropout, cpu_cache=cpu_cache), ) # # output blocks @@ -730,7 +764,8 @@ class Decoder3d(nn.Module): attn_scales=[], temperal_upsample=[False, True, True], dropout=0.0, - pruning_rate=0.0): + pruning_rate=0.0, + cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim @@ -738,6 +773,7 @@ class Decoder3d(nn.Module): 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]] @@ -748,9 +784,9 @@ class Decoder3d(nn.Module): self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks - self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache), AttentionBlock(dims[0]), - ResidualBlock(dims[0], dims[0], dropout)) + ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache)) # upsample blocks upsamples = [] @@ -759,7 +795,7 @@ class Decoder3d(nn.Module): 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)) + 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 @@ -838,7 +874,8 @@ class Decoder3d_38(nn.Module): attn_scales=[], temperal_upsample=[False, True, True], dropout=0.0, - pruning_rate=0.0): + pruning_rate=0.0, + cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim @@ -855,9 +892,9 @@ class Decoder3d_38(nn.Module): self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) # middle blocks - self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache), AttentionBlock(dims[0]), - ResidualBlock(dims[0], dims[0], dropout)) + ResidualBlock(dims[0], dims[0], dropout, cpu_cache=cpu_cache)) # upsample blocks upsamples = [] @@ -951,7 +988,8 @@ class VideoVAE_(nn.Module): dropout=0.0, mean=None, inv_std=None, - pruning_rate=0.0): + pruning_rate=0.0, + cpu_cache=False): super().__init__() self.dim = dim self.z_dim = z_dim @@ -965,11 +1003,11 @@ class VideoVAE_(nn.Module): # modules self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, - attn_scales, self.temperal_downsample, dropout, pruning_rate) + 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) + attn_scales, self.temperal_upsample, dropout, pruning_rate, cpu_cache=cpu_cache) def forward(self, x): mu, log_var = self.encode(x) @@ -1016,10 +1054,15 @@ class VideoVAE_(nn.Module): 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: + pass - for i in range(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, :, :], @@ -1042,6 +1085,12 @@ class VideoVAE_(nn.Module): std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) eps = torch.randn_like(std) return mu + std * eps + 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: + pass return mu @@ -1065,12 +1114,12 @@ class VideoVAE_(nn.Module): out_ = self.decoder(x[:, :, -1, :, :].unsqueeze(2), feat_cache=None, feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) # may add tensor offload + 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) # may add tensor offload + out = torch.cat([out, out_], 2) self.clear_cache() return out @@ -1079,9 +1128,14 @@ class VideoVAE_(nn.Module): 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: + pass x = self.conv2(z) for i in range(iter_): self._conv_idx = [0] @@ -1093,13 +1147,19 @@ class VideoVAE_(nn.Module): out_ = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) - out = torch.cat([out, out_], 2) # may add tensor offload + out = torch.cat([out, out_], 2) if pbar: pbar.update(1) if pbar: pbar.update_absolute(0) self.clear_cache() + 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: + pass return out def reparameterize(self, mu, log_var): @@ -1126,10 +1186,11 @@ class VideoVAE_(nn.Module): class WanVideoVAE(nn.Module): - def __init__(self, z_dim=16, dtype=torch.float32, pruning_rate=0.0): + def __init__(self, z_dim=16, dtype=torch.float32, pruning_rate=0.0, cpu_cache=False): super().__init__() self.dtype = dtype + self.cpu_cache = cpu_cache mean = [ -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, @@ -1144,7 +1205,7 @@ class WanVideoVAE(nn.Module): 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).eval().requires_grad_(False) + self.model = VideoVAE_(z_dim=z_dim, mean=self.mean, inv_std=self.inv_std, pruning_rate=pruning_rate, cpu_cache=self.cpu_cache).eval().requires_grad_(False) self.upsampling_factor = 8 @@ -1363,7 +1424,8 @@ class VideoVAE38_(VideoVAE_): dtype=torch.bfloat16, mean=None, inv_std=None, - pruning_rate=0.0): + pruning_rate=0.0, + cpu_cache=False): super(VideoVAE_, self).__init__() self.dim = dim self.z_dim = z_dim @@ -1375,24 +1437,29 @@ class VideoVAE38_(VideoVAE_): self.dtype = dtype self.mean = mean self.inv_std = inv_std + self.cpu_cache = cpu_cache # modules self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, - attn_scales, self.temperal_downsample, dropout, pruning_rate) + 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) - + 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: + pass x = patchify(x, patch_size=2) t = x.shape[2] iter_ = 1 + (t - 1) // 4 if pbar: pbar = ProgressBar(iter_) - for i in range(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, :, :], @@ -1408,18 +1475,29 @@ class VideoVAE38_(VideoVAE_): mu = self.conv1(out).chunk(2, dim=1)[0] mu = (mu - self.mean.to(mu)) * self.inv_std.to(mu) self.clear_cache() + 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: + 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: + 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 range(iter_): + 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, :, :], @@ -1435,12 +1513,18 @@ class VideoVAE38_(VideoVAE_): pbar.update(1) out = unpatchify(out, patch_size=2) self.clear_cache() + 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: + pass return out class WanVideoVAE38(WanVideoVAE): - def __init__(self, z_dim=48, dim=160, dtype=torch.bfloat16, pruning_rate=0.0): + def __init__(self, z_dim=48, dim=160, dtype=torch.bfloat16, pruning_rate=0.0, cpu_cache=False): super(WanVideoVAE, self).__init__() mean = [ @@ -1463,7 +1547,8 @@ class WanVideoVAE38(WanVideoVAE): 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 # 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).eval().requires_grad_(False) + 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).eval().requires_grad_(False) self.upsampling_factor = 16 \ No newline at end of file