non_blocking teacache offload

nodes.py

@@ -2277,10 +2277,6 @@ class WanVideoSampler:
         else:
             transformer.slg_blocks = None
 
-        mm.unload_all_models()
-        mm.soft_empty_cache()
-        gc.collect()
-
         self.teacache_state = [None, None]
         self.teacache_state_source = [None, None]
         self.teacache_states_context = []
@@ -2429,11 +2425,6 @@ class WanVideoSampler:
                     noise_pred_uncond=noise_pred_uncond.to(intermediate_device)
 
                 return noise_pred_uncond + cfg_scale * (noise_pred_cond - noise_pred_uncond), [teacache_state_cond]
-        
-        try:
-            torch.cuda.reset_peak_memory_stats(device)
-        except:
-            pass
 
         log.info(f"Sampling {(latent_video_length-1) * 4 + 1} frames at {latent.shape[3]*8}x{latent.shape[2]*8} with {steps} steps")
 
@@ -2456,7 +2447,17 @@ class WanVideoSampler:
             latent_shift_start_percent = loop_args["start_percent"]
             latent_shift_end_percent = loop_args["end_percent"]
             shift_idx = 0
-        #main loop start
+
+        #clear memory before sampling
+        mm.unload_all_models()
+        mm.soft_empty_cache()
+        gc.collect()
+        try:
+            torch.cuda.reset_peak_memory_stats(device)
+        except:
+            pass
+
+        #region main loop start
         for idx, t in enumerate(tqdm(timesteps)):    
             if flowedit_args is not None:
                 if idx < skip_steps:

wanvideo/modules/model.py

@@ -1035,7 +1035,7 @@ class WanModel(ModelMixin, ConfigMixin):
 
         if not self.enable_teacache or (self.enable_teacache and should_calc):
             if self.enable_teacache:
-                original_x = x.clone().to(self.teacache_cache_device)
+                original_x = x.clone().to(self.teacache_cache_device, non_blocking=self.use_non_blocking)
 
             # arguments
             kwargs = dict(
@@ -1072,8 +1072,8 @@ class WanModel(ModelMixin, ConfigMixin):
                 self.teacache_state.update(
                     pred_id,
                     previous_residual=(x.to(original_x.device) - original_x),
-                    accumulated_rel_l1_distance=accumulated_rel_l1_distance.to(self.teacache_cache_device),
-                    previous_modulated_input=previous_modulated_input.to(self.teacache_cache_device)
+                    accumulated_rel_l1_distance=accumulated_rel_l1_distance.to(self.teacache_cache_device, non_blocking=self.use_non_blocking),
+                    previous_modulated_input=previous_modulated_input.to(self.teacache_cache_device, non_blocking=self.use_non_blocking)
                 )
 
         x = self.head(x, e)