cache-dit 1.0.3__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/caching/cache_blocks/pattern_utils.py

@@ -0,0 +1,86 @@
+import torch
+
+from typing import Any
+from cache_dit.caching import CachedContext
+from cache_dit.caching import CachedContextManager
+from cache_dit.caching import PrunedContextManager
+
+
+def apply_stats(
+    module: torch.nn.Module | Any,
+    cache_context: CachedContext | str = None,
+    context_manager: CachedContextManager | PrunedContextManager = None,
+):
+    # Patch the cached stats to the module, the cached stats
+    # will be reset for each calling of pipe.__call__(**kwargs).
+    if module is None or context_manager is None:
+        return
+
+    if cache_context is not None:
+        context_manager.set_context(cache_context)
+
+    # Cache stats for Dual Block Cache
+    module._cached_steps = context_manager.get_cached_steps()
+    module._residual_diffs = context_manager.get_residual_diffs()
+    module._cfg_cached_steps = context_manager.get_cfg_cached_steps()
+    module._cfg_residual_diffs = context_manager.get_cfg_residual_diffs()
+    # Pruned stats for Dynamic Block Prune
+    if not isinstance(context_manager, PrunedContextManager):
+        return
+    module._pruned_steps = context_manager.get_pruned_steps()
+    module._cfg_pruned_steps = context_manager.get_cfg_pruned_steps()
+    module._pruned_blocks = context_manager.get_pruned_blocks()
+    module._cfg_pruned_blocks = context_manager.get_cfg_pruned_blocks()
+    module._actual_blocks = context_manager.get_actual_blocks()
+    module._cfg_actual_blocks = context_manager.get_cfg_actual_blocks()
+    # Caculate pruned ratio
+    if len(module._pruned_blocks) > 0 and sum(module._actual_blocks) > 0:
+        module._pruned_ratio = sum(module._pruned_blocks) / sum(
+            module._actual_blocks
+        )
+    else:
+        module._pruned_ratio = None
+    if (
+        len(module._cfg_pruned_blocks) > 0
+        and sum(module._cfg_actual_blocks) > 0
+    ):
+        module._cfg_pruned_ratio = sum(module._cfg_pruned_blocks) / sum(
+            module._cfg_actual_blocks
+        )
+    else:
+        module._cfg_pruned_ratio = None
+
+
+def remove_stats(
+    module: torch.nn.Module | Any,
+):
+    if module is None:
+        return
+
+    # Dual Block Cache
+    if hasattr(module, "_cached_steps"):
+        del module._cached_steps
+    if hasattr(module, "_residual_diffs"):
+        del module._residual_diffs
+    if hasattr(module, "_cfg_cached_steps"):
+        del module._cfg_cached_steps
+    if hasattr(module, "_cfg_residual_diffs"):
+        del module._cfg_residual_diffs
+
+    # Dynamic Block Prune
+    if hasattr(module, "_pruned_steps"):
+        del module._pruned_steps
+    if hasattr(module, "_cfg_pruned_steps"):
+        del module._cfg_pruned_steps
+    if hasattr(module, "_pruned_blocks"):
+        del module._pruned_blocks
+    if hasattr(module, "_cfg_pruned_blocks"):
+        del module._cfg_pruned_blocks
+    if hasattr(module, "_actual_blocks"):
+        del module._actual_blocks
+    if hasattr(module, "_cfg_actual_blocks"):
+        del module._cfg_actual_blocks
+    if hasattr(module, "_pruned_ratio"):
+        del module._pruned_ratio
+    if hasattr(module, "_cfg_pruned_ratio"):
+        del module._cfg_pruned_ratio

cache_dit/caching/cache_contexts/__init__.py

@@ -0,0 +1,28 @@
+from cache_dit.caching.cache_contexts.calibrators import (
+    Calibrator,
+    CalibratorBase,
+    CalibratorConfig,
+    TaylorSeerCalibratorConfig,
+    FoCaCalibratorConfig,
+)
+from cache_dit.caching.cache_contexts.cache_config import (
+    BasicCacheConfig,
+    DBCacheConfig,
+)
+from cache_dit.caching.cache_contexts.cache_context import (
+    CachedContext,
+)
+from cache_dit.caching.cache_contexts.cache_manager import (
+    CachedContextManager,
+    ContextNotExistError,
+)
+from cache_dit.caching.cache_contexts.prune_config import DBPruneConfig
+from cache_dit.caching.cache_contexts.prune_context import (
+    PrunedContext,
+)
+from cache_dit.caching.cache_contexts.prune_manager import (
+    PrunedContextManager,
+)
+from cache_dit.caching.cache_contexts.context_manager import (
+    ContextManager,
+)

cache_dit/caching/cache_contexts/cache_config.py

@@ -0,0 +1,120 @@
+import torch
+import dataclasses
+from typing import Optional, Union
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass
+class BasicCacheConfig:
+    # Default: Dual Block Cache with Flexible FnBn configuration.
+    cache_type: CacheType = CacheType.DBCache  # DBCache, DBPrune, NONE
+
+    # Fn_compute_blocks: (`int`, *required*, defaults to 8):
+    #     Specifies that `DBCache` uses the **first n** Transformer blocks to fit the information
+    #     at time step t, enabling the calculation of a more stable L1 diff and delivering more
+    #     accurate information to subsequent blocks. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
+    #     for more details of DBCache.
+    Fn_compute_blocks: int = 8
+    # Bn_compute_blocks: (`int`, *required*, defaults to 0):
+    #     Further fuses approximate information in the **last n** Transformer blocks to enhance
+    #     prediction accuracy. These blocks act as an auto-scaler for approximate hidden states
+    #     that use residual cache.
+    Bn_compute_blocks: int = 0
+    # residual_diff_threshold (`float`, *required*, defaults to 0.08):
+    #     the value of residual diff threshold, a higher value leads to faster performance at the
+    #     cost of lower precision.
+    residual_diff_threshold: Union[torch.Tensor, float] = 0.08
+    # max_warmup_steps (`int`, *required*, defaults to 8):
+    #     DBCache does not apply the caching strategy when the number of running steps is less than
+    #     or equal to this value, ensuring the model sufficiently learns basic features during warmup.
+    max_warmup_steps: int = 8  # DON'T Cache in warmup steps
+    # warmup_interval (`int`, *required*, defaults to 1):
+    #     Skip interval in warmup steps, e.g., when warmup_interval is 2, only 0, 2, 4, ... steps
+    #     in warmup steps will be computed, others will use dynamic cache.
+    warmup_interval: int = 1  # skip interval in warmup steps
+    # max_cached_steps (`int`, *required*, defaults to -1):
+    #     DBCache disables the caching strategy when the previous cached steps exceed this value to
+    #     prevent precision degradation.
+    max_cached_steps: int = -1  # for both CFG and non-CFG
+    # max_continuous_cached_steps (`int`, *required*, defaults to -1):
+    #     DBCache disables the caching strategy when the previous continous cached steps exceed this value to
+    #     prevent precision degradation.
+    max_continuous_cached_steps: int = -1  # the max continuous cached steps
+    # enable_separate_cfg (`bool`, *required*,  defaults to None):
+    #     Whether to do separate cfg or not, such as Wan 2.1, Qwen-Image. For model that fused CFG
+    #     and non-CFG into single forward step, should set enable_separate_cfg as False, for example:
+    #     CogVideoX, HunyuanVideo, Mochi, etc.
+    enable_separate_cfg: Optional[bool] = None
+    # cfg_compute_first (`bool`, *required*,  defaults to False):
+    #     Compute cfg forward first or not, default False, namely, 0, 2, 4, ..., -> non-CFG step;
+    #     1, 3, 5, ... -> CFG step.
+    cfg_compute_first: bool = False
+    # cfg_diff_compute_separate (`bool`, *required*,  defaults to True):
+    #     Compute separate diff values for CFG and non-CFG step, default True. If False, we will
+    #     use the computed diff from current non-CFG transformer step for current CFG step.
+    cfg_diff_compute_separate: bool = True
+    # num_inference_steps (`int`, *optional*, defaults to None):
+    #     num_inference_steps for DiffusionPipeline, used to adjust some internal settings
+    #     for better caching performance. For example, we will refresh the cache once the
+    #     executed steps exceed num_inference_steps if num_inference_steps is provided.
+    num_inference_steps: Optional[int] = None
+
+    def update(self, **kwargs) -> "BasicCacheConfig":
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                if value is not None:
+                    setattr(self, key, value)
+        return self
+
+    def empty(self, **kwargs) -> "BasicCacheConfig":
+        # Set all fields to None
+        for field in dataclasses.fields(self):
+            if hasattr(self, field.name):
+                setattr(self, field.name, None)
+        if kwargs:
+            self.update(**kwargs)
+        return self
+
+    def reset(self, **kwargs) -> "BasicCacheConfig":
+        return self.empty(**kwargs)
+
+    def as_dict(self) -> dict:
+        return dataclasses.asdict(self)
+
+    def strify(self) -> str:
+        return (
+            f"{self.cache_type}_"
+            f"F{self.Fn_compute_blocks}"
+            f"B{self.Bn_compute_blocks}_"
+            f"W{self.max_warmup_steps}"
+            f"I{self.warmup_interval}"
+            f"M{max(0, self.max_cached_steps)}"
+            f"MC{max(0, self.max_continuous_cached_steps)}_"
+            f"R{self.residual_diff_threshold}"
+        )
+
+
+@dataclasses.dataclass
+class ExtraCacheConfig:
+    # Some other not very important settings for Dual Block Cache.
+    # NOTE: These flags maybe deprecated in the future and users
+    # should never use these extra configurations in their cases.
+
+    # l1_hidden_states_diff_threshold (`float`, *optional*, defaults to None):
+    #     The hidden states diff threshold for DBCache if use hidden_states as
+    #     cache (not residual).
+    l1_hidden_states_diff_threshold: float = None
+    # important_condition_threshold (`float`, *optional*, defaults to 0.0):
+    #     Only select the most important tokens while calculating the l1 diff.
+    important_condition_threshold: float = 0.0
+    # downsample_factor (`int`, *optional*, defaults to 1):
+    #     Downsample factor for Fn buffer, in order the save GPU memory.
+    downsample_factor: int = 1
+
+
+@dataclasses.dataclass
+class DBCacheConfig(BasicCacheConfig):
+    pass  # Just an alias for BasicCacheConfig

cache_dit/{cache_factory → caching}/cache_contexts/cache_context.py

@@ -1,3 +1,5 @@
+# The cache context codebase is adapted from FBCache. Over time its codebase
+# diverged a lot, and context API is no longer compatible with FBCache.
 import logging
 import dataclasses
 from collections import defaultdict
@@ -5,7 +7,12 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union, Tuple
 
 import torch
 
-from cache_dit.cache_factory.cache_contexts.calibrators import (
+from cache_dit.caching.cache_contexts.cache_config import (
+    BasicCacheConfig,
+    ExtraCacheConfig,
+    DBCacheConfig,
+)
+from cache_dit.caching.cache_contexts.calibrators import (
     Calibrator,
     CalibratorBase,
     CalibratorConfig,
@@ -15,101 +22,16 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-@dataclasses.dataclass
-class BasicCacheConfig:
-    # Dual Block Cache with Flexible FnBn configuration.
-
-    # Fn_compute_blocks: (`int`, *required*, defaults to 8):
-    #     Specifies that `DBCache` uses the **first n** Transformer blocks to fit the information
-    #     at time step t, enabling the calculation of a more stable L1 diff and delivering more
-    #     accurate information to subsequent blocks. Please check https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md
-    #     for more details of DBCache.
-    Fn_compute_blocks: int = 8
-    # Bn_compute_blocks: (`int`, *required*, defaults to 0):
-    #     Further fuses approximate information in the **last n** Transformer blocks to enhance
-    #     prediction accuracy. These blocks act as an auto-scaler for approximate hidden states
-    #     that use residual cache.
-    Bn_compute_blocks: int = 0
-    # residual_diff_threshold (`float`, *required*, defaults to 0.08):
-    #     the value of residual diff threshold, a higher value leads to faster performance at the
-    #     cost of lower precision.
-    residual_diff_threshold: Union[torch.Tensor, float] = 0.08
-    # max_warmup_steps (`int`, *required*, defaults to 8):
-    #     DBCache does not apply the caching strategy when the number of running steps is less than
-    #     or equal to this value, ensuring the model sufficiently learns basic features during warmup.
-    max_warmup_steps: int = 8  # DON'T Cache in warmup steps
-    # warmup_interval (`int`, *required*, defaults to 1):
-    #     Skip interval in warmup steps, e.g., when warmup_interval is 2, only 0, 2, 4, ... steps
-    #     in warmup steps will be computed, others will use dynamic cache.
-    warmup_interval: int = 1  # skip interval in warmup steps
-    # max_cached_steps (`int`, *required*, defaults to -1):
-    #     DBCache disables the caching strategy when the previous cached steps exceed this value to
-    #     prevent precision degradation.
-    max_cached_steps: int = -1  # for both CFG and non-CFG
-    # max_continuous_cached_steps (`int`, *required*, defaults to -1):
-    #     DBCache disables the caching strategy when the previous continous cached steps exceed this value to
-    #     prevent precision degradation.
-    max_continuous_cached_steps: int = -1  # the max continuous cached steps
-    # enable_separate_cfg (`bool`, *required*,  defaults to None):
-    #     Whether to do separate cfg or not, such as Wan 2.1, Qwen-Image. For model that fused CFG
-    #     and non-CFG into single forward step, should set enable_separate_cfg as False, for example:
-    #     CogVideoX, HunyuanVideo, Mochi, etc.
-    enable_separate_cfg: Optional[bool] = None
-    # cfg_compute_first (`bool`, *required*,  defaults to False):
-    #     Compute cfg forward first or not, default False, namely, 0, 2, 4, ..., -> non-CFG step;
-    #     1, 3, 5, ... -> CFG step.
-    cfg_compute_first: bool = False
-    # cfg_diff_compute_separate (`bool`, *required*,  defaults to True):
-    #     Compute separate diff values for CFG and non-CFG step, default True. If False, we will
-    #     use the computed diff from current non-CFG transformer step for current CFG step.
-    cfg_diff_compute_separate: bool = True
-
-    def update(self, **kwargs) -> "BasicCacheConfig":
-        for key, value in kwargs.items():
-            if hasattr(self, key):
-                setattr(self, key, value)
-        return self
-
-    def strify(self) -> str:
-        return (
-            f"DBCACHE_F{self.Fn_compute_blocks}"
-            f"B{self.Bn_compute_blocks}_"
-            f"W{self.max_warmup_steps}"
-            f"I{self.warmup_interval}"
-            f"M{max(0, self.max_cached_steps)}"
-            f"MC{max(0, self.max_continuous_cached_steps)}_"
-            f"R{self.residual_diff_threshold}"
-        )
-
-
-@dataclasses.dataclass
-class ExtraCacheConfig:
-    # Some other not very important settings for Dual Block Cache.
-    # NOTE: These flags maybe deprecated in the future and users
-    # should never use these extra configurations in their cases.
-
-    # l1_hidden_states_diff_threshold (`float`, *optional*, defaults to None):
-    #     The hidden states diff threshold for DBCache if use hidden_states as
-    #     cache (not residual).
-    l1_hidden_states_diff_threshold: float = None
-    # important_condition_threshold (`float`, *optional*, defaults to 0.0):
-    #     Only select the most important tokens while calculating the l1 diff.
-    important_condition_threshold: float = 0.0
-    # downsample_factor (`int`, *optional*, defaults to 1):
-    #     Downsample factor for Fn buffer, in order the save GPU memory.
-    downsample_factor: int = 1
-    # num_inference_steps (`int`, *optional*, defaults to -1):
-    #     num_inference_steps for DiffusionPipeline, for future use.
-    num_inference_steps: int = -1
-
-
 @dataclasses.dataclass
 class CachedContext:
     name: str = "default"
     # Buffer for storing the residuals and other tensors
     buffers: Dict[str, Any] = dataclasses.field(default_factory=dict)
     # Basic Dual Block Cache Config
-    cache_config: BasicCacheConfig = dataclasses.field(
+    cache_config: Union[
+        BasicCacheConfig,
+        DBCacheConfig,
+    ] = dataclasses.field(
         default_factory=BasicCacheConfig,
     )
     # NOTE: Users should never use these extra configurations.
@@ -131,14 +53,14 @@ class CachedContext:
     # be double of executed_steps.
     transformer_executed_steps: int = 0
 
-    # CFG & non-CFG cached steps
+    # CFG & non-CFG cached/pruned steps
     cached_steps: List[int] = dataclasses.field(default_factory=list)
-    residual_diffs: DefaultDict[str, float] = dataclasses.field(
+    residual_diffs: DefaultDict[str, float | list] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
     continuous_cached_steps: int = 0
     cfg_cached_steps: List[int] = dataclasses.field(default_factory=list)
-    cfg_residual_diffs: DefaultDict[str, float] = dataclasses.field(
+    cfg_residual_diffs: DefaultDict[str, float | list] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
     cfg_continuous_cached_steps: int = 0
@@ -286,7 +208,9 @@ class CachedContext:
     def get_cfg_calibrators(self) -> Tuple[CalibratorBase, CalibratorBase]:
         return self.cfg_calibrator, self.cfg_encoder_calibrator
 
-    def add_residual_diff(self, diff):
+    def add_residual_diff(self, diff: float | torch.Tensor):
+        if isinstance(diff, torch.Tensor):
+            diff = diff.item()
         # step: executed_steps - 1, not transformer_steps - 1
         step = str(self.get_current_step())
         # Only add the diff if it is not already recorded for this step

cache_dit/{cache_factory → caching}/cache_contexts/cache_manager.py

@@ -5,8 +5,9 @@ from typing import Dict, Optional, Tuple, Union, List
 import torch
 import torch.distributed as dist
 
-from cache_dit.cache_factory.cache_contexts.calibrators import CalibratorBase
-from cache_dit.cache_factory.cache_contexts.cache_context import (
+from cache_dit.caching.cache_contexts.calibrators import CalibratorBase
+from cache_dit.caching.cache_contexts.cache_context import (
+    BasicCacheConfig,
     CachedContext,
 )
 from cache_dit.logger import init_logger
@@ -14,36 +15,156 @@ from cache_dit.logger import init_logger
 logger = init_logger(__name__)
 
 
-class CacheNotExistError(Exception):
+class ContextNotExistError(Exception):
     pass
 
 
 class CachedContextManager:
     # Each Pipeline should have it's own context manager instance.
 
-    def __init__(self, name: str = None):
+    def __init__(self, name: str = None, persistent_context: bool = False):
         self.name = name
         self._current_context: CachedContext = None
         self._cached_context_manager: Dict[str, CachedContext] = {}
+        # Whether to create new context automatically when setting
+        # a non-exist context name. Persistent context is useful when
+        # the pipeline class is not provided and users want to use
+        # cache-dit in a transformer-only way.
+        self._persistent_context = persistent_context
+        self._current_step_refreshed: bool = False
+
+    @property
+    def persistent_context(self) -> bool:
+        return self._persistent_context
+
+    @property
+    def current_context(self) -> CachedContext:
+        return self._current_context
+
+    @property
+    @torch.compiler.disable
+    def current_step_refreshed(self) -> bool:
+        return self._current_step_refreshed
 
+    @torch.compiler.disable
+    def is_pre_refreshed(self) -> bool:
+        _context = self._current_context
+        if _context is None:
+            return False
+
+        num_inference_steps = _context.cache_config.num_inference_steps
+        if num_inference_steps is not None:
+            current_step = _context.get_current_step()  # e.g, 0~49,50~99,...
+            return current_step == num_inference_steps - 1
+        return False
+
+    @torch.compiler.disable
     def new_context(self, *args, **kwargs) -> CachedContext:
+        if self._persistent_context:
+            cache_config: BasicCacheConfig = kwargs.get("cache_config", None)
+            assert (
+                cache_config is not None
+                and cache_config.num_inference_steps is not None
+            ), (
+                "When persistent_context is True, num_inference_steps "
+                "must be set in cache_config for proper cache refreshing."
+                f"\nkwargs: {kwargs}"
+            )
         _context = CachedContext(*args, **kwargs)
+        # NOTE: Patch args and kwargs for implicit refresh.
+        _context._init_args = args  # maybe empty tuple: ()
+        _context._init_kwargs = kwargs  # maybe empty dict: {}
         self._cached_context_manager[_context.name] = _context
         return _context
 
-    def set_context(self, cached_context: CachedContext | str) -> CachedContext:
+    @torch.compiler.disable
+    def maybe_refresh(
+        self,
+        cached_context: Optional[CachedContext | str] = None,
+    ) -> bool:
+        if cached_context is None:
+            _context = self._current_context
+            assert _context is not None, "Current context is not set!"
+
+        if isinstance(cached_context, CachedContext):
+            _context = cached_context
+        else:
+            if cached_context not in self._cached_context_manager:
+                raise ContextNotExistError("Context not exist!")
+            _context = self._cached_context_manager[cached_context]
+
+        if self._persistent_context:
+            assert _context.cache_config.num_inference_steps is not None, (
+                "When persistent_context is True, num_inference_steps must be set "
+                "in cache_config for proper cache refreshing."
+            )
+
+        num_inference_steps = _context.cache_config.num_inference_steps
+        if num_inference_steps is not None:
+            current_step = _context.get_current_step()  # e.g, 0~49,50~99,...
+            # Another round of inference, need to refresh cache context.
+            if current_step >= num_inference_steps:
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        f"Refreshing cache context '{_context.name}' "
+                        f"as current step: {current_step} >= "
+                        f"num_inference_steps: {num_inference_steps}."
+                    )
+                return True
+        return False
+
+    @torch.compiler.disable
+    def set_context(
+        self,
+        cached_context: CachedContext | str,
+        *args,
+        **kwargs,
+    ) -> CachedContext:
         if isinstance(cached_context, CachedContext):
             self._current_context = cached_context
         else:
             if cached_context not in self._cached_context_manager:
-                raise CacheNotExistError("Context not exist!")
-            self._current_context = self._cached_context_manager[cached_context]
+                if not self._persistent_context:
+                    raise ContextNotExistError(
+                        "Context not exist and persistent_context is False. Please "
+                        "create new context first or set persistent_context=True."
+                    )
+                else:
+                    # Create new context if not exist
+                    if any((bool(args), bool(kwargs))):
+                        kwargs["name"] = cached_context
+                        self._current_context = self.new_context(
+                            *args, **kwargs
+                        )
+                    else:
+                        raise ValueError(
+                            "To create new context, please provide args and kwargs."
+                        )
+            else:
+                self._current_context = self._cached_context_manager[
+                    cached_context
+                ]
+
+        if self.maybe_refresh(self._current_context):
+            if not any((bool(args), bool(kwargs))):
+                assert hasattr(self._current_context, "_init_args")
+                assert hasattr(self._current_context, "_init_kwargs")
+                args = self._current_context._init_args
+                kwargs = self._current_context._init_kwargs
+
+            self._current_context = self.reset_context(
+                self._current_context, *args, **kwargs
+            )
+            self._current_step_refreshed = True
+        else:
+            self._current_step_refreshed = False
+
         return self._current_context
 
     def get_context(self, name: str = None) -> CachedContext:
         if name is not None:
             if name not in self._cached_context_manager:
-                raise CacheNotExistError("Context not exist!")
+                raise ContextNotExistError("Context not exist!")
             return self._cached_context_manager[name]
         return self._current_context
 
@@ -482,7 +603,7 @@ class CachedContextManager:
 
             if calibrator is not None:
                 # Use calibrator to update the buffer
-                calibrator.update(buffer)
+                calibrator.update(buffer, name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
@@ -513,7 +634,7 @@ class CachedContextManager:
                 calibrator, _ = self.get_calibrator()
 
             if calibrator is not None:
-                return calibrator.approximate()
+                return calibrator.approximate(name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
@@ -551,7 +672,7 @@ class CachedContextManager:
 
             if encoder_calibrator is not None:
                 # Use CalibratorBase to update the buffer
-                encoder_calibrator.update(buffer)
+                encoder_calibrator.update(buffer, name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(
@@ -582,7 +703,7 @@ class CachedContextManager:
 
             if encoder_calibrator is not None:
                 # Use calibrator to approximate the value
-                return encoder_calibrator.approximate()
+                return encoder_calibrator.approximate(name=prefix)
             else:
                 if logger.isEnabledFor(logging.DEBUG):
                     logger.debug(

cache_dit/{cache_factory → caching}/cache_contexts/calibrators/__init__.py

@@ -1,10 +1,10 @@
-from cache_dit.cache_factory.cache_contexts.calibrators.base import (
+from cache_dit.caching.cache_contexts.calibrators.base import (
     CalibratorBase,
 )
-from cache_dit.cache_factory.cache_contexts.calibrators.taylorseer import (
+from cache_dit.caching.cache_contexts.calibrators.taylorseer import (
     TaylorSeerCalibrator,
 )
-from cache_dit.cache_factory.cache_contexts.calibrators.foca import (
+from cache_dit.caching.cache_contexts.calibrators.foca import (
     FoCaCalibrator,
 )
 
@@ -45,6 +45,28 @@ class CalibratorConfig:
     def to_kwargs(self) -> Dict:
         return self.calibrator_kwargs.copy()
 
+    def as_dict(self) -> dict:
+        return dataclasses.asdict(self)
+
+    def update(self, **kwargs) -> "CalibratorConfig":
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                if value is not None:
+                    setattr(self, key, value)
+        return self
+
+    def empty(self, **kwargs) -> "CalibratorConfig":
+        # Set all fields to None
+        for field in dataclasses.fields(self):
+            if hasattr(self, field.name):
+                setattr(self, field.name, None)
+        if kwargs:
+            self.update(**kwargs)
+        return self
+
+    def reset(self, **kwargs) -> "CalibratorConfig":
+        return self.empty(**kwargs)
+
 
 @dataclasses.dataclass
 class TaylorSeerCalibratorConfig(CalibratorConfig):

cache_dit/{cache_factory → caching}/cache_contexts/calibrators/foca.py

@@ -1,4 +1,4 @@
-from cache_dit.cache_factory.cache_contexts.calibrators.base import (
+from cache_dit.caching.cache_contexts.calibrators.base import (
     CalibratorBase,
 )