tea_cache

app.py

@@ -258,23 +258,23 @@ def calculate_time_required(max_duration_s, rc_bool):
     if max_duration_s == -1:
         return 75
     if max_duration_s == 2:
-        return 110
+        return 90
     elif max_duration_s == 3:
-        return 170
+        return 150
     elif max_duration_s == 4:
-        return 180
+        return 150
     elif max_duration_s == 5:
-        return 190
+        return 150
     elif max_duration_s == 6:
-        return 250
+        return 210
     elif max_duration_s == 7:
-        return 260
+        return 220
     elif max_duration_s == 8:
-        return 320
+        return 260
     elif max_duration_s == 9:
-        return 330
+        return 270
     elif max_duration_s == 10:
-        return 340
+        return 270
     
 def get_display_time_required(max_duration_s, rc_bool):
     # the 30 seconds extra is just for saftey in case of a unexpected slow down
@@ -484,6 +484,8 @@ def get_sam_mask(prompt_state, frames, frame_idx, input_points, input_labels):
 def animate_scene(input_video, max_duration_s, edited_frame, rc_str,
                   prompt_state=None, session_id=None,
                   progress=gr.Progress(track_tqdm=True)):
+    
+    start_preprocess = time.time()
 
     # Build dicts keyed by frame index, or empty if none
     if not prompt_state or "selections" not in prompt_state:
@@ -546,6 +548,9 @@ def animate_scene(input_video, max_duration_s, edited_frame, rc_str,
     else:
         final_video_path = output_video_path
 
+    preprocess_time = time.time() - start_preprocess
+    print(f"total time for duration {max_duration_s} took {preprocess_time:.2f} seconds")
+
     print(f"task for {session_id} finalised")
     return final_video_path, pose_video, bg_video, mask_video, face_video
 

wan/animate.py

@@ -611,6 +611,15 @@ class WanAnimate:
                             "face_pixel_values": face_pixel_values_uncond,
                         }
 
+                    
+                    tea_cache_posi = {
+                        "tea_cache": TeaCache(
+                            num_inference_steps=sampling_steps,
+                            rel_l1_thresh=0.3,
+                            timestep_scale=float(self.num_train_timesteps),
+                        )
+                    }
+
                     for i, t in enumerate(timesteps):
                         latent_model_input = latents
                         timestep = [t]
@@ -618,7 +627,7 @@ class WanAnimate:
                         timestep = torch.stack(timestep)
 
                         noise_pred_cond = TensorList(
-                            self.noise_model(TensorList(latent_model_input), t=timestep, **arg_c)
+                            self.noise_model(TensorList(latent_model_input), t=timestep, **tea_cache_posi, **arg_c)
                         )
 
                         if guide_scale > 1:
@@ -661,3 +670,53 @@ class WanAnimate:
 
         videos = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len]
         return videos[0] if self.rank == 0 else None
+
+class TeaCache:
+    def __init__(self, num_inference_steps, rel_l1_thresh, timestep_scale=1000.0):
+        self.num_inference_steps = num_inference_steps
+        self.rel_l1_thresh = rel_l1_thresh
+        self.timestep_scale = timestep_scale
+
+        self.step = 0
+        self.accumulated_rel_l1_distance = 0.0
+        self.previous_modulated_input = None
+        self.previous_residual = None
+        self.previous_hidden_states = None
+
+    def check(self, x, t_scalar):
+        # scheduler-agnostic timestep-aware weighting
+        # works for UniPC / flow schedulers that only expose t
+        w = float(t_scalar) / max(self.timestep_scale, 1.0)
+        modulated_input = x * w
+
+        if self.step == 0 or self.step == self.num_inference_steps - 1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0.0
+        else:
+            diff = torch.abs(modulated_input - self.previous_modulated_input)
+            ref = torch.abs(self.previous_modulated_input)
+            relative_l1 = diff.sum() / (ref.sum() + 1e-8)
+
+            self.accumulated_rel_l1_distance += relative_l1.item()
+            should_calc = self.accumulated_rel_l1_distance >= self.rel_l1_thresh
+
+            if should_calc:
+                self.accumulated_rel_l1_distance = 0.0
+
+        self.previous_modulated_input = modulated_input.detach()
+
+        self.step += 1
+        if self.step == self.num_inference_steps:
+            self.step = 0
+
+        if should_calc:
+            self.previous_hidden_states = x.clone()
+
+        return not should_calc
+
+    def store(self, hidden_states):
+        self.previous_residual = hidden_states - self.previous_hidden_states
+        self.previous_hidden_states = None
+
+    def update(self, hidden_states):
+        return hidden_states + self.previous_residual

wan/modules/animate/model_animate.py

@@ -375,6 +375,7 @@ class WanAnimateModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         x,
         t,
         clip_fea,
+        tea_cache,
         context,
         seq_len,
         y=None,
@@ -391,14 +392,17 @@ class WanAnimateModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
         x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
 
+        device = x[0].device
         grid_sizes = torch.stack(
-            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+            [torch.tensor(u.shape[2:], dtype=torch.long, device=device) for u in x]
+        )
         x = [u.flatten(2).transpose(1, 2) for u in x]
-        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long, device=device)
         assert seq_lens.max() <= seq_len
+
         x = torch.cat([
-            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
-                      dim=1) for u in x
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+            for u in x
         ])
 
         # time embeddings
@@ -413,36 +417,64 @@ class WanAnimateModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         context_lens = None
         context = self.text_embedding(
             torch.stack([
-                torch.cat(
-                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
                 for u in context
-            ]))
+            ])
+        )
 
         if self.use_img_emb:
-            context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+            context_clip = self.img_emb(clip_fea)
             context = torch.concat([context_clip, context], dim=1)
 
-        # arguments
         kwargs = dict(
             e=e0,
             seq_lens=seq_lens,
             grid_sizes=grid_sizes,
             freqs=freqs,
             context=context,
-            context_lens=context_lens)
+            context_lens=context_lens,
+        )
 
-        if self.use_context_parallel:
-            x = torch.chunk(x, get_world_size(), dim=1)[get_rank()]
+        # tea cache check must happen after e0 exists
+        ori_x_len = x.shape[1]
+        if tea_cache is not None:
+            t_scalar = int(t.flatten()[0].item())
+            tea_cache_update = tea_cache.check(x, t_scalar)
+        else:
+            tea_cache_update = False
+
+        if tea_cache_update:
+            x = tea_cache.update(x)
+        else:
+            if self.use_context_parallel:
+                world_size = get_world_size()
+                rank = get_rank()
+
+                pad_size = 0
+                if x.shape[1] % world_size != 0:
+                    pad_size = world_size - (x.shape[1] % world_size)
+                    x = torch.cat(
+                        [x, x.new_zeros(x.shape[0], pad_size, x.shape[2])],
+                        dim=1
+                    )
+
+                x = torch.chunk(x, world_size, dim=1)[rank]
+
+            for idx, block in enumerate(self.blocks):
+                x = block(x, **kwargs)
+                x = self.after_transformer_block(idx, x, motion_vec)
 
-        for idx, block in enumerate(self.blocks):
-            x = block(x, **kwargs)
-            x = self.after_transformer_block(idx, x, motion_vec)
+            if tea_cache is not None:
+                x_cache = gather_forward(x, dim=1) if self.use_context_parallel else x
+                x_cache = x_cache[:, :ori_x_len]
+                tea_cache.store(x_cache)
 
         # head
         x = self.head(x, e)
 
         if self.use_context_parallel:
             x = gather_forward(x, dim=1)
+            x = x[:, :ori_x_len]
 
         # unpatchify
         x = self.unpatchify(x, grid_sizes)