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

cache_dit/__init__.py

@@ -18,6 +18,7 @@ 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

cache_dit/_version.py

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

cache_dit/cache_factory/__init__.py

@@ -3,25 +3,22 @@ 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 CachedContextV2
-from cache_dit.cache_factory.cache_contexts import CachedContextManagerV2
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig  # no V1
+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
-from cache_dit.cache_factory.cache_adapters import CachedAdapterV2
 
 from cache_dit.cache_factory.cache_interface import enable_cache
 from cache_dit.cache_factory.cache_interface import disable_cache

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -7,73 +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.cache_contexts import CalibratorConfig
+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,
-        # New param only for v2 API
-        calibrator_config: Optional[CalibratorConfig] = 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
-        )
-        # V1 only supports the Taylorseer calibrator. We have decided to
-        # keep this code for API compatibility reasons.
-        if calibrator_config is None:
-            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)
-        else:
-            self._maybe_update_param("calibrator_config", calibrator_config)
-
-    def _maybe_update_param(self, key: str, value: Any):
-        if value is not None:
-            self._context_kwargs[key] = value
-
-
 @dataclasses.dataclass
 class BlockAdapter:
 
@@ -123,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

@@ -1,2 +1 @@
 from cache_dit.cache_factory.cache_adapters.cache_adapter import CachedAdapter
-from cache_dit.cache_factory.cache_adapters.v2 import CachedAdapterV2

cache_dit/cache_factory/cache_adapters/cache_adapter.py

@@ -13,6 +13,8 @@ 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,
@@ -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,12 +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.cache_manager import (
-    CachedContextManager,
-)
-from cache_dit.cache_factory.cache_contexts.v2 import (
-    CachedContextV2,
-    CachedContextManagerV2,
+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,
+)