| 308 | |
| 309 | |
| 310 | class TeaCache: |
| 311 | def __init__(self, num_inference_steps, rel_l1_thresh, model_id): |
| 312 | self.num_inference_steps = num_inference_steps |
| 313 | self.step = 0 |
| 314 | self.accumulated_rel_l1_distance = 0 |
| 315 | self.previous_modulated_input = None |
| 316 | self.rel_l1_thresh = rel_l1_thresh |
| 317 | self.previous_residual = None |
| 318 | self.previous_hidden_states = None |
| 319 | |
| 320 | self.coefficients_dict = { |
| 321 | "Wan2.1-T2V-1.3B": [-5.21862437e+04, 9.23041404e+03, -5.28275948e+02, 1.36987616e+01, -4.99875664e-02], |
| 322 | "Wan2.1-T2V-14B": [-3.03318725e+05, 4.90537029e+04, -2.65530556e+03, 5.87365115e+01, -3.15583525e-01], |
| 323 | "Wan2.1-I2V-14B-480P": [2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01], |
| 324 | "Wan2.1-I2V-14B-720P": [ 8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02], |
| 325 | } |
| 326 | if model_id not in self.coefficients_dict: |
| 327 | supported_model_ids = ", ".join([i for i in self.coefficients_dict]) |
| 328 | raise ValueError(f"{model_id} is not a supported TeaCache model id. Please choose a valid model id in ({supported_model_ids}).") |
| 329 | self.coefficients = self.coefficients_dict[model_id] |
| 330 | |
| 331 | def check(self, dit: WanModel, x, t_mod): |
| 332 | modulated_inp = t_mod.clone() |
| 333 | if self.step == 0 or self.step == self.num_inference_steps - 1: |
| 334 | should_calc = True |
| 335 | self.accumulated_rel_l1_distance = 0 |
| 336 | else: |
| 337 | coefficients = self.coefficients |
| 338 | rescale_func = np.poly1d(coefficients) |
| 339 | self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()) |
| 340 | if self.accumulated_rel_l1_distance < self.rel_l1_thresh: |
| 341 | should_calc = False |
| 342 | else: |
| 343 | should_calc = True |
| 344 | self.accumulated_rel_l1_distance = 0 |
| 345 | self.previous_modulated_input = modulated_inp |
| 346 | self.step += 1 |
| 347 | if self.step == self.num_inference_steps: |
| 348 | self.step = 0 |
| 349 | if should_calc: |
| 350 | self.previous_hidden_states = x.clone() |
| 351 | return not should_calc |
| 352 | |
| 353 | def store(self, hidden_states): |
| 354 | self.previous_residual = hidden_states - self.previous_hidden_states |
| 355 | self.previous_hidden_states = None |
| 356 | |
| 357 | def update(self, hidden_states): |
| 358 | hidden_states = hidden_states + self.previous_residual |
| 359 | return hidden_states |
| 360 | |
| 361 | |
| 362 | |