| from __future__ import annotations |
|
|
| from methods.cache_strategy.common import WorldCacheConfig |
| from methods.cache_strategy.WorldCache.ctrl_world_runtime import CtrlWorldWorldCacheRuntime |
|
|
|
|
| def enable_worldcache( |
| unet, |
| *, |
| num_steps: int, |
| rel_l1_thresh: float = 0.03, |
| ret_ratio: float = 0.4, |
| probe_depth: int = 3, |
| motion_sensitivity: float = 5.0, |
| hf_enabled: bool = False, |
| hf_thresh: float = 0.01, |
| saliency_enabled: bool = False, |
| saliency_weight: float = 5.0, |
| osi_enabled: bool = False, |
| dynamic_decay: bool = False, |
| ) -> None: |
| config = WorldCacheConfig( |
| num_steps=num_steps, |
| rel_l1_thresh=rel_l1_thresh, |
| ret_ratio=ret_ratio, |
| probe_depth=probe_depth, |
| motion_sensitivity=motion_sensitivity, |
| flow_enabled=False, |
| flow_scale=0.5, |
| hf_enabled=hf_enabled, |
| hf_thresh=hf_thresh, |
| saliency_enabled=saliency_enabled, |
| saliency_weight=saliency_weight, |
| osi_enabled=osi_enabled, |
| dynamic_decay=dynamic_decay, |
| aduc_enabled=False, |
| aduc_start=0.5, |
| parallel_cfg=False, |
| ) |
| unet._ctrl_cache_runtime = CtrlWorldWorldCacheRuntime(config=config, total_steps=num_steps) |
| print( |
| "[WorldCache] Enabled on Ctrl-World " |
| f"(steps={num_steps}, rel_l1_thresh={rel_l1_thresh}, ret_ratio={ret_ratio}, " |
| f"probe_depth={probe_depth}, motion_sensitivity={motion_sensitivity}, " |
| f"hf={hf_enabled}, saliency={saliency_enabled}, osi={osi_enabled}, decay={dynamic_decay})" |
| ) |
|
|
|
|
| def disable_worldcache(unet) -> None: |
| if hasattr(unet, "_ctrl_cache_runtime"): |
| delattr(unet, "_ctrl_cache_runtime") |
| print("[WorldCache] Disabled on Ctrl-World.") |
|
|
