| 466 | |
| 467 | class TeaCache: |
| 468 | def __init__(self, num_inference_steps, rel_l1_thresh, model_id): |
| 469 | self.num_inference_steps = num_inference_steps |
| 470 | self.step = 0 |
| 471 | self.accumulated_rel_l1_distance = 0 |
| 472 | self.previous_modulated_input = None |
| 473 | self.rel_l1_thresh = rel_l1_thresh |
| 474 | self.previous_residual = None |
| 475 | self.previous_hidden_states = None |
| 476 | |
| 477 | self.coefficients_dict = { |
| 478 | "Wan2.1-T2V-1.3B": [-5.21862437e+04, 9.23041404e+03, -5.28275948e+02, 1.36987616e+01, -4.99875664e-02], |
| 479 | "Wan2.1-T2V-14B": [-3.03318725e+05, 4.90537029e+04, -2.65530556e+03, 5.87365115e+01, -3.15583525e-01], |
| 480 | "Wan2.1-I2V-14B-480P": [2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01], |
| 481 | "Wan2.1-I2V-14B-720P": [ 8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02], |
| 482 | } |
| 483 | if model_id not in self.coefficients_dict: |
| 484 | supported_model_ids = ", ".join([i for i in self.coefficients_dict]) |
| 485 | raise ValueError(f"{model_id} is not a supported TeaCache model id. Please choose a valid model id in ({supported_model_ids}).") |
| 486 | self.coefficients = self.coefficients_dict[model_id] |
| 487 | |
| 488 | def check(self, dit: WanModel, x, t_mod): |
| 489 | modulated_inp = t_mod.clone() |