| 377 | state.calibration_ratios.append(ratio) |
| 378 | |
| 379 | def _advance_step(self, state: MagCacheState): |
| 380 | state.step_index += 1 |
| 381 | if state.step_index >= self.config.num_inference_steps: |
| 382 | # End of inference loop |
| 383 | if self.config.calibrate: |
| 384 | print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):") |
| 385 | print(f"{state.calibration_ratios}\n") |
| 386 | logger.info(f"MagCache Calibration Results: {state.calibration_ratios}") |
| 387 | |
| 388 | # Reset state |
| 389 | state.step_index = 0 |
| 390 | state.accumulated_ratio = 1.0 |
| 391 | state.accumulated_steps = 0 |
| 392 | state.accumulated_err = 0.0 |
| 393 | state.previous_residual = None |
| 394 | state.calibration_ratios = [] |
| 395 | |
| 396 | |
| 397 | def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None: |