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

cache_dit/__init__.py

@@ -10,9 +10,18 @@ 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.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.21'
+__version_tuple__ = version_tuple = (0, 2, 21)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -1,7 +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_adapters import BlockAdapterParams
 from cache_dit.cache_factory.cache_adapters import UnifiedCacheAdapter
 from cache_dit.cache_factory.utils import load_cache_options_from_yaml
 
@@ -39,24 +40,26 @@ def block_range(
 
 
 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,
+    pipe_or_adapter: DiffusionPipeline | BlockAdapterParams,
+    forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
     **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,
-    )
+    if isinstance(pipe_or_adapter, BlockAdapterParams):
+        return UnifiedCacheAdapter.apply(
+            pipe=None,
+            adapter_params=pipe_or_adapter,
+            forward_pattern=forward_pattern,
+            **cache_options_kwargs,
+        )
+    elif isinstance(pipe_or_adapter, DiffusionPipeline):
+        return UnifiedCacheAdapter.apply(
+            pipe=pipe_or_adapter,
+            adapter_params=None,
+            forward_pattern=forward_pattern,
+            **cache_options_kwargs,
+        )
+    else:
+        raise ValueError(
+            "Please pass DiffusionPipeline or BlockAdapterParams"
+            "(BlockAdapter) for the 1 position param: pipe_or_adapter"
+        )