diff --git a/nodes.py b/nodes.py index ed72325..834d335 100644 --- a/nodes.py +++ b/nodes.py @@ -2038,6 +2038,8 @@ class WanVideoSampler: block_swap_args["offload_txt_emb"], block_swap_args["offload_img_emb"], vace_blocks_to_swap = block_swap_args.get("vace_blocks_to_swap", None), + prefetch_blocks = block_swap_args.get("prefetch_blocks", 0), + block_swap_debug = block_swap_args.get("block_swap_debug", False), ) elif model["auto_cpu_offload"]: for module in transformer.modules(): diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 0e7b7cc..38e6663 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -274,6 +274,8 @@ class WanVideoBlockSwap: "optional": { "use_non_blocking": ("BOOLEAN", {"default": False, "tooltip": "Use non-blocking memory transfer for offloading, reserves more RAM but is faster"}), "vace_blocks_to_swap": ("INT", {"default": 0, "min": 0, "max": 15, "step": 1, "tooltip": "Number of VACE blocks to swap, the VACE model has 15 blocks"}), + "prefetch_blocks": ("INT", {"default": 0, "min": 0, "max": 40, "step": 1, "tooltip": "Number of blocks to prefetch ahead, can speed up processing but increases memory usage. 1 is usually enough to offset speed loss from block swapping, use the debug option to confirm it for your system"}), + "block_swap_debug": ("BOOLEAN", {"default": False, "tooltip": "Enable debug logging for block swapping"}), }, } RETURN_TYPES = ("BLOCKSWAPARGS",) diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 47df295..f4d77df 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn from einops import repeat, rearrange from ...enhance_a_video.enhance import get_feta_scores - +import time try: from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask create_block_mask = torch.compile(create_block_mask) @@ -29,7 +29,7 @@ from ...cache_methods.cache_methods import TeaCacheState, MagCacheState, EasyCac from ...multitalk.multitalk import get_attn_map_with_target from ...echoshot.echoshot import rope_apply_z, rope_apply_c, rope_apply_echoshot -from comfy.model_management import get_torch_device, get_autocast_device +from comfy.model_management import get_torch_device, get_autocast_device, get_offload_stream from comfy.ldm.flux.math import apply_rope as apply_rope_comfy def apply_rope_comfy_chunked(xq, xk, freqs_cis, num_chunks=4): @@ -1094,6 +1094,8 @@ class WanModel(torch.nn.Module): self.slg_end_percent = 1.0 self.use_non_blocking = False + self.prefetch_blocks = 0 + self.block_swap_debug = False self.video_attention_split_steps = [] self.lora_scheduling_enabled = False @@ -1234,9 +1236,11 @@ class WanModel(torch.nn.Module): return block_mask - def block_swap(self, blocks_to_swap, offload_txt_emb=False, offload_img_emb=False, vace_blocks_to_swap=None): + def block_swap(self, blocks_to_swap, offload_txt_emb=False, offload_img_emb=False, vace_blocks_to_swap=None, prefetch_blocks=0, block_swap_debug=False): log.info(f"Swapping {blocks_to_swap + 1} transformer blocks") self.blocks_to_swap = blocks_to_swap + self.prefetch_blocks = prefetch_blocks + self.block_swap_debug = block_swap_debug self.offload_img_emb = offload_img_emb self.offload_txt_emb = offload_txt_emb @@ -1811,15 +1815,47 @@ class WanModel(torch.nn.Module): device=self.offload_device) self.controlnet.to(self.offload_device) + # Asynchronous block offloading with CUDA streams and events + cuda_stream = mm.get_offload_stream(device) + events = [torch.cuda.Event() for _ in self.blocks] + for b, block in enumerate(self.blocks): + # Prefetch blocks if enabled + if self.prefetch_blocks > 0: + for prefetch_offset in range(1, self.prefetch_blocks + 1): + prefetch_idx = b + prefetch_offset + if prefetch_idx < len(self.blocks) and self.blocks_to_swap >= 0 and prefetch_idx <= self.blocks_to_swap: + with torch.cuda.stream(cuda_stream): + self.blocks[prefetch_idx].to(self.main_device, non_blocking=self.use_non_blocking) + events[prefetch_idx].record(cuda_stream) + if self.block_swap_debug: + transfer_start = time.perf_counter() + # Wait for block to be ready + if b <= self.blocks_to_swap and self.blocks_to_swap >= 0: + if self.prefetch_blocks > 0: + if not events[b].query(): + events[b].synchronize() + block.to(self.main_device) + if self.block_swap_debug: + transfer_end = time.perf_counter() + transfer_time = transfer_end - transfer_start + compute_start = time.perf_counter() #skip layer guidance if self.slg_blocks is not None: if b in self.slg_blocks and is_uncond: if self.slg_start_percent <= current_step_percentage <= self.slg_end_percent: continue - if b <= self.blocks_to_swap and self.blocks_to_swap >= 0: - block.to(self.main_device) x = block(x, **kwargs) + if self.block_swap_debug: + compute_end = time.perf_counter() + compute_time = compute_end - compute_start + to_cpu_transfer_start = time.perf_counter() + if b <= self.blocks_to_swap and self.blocks_to_swap >= 0: + block.to(self.offload_device, non_blocking=self.use_non_blocking) + if self.block_swap_debug: + to_cpu_transfer_end = time.perf_counter() + to_cpu_transfer_time = to_cpu_transfer_end - to_cpu_transfer_start + log.info(f"Block {b}: transfer_time={transfer_time:.4f}s, compute_time={compute_time:.4f}s, to_cpu_transfer_time={to_cpu_transfer_time:.4f}s") #uni3c controlnet if pdc_controlnet_states is not None and b < len(pdc_controlnet_states):