| 620 | self.previous_hidden_states = None |
| 621 | |
| 622 | def check(self, dit: FluxDiT, hidden_states, conditioning): |
| 623 | inp = hidden_states.clone() |
| 624 | temb_ = conditioning.clone() |
| 625 | modulated_inp, _, _, _, _ = dit.blocks[0].norm1_a(inp, emb=temb_) |
| 626 | if self.step == 0 or self.step == self.num_inference_steps - 1: |
| 627 | should_calc = True |
| 628 | self.accumulated_rel_l1_distance = 0 |
| 629 | else: |
| 630 | coefficients = [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01] |
| 631 | rescale_func = np.poly1d(coefficients) |
| 632 | self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item()) |
| 633 | if self.accumulated_rel_l1_distance < self.rel_l1_thresh: |
| 634 | should_calc = False |
| 635 | else: |
| 636 | should_calc = True |
| 637 | self.accumulated_rel_l1_distance = 0 |
| 638 | self.previous_modulated_input = modulated_inp |
| 639 | self.step += 1 |
| 640 | if self.step == self.num_inference_steps: |
| 641 | self.step = 0 |
| 642 | if should_calc: |
| 643 | self.previous_hidden_states = hidden_states.clone() |
| 644 | return not should_calc |
| 645 | |
| 646 | def store(self, hidden_states): |
| 647 | self.previous_residual = hidden_states - self.previous_hidden_states |