cache-dit 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

cache_dit/__init__.py

@@ -18,6 +18,10 @@ from cache_dit.cache_factory import BlockAdapter
 from cache_dit.cache_factory import ParamsModifier
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory import PatchFunctor
+from cache_dit.cache_factory import BasicCacheConfig
+from cache_dit.cache_factory import CalibratorConfig
+from cache_dit.cache_factory import TaylorSeerCalibratorConfig
+from cache_dit.cache_factory import FoCaCalibratorConfig
 from cache_dit.cache_factory import supported_pipelines
 from cache_dit.cache_factory import get_adapter
 from cache_dit.compile import set_compile_configs

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.0'
-__version_tuple__ = version_tuple = (0, 3, 0)
+__version__ = version = '0.3.2'
+__version_tuple__ = version_tuple = (0, 3, 2)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -3,15 +3,19 @@ from cache_dit.cache_factory.cache_types import cache_type
 from cache_dit.cache_factory.cache_types import block_range
 
 from cache_dit.cache_factory.forward_pattern import ForwardPattern
-
+from cache_dit.cache_factory.params_modifier import ParamsModifier
 from cache_dit.cache_factory.patch_functors import PatchFunctor
 
 from cache_dit.cache_factory.block_adapters import BlockAdapter
-from cache_dit.cache_factory.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 
 from cache_dit.cache_factory.cache_contexts import CachedContext
+from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 from cache_dit.cache_factory.cache_contexts import CachedContextManager
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+from cache_dit.cache_factory.cache_contexts import TaylorSeerCalibratorConfig
+from cache_dit.cache_factory.cache_contexts import FoCaCalibratorConfig
+
 from cache_dit.cache_factory.cache_blocks import CachedBlocks
 
 from cache_dit.cache_factory.cache_adapters import CachedAdapter

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -7,63 +7,15 @@ from collections.abc import Iterable
 from typing import Any, Tuple, List, Optional, Union
 
 from diffusers import DiffusionPipeline
-from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.patch_functors import PatchFunctor
+from cache_dit.cache_factory.forward_pattern import ForwardPattern
+from cache_dit.cache_factory.params_modifier import ParamsModifier
 
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class ParamsModifier:
-    def __init__(
-        self,
-        # Cache context kwargs
-        Fn_compute_blocks: Optional[int] = None,
-        Bn_compute_blocks: Optional[int] = None,
-        max_warmup_steps: Optional[int] = None,
-        max_cached_steps: Optional[int] = None,
-        max_continuous_cached_steps: Optional[int] = None,
-        residual_diff_threshold: Optional[float] = None,
-        # Cache CFG or not
-        enable_separate_cfg: Optional[bool] = None,
-        cfg_compute_first: Optional[bool] = None,
-        cfg_diff_compute_separate: Optional[bool] = None,
-        # Hybird TaylorSeer
-        enable_taylorseer: Optional[bool] = None,
-        enable_encoder_taylorseer: Optional[bool] = None,
-        taylorseer_cache_type: Optional[str] = None,
-        taylorseer_order: Optional[int] = None,
-        **other_cache_context_kwargs,
-    ):
-        self._context_kwargs = other_cache_context_kwargs.copy()
-        self._maybe_update_param("Fn_compute_blocks", Fn_compute_blocks)
-        self._maybe_update_param("Bn_compute_blocks", Bn_compute_blocks)
-        self._maybe_update_param("max_warmup_steps", max_warmup_steps)
-        self._maybe_update_param("max_cached_steps", max_cached_steps)
-        self._maybe_update_param(
-            "max_continuous_cached_steps", max_continuous_cached_steps
-        )
-        self._maybe_update_param(
-            "residual_diff_threshold", residual_diff_threshold
-        )
-        self._maybe_update_param("enable_separate_cfg", enable_separate_cfg)
-        self._maybe_update_param("cfg_compute_first", cfg_compute_first)
-        self._maybe_update_param(
-            "cfg_diff_compute_separate", cfg_diff_compute_separate
-        )
-        self._maybe_update_param("enable_taylorseer", enable_taylorseer)
-        self._maybe_update_param(
-            "enable_encoder_taylorseer", enable_encoder_taylorseer
-        )
-        self._maybe_update_param("taylorseer_cache_type", taylorseer_cache_type)
-        self._maybe_update_param("taylorseer_order", taylorseer_order)
-
-    def _maybe_update_param(self, key: str, value: Any):
-        if value is not None:
-            self._context_kwargs[key] = value
-
-
 @dataclasses.dataclass
 class BlockAdapter:
 
@@ -113,10 +65,12 @@ class BlockAdapter:
     ] = None
 
     # modify cache context params for specific blocks.
-    params_modifiers: Union[
-        ParamsModifier,
-        List[ParamsModifier],
-        List[List[ParamsModifier]],
+    params_modifiers: Optional[
+        Union[
+            ParamsModifier,
+            List[ParamsModifier],
+            List[List[ParamsModifier]],
+        ]
     ] = None
 
     check_forward_pattern: bool = True

cache_dit/cache_factory/cache_adapters/__init__.py

@@ -0,0 +1 @@
+from cache_dit.cache_factory.cache_adapters.cache_adapter import CachedAdapter

cache_dit/cache_factory/{cache_adapters.py → cache_adapters/cache_adapter.py}

@@ -8,12 +8,14 @@ from typing import Dict, List, Tuple, Any, Union, Callable
 
 from diffusers import DiffusionPipeline
 
-from cache_dit.cache_factory import CacheType
-from cache_dit.cache_factory import BlockAdapter
-from cache_dit.cache_factory import ParamsModifier
-from cache_dit.cache_factory import BlockAdapterRegistry
-from cache_dit.cache_factory import CachedContextManager
-from cache_dit.cache_factory import CachedBlocks
+from cache_dit.cache_factory.cache_types import CacheType
+from cache_dit.cache_factory.block_adapters import BlockAdapter
+from cache_dit.cache_factory.block_adapters import ParamsModifier
+from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
+from cache_dit.cache_factory.cache_contexts import CachedContextManager
+from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+from cache_dit.cache_factory.cache_blocks import CachedBlocks
 from cache_dit.cache_factory.cache_blocks.utils import (
     patch_cached_stats,
     remove_cached_stats,
@@ -55,6 +57,12 @@ class CachedAdapter:
                 block_adapter = BlockAdapterRegistry.get_adapter(
                     pipe_or_adapter
                 )
+                if params_modifiers := cache_context_kwargs.pop(
+                    "params_modifiers",
+                    None,
+                ):
+                    block_adapter.params_modifiers = params_modifiers
+
                 return cls.cachify(
                     block_adapter,
                     **cache_context_kwargs,
@@ -69,6 +77,12 @@ class CachedAdapter:
             logger.info(
                 "Adapting Cache Acceleration using custom BlockAdapter!"
             )
+            if pipe_or_adapter.params_modifiers is None:
+                if params_modifiers := cache_context_kwargs.pop(
+                    "params_modifiers", None
+                ):
+                    pipe_or_adapter.params_modifiers = params_modifiers
+
             return cls.cachify(
                 pipe_or_adapter,
                 **cache_context_kwargs,
@@ -114,33 +128,36 @@ class CachedAdapter:
         **cache_context_kwargs,
     ):
         # Check cache_context_kwargs
-        if cache_context_kwargs["enable_separate_cfg"] is None:
+        cache_config: BasicCacheConfig = cache_context_kwargs[
+            "cache_config"
+        ]  # ref
+        assert cache_config is not None, "cache_config can not be None."
+        if cache_config.enable_separate_cfg is None:
             # Check cfg for some specific case if users don't set it as True
             if BlockAdapterRegistry.has_separate_cfg(block_adapter):
-                cache_context_kwargs["enable_separate_cfg"] = True
+                cache_config.enable_separate_cfg = True
                 logger.info(
                     f"Use custom 'enable_separate_cfg' from BlockAdapter: True. "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
             else:
-                cache_context_kwargs["enable_separate_cfg"] = (
+                cache_config.enable_separate_cfg = (
                     BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
                 )
                 logger.info(
                     f"Use default 'enable_separate_cfg' from block adapter "
-                    f"register: {cache_context_kwargs['enable_separate_cfg']}, "
+                    f"register: {cache_config.enable_separate_cfg}, "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
         else:
             logger.info(
                 f"Use custom 'enable_separate_cfg' from cache context "
-                f"kwargs: {cache_context_kwargs['enable_separate_cfg']}. "
+                f"kwargs: {cache_config.enable_separate_cfg}. "
                 f"Pipeline: {block_adapter.pipe.__class__.__name__}."
             )
 
-        if (
-            cache_type := cache_context_kwargs.pop("cache_type", None)
-        ) is not None:
+        cache_type = cache_context_kwargs.pop("cache_type", None)
+        if cache_type is not None:
             assert (
                 cache_type == CacheType.DBCache
             ), "Custom cache setting only support for DBCache now!"
@@ -176,7 +193,7 @@ class CachedAdapter:
         block_adapter.pipe._cache_manager = cache_manager  # instance level
 
         flatten_contexts, contexts_kwargs = cls.modify_context_params(
-            block_adapter, cache_manager, **cache_context_kwargs
+            block_adapter, **cache_context_kwargs
         )
 
         original_call = block_adapter.pipe.__class__.__call__
@@ -212,7 +229,6 @@ class CachedAdapter:
     def modify_context_params(
         cls,
         block_adapter: BlockAdapter,
-        cache_manager: CachedContextManager,
         **cache_context_kwargs,
     ) -> Tuple[List[str], List[Dict[str, Any]]]:
 
@@ -230,6 +246,8 @@ class CachedAdapter:
             contexts_kwargs[i]["name"] = flatten_contexts[i]
 
         if block_adapter.params_modifiers is None:
+            for i in range(len(contexts_kwargs)):
+                cls._config_messages(**contexts_kwargs[i])
             return flatten_contexts, contexts_kwargs
 
         flatten_modifiers: List[ParamsModifier] = BlockAdapter.flatten(
@@ -242,12 +260,26 @@ class CachedAdapter:
             contexts_kwargs[i].update(
                 flatten_modifiers[i]._context_kwargs,
             )
-            contexts_kwargs[i], _ = cache_manager.collect_cache_kwargs(
-                default_attrs={}, **contexts_kwargs[i]
-            )
+            cls._config_messages(**contexts_kwargs[i])
 
         return flatten_contexts, contexts_kwargs
 
+    @classmethod
+    def _config_messages(cls, **contexts_kwargs):
+        cache_config: BasicCacheConfig = contexts_kwargs.get(
+            "cache_config", None
+        )
+        calibrator_config: CalibratorConfig = contexts_kwargs.get(
+            "calibrator_config", None
+        )
+        if cache_config is not None:
+            message = f"Collected Cache Config: {cache_config.strify()}"
+            if calibrator_config is not None:
+                message += f", Calibrator Config: {calibrator_config.strify(details=True)}"
+            else:
+                message += ", Calibrator Config: None"
+            logger.info(message)
+
     @classmethod
     def mock_blocks(
         cls,
@@ -335,7 +367,8 @@ class CachedAdapter:
         total_cached_blocks: List[Dict[str, torch.nn.ModuleList]] = []
         assert hasattr(block_adapter.pipe, "_cache_manager")
         assert isinstance(
-            block_adapter.pipe._cache_manager, CachedContextManager
+            block_adapter.pipe._cache_manager,
+            CachedContextManager,
         )
 
         for i in range(len(block_adapter.transformer)):

cache_dit/cache_factory/cache_contexts/__init__.py

@@ -1,5 +1,14 @@
-# namespace alias: for _CachedContext and many others' cache context funcs.
-from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
+from cache_dit.cache_factory.cache_contexts.calibrators import (
+    Calibrator,
+    CalibratorBase,
+    CalibratorConfig,
+    TaylorSeerCalibratorConfig,
+    FoCaCalibratorConfig,
+)
+from cache_dit.cache_factory.cache_contexts.cache_context import (
+    CachedContext,
+    BasicCacheConfig,
+)
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
     CachedContextManager,
 )