| from __future__ import annotations |
|
|
| from methods.cache_strategy.FasterCache.processor import SVDFasterCacheAttnProcessor |
|
|
|
|
| def enable_fastercache( |
| unet, |
| *, |
| start_step: int = 0, |
| model_interval: int = 5, |
| block_interval: int = 3, |
| first_layers_fp: int = 2, |
| ) -> None: |
| runtime_state = { |
| "start_step": start_step, |
| "model_interval": model_interval, |
| "block_interval": block_interval, |
| "first_layers_fp": first_layers_fp, |
| "current_step_idx": -1, |
| "resolved_start_step": 0, |
| "delta_lf": None, |
| "delta_hf": None, |
| "block_histories": {}, |
| "block_alpha": 0.3, |
| } |
| processors = {} |
| for layer_idx, name in enumerate(unet.attn_processors.keys()): |
| processors[name] = SVDFasterCacheAttnProcessor( |
| runtime_state=runtime_state, |
| layer_idx=layer_idx, |
| first_layers_fp=first_layers_fp, |
| ) |
|
|
| unet._fastercache_state = runtime_state |
| unet.set_attn_processor(processors) |
| print( |
| "[FasterCache] Enabled on " |
| f"{len(processors)} attention layers (start_step={start_step}, " |
| f"model_interval={model_interval}, block_interval={block_interval})" |
| ) |
|
|
|
|
| def disable_fastercache(unet) -> None: |
| if hasattr(unet, "_fastercache_state"): |
| delattr(unet, "_fastercache_state") |
| unet.set_default_attn_processor() |
| print("[FasterCache] Disabled - restored default attention processors.") |
|
|
