cache-dit 0.2.26__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
     }
@@ -1057,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
 
@@ -1104,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,23 +1,18 @@
 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
 
 logger = init_logger(__name__)
 
 
-def supported_pipelines() -> Tuple[int, List[str]]:
-    return UnifiedCacheAdapter.supported_pipelines()
-
-
 def enable_cache(
-    # BlockAdapter & forward pattern
+    # DiffusionPipeline or BlockAdapter
     pipe_or_adapter: DiffusionPipeline | BlockAdapter | Any,
-    forward_pattern: ForwardPattern = ForwardPattern.Pattern_0,
     # Cache context kwargs
     Fn_compute_blocks: int = 8,
     Bn_compute_blocks: int = 0,
@@ -34,7 +29,7 @@ def enable_cache(
     enable_encoder_taylorseer: bool = False,
     taylorseer_cache_type: str = "residual",
     taylorseer_order: int = 2,
-    **other_cache_kwargs,
+    **other_cache_context_kwargs,
 ) -> DiffusionPipeline | Any:
     r"""
     Unified Cache API for  almost Any Diffusion Transformers (with Transformer Blocks
@@ -48,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
@@ -111,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
@@ -141,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/patch_functors/functor_chroma.py

@@ -30,6 +30,9 @@ class ChromaPatchFunctor(PatchFunctor):
         blocks: torch.nn.ModuleList = None,
         **kwargs,
     ) -> ChromaTransformer2DModel:
+        if getattr(transformer, "_is_patched", False):
+            return transformer
+
         if blocks is None:
             blocks = transformer.single_transformer_blocks
 

cache_dit/cache_factory/patch_functors/functor_flux.py

@@ -30,6 +30,10 @@ class FluxPatchFunctor(PatchFunctor):
         blocks: torch.nn.ModuleList = None,
         **kwargs,
     ) -> FluxTransformer2DModel:
+
+        if getattr(transformer, "_is_patched", False):
+            return transformer
+
         if blocks is None:
             blocks = transformer.single_transformer_blocks
 

cache_dit/quantize/quantize_ao.py

@@ -179,6 +179,7 @@ def quantize_ao(
     force_empty_cache()
 
     logger.info(
+        f"Quantized        Method: {quant_type:>5}\n"
         f"Quantized Linear Layers: {num_quant_linear:>5}\n"
         f"Skipped   Linear Layers: {num_skip_linear:>5}\n"
         f"Total     Linear Layers: {num_linear_layers:>5}\n"

cache_dit/utils.py

@@ -30,27 +30,32 @@ class CacheStats:
 
 
 def summary(
-    pipe_or_transformer: DiffusionPipeline | torch.nn.Module | Any,
+    pipe_or_module: DiffusionPipeline | torch.nn.Module | Any,
     details: bool = False,
     logging: bool = True,
 ) -> CacheStats:
     cache_stats = CacheStats()
-    cls_name = pipe_or_transformer.__class__.__name__
-    if not isinstance(pipe_or_transformer, torch.nn.Module):
-        assert hasattr(pipe_or_transformer, "transformer")
-        transformer = pipe_or_transformer.transformer
+
+    if not isinstance(pipe_or_module, torch.nn.Module):
+        assert hasattr(pipe_or_module, "transformer")
+        module = pipe_or_module.transformer
+        cls_name = module.__class__.__name__
     else:
-        transformer = pipe_or_transformer
+        module = pipe_or_module
+
+    cls_name = module.__class__.__name__
+    if isinstance(module, torch.nn.ModuleList):
+        cls_name = module[0].__class__.__name__
 
-    if hasattr(transformer, "_cache_context_kwargs"):
-        cache_options = transformer._cache_context_kwargs
+    if hasattr(module, "_cache_context_kwargs"):
+        cache_options = module._cache_context_kwargs
         cache_stats.cache_options = cache_options
         if logging:
             print(f"\n🤗Cache Options: {cls_name}\n\n{cache_options}")
 
-    if hasattr(transformer, "_cached_steps"):
-        cached_steps: list[int] = transformer._cached_steps
-        residual_diffs: dict[str, float] = dict(transformer._residual_diffs)
+    if hasattr(module, "_cached_steps"):
+        cached_steps: list[int] = module._cached_steps
+        residual_diffs: dict[str, float] = dict(module._residual_diffs)
         cache_stats.cached_steps = cached_steps
         cache_stats.residual_diffs = residual_diffs
 
@@ -91,11 +96,9 @@ def summary(
                     compact=True,
                 )
 
-    if hasattr(transformer, "_cfg_cached_steps"):
-        cfg_cached_steps: list[int] = transformer._cfg_cached_steps
-        cfg_residual_diffs: dict[str, float] = dict(
-            transformer._cfg_residual_diffs
-        )
+    if hasattr(module, "_cfg_cached_steps"):
+        cfg_cached_steps: list[int] = module._cfg_cached_steps
+        cfg_residual_diffs: dict[str, float] = dict(module._cfg_residual_diffs)
         cache_stats.cfg_cached_steps = cfg_cached_steps
         cache_stats.cfg_residual_diffs = cfg_residual_diffs