cache-dit 0.2.16__py3-none-any.whl → 0.2.17__py3-none-any.whl

cache_dit/__init__.py

@@ -3,3 +3,15 @@ try:
 except ImportError:
     __version__ = "unknown version"
     version_tuple = (0, 0, "unknown version")
+
+from cache_dit.cache_factory import load_options
+from cache_dit.cache_factory import enable_cache
+from cache_dit.cache_factory import cache_type
+from cache_dit.cache_factory import default_options
+from cache_dit.cache_factory import block_range
+from cache_dit.cache_factory import CacheType
+from cache_dit.compile import set_compile_configs
+from cache_dit.logger import init_logger
+
+NONE = CacheType.NONE
+DBCache = CacheType.DBCache

cache_dit/_version.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     from typing import Union
 
     VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
 else:
     VERSION_TUPLE = object
+    COMMIT_ID = object
 
 version: str
 __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.16'
-__version_tuple__ = version_tuple = (0, 2, 16)
+__version__ = version = '0.2.17'
+__version_tuple__ = version_tuple = (0, 2, 17)
+
+__commit_id__ = commit_id = None

cache_dit/cache_factory/.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+deprecated

cache_dit/cache_factory/__init__.py

@@ -1,4 +1,53 @@
-from cache_dit.cache_factory.adapters import CacheType
-from cache_dit.cache_factory.adapters import apply_cache_on_pipe
-from cache_dit.cache_factory.adapters import apply_cache_on_transformer
+import torch
+from typing import Dict, List
+from diffusers import DiffusionPipeline
+from cache_dit.cache_factory.cache_adapters import CacheType
+from cache_dit.cache_factory.cache_adapters import UnifiedCacheAdapter
 from cache_dit.cache_factory.utils import load_cache_options_from_yaml
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def load_options(path: str):
+    """cache_dit.load_options(cache_config.yaml)"""
+    return load_cache_options_from_yaml(path)
+
+
+def cache_type(type_hint: "CacheType | str") -> CacheType:
+    return CacheType.type(cache_type=type_hint)
+
+
+def default_options(cache_type: CacheType = None) -> Dict:
+    if cache_type is None:
+        return CacheType.default_options(CacheType.DBCache)
+    return CacheType.default_options(cache_type)
+
+
+def block_range(start: int, end: int, step: int = 1) -> List[int]:
+    return CacheType.block_range(start, end, step)
+
+
+def enable_cache(
+    pipe: DiffusionPipeline,
+    *,
+    transformer: torch.nn.Module = None,
+    blocks: torch.nn.ModuleList = None,
+    # transformer_blocks, blocks, etc.
+    blocks_name: str = None,
+    dummy_blocks_names: list[str] = [],
+    return_hidden_states_first: bool = True,
+    return_hidden_states_only: bool = False,
+    **cache_options_kwargs,
+) -> DiffusionPipeline:
+    return UnifiedCacheAdapter.apply(
+        pipe,
+        transformer=transformer,
+        blocks=blocks,
+        blocks_name=blocks_name,
+        dummy_blocks_names=dummy_blocks_names,
+        return_hidden_states_first=return_hidden_states_first,
+        return_hidden_states_only=return_hidden_states_only,
+        **cache_options_kwargs,
+    )