cache-dit 0.2.25__py3-none-any.whl → 0.2.27__py3-none-any.whl

cache_dit/cache_factory/{cache_context.py → cache_contexts/cache_context.py}

@@ -7,14 +7,15 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union, Tuple
 import torch
 import torch.distributed as dist
 
-from cache_dit.cache_factory.taylorseer import TaylorSeer
+from cache_dit.cache_factory.cache_contexts.taylorseer import TaylorSeer
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
 @dataclasses.dataclass
-class DBCacheContext:
+class _CachedContext:  # Internal CachedContext Impl class
+    name: str = "default"
     # Dual Block Cache
     # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
     Fn_compute_blocks: int = 1
@@ -99,6 +100,8 @@ class DBCacheContext:
 
     @torch.compiler.disable
     def __post_init__(self):
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.info(f"Created _CacheContext: {self.name}")
         # Some checks for settings
         if self.do_separate_cfg:
             assert self.enable_alter_cache is False, (
@@ -329,26 +332,60 @@ class DBCacheContext:
 
 
 # TODO: Support context manager for different cache_context
+_current_cache_context: _CachedContext = None
+
+_cache_context_manager: Dict[str, _CachedContext] = {}
 
 
 def create_cache_context(*args, **kwargs):
-    return DBCacheContext(*args, **kwargs)
+    global _cache_context_manager
+    _context = _CachedContext(*args, **kwargs)
+    _cache_context_manager[_context.name] = _context
+    return _context
 
 
-def get_current_cache_context():
+def get_cache_context():
     return _current_cache_context
 
 
-def set_current_cache_context(cache_context=None):
-    global _current_cache_context
-    _current_cache_context = cache_context
+def set_cache_context(cache_context: _CachedContext | str):
+    global _current_cache_context, _cache_context_manager
+    if isinstance(cache_context, _CachedContext):
+        _current_cache_context = cache_context
+    else:
+        _current_cache_context = _cache_context_manager[cache_context]
+
+
+def reset_cache_context(cache_context: _CachedContext | str, *args, **kwargs):
+    global _cache_context_manager
+    if isinstance(cache_context, _CachedContext):
+        old_context_name = cache_context.name
+        if cache_context.name in _cache_context_manager:
+            del _cache_context_manager[cache_context.name]
+        # force use old_context name
+        kwargs["name"] = old_context_name
+        _context = _CachedContext(*args, **kwargs)
+        _cache_context_manager[_context.name] = _context
+    else:
+        old_context_name = cache_context
+        if cache_context in _cache_context_manager:
+            del _cache_context_manager[cache_context]
+        # force use old_context name
+        kwargs["name"] = old_context_name
+        _context = _CachedContext(*args, **kwargs)
+        _cache_context_manager[_context.name] = _context
+
+    return _context
 
 
 @contextlib.contextmanager
-def cache_context(cache_context):
-    global _current_cache_context
+def cache_context(cache_context: _CachedContext | str):
+    global _current_cache_context, _cache_context_manager
     old_cache_context = _current_cache_context
-    _current_cache_context = cache_context
+    if isinstance(cache_context, _CachedContext):
+        _current_cache_context = cache_context
+    else:
+        _current_cache_context = _cache_context_manager[cache_context]
     try:
         yield
     finally:
@@ -357,49 +394,49 @@ def cache_context(cache_context):
 
 @torch.compiler.disable
 def get_residual_diff_threshold():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_residual_diff_threshold()
 
 
 @torch.compiler.disable
 def get_buffer(name):
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_buffer(name)
 
 
 @torch.compiler.disable
 def set_buffer(name, buffer):
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     cache_context.set_buffer(name, buffer)
 
 
 @torch.compiler.disable
 def remove_buffer(name):
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     cache_context.remove_buffer(name)
 
 
 @torch.compiler.disable
 def mark_step_begin():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     cache_context.mark_step_begin()
 
 
 @torch.compiler.disable
 def get_current_step():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_current_step()
 
 
 @torch.compiler.disable
 def get_current_step_residual_diff():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     step = str(get_current_step())
     residual_diffs = get_residual_diffs()
@@ -410,7 +447,7 @@ def get_current_step_residual_diff():
 
 @torch.compiler.disable
 def get_current_step_cfg_residual_diff():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     step = str(get_current_step())
     cfg_residual_diffs = get_cfg_residual_diffs()
@@ -421,110 +458,110 @@ def get_current_step_cfg_residual_diff():
 
 @torch.compiler.disable
 def get_current_transformer_step():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_current_transformer_step()
 
 
 @torch.compiler.disable
 def get_cached_steps():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_cached_steps()
 
 
 @torch.compiler.disable
 def get_cfg_cached_steps():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_cfg_cached_steps()
 
 
 @torch.compiler.disable
 def get_max_cached_steps():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.max_cached_steps
 
 
 @torch.compiler.disable
 def get_max_continuous_cached_steps():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.max_continuous_cached_steps
 
 
 @torch.compiler.disable
 def get_continuous_cached_steps():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.continuous_cached_steps
 
 
 @torch.compiler.disable
 def get_cfg_continuous_cached_steps():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.cfg_continuous_cached_steps
 
 
 @torch.compiler.disable
 def add_cached_step():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     cache_context.add_cached_step()
 
 
 @torch.compiler.disable
 def add_residual_diff(diff):
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     cache_context.add_residual_diff(diff)
 
 
 @torch.compiler.disable
 def get_residual_diffs():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_residual_diffs()
 
 
 @torch.compiler.disable
 def get_cfg_residual_diffs():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_cfg_residual_diffs()
 
 
 @torch.compiler.disable
 def is_taylorseer_enabled():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.enable_taylorseer
 
 
 @torch.compiler.disable
 def is_encoder_taylorseer_enabled():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.enable_encoder_taylorseer
 
 
 def get_taylorseers() -> Tuple[TaylorSeer, TaylorSeer]:
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_taylorseers()
 
 
 def get_cfg_taylorseers() -> Tuple[TaylorSeer, TaylorSeer]:
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.get_cfg_taylorseers()
 
 
 @torch.compiler.disable
 def is_taylorseer_cache_residual():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.taylorseer_cache_type == "residual"
 
@@ -547,28 +584,28 @@ def is_encoder_cache_residual():
 
 @torch.compiler.disable
 def is_alter_cache_enabled():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.enable_alter_cache
 
 
 @torch.compiler.disable
 def is_alter_cache():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.is_alter_cache
 
 
 @torch.compiler.disable
 def is_in_warmup():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.is_in_warmup()
 
 
 @torch.compiler.disable
 def is_l1_diff_enabled():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return (
         cache_context.l1_hidden_states_diff_threshold is not None
@@ -578,21 +615,21 @@ def is_l1_diff_enabled():
 
 @torch.compiler.disable
 def get_important_condition_threshold():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.important_condition_threshold
 
 
 @torch.compiler.disable
 def non_compute_blocks_diff_threshold():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.non_compute_blocks_diff_threshold
 
 
 @torch.compiler.disable
 def Fn_compute_blocks():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     assert (
         cache_context.Fn_compute_blocks >= 1
@@ -612,7 +649,7 @@ def Fn_compute_blocks():
 
 @torch.compiler.disable
 def Fn_compute_blocks_ids():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     assert (
         len(cache_context.Fn_compute_blocks_ids)
@@ -627,7 +664,7 @@ def Fn_compute_blocks_ids():
 
 @torch.compiler.disable
 def Bn_compute_blocks():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     assert (
         cache_context.Bn_compute_blocks >= 0
@@ -647,7 +684,7 @@ def Bn_compute_blocks():
 
 @torch.compiler.disable
 def Bn_compute_blocks_ids():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     assert (
         len(cache_context.Bn_compute_blocks_ids)
@@ -662,44 +699,41 @@ def Bn_compute_blocks_ids():
 
 @torch.compiler.disable
 def do_separate_cfg():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.do_separate_cfg
 
 
 @torch.compiler.disable
 def is_separate_cfg_step():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.is_separate_cfg_step()
 
 
 @torch.compiler.disable
 def cfg_diff_compute_separate():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.cfg_diff_compute_separate
 
 
-_current_cache_context: DBCacheContext = None
-
-
 def collect_cache_kwargs(default_attrs: dict, **kwargs):
     # NOTE: This API will split kwargs into cache_kwargs and other_kwargs
     # default_attrs: specific settings for different pipelines
-    cache_attrs = dataclasses.fields(DBCacheContext)
+    cache_attrs = dataclasses.fields(_CachedContext)
     cache_attrs = [
         attr
         for attr in cache_attrs
         if hasattr(
-            DBCacheContext,
+            _CachedContext,
             attr.name,
         )
     ]
     cache_kwargs = {
         attr.name: kwargs.pop(
             attr.name,
-            getattr(DBCacheContext, attr.name),
+            getattr(_CachedContext, attr.name),
         )
         for attr in cache_attrs
     }
@@ -941,7 +975,11 @@ def get_Bn_buffer(prefix: str = "Bn"):
 
 
 @torch.compiler.disable
-def set_Bn_encoder_buffer(buffer: torch.Tensor, prefix: str = "Bn"):
+def set_Bn_encoder_buffer(buffer: torch.Tensor | None, prefix: str = "Bn"):
+    # DON'T set None Buffer
+    if buffer is None:
+        return
+
     # This buffer is use for encoder hidden states approximation.
     if is_encoder_taylorseer_enabled():
         # taylorseer, encoder_taylorseer
@@ -1053,7 +1091,7 @@ def apply_hidden_states_residual(
 
 @torch.compiler.disable
 def get_downsample_factor():
-    cache_context = get_current_cache_context()
+    cache_context = get_cache_context()
     assert cache_context is not None, "cache_context must be set before"
     return cache_context.downsample_factor
 
@@ -1100,7 +1138,7 @@ def get_can_use_cache(
                 "can not use cache."
             )
         # reset continuous cached steps stats
-        cache_context = get_current_cache_context()
+        cache_context = get_cache_context()
         if not is_separate_cfg_step():
             cache_context.continuous_cached_steps = 0
         else:

cache_dit/cache_factory/cache_interface.py

@@ -1,8 +1,9 @@
+from typing import Any, Tuple, List
 from diffusers import DiffusionPipeline
-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 BlockAdapter
-from cache_dit.cache_factory.cache_adapters import UnifiedCacheAdapter
+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.logger import init_logger
 
@@ -10,9 +11,8 @@ logger = init_logger(__name__)
 
 
 def enable_cache(
-    # BlockAdapter & forward pattern
-    pipe_or_adapter: DiffusionPipeline | BlockAdapter,
-    forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
+    # DiffusionPipeline or BlockAdapter
+    pipe_or_adapter: DiffusionPipeline | BlockAdapter | Any,
     # Cache context kwargs
     Fn_compute_blocks: int = 8,
     Bn_compute_blocks: int = 0,
@@ -29,8 +29,8 @@ def enable_cache(
     enable_encoder_taylorseer: bool = False,
     taylorseer_cache_type: str = "residual",
     taylorseer_order: int = 2,
-    **other_cache_kwargs,
-) -> DiffusionPipeline:
+    **other_cache_context_kwargs,
+) -> DiffusionPipeline | Any:
     r"""
     Unified Cache API for  almost Any Diffusion Transformers (with Transformer Blocks
     that match the specific Input and Output patterns).
@@ -43,9 +43,6 @@ def enable_cache(
             The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined).
             For example: cache_dit.enable_cache(FluxPipeline(...)). Please check https://github.com/vipshop/cache-dit/blob/main/docs/BlockAdapter.md
             for the usgae of BlockAdapter.
-        forward_pattern (`ForwardPattern`, *required*, defaults to `ForwardPattern.Pattern_0`):
-            The forward pattern of Transformer block, please check https://github.com/vipshop/cache-dit/tree/main?tab=readme-ov-file#forward-pattern-matching
-            for more details.
         Fn_compute_blocks (`int`, *required*, defaults to 8):
             Specifies that `DBCache` uses the **first n** Transformer blocks to fit the information
             at time step t, enabling the calculation of a more stable L1 diff and delivering more
@@ -106,7 +103,7 @@ def enable_cache(
     """
 
     # Collect cache context kwargs
-    cache_context_kwargs = other_cache_kwargs.copy()
+    cache_context_kwargs = other_cache_context_kwargs.copy()
     cache_context_kwargs["cache_type"] = CacheType.DBCache
     cache_context_kwargs["Fn_compute_blocks"] = Fn_compute_blocks
     cache_context_kwargs["Bn_compute_blocks"] = Bn_compute_blocks
@@ -136,21 +133,32 @@ def enable_cache(
         }
 
     if isinstance(pipe_or_adapter, BlockAdapter):
-        return UnifiedCacheAdapter.apply(
+        return CachedAdapter.apply(
             pipe=None,
             block_adapter=pipe_or_adapter,
-            forward_pattern=forward_pattern,
             **cache_context_kwargs,
         )
     elif isinstance(pipe_or_adapter, DiffusionPipeline):
-        return UnifiedCacheAdapter.apply(
+        return CachedAdapter.apply(
             pipe=pipe_or_adapter,
             block_adapter=None,
-            forward_pattern=forward_pattern,
             **cache_context_kwargs,
         )
     else:
         raise ValueError(
+            f"type: {type(pipe_or_adapter)} is not valid, "
             "Please pass DiffusionPipeline or BlockAdapter"
             "for the 1's position param: pipe_or_adapter"
         )
+
+
+def supported_pipelines(
+    **kwargs,
+) -> Tuple[int, List[str]]:
+    return BlockAdapterRegistry.supported_pipelines(**kwargs)
+
+
+def get_adapter(
+    pipe: DiffusionPipeline | str | Any,
+) -> BlockAdapter:
+    return BlockAdapterRegistry.get_adapter(pipe)

cache_dit/cache_factory/forward_pattern.py

@@ -19,39 +19,57 @@ class ForwardPattern(Enum):
         self.Supported = Supported
 
     Pattern_0 = (
-        True,
-        False,
-        False,
-        ("hidden_states", "encoder_hidden_states"),
-        ("hidden_states", "encoder_hidden_states"),
-        True,
+        True,  # Return_H_First
+        False, # Return_H_Only
+        False, # Forward_H_only
+        ("hidden_states", "encoder_hidden_states"),  # In
+        ("hidden_states", "encoder_hidden_states"),  # Out
+        True,  # Supported
     )
 
     Pattern_1 = (
-        False,
-        False,
-        False,
-        ("hidden_states", "encoder_hidden_states"),
-        ("encoder_hidden_states", "hidden_states"),
-        True,
+        False, # Return_H_First
+        False, # Return_H_Only
+        False, # Forward_H_only
+        ("hidden_states", "encoder_hidden_states"),  # In
+        ("encoder_hidden_states", "hidden_states"),  # Out
+        True,  # Supported
     )
 
     Pattern_2 = (
-        False,
-        True,
-        False,
-        ("hidden_states", "encoder_hidden_states"),
-        ("hidden_states",),
-        True,
+        False, # Return_H_First
+        True,  # Return_H_Only
+        False, # Forward_H_only
+        ("hidden_states", "encoder_hidden_states"),  # In
+        ("hidden_states",),                          # Out
+        True,  # Supported
     )
 
     Pattern_3 = (
-        False,
-        True,
-        False,
-        ("hidden_states",),
-        ("hidden_states",),
-        False,
+        False, # Return_H_First
+        True,  # Return_H_Only
+        True,  # Forward_H_only
+        ("hidden_states",),  # In
+        ("hidden_states",),  # Out
+        True,  # Supported
+    )
+
+    Pattern_4 = (
+        True,  # Return_H_First
+        False, # Return_H_Only
+        True,  # Forward_H_only
+        ("hidden_states",),                          # In
+        ("hidden_states", "encoder_hidden_states"),  # Out
+        True,  # Supported
+    )
+
+    Pattern_5 = (
+        False, # Return_H_First
+        False, # Return_H_Only
+        True,  # Forward_H_only
+        ("hidden_states",),                          # In
+        ("encoder_hidden_states", "hidden_states"),  # Out
+        True,  # Supported
     )
 
     @staticmethod
@@ -60,4 +78,7 @@ class ForwardPattern(Enum):
             ForwardPattern.Pattern_0,
             ForwardPattern.Pattern_1,
             ForwardPattern.Pattern_2,
+            ForwardPattern.Pattern_3,
+            ForwardPattern.Pattern_4,
+            ForwardPattern.Pattern_5,
         ]

cache_dit/cache_factory/patch_functors/__init__.py

@@ -0,0 +1,5 @@
+from cache_dit.cache_factory.patch_functors.functor_base import PatchFunctor
+from cache_dit.cache_factory.patch_functors.functor_flux import FluxPatchFunctor
+from cache_dit.cache_factory.patch_functors.functor_chroma import (
+    ChromaPatchFunctor,
+)

cache_dit/cache_factory/patch_functors/functor_base.py

@@ -0,0 +1,18 @@
+import torch
+from abc import abstractmethod
+
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class PatchFunctor:
+
+    @abstractmethod
+    def apply(
+        self,
+        transformer: torch.nn.Module,
+        *args,
+        **kwargs,
+    ) -> torch.nn.Module:
+        raise NotImplementedError("apply method is not implemented.")