Timestep Embedding Aware Cache, a training-free caching approach that estimates and leverages the fluctuating differences among model outputs across timesteps, thereby accelerating the inference. Please refer to: 1. https://github.com/ali-vilab/TeaCache. 2. Liu, Feng, et al. "Ti
| 23 | |
| 24 | |
| 25 | class TeaCache(): |
| 26 | """ |
| 27 | Timestep Embedding Aware Cache, a training-free caching approach that estimates and leverages |
| 28 | the fluctuating differences among model outputs across timesteps, thereby accelerating the inference. |
| 29 | Please refer to: |
| 30 | 1. https://github.com/ali-vilab/TeaCache. |
| 31 | 2. Liu, Feng, et al. "Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model." arXiv preprint arXiv:2411.19108 (2024). |
| 32 | """ |
| 33 | def __init__( |
| 34 | self, |
| 35 | coefficients: list[float], |
| 36 | num_steps: int, |
| 37 | rel_l1_thresh: float = 0.0, |
| 38 | num_skip_start_steps: int = 0, |
| 39 | offload: bool = True, |
| 40 | ): |
| 41 | if num_steps < 1: |
| 42 | raise ValueError(f"`num_steps` must be greater than 0 but is {num_steps}.") |
| 43 | if rel_l1_thresh < 0: |
| 44 | raise ValueError(f"`rel_l1_thresh` must be greater than or equal to 0 but is {rel_l1_thresh}.") |
| 45 | if num_skip_start_steps < 0 or num_skip_start_steps > num_steps: |
| 46 | raise ValueError( |
| 47 | "`num_skip_start_steps` must be great than or equal to 0 and " |
| 48 | f"less than or equal to `num_steps={num_steps}` but is {num_skip_start_steps}." |
| 49 | ) |
| 50 | self.coefficients = coefficients |
| 51 | self.num_steps = num_steps |
| 52 | self.rel_l1_thresh = rel_l1_thresh |
| 53 | self.num_skip_start_steps = num_skip_start_steps |
| 54 | self.offload = offload |
| 55 | self.rescale_func = np.poly1d(self.coefficients) |
| 56 | |
| 57 | self.cnt = 0 |
| 58 | self.should_calc = True |
| 59 | self.accumulated_rel_l1_distance = 0 |
| 60 | self.previous_modulated_input = None |
| 61 | # Some pipelines concatenate the unconditional and text guide in forward. |
| 62 | self.previous_residual = None |
| 63 | # Some pipelines perform forward propagation separately on the unconditional and text guide. |
| 64 | self.previous_residual_cond = None |
| 65 | self.previous_residual_uncond = None |
| 66 | |
| 67 | @staticmethod |
| 68 | def compute_rel_l1_distance(prev: torch.Tensor, cur: torch.Tensor) -> torch.Tensor: |
| 69 | rel_l1_distance = (torch.abs(cur - prev).mean()) / torch.abs(prev).mean() |
| 70 | |
| 71 | return rel_l1_distance.cpu().item() |
| 72 | |
| 73 | def reset(self): |
| 74 | self.cnt = 0 |
| 75 | self.should_calc = True |
| 76 | self.accumulated_rel_l1_distance = 0 |
| 77 | self.previous_modulated_input = None |
| 78 | self.previous_residual = None |
| 79 | self.previous_residual_cond = None |
| 80 | self.previous_residual_uncond = None |