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