File size: 948 Bytes
ec0a9aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | from __future__ import annotations
from methods.cache_strategy.common import DiCacheConfig
from methods.cache_strategy.DiCache.ctrl_world_runtime import CtrlWorldDiCacheRuntime
def enable_dicache(
unet,
*,
num_steps: int,
rel_l1_thresh: float = 0.08,
ret_ratio: float = 0.2,
probe_depth: int = 2,
) -> None:
config = DiCacheConfig(
num_steps=num_steps,
rel_l1_thresh=rel_l1_thresh,
ret_ratio=ret_ratio,
probe_depth=probe_depth,
)
unet._ctrl_cache_runtime = CtrlWorldDiCacheRuntime(config=config, total_steps=num_steps)
print(
"[DiCache] Enabled on Ctrl-World "
f"(steps={num_steps}, rel_l1_thresh={rel_l1_thresh}, "
f"ret_ratio={ret_ratio}, probe_depth={probe_depth})"
)
def disable_dicache(unet) -> None:
if hasattr(unet, "_ctrl_cache_runtime"):
delattr(unet, "_ctrl_cache_runtime")
print("[DiCache] Disabled on Ctrl-World.")
|