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

cache_dit/cache_factory/cache_contexts/prune_manager.py

@@ -0,0 +1,154 @@
+import torch
+import functools
+from typing import Dict, List, Tuple, Union
+
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CachedContextManager,
+)
+from cache_dit.cache_factory.cache_contexts.prune_context import (
+    PrunedContext,
+)
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class PrunedContextManager(CachedContextManager):
+    # Reuse CachedContextManager for Dynamic Block Prune
+
+    def __init__(self, name: str = None):
+        super().__init__(name)
+        # Overwrite for Dynamic Block Prune
+        self._current_context: PrunedContext = None
+        self._cached_context_manager: Dict[str, PrunedContext] = {}
+
+    # Overwrite for Dynamic Block Prune
+    def new_context(self, *args, **kwargs) -> PrunedContext:
+        _context = PrunedContext(*args, **kwargs)
+        self._cached_context_manager[_context.name] = _context
+        return _context
+
+    def set_context(
+        self,
+        cached_context: PrunedContext | str,
+    ) -> PrunedContext:
+        return super().set_context(cached_context)
+
+    def get_context(self, name: str = None) -> PrunedContext:
+        return super().get_context(name)
+
+    def reset_context(
+        self,
+        cached_context: PrunedContext | str,
+        *args,
+        **kwargs,
+    ) -> PrunedContext:
+        return super().reset_context(cached_context, *args, **kwargs)
+
+    # Specially for Dynamic Block Prune
+    @torch.compiler.disable
+    def add_pruned_step(self):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_pruned_step()
+
+    @torch.compiler.disable
+    def add_pruned_block(self, num_blocks):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_pruned_block(num_blocks)
+
+    @torch.compiler.disable
+    def add_actual_block(self, num_blocks):
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        cached_context.add_actual_block(num_blocks)
+
+    @torch.compiler.disable
+    def get_pruned_steps(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_pruned_steps()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_steps(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_pruned_steps()
+
+    @torch.compiler.disable
+    def get_pruned_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_pruned_blocks()
+
+    @torch.compiler.disable
+    def get_actual_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_actual_blocks()
+
+    @torch.compiler.disable
+    def get_cfg_pruned_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_pruned_blocks()
+
+    @torch.compiler.disable
+    def get_cfg_actual_blocks(self) -> List[int]:
+        cached_context = self.get_context()
+        assert cached_context is not None, "cached_context must be set before"
+        return cached_context.get_cfg_actual_blocks()
+
+    @torch.compiler.disable
+    @functools.lru_cache(maxsize=8)
+    def get_non_prune_blocks_ids(self, num_blocks: int) -> List[int]:
+        assert num_blocks is not None, "num_blocks must be provided"
+        assert num_blocks > 0, "num_blocks must be greater than 0"
+        # Get the non-prune block ids for current context
+        # Never prune the first `Fn` and last `Bn` blocks.
+        Fn_compute_blocks_ids = list(
+            range(
+                self.Fn_compute_blocks()
+                if self.Fn_compute_blocks() < num_blocks
+                else num_blocks
+            )
+        )
+
+        Bn_compute_blocks_ids = list(
+            range(
+                num_blocks
+                - (
+                    self.Bn_compute_blocks()
+                    if self.Bn_compute_blocks() < num_blocks
+                    else num_blocks
+                ),
+                num_blocks,
+            )
+        )
+        context = self.get_context()
+        assert context is not None, "cached_context must be set before"
+
+        non_prune_blocks_ids = list(
+            set(
+                Fn_compute_blocks_ids
+                + Bn_compute_blocks_ids
+                + context.cache_config.non_prune_block_ids
+            )
+        )
+        non_prune_blocks_ids = [
+            d for d in non_prune_blocks_ids if d < num_blocks
+        ]
+        return sorted(non_prune_blocks_ids)
+
+    @torch.compiler.disable
+    def can_prune(self, *args, **kwargs) -> bool:
+        # Directly reuse can_cache for Dynamic Block Prune
+        return self.can_cache(*args, **kwargs)
+
+    @torch.compiler.disable
+    def apply_prune(
+        self, *args, **kwargs
+    ) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
+        # Directly reuse apply_cache for Dynamic Block Prune
+        return self.apply_cache(*args, **kwargs)

cache_dit/cache_factory/cache_interface.py

@@ -5,6 +5,8 @@ from cache_dit.cache_factory.block_adapters import BlockAdapter
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
 from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+from cache_dit.cache_factory.cache_contexts import DBCacheConfig
+from cache_dit.cache_factory.cache_contexts import DBPruneConfig
 from cache_dit.cache_factory.cache_contexts import CalibratorConfig
 from cache_dit.cache_factory.params_modifier import ParamsModifier
 
@@ -19,8 +21,12 @@ def enable_cache(
         DiffusionPipeline,
         BlockAdapter,
     ],
-    # Basic DBCache config: BasicCacheConfig
-    cache_config: BasicCacheConfig = BasicCacheConfig(),
+    # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
+    cache_config: Union[
+        BasicCacheConfig,
+        DBCacheConfig,
+        DBPruneConfig,
+    ] = DBCacheConfig(),
     # Calibrator config: TaylorSeerCalibratorConfig, etc.
     calibrator_config: Optional[CalibratorConfig] = None,
     # Modify cache context params for specific blocks.
@@ -136,14 +142,14 @@ def enable_cache(
     >>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.
     """
     # Collect cache context kwargs
-    cache_context_kwargs = {}
-    if (cache_type := cache_context_kwargs.pop("cache_type", None)) is not None:
+    context_kwargs = {}
+    if (cache_type := context_kwargs.get("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
     # WARNING: Deprecated cache config params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
-    deprecated_cache_kwargs = {
+    deprecated_kwargs = {
         "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
         "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
         "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -159,23 +165,23 @@ def enable_cache(
         ),
     }
 
-    deprecated_cache_kwargs = {
-        k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+    deprecated_kwargs = {
+        k: v for k, v in deprecated_kwargs.items() if v is not None
     }
 
-    if deprecated_cache_kwargs:
+    if deprecated_kwargs:
         logger.warning(
             "Manually settup DBCache context without BasicCacheConfig is "
             "deprecated and will be removed in the future, please use "
             "`cache_config` parameter instead!"
         )
         if cache_config is not None:
-            cache_config.update(**deprecated_cache_kwargs)
+            cache_config.update(**deprecated_kwargs)
         else:
-            cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+            cache_config = BasicCacheConfig(**deprecated_kwargs)
 
     if cache_config is not None:
-        cache_context_kwargs["cache_config"] = cache_config
+        context_kwargs["cache_config"] = cache_config
 
     # WARNING: Deprecated taylorseer params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
@@ -202,15 +208,15 @@ def enable_cache(
         )
 
     if calibrator_config is not None:
-        cache_context_kwargs["calibrator_config"] = calibrator_config
+        context_kwargs["calibrator_config"] = calibrator_config
 
     if params_modifiers is not None:
-        cache_context_kwargs["params_modifiers"] = params_modifiers
+        context_kwargs["params_modifiers"] = params_modifiers
 
     if isinstance(pipe_or_adapter, (DiffusionPipeline, BlockAdapter)):
         return CachedAdapter.apply(
             pipe_or_adapter,
-            **cache_context_kwargs,
+            **context_kwargs,
         )
     else:
         raise ValueError(

cache_dit/cache_factory/cache_types.py

@@ -6,7 +6,8 @@ logger = init_logger(__name__)
 
 class CacheType(Enum):
     NONE = "NONE"
-    DBCache = "Dual_Block_Cache"
+    DBCache = "DBCache"  # "Dual_Block_Cache"
+    DBPrune = "DBPrune"  # "Dynamic_Block_Prune"
 
     @staticmethod
     def type(type_hint: "CacheType | str") -> "CacheType":
@@ -14,6 +15,9 @@ class CacheType(Enum):
             return type_hint
         return cache_type(type_hint)
 
+    def __str__(self) -> str:
+        return self.value
+
 
 def cache_type(type_hint: "CacheType | str") -> "CacheType":
     if type_hint is None:
@@ -21,7 +25,6 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
 
     if isinstance(type_hint, CacheType):
         return type_hint
-
     elif type_hint.upper() in (
         "DUAL_BLOCK_CACHE",
         "DB_CACHE",
@@ -29,6 +32,20 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
         "DB",
     ):
         return CacheType.DBCache
+    elif type_hint.upper() in (
+        "DYNAMIC_BLOCK_PRUNE",
+        "DB_PRUNE",
+        "DBPRUNE",
+        "DBP",
+    ):
+        return CacheType.DBPrune
+    elif type_hint.upper() in (
+        "NONE",
+        "NO_CACHE",
+        "NOCACHE",
+        "NC",
+    ):
+        return CacheType.NONE
     return CacheType.NONE
 
 

cache_dit/cache_factory/params_modifier.py

@@ -11,7 +11,7 @@ logger = init_logger(__name__)
 class ParamsModifier:
     def __init__(
         self,
-        # Basic DBCache config: BasicCacheConfig
+        # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
         cache_config: BasicCacheConfig = None,
         # Calibrator config: TaylorSeerCalibratorConfig, etc.
         calibrator_config: Optional[CalibratorConfig] = None,
@@ -22,7 +22,7 @@ class ParamsModifier:
 
         # WARNING: Deprecated cache config params. These parameters are now retained
         # for backward compatibility but will be removed in the future.
-        deprecated_cache_kwargs = {
+        deprecated_kwargs = {
             "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
             "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
             "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -40,20 +40,20 @@ class ParamsModifier:
             ),
         }
 
-        deprecated_cache_kwargs = {
-            k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+        deprecated_kwargs = {
+            k: v for k, v in deprecated_kwargs.items() if v is not None
         }
 
-        if deprecated_cache_kwargs:
+        if deprecated_kwargs:
             logger.warning(
                 "Manually settup DBCache context without BasicCacheConfig is "
                 "deprecated and will be removed in the future, please use "
                 "`cache_config` parameter instead!"
             )
             if cache_config is not None:
-                cache_config.update(**deprecated_cache_kwargs)
+                cache_config.update(**deprecated_kwargs)
             else:
-                cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+                cache_config = BasicCacheConfig(**deprecated_kwargs)
 
         if cache_config is not None:
             self._context_kwargs["cache_config"] = cache_config

cache_dit/cache_factory/utils.py

@@ -7,10 +7,6 @@ def load_cache_options_from_yaml(yaml_file_path):
             kwargs: dict = yaml.safe_load(f)
 
         required_keys = [
-            "max_warmup_steps",
-            "max_cached_steps",
-            "Fn_compute_blocks",
-            "Bn_compute_blocks",
             "residual_diff_threshold",
         ]
         for key in required_keys:
@@ -38,10 +34,25 @@ def load_cache_options_from_yaml(yaml_file_path):
                 )
             )
 
-        from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+        if "cache_type" not in kwargs:
+            from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 
-        cache_context_kwargs["cache_config"] = BasicCacheConfig()
-        cache_context_kwargs["cache_config"].update(**kwargs)
+            cache_context_kwargs["cache_config"] = BasicCacheConfig()
+            cache_context_kwargs["cache_config"].update(**kwargs)
+        else:
+            cache_type = kwargs.pop("cache_type")
+            if cache_type == "DBCache":
+                from cache_dit.cache_factory.cache_contexts import DBCacheConfig
+
+                cache_context_kwargs["cache_config"] = DBCacheConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            elif cache_type == "DBPrune":
+                from cache_dit.cache_factory.cache_contexts import DBPruneConfig
+
+                cache_context_kwargs["cache_config"] = DBPruneConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            else:
+                raise ValueError(f"Unsupported cache_type: {cache_type}.")
 
         return cache_context_kwargs