cache-dit 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

cache_dit/__init__.py

@@ -7,12 +7,21 @@ except ImportError:
 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.cache_factory import ForwardPattern
+from cache_dit.cache_factory import BlockAdapterParams
 from cache_dit.compile import set_compile_configs
 from cache_dit.utils import summary
+from cache_dit.utils import strify
 from cache_dit.logger import init_logger
 
 NONE = CacheType.NONE
 DBCache = CacheType.DBCache
+
+BlockAdapter = BlockAdapterParams
+
+Forward_Pattern_0 = ForwardPattern.Pattern_0
+Forward_Pattern_1 = ForwardPattern.Pattern_1
+Forward_Pattern_2 = ForwardPattern.Pattern_2
+Forward_Pattern_3 = ForwardPattern.Pattern_3

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.20'
-__version_tuple__ = version_tuple = (0, 2, 20)
+__version__ = version = '0.2.22'
+__version_tuple__ = version_tuple = (0, 2, 22)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -1,62 +1,8 @@
-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.forward_pattern import ForwardPattern
+from cache_dit.cache_factory.cache_types import CacheType
+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.cache_adapters import BlockAdapterParams
 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):
-    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 = CacheType.DBCache,
-) -> Dict:
-    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,
-    )
+from cache_dit.cache_factory.cache_interface import enable_cache
+from cache_dit.cache_factory.utils import load_options