cache-dit 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

cache_dit/__init__.py

@@ -18,6 +18,7 @@ from cache_dit.cache_factory import BlockAdapter
 from cache_dit.cache_factory import ParamsModifier
 from cache_dit.cache_factory import ForwardPattern
 from cache_dit.cache_factory import PatchFunctor
+from cache_dit.cache_factory import BasicCacheConfig
 from cache_dit.cache_factory import CalibratorConfig
 from cache_dit.cache_factory import TaylorSeerCalibratorConfig
 from cache_dit.cache_factory import FoCaCalibratorConfig

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.1'
-__version_tuple__ = version_tuple = (0, 3, 1)
+__version__ = version = '0.3.3'
+__version_tuple__ = version_tuple = (0, 3, 3)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/__init__.py

@@ -3,25 +3,22 @@ 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.forward_pattern import ForwardPattern
-
+from cache_dit.cache_factory.params_modifier import ParamsModifier
 from cache_dit.cache_factory.patch_functors import PatchFunctor
 
 from cache_dit.cache_factory.block_adapters import BlockAdapter
-from cache_dit.cache_factory.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 
 from cache_dit.cache_factory.cache_contexts import CachedContext
+from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
 from cache_dit.cache_factory.cache_contexts import CachedContextManager
-from cache_dit.cache_factory.cache_contexts import CachedContextV2
-from cache_dit.cache_factory.cache_contexts import CachedContextManagerV2
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig  # no V1
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig
 from cache_dit.cache_factory.cache_contexts import TaylorSeerCalibratorConfig
 from cache_dit.cache_factory.cache_contexts import FoCaCalibratorConfig
 
 from cache_dit.cache_factory.cache_blocks import CachedBlocks
 
 from cache_dit.cache_factory.cache_adapters import CachedAdapter
-from cache_dit.cache_factory.cache_adapters import CachedAdapterV2
 
 from cache_dit.cache_factory.cache_interface import enable_cache
 from cache_dit.cache_factory.cache_interface import disable_cache

cache_dit/cache_factory/block_adapters/block_adapters.py

@@ -7,73 +7,15 @@ from collections.abc import Iterable
 from typing import Any, Tuple, List, Optional, Union
 
 from diffusers import DiffusionPipeline
-from cache_dit.cache_factory.forward_pattern import ForwardPattern
 from cache_dit.cache_factory.patch_functors import PatchFunctor
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+from cache_dit.cache_factory.forward_pattern import ForwardPattern
+from cache_dit.cache_factory.params_modifier import ParamsModifier
 
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-class ParamsModifier:
-    def __init__(
-        self,
-        # Cache context kwargs
-        Fn_compute_blocks: Optional[int] = None,
-        Bn_compute_blocks: Optional[int] = None,
-        max_warmup_steps: Optional[int] = None,
-        max_cached_steps: Optional[int] = None,
-        max_continuous_cached_steps: Optional[int] = None,
-        residual_diff_threshold: Optional[float] = None,
-        # Cache CFG or not
-        enable_separate_cfg: Optional[bool] = None,
-        cfg_compute_first: Optional[bool] = None,
-        cfg_diff_compute_separate: Optional[bool] = None,
-        # Hybird TaylorSeer
-        enable_taylorseer: Optional[bool] = None,
-        enable_encoder_taylorseer: Optional[bool] = None,
-        taylorseer_cache_type: Optional[str] = None,
-        taylorseer_order: Optional[int] = None,
-        # New param only for v2 API
-        calibrator_config: Optional[CalibratorConfig] = None,
-        **other_cache_context_kwargs,
-    ):
-        self._context_kwargs = other_cache_context_kwargs.copy()
-        self._maybe_update_param("Fn_compute_blocks", Fn_compute_blocks)
-        self._maybe_update_param("Bn_compute_blocks", Bn_compute_blocks)
-        self._maybe_update_param("max_warmup_steps", max_warmup_steps)
-        self._maybe_update_param("max_cached_steps", max_cached_steps)
-        self._maybe_update_param(
-            "max_continuous_cached_steps", max_continuous_cached_steps
-        )
-        self._maybe_update_param(
-            "residual_diff_threshold", residual_diff_threshold
-        )
-        self._maybe_update_param("enable_separate_cfg", enable_separate_cfg)
-        self._maybe_update_param("cfg_compute_first", cfg_compute_first)
-        self._maybe_update_param(
-            "cfg_diff_compute_separate", cfg_diff_compute_separate
-        )
-        # V1 only supports the Taylorseer calibrator. We have decided to
-        # keep this code for API compatibility reasons.
-        if calibrator_config is None:
-            self._maybe_update_param("enable_taylorseer", enable_taylorseer)
-            self._maybe_update_param(
-                "enable_encoder_taylorseer", enable_encoder_taylorseer
-            )
-            self._maybe_update_param(
-                "taylorseer_cache_type", taylorseer_cache_type
-            )
-            self._maybe_update_param("taylorseer_order", taylorseer_order)
-        else:
-            self._maybe_update_param("calibrator_config", calibrator_config)
-
-    def _maybe_update_param(self, key: str, value: Any):
-        if value is not None:
-            self._context_kwargs[key] = value
-
-
 @dataclasses.dataclass
 class BlockAdapter:
 
@@ -123,10 +65,12 @@ class BlockAdapter:
     ] = None
 
     # modify cache context params for specific blocks.
-    params_modifiers: Union[
-        ParamsModifier,
-        List[ParamsModifier],
-        List[List[ParamsModifier]],
+    params_modifiers: Optional[
+        Union[
+            ParamsModifier,
+            List[ParamsModifier],
+            List[List[ParamsModifier]],
+        ]
     ] = None
 
     check_forward_pattern: bool = True
@@ -169,6 +113,19 @@ class BlockAdapter:
         if any((self.pipe is not None, self.transformer is not None)):
             self.maybe_fill_attrs()
             self.maybe_patchify()
+            self.maybe_skip_checks()
+
+    def maybe_skip_checks(self):
+        if getattr(self.transformer, "_hf_hook", None) is not None:
+            logger.warning("_hf_hook is not None, force skip pattern check!")
+            self.check_forward_pattern = False
+            self.check_num_outputs = False
+        elif getattr(self.transformer, "_diffusers_hook", None) is not None:
+            logger.warning(
+                "_diffusers_hook is not None, force skip pattern check!"
+            )
+            self.check_forward_pattern = False
+            self.check_num_outputs = False
 
     def maybe_fill_attrs(self):
         # NOTE: This func should be call before normalize.

cache_dit/cache_factory/cache_adapters/__init__.py

@@ -1,2 +1 @@
 from cache_dit.cache_factory.cache_adapters.cache_adapter import CachedAdapter
-from cache_dit.cache_factory.cache_adapters.v2 import CachedAdapterV2

cache_dit/cache_factory/cache_adapters/cache_adapter.py

@@ -1,10 +1,8 @@
 import torch
-
 import unittest
 import functools
-
 from contextlib import ExitStack
-from typing import Dict, List, Tuple, Any, Union, Callable
+from typing import Dict, List, Tuple, Any, Union, Callable, Optional
 
 from diffusers import DiffusionPipeline
 
@@ -13,8 +11,10 @@ from cache_dit.cache_factory.block_adapters import BlockAdapter
 from cache_dit.cache_factory.block_adapters import ParamsModifier
 from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
 from cache_dit.cache_factory.cache_contexts import CachedContextManager
+from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+from cache_dit.cache_factory.cache_contexts import CalibratorConfig
 from cache_dit.cache_factory.cache_blocks import CachedBlocks
-from cache_dit.cache_factory.cache_blocks.utils import (
+from cache_dit.cache_factory.cache_blocks import (
     patch_cached_stats,
     remove_cached_stats,
 )
@@ -55,6 +55,12 @@ class CachedAdapter:
                 block_adapter = BlockAdapterRegistry.get_adapter(
                     pipe_or_adapter
                 )
+                if params_modifiers := cache_context_kwargs.pop(
+                    "params_modifiers",
+                    None,
+                ):
+                    block_adapter.params_modifiers = params_modifiers
+
                 return cls.cachify(
                     block_adapter,
                     **cache_context_kwargs,
@@ -69,6 +75,12 @@ class CachedAdapter:
             logger.info(
                 "Adapting Cache Acceleration using custom BlockAdapter!"
             )
+            if pipe_or_adapter.params_modifiers is None:
+                if params_modifiers := cache_context_kwargs.pop(
+                    "params_modifiers", None
+                ):
+                    pipe_or_adapter.params_modifiers = params_modifiers
+
             return cls.cachify(
                 pipe_or_adapter,
                 **cache_context_kwargs,
@@ -114,33 +126,36 @@ class CachedAdapter:
         **cache_context_kwargs,
     ):
         # Check cache_context_kwargs
-        if cache_context_kwargs["enable_separate_cfg"] is None:
+        cache_config: BasicCacheConfig = cache_context_kwargs[
+            "cache_config"
+        ]  # ref
+        assert cache_config is not None, "cache_config can not be None."
+        if cache_config.enable_separate_cfg is None:
             # Check cfg for some specific case if users don't set it as True
             if BlockAdapterRegistry.has_separate_cfg(block_adapter):
-                cache_context_kwargs["enable_separate_cfg"] = True
+                cache_config.enable_separate_cfg = True
                 logger.info(
                     f"Use custom 'enable_separate_cfg' from BlockAdapter: True. "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
             else:
-                cache_context_kwargs["enable_separate_cfg"] = (
+                cache_config.enable_separate_cfg = (
                     BlockAdapterRegistry.has_separate_cfg(block_adapter.pipe)
                 )
                 logger.info(
                     f"Use default 'enable_separate_cfg' from block adapter "
-                    f"register: {cache_context_kwargs['enable_separate_cfg']}, "
+                    f"register: {cache_config.enable_separate_cfg}, "
                     f"Pipeline: {block_adapter.pipe.__class__.__name__}."
                 )
         else:
             logger.info(
                 f"Use custom 'enable_separate_cfg' from cache context "
-                f"kwargs: {cache_context_kwargs['enable_separate_cfg']}. "
+                f"kwargs: {cache_config.enable_separate_cfg}. "
                 f"Pipeline: {block_adapter.pipe.__class__.__name__}."
             )
 
-        if (
-            cache_type := cache_context_kwargs.pop("cache_type", None)
-        ) is not None:
+        cache_type = cache_context_kwargs.pop("cache_type", None)
+        if cache_type is not None:
             assert (
                 cache_type == CacheType.DBCache
             ), "Custom cache setting only support for DBCache now!"
@@ -176,7 +191,7 @@ class CachedAdapter:
         block_adapter.pipe._cache_manager = cache_manager  # instance level
 
         flatten_contexts, contexts_kwargs = cls.modify_context_params(
-            block_adapter, cache_manager, **cache_context_kwargs
+            block_adapter, **cache_context_kwargs
         )
 
         original_call = block_adapter.pipe.__class__.__call__
@@ -212,7 +227,6 @@ class CachedAdapter:
     def modify_context_params(
         cls,
         block_adapter: BlockAdapter,
-        cache_manager: CachedContextManager,
         **cache_context_kwargs,
     ) -> Tuple[List[str], List[Dict[str, Any]]]:
 
@@ -230,6 +244,8 @@ class CachedAdapter:
             contexts_kwargs[i]["name"] = flatten_contexts[i]
 
         if block_adapter.params_modifiers is None:
+            for i in range(len(contexts_kwargs)):
+                cls._config_messages(**contexts_kwargs[i])
             return flatten_contexts, contexts_kwargs
 
         flatten_modifiers: List[ParamsModifier] = BlockAdapter.flatten(
@@ -242,12 +258,26 @@ class CachedAdapter:
             contexts_kwargs[i].update(
                 flatten_modifiers[i]._context_kwargs,
             )
-            contexts_kwargs[i], _ = cache_manager.collect_cache_kwargs(
-                default_attrs={}, **contexts_kwargs[i]
-            )
+            cls._config_messages(**contexts_kwargs[i])
 
         return flatten_contexts, contexts_kwargs
 
+    @classmethod
+    def _config_messages(cls, **contexts_kwargs):
+        cache_config: BasicCacheConfig = contexts_kwargs.get(
+            "cache_config", None
+        )
+        calibrator_config: CalibratorConfig = contexts_kwargs.get(
+            "calibrator_config", None
+        )
+        if cache_config is not None:
+            message = f"Collected Cache Config: {cache_config.strify()}"
+            if calibrator_config is not None:
+                message += f", Calibrator Config: {calibrator_config.strify(details=True)}"
+            else:
+                message += ", Calibrator Config: None"
+            logger.info(message)
+
     @classmethod
     def mock_blocks(
         cls,
@@ -298,7 +328,19 @@ class CachedAdapter:
 
         assert isinstance(dummy_blocks_names, list)
 
-        @functools.wraps(original_forward)
+        from accelerate import hooks
+
+        _hf_hook: Optional[hooks.ModelHook] = None
+
+        if getattr(transformer, "_hf_hook", None) is not None:
+            _hf_hook = transformer._hf_hook  # hooks from accelerate.hooks
+
+        # TODO: remove group offload hooks the re-apply after cache applied.
+        # hooks = _diffusers_hook.hooks.copy(); _diffusers_hook.hooks.clear()
+        # re-apply hooks to transformer after cache applied.
+        # from diffusers.hooks.hooks import HookFunctionReference, HookRegistry
+        # from diffusers.hooks.group_offloading import apply_group_offloading
+
         def new_forward(self, *args, **kwargs):
             with ExitStack() as stack:
                 for name, context_name in zip(
@@ -316,9 +358,27 @@ class CachedAdapter:
                             self, dummy_name, dummy_blocks
                         )
                     )
-                return original_forward(*args, **kwargs)
+                outputs = original_forward(*args, **kwargs)
+            return outputs
+
+        def new_forward_with_hf_hook(self, *args, **kwargs):
+            # Compatible with model cpu offload
+            if _hf_hook is not None and hasattr(_hf_hook, "pre_forward"):
+                args, kwargs = _hf_hook.pre_forward(self, *args, **kwargs)
+
+            outputs = new_forward(self, *args, **kwargs)
+
+            if _hf_hook is not None and hasattr(_hf_hook, "post_forward"):
+                outputs = _hf_hook.post_forward(self, outputs)
+
+            return outputs
+
+        # NOTE: Still can't fully compatible with group offloading
+        transformer.forward = functools.update_wrapper(
+            functools.partial(new_forward_with_hf_hook, transformer),
+            new_forward_with_hf_hook,
+        )
 
-        transformer.forward = new_forward.__get__(transformer)
         transformer._original_forward = original_forward
         transformer._is_cached = True
 
@@ -335,7 +395,8 @@ class CachedAdapter:
         total_cached_blocks: List[Dict[str, torch.nn.ModuleList]] = []
         assert hasattr(block_adapter.pipe, "_cache_manager")
         assert isinstance(
-            block_adapter.pipe._cache_manager, CachedContextManager
+            block_adapter.pipe._cache_manager,
+            CachedContextManager,
         )
 
         for i in range(len(block_adapter.transformer)):

cache_dit/cache_factory/cache_blocks/__init__.py

@@ -12,6 +12,10 @@ from cache_dit.cache_factory.cache_blocks.pattern_0_1_2 import (
 from cache_dit.cache_factory.cache_blocks.pattern_3_4_5 import (
     CachedBlocks_Pattern_3_4_5,
 )
+from cache_dit.cache_factory.cache_blocks.pattern_utils import (
+    patch_cached_stats,
+    remove_cached_stats,
+)
 
 from cache_dit.logger import init_logger
 

cache_dit/cache_factory/cache_blocks/offload_utils.py

@@ -0,0 +1,115 @@
+import torch
+import asyncio
+import logging
+from contextlib import contextmanager
+from typing import Generator, Optional, List
+from diffusers.hooks.group_offloading import _is_group_offload_enabled
+from cache_dit.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@torch.compiler.disable
+@contextmanager
+def maybe_onload(
+    block: torch.nn.Module,
+    reference_tensor: torch.Tensor,
+    pending_tasks: List[asyncio.Task] = [],
+) -> Generator:
+
+    if not _is_group_offload_enabled(block):
+        yield block
+        return
+
+    original_devices: Optional[List[torch.device]] = None
+    if hasattr(block, "parameters"):
+        params = list(block.parameters())
+        if params:
+            original_devices = [param.data.device for param in params]
+
+    target_device: torch.device = reference_tensor.device
+    move_task: Optional[asyncio.Task] = None
+    need_restore: bool = False
+
+    try:
+        if original_devices is not None:
+            unique_devices = list(set(original_devices))
+            if len(unique_devices) > 1 or unique_devices[0] != target_device:
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        f"Onloading from {unique_devices} to {target_device}"
+                    )
+
+                has_meta_params = any(
+                    dev.type == "meta" for dev in original_devices
+                )
+                if has_meta_params:  # compatible with sequential cpu offload
+                    block = block.to_empty(device=target_device)
+                else:
+                    block = block.to(target_device, non_blocking=False)
+                need_restore = True
+        yield block
+    finally:
+        if need_restore and original_devices:
+
+            async def restore_device():
+                for param, original_device in zip(
+                    block.parameters(), original_devices
+                ):
+                    param.data = await asyncio.to_thread(
+                        lambda p, d: p.to(d, non_blocking=True),
+                        param.data,  # type: torch.Tensor
+                        original_device,  # type: torch.device
+                    )  # type: ignore[assignment]
+
+            loop = get_event_loop()
+            move_task = loop.create_task(restore_device())
+            if move_task:
+                pending_tasks.append(move_task)
+
+
+def get_event_loop() -> asyncio.AbstractEventLoop:
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+    if not loop.is_running():
+
+        def run_loop() -> None:
+            asyncio.set_event_loop(loop)
+            loop.run_forever()
+
+        import threading
+
+        if not any(t.name == "_my_loop" for t in threading.enumerate()):
+            threading.Thread(
+                target=run_loop, name="_my_loop", daemon=True
+            ).start()
+
+    return loop
+
+
+@torch.compiler.disable
+def maybe_offload(
+    pending_tasks: List[asyncio.Task],
+) -> None:
+    if not pending_tasks:
+        return
+
+    loop = get_event_loop()
+
+    async def gather_tasks():
+        return await asyncio.gather(*pending_tasks)
+
+    future = asyncio.run_coroutine_threadsafe(gather_tasks(), loop)
+    try:
+        future.result(timeout=30.0)
+    except Exception as e:
+        logger.error(f"May Offload Error: {e}")
+
+    pending_tasks.clear()

cache_dit/cache_factory/cache_blocks/pattern_base.py

@@ -1,7 +1,9 @@
 import inspect
+import asyncio
 import torch
 import torch.distributed as dist
 
+from typing import List
 from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
 from cache_dit.cache_factory.cache_contexts.cache_manager import (
     CachedContextManager,
@@ -45,6 +47,7 @@ class CachedBlocks_Pattern_Base(torch.nn.Module):
         self.cache_prefix = cache_prefix
         self.cache_context = cache_context
         self.cache_manager = cache_manager
+        self.pending_tasks: List[asyncio.Task] = []
 
         self._check_forward_pattern()
         logger.info(

cache_dit/cache_factory/cache_contexts/__init__.py

@@ -1,12 +1,14 @@
-# namespace alias: for _CachedContext and many others' cache context funcs.
-from cache_dit.cache_factory.cache_contexts.cache_context import CachedContext
-from cache_dit.cache_factory.cache_contexts.cache_manager import (
-    CachedContextManager,
-)
-from cache_dit.cache_factory.cache_contexts.v2 import (
-    CachedContextV2,
-    CachedContextManagerV2,
+from cache_dit.cache_factory.cache_contexts.calibrators import (
+    Calibrator,
+    CalibratorBase,
     CalibratorConfig,
     TaylorSeerCalibratorConfig,
     FoCaCalibratorConfig,
 )
+from cache_dit.cache_factory.cache_contexts.cache_context import (
+    CachedContext,
+    BasicCacheConfig,
+)
+from cache_dit.cache_factory.cache_contexts.cache_manager import (
+    CachedContextManager,
+)