cache-dit 0.1.0__py3-none-any.whl

cache_dit/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)

cache_dit/cache_factory/__init__.py

@@ -0,0 +1,166 @@
+from enum import Enum
+
+from diffusers import DiffusionPipeline
+
+from cache_dit.cache_factory.dual_block_cache.diffusers_adapters import (
+    apply_db_cache_on_pipe,
+)
+from cache_dit.cache_factory.first_block_cache.diffusers_adapters import (
+    apply_fb_cache_on_pipe,
+)
+from cache_dit.cache_factory.dynamic_block_prune.diffusers_adapters import (
+    apply_db_prune_on_pipe,
+)
+from cache_dit.logger import init_logger
+
+
+logger = init_logger(__name__)
+
+
+class CacheType(Enum):
+    NONE = "NONE"
+    FBCache = "First_Block_Cache"
+    DBCache = "Dual_Block_Cache"
+    DBPrune = "Dynamic_Block_Prune"
+
+    @staticmethod
+    def type(cache_type: "CacheType | str") -> "CacheType":
+        if isinstance(cache_type, CacheType):
+            return cache_type
+        return CacheType.cache_type(cache_type)
+
+    @staticmethod
+    def cache_type(cache_type: "CacheType | str") -> "CacheType":
+        if cache_type is None:
+            return CacheType.NONE
+
+        if isinstance(cache_type, CacheType):
+            return cache_type
+        if cache_type.lower() in (
+            "first_block_cache",
+            "fb_cache",
+            "fbcache",
+            "fb",
+        ):
+            return CacheType.FBCache
+        elif cache_type.lower() in (
+            "dual_block_cache",
+            "db_cache",
+            "dbcache",
+            "db",
+        ):
+            return CacheType.DBCache
+        elif cache_type.lower() in (
+            "dynamic_block_prune",
+            "db_prune",
+            "dbprune",
+            "dbp",
+        ):
+            return CacheType.DBPrune
+        elif cache_type.lower() in (
+            "none_cache",
+            "nonecache",
+            "no_cache",
+            "nocache",
+            "none",
+            "no",
+        ):
+            return CacheType.NONE
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type}")
+
+    @staticmethod
+    def range(start: int, end: int, step: int = 1) -> list[int]:
+        if start > end or end <= 0 or step <= 1:
+            return []
+        # Always compute 0 and end - 1 blocks for DB Cache
+        return list(
+            sorted(set([0] + list(range(start, end, step)) + [end - 1]))
+        )
+
+    @staticmethod
+    def default_options(cache_type: "CacheType | str") -> dict:
+        _no_options = {
+            "cache_type": CacheType.NONE,
+        }
+
+        _fb_options = {
+            "cache_type": CacheType.FBCache,
+            "residual_diff_threshold": 0.08,
+            "warmup_steps": 8,
+            "max_cached_steps": 8,
+        }
+
+        _Fn_compute_blocks = 8
+        _Bn_compute_blocks = 8
+
+        _db_options = {
+            "cache_type": CacheType.DBCache,
+            "residual_diff_threshold": 0.12,
+            "warmup_steps": 8,
+            "max_cached_steps": -1,  # -1 means no limit
+            # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
+            "Fn_compute_blocks": _Fn_compute_blocks,
+            "Bn_compute_blocks": _Bn_compute_blocks,
+            "max_Fn_compute_blocks": 16,
+            "max_Bn_compute_blocks": 16,
+            "Fn_compute_blocks_ids": [],  # 0, 1, 2, ..., 7, etc.
+            "Bn_compute_blocks_ids": [],  # 0, 1, 2, ..., 7, etc.
+        }
+
+        _dbp_options = {
+            "cache_type": CacheType.DBPrune,
+            "residual_diff_threshold": 0.08,
+            "Fn_compute_blocks": _Fn_compute_blocks,
+            "Bn_compute_blocks": _Bn_compute_blocks,
+            "warmup_steps": 8,
+            "max_pruned_steps": -1,  # -1 means no limit
+        }
+
+        if cache_type == CacheType.FBCache:
+            return _fb_options
+        elif cache_type == CacheType.DBCache:
+            return _db_options
+        elif cache_type == CacheType.DBPrune:
+            return _dbp_options
+        elif cache_type == CacheType.NONE:
+            return _no_options
+        else:
+            raise ValueError(f"Unknown cache type: {cache_type}")
+
+
+def apply_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
+    assert isinstance(pipe, DiffusionPipeline)
+
+    if hasattr(pipe, "_is_cached") and pipe._is_cached:
+        return pipe
+
+    if hasattr(pipe, "_is_pruned") and pipe._is_pruned:
+        return pipe
+
+    cache_type = kwargs.pop("cache_type", None)
+    if cache_type is None:
+        logger.warning(
+            "No cache type specified, we will use DBCache by default. "
+            "Please specify the cache_type explicitly if you want to "
+            "use a different cache type."
+        )
+        # Force to use DBCache with default cache options
+        return apply_db_cache_on_pipe(
+            pipe,
+            **CacheType.default_options(CacheType.DBCache),
+        )
+
+    cache_type = CacheType.type(cache_type)
+
+    if cache_type == CacheType.FBCache:
+        return apply_fb_cache_on_pipe(pipe, *args, **kwargs)
+    elif cache_type == CacheType.DBCache:
+        return apply_db_cache_on_pipe(pipe, *args, **kwargs)
+    elif cache_type == CacheType.DBPrune:
+        return apply_db_prune_on_pipe(pipe, *args, **kwargs)
+    elif cache_type == CacheType.NONE:
+        logger.warning("Cache type is NONE, no caching will be applied.")
+        return pipe
+    else:
+        raise ValueError(f"Unknown cache type: {cache_type}")