cache-dit 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

cache_dit/__init__.py

@@ -0,0 +1,5 @@
+try:
+    from ._version import version as __version__, version_tuple
+except ImportError:
+    __version__ = "unknown version"
+    version_tuple = (0, 0, "unknown version")

cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.12'
-__version_tuple__ = version_tuple = (0, 2, 12)
+__version__ = version = '0.2.14'
+__version_tuple__ = version_tuple = (0, 2, 14)

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -8,8 +8,9 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union
 
 import torch
 
-import cache_dit.primitives as DP
+import cache_dit.primitives as primitives
 from cache_dit.cache_factory.taylorseer import TaylorSeer
+from cache_dit.utils import is_diffusers_at_least_0_3_5
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -772,8 +773,8 @@ def are_two_tensors_similar(
             mean_t1 = t1.abs().mean()
 
         if parallelized:
-            mean_diff = DP.all_reduce_sync(mean_diff, "avg")
-            mean_t1 = DP.all_reduce_sync(mean_t1, "avg")
+            mean_diff = primitives.all_reduce_sync(mean_diff, "avg")
+            mean_t1 = primitives.all_reduce_sync(mean_t1, "avg")
 
         # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
         # Futher, if we assume that (H(t,  0) - H(t-1,0)) ~ 0, then,
@@ -1365,6 +1366,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
     ):
         original_hidden_states = hidden_states
         original_encoder_hidden_states = encoder_hidden_states
+        # This condition branch is mainly for FLUX series.
         if self.single_transformer_blocks is not None:
             for block in self.transformer_blocks[Fn_compute_blocks() :]:
                 hidden_states = block(
@@ -1381,22 +1383,32 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                             hidden_states,
                         )
 
-            hidden_states = torch.cat(
-                [encoder_hidden_states, hidden_states], dim=1
-            )
-            for block in self._Mn_single_transformer_blocks():
-                hidden_states = block(
-                    hidden_states,
-                    *args,
-                    **kwargs,
+            # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L380
+            if is_diffusers_at_least_0_3_5():
+                for block in self._Mn_single_transformer_blocks():
+                    encoder_hidden_states, hidden_states = block(
+                        hidden_states,
+                        encoder_hidden_states,
+                        *args,
+                        **kwargs,
+                    )
+            else:
+                hidden_states = torch.cat(
+                    [encoder_hidden_states, hidden_states], dim=1
+                )
+                for block in self._Mn_single_transformer_blocks():
+                    hidden_states = block(
+                        hidden_states,
+                        *args,
+                        **kwargs,
+                    )
+                encoder_hidden_states, hidden_states = hidden_states.split(
+                    [
+                        encoder_hidden_states.shape[1],
+                        hidden_states.shape[1] - encoder_hidden_states.shape[1],
+                    ],
+                    dim=1,
                 )
-            encoder_hidden_states, hidden_states = hidden_states.split(
-                [
-                    encoder_hidden_states.shape[1],
-                    hidden_states.shape[1] - encoder_hidden_states.shape[1],
-                ],
-                dim=1,
-            )
         else:
             for block in self._Mn_transformer_blocks():
                 hidden_states = block(
@@ -1789,43 +1801,71 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
         original_hidden_states = hidden_states
         original_encoder_hidden_states = encoder_hidden_states
+        # This condition branch is mainly for FLUX series.
         if self.single_transformer_blocks is not None:
             assert Bn_compute_blocks() <= len(self.single_transformer_blocks), (
                 f"Bn_compute_blocks {Bn_compute_blocks()} must be less than "
                 f"the number of single transformer blocks {len(self.single_transformer_blocks)}"
             )
-
-            hidden_states = torch.cat(
-                [encoder_hidden_states, hidden_states], dim=1
-            )
-            if len(Bn_compute_blocks_ids()) > 0:
-                for i, block in enumerate(self._Bn_single_transformer_blocks()):
-                    hidden_states = (
-                        self._compute_and_cache_single_transformer_block(
-                            i,
-                            original_hidden_states,
-                            original_encoder_hidden_states,
-                            block,
+            if is_diffusers_at_least_0_3_5():
+                if len(Bn_compute_blocks_ids()) > 0:
+                    # NOTE: Reuse _compute_and_cache_transformer_block here.
+                    for i, block in enumerate(
+                        self._Bn_single_transformer_blocks()
+                    ):
+                        hidden_states, encoder_hidden_states = (
+                            self._compute_and_cache_transformer_block(
+                                i,
+                                block,
+                                hidden_states,
+                                encoder_hidden_states,
+                                *args,
+                                **kwargs,
+                            )
+                        )
+                else:
+                    # Compute all Bn blocks if no specific Bn compute blocks ids are set.
+                    for block in self._Bn_single_transformer_blocks():
+                        encoder_hidden_states, hidden_states = block(
                             hidden_states,
+                            encoder_hidden_states,
                             *args,
                             **kwargs,
                         )
-                    )
             else:
-                # Compute all Bn blocks if no specific Bn compute blocks ids are set.
-                for block in self._Bn_single_transformer_blocks():
-                    hidden_states = block(
-                        hidden_states,
-                        *args,
-                        **kwargs,
-                    )
-            encoder_hidden_states, hidden_states = hidden_states.split(
-                [
-                    encoder_hidden_states.shape[1],
-                    hidden_states.shape[1] - encoder_hidden_states.shape[1],
-                ],
-                dim=1,
-            )
+                hidden_states = torch.cat(
+                    [encoder_hidden_states, hidden_states], dim=1
+                )
+                if len(Bn_compute_blocks_ids()) > 0:
+                    for i, block in enumerate(
+                        self._Bn_single_transformer_blocks()
+                    ):
+                        hidden_states = (
+                            self._compute_and_cache_single_transformer_block(
+                                i,
+                                original_hidden_states,
+                                original_encoder_hidden_states,
+                                block,
+                                hidden_states,
+                                *args,
+                                **kwargs,
+                            )
+                        )
+                else:
+                    # Compute all Bn blocks if no specific Bn compute blocks ids are set.
+                    for block in self._Bn_single_transformer_blocks():
+                        hidden_states = block(
+                            hidden_states,
+                            *args,
+                            **kwargs,
+                        )
+                encoder_hidden_states, hidden_states = hidden_states.split(
+                    [
+                        encoder_hidden_states.shape[1],
+                        hidden_states.shape[1] - encoder_hidden_states.shape[1],
+                    ],
+                    dim=1,
+                )
         else:
             assert Bn_compute_blocks() <= len(self.transformer_blocks), (
                 f"Bn_compute_blocks {Bn_compute_blocks()} must be less than "

cache_dit/cache_factory/dynamic_block_prune/prune_context.py

@@ -7,7 +7,8 @@ from typing import Any, Dict, List, Optional, Union
 
 import torch
 
-import cache_dit.primitives as DP
+import cache_dit.primitives as primitives
+from cache_dit.utils import is_diffusers_at_least_0_3_5
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -446,8 +447,8 @@ def are_two_tensors_similar(
         mean_t1 = t1.abs().mean()
 
     if parallelized:
-        mean_diff = DP.all_reduce_sync(mean_diff, "avg")
-        mean_t1 = DP.all_reduce_sync(mean_t1, "avg")
+        mean_diff = primitives.all_reduce_sync(mean_diff, "avg")
+        mean_t1 = primitives.all_reduce_sync(mean_t1, "avg")
 
     # D = (t1 - t2) / t1 = 1 - (t2 / t1), if D = 0, then t1 = t2.
     # Futher, if we assume that (H(t,  0) - H(t-1,0)) ~ 0, then,
@@ -936,27 +937,44 @@ class DBPrunedTransformerBlocks(torch.nn.Module):
             )
 
         if self.single_transformer_blocks is not None:
-            hidden_states = torch.cat(
-                [encoder_hidden_states, hidden_states], dim=1
-            )
-            for j, block in enumerate(self.single_transformer_blocks):
-                hidden_states = self._compute_or_prune_single_transformer_block(
-                    j + len(self.transformer_blocks),
-                    original_hidden_states,
-                    original_encoder_hidden_states,
-                    block,
-                    hidden_states,
-                    *args,
-                    **kwargs,
+            # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py#L380
+            if is_diffusers_at_least_0_3_5():
+                for j, block in enumerate(self.single_transformer_blocks):
+                    # NOTE: Reuse _compute_or_prune_transformer_block here.
+                    hidden_states, encoder_hidden_states = (
+                        self._compute_or_prune_transformer_block(
+                            j + len(self.transformer_blocks),
+                            block,
+                            hidden_states,
+                            encoder_hidden_states,
+                            *args,
+                            **kwargs,
+                        )
+                    )
+            else:
+                hidden_states = torch.cat(
+                    [encoder_hidden_states, hidden_states], dim=1
                 )
+                for j, block in enumerate(self.single_transformer_blocks):
+                    hidden_states = (
+                        self._compute_or_prune_single_transformer_block(
+                            j + len(self.transformer_blocks),
+                            original_hidden_states,
+                            original_encoder_hidden_states,
+                            block,
+                            hidden_states,
+                            *args,
+                            **kwargs,
+                        )
+                    )
 
-            encoder_hidden_states, hidden_states = hidden_states.split(
-                [
-                    encoder_hidden_states.shape[1],
-                    hidden_states.shape[1] - encoder_hidden_states.shape[1],
-                ],
-                dim=1,
-            )
+                encoder_hidden_states, hidden_states = hidden_states.split(
+                    [
+                        encoder_hidden_states.shape[1],
+                        hidden_states.shape[1] - encoder_hidden_states.shape[1],
+                    ],
+                    dim=1,
+                )
 
         hidden_states = (
             hidden_states.reshape(-1)

cache_dit/cache_factory/first_block_cache/cache_context.py

@@ -7,7 +7,7 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union
 
 import torch
 
-import cache_dit.primitives as DP
+import cache_dit.primitives as primitives
 from cache_dit.cache_factory.taylorseer import TaylorSeer
 from cache_dit.logger import init_logger
 
@@ -348,8 +348,8 @@ def are_two_tensors_similar(
     mean_diff = (t1 - t2).abs().mean()
     mean_t1 = t1.abs().mean()
     if parallelized:
-        mean_diff = DP.all_reduce_sync(mean_diff, "avg")
-        mean_t1 = DP.all_reduce_sync(mean_t1, "avg")
+        mean_diff = primitives.all_reduce_sync(mean_diff, "avg")
+        mean_t1 = primitives.all_reduce_sync(mean_t1, "avg")
     diff = (mean_diff / mean_t1).item()
 
     add_residual_diff(diff)

cache_dit/metrics/lpips.py

@@ -0,0 +1,43 @@
+import builtins as __builtin__
+import contextlib
+import warnings
+
+import lpips
+import torch
+
+warnings.filterwarnings("ignore")
+
+lpips_loss_fn_vgg = None
+lpips_loss_fn_alex = None
+
+
+def dummy_print(*args, **kwargs):
+    pass
+
+
+@contextlib.contextmanager
+def disable_print():
+    origin_print = __builtin__.print
+    __builtin__.print = dummy_print
+    yield
+    __builtin__.print = origin_print
+
+
+def compute_lpips_img(img0, img1, net: str = "alex"):
+    global lpips_loss_fn_vgg
+    global lpips_loss_fn_alex
+    if net.lower() == "alex":
+        if lpips_loss_fn_alex is None:
+            with disable_print():
+                lpips_loss_fn_alex = lpips.LPIPS(net="alex")
+        loss_fn = lpips_loss_fn_alex
+    elif net.lower() == "vgg":
+        if lpips_loss_fn_vgg is None:
+            with disable_print():
+                lpips_loss_fn_vgg = lpips.LPIPS(net="vgg")
+        loss_fn = lpips_loss_fn_vgg
+    else:
+        assert False, f"unsupport net {net}"
+
+    with torch.no_grad():
+        return loss_fn(img0, img1).item()

cache_dit/metrics/metrics.py

@@ -14,6 +14,7 @@ from cache_dit.metrics.config import get_metrics_verbose
 from cache_dit.metrics.config import _IMAGE_EXTENSIONS
 from cache_dit.metrics.config import _VIDEO_EXTENSIONS
 from cache_dit.logger import init_logger
+from cache_dit.metrics.lpips import compute_lpips_img
 
 logger = init_logger(__name__)
 
@@ -21,6 +22,35 @@ logger = init_logger(__name__)
 DISABLE_VERBOSE = not get_metrics_verbose()
 
 
+def compute_lpips_file(
+    image_true: np.ndarray | str,
+    image_test: np.ndarray | str,
+) -> float:
+    import torch
+    from PIL import Image
+    from torchvision.transforms.v2.functional import (
+        convert_image_dtype,
+        normalize,
+        pil_to_tensor,
+    )
+
+    def load_img_as_tensor(path):
+        pil = Image.open(path)
+        img = pil_to_tensor(pil)
+        img = convert_image_dtype(img, dtype=torch.float32)
+        img = normalize(img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        return img
+
+    if isinstance(image_true, str):
+        image_true = load_img_as_tensor(image_true)
+    if isinstance(image_test, str):
+        image_test = load_img_as_tensor(image_test)
+    return compute_lpips_img(
+        image_true,
+        image_test,
+    )
+
+
 def compute_psnr_file(
     image_true: np.ndarray | str,
     image_test: np.ndarray | str,
@@ -305,6 +335,11 @@ def compute_video_metric(
         return None, None
 
 
+compute_lpips = partial(
+    compute_dir_metric,
+    compute_file_func=compute_lpips_file,
+)
+
 compute_psnr = partial(
     compute_dir_metric,
     compute_file_func=compute_psnr_file,
@@ -320,6 +355,10 @@ compute_mse = partial(
     compute_file_func=compute_mse_file,
 )
 
+compute_video_lpips = partial(
+    compute_video_metric,
+    compute_frame_func=compute_lpips_file,
+)
 compute_video_psnr = partial(
     compute_video_metric,
     compute_frame_func=compute_psnr_file,
@@ -335,6 +374,7 @@ compute_video_mse = partial(
 
 
 METRICS_CHOICES = [
+    "lpips",
     "psnr",
     "ssim",
     "mse",
@@ -522,6 +562,9 @@ def entrypoint():
                 METRICS_META[msg] = value
                 logger.info(msg)
 
+            if metric == "lpips" or metric == "all":
+                img_lpips, n = compute_lpips(img_true, img_test)
+                _logging_msg(img_lpips, "lpips", n)
             if metric == "psnr" or metric == "all":
                 img_psnr, n = compute_psnr(img_true, img_test)
                 _logging_msg(img_psnr, "psnr", n)
@@ -558,6 +601,9 @@ def entrypoint():
                 METRICS_META[msg] = value
                 logger.info(msg)
 
+            if metric == "lpips" or metric == "all":
+                video_lpips, n = compute_video_lpips(video_true, video_test)
+                _logging_msg(video_lpips, "lpips", n)
             if metric == "psnr" or metric == "all":
                 video_psnr, n = compute_video_psnr(video_true, video_test)
                 _logging_msg(video_psnr, "psnr", n)

cache_dit/utils.py

@@ -0,0 +1,7 @@
+import torch
+import diffusers
+
+
+@torch.compiler.disable
+def is_diffusers_at_least_0_3_5() -> bool:
+    return diffusers.__version__ >= "0.35.0"

{cache_dit-0.2.12.dist-info → cache_dit-0.2.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.12
+Version: 0.2.14
 Summary: 🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration Toolbox for Diffusion Transformers
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -16,6 +16,7 @@ Requires-Dist: transformers>=4.51.3
 Requires-Dist: diffusers>=0.33.1
 Requires-Dist: scikit-image
 Requires-Dist: scipy
+Requires-Dist: lpips==0.1.4
 Provides-Extra: all
 Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
@@ -52,7 +53,7 @@ Dynamic: requires-python
       <img src=https://img.shields.io/badge/Release-v0.2-brightgreen.svg >
  </div>
   <p align="center">
-    DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT offers <br>a set of training-free cache accelerators for DiT: <b>🔥<a href="#dbcache">DBCache</a>, <a href="#dbprune">DBPrune</a>, <a href="#taylorseer">TaylorSeer</a>, <a href="#fbcache">FBCache</a></b>, etc🔥
+    DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. <br> CacheDiT offers a set of training-free cache accelerators for Diffusion Transformers: <br> <b>🔥<a href="#dbcache">DBCache</a>, <a href="#dbprune">DBPrune</a>, <a href="#taylorseer">Hybrid TaylorSeer</a>, <a href="#cfg">Hybrid Cache CFG</a>, <a href="#fbcache">FBCache</a></b>, etc🔥
   </p>
 </div>
 
@@ -66,95 +67,6 @@ Dynamic: requires-python
 - [2025-07-18] 🎉First caching mechanism in **[🤗huggingface/flux-fast](https://github.com/huggingface/flux-fast)** with **[cache-dit](https://github.com/vipshop/cache-dit)**, also check the **[PR](https://github.com/huggingface/flux-fast/pull/13)**. 
 - [2025-07-13] **[🤗flux-faster](https://github.com/xlite-dev/flux-faster)** is released! A forked version of **[🤗huggingface/flux-fast](https://github.com/huggingface/flux-fast)** that **makes flux-fast even faster** with **[cache-dit](https://github.com/vipshop/cache-dit)**, **3.3x** speedup on NVIDIA L20.
 
-## 🤗 Introduction 
-
-<div align="center">
-  <p align="center">
-    <h3>🔥DBCache: Dual Block Caching for Diffusion Transformers</h3>
-  </p>
-</div> 
-
-**DBCache**: **Dual Block Caching** for Diffusion Transformers. We have enhanced `FBCache` into a more general and customizable cache algorithm, namely `DBCache`, enabling it to achieve fully `UNet-style` cache acceleration for DiT models. Different configurations of compute blocks (**F8B12**, etc.) can be customized in DBCache. Moreover, it can be entirely **training**-**free**. DBCache can strike a perfect **balance** between performance and precision!
-
-<div align="center">
-  <p align="center">
-    DBCache, <b> L20x1 </b>, Steps: 28, "A cat holding a sign that says hello world with complex background"
-  </p>
-</div>
-
-|Baseline(L20x1)|F1B0 (0.08)|F1B0 (0.20)|F8B8 (0.15)|F12B12 (0.20)|F16B16 (0.20)|
-|:---:|:---:|:---:|:---:|:---:|:---:|
-|24.85s|15.59s|8.58s|15.41s|15.11s|17.74s|
-|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.08_S11.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.2_S19.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F8B8S1_R0.15_S15.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F12B12S4_R0.2_S16.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F16B16S4_R0.2_S13.png width=105px>|
-|**Baseline(L20x1)**|**F1B0 (0.08)**|**F8B8 (0.12)**|**F8B12 (0.12)**|**F8B16 (0.20)**|**F8B20 (0.20)**|
-|27.85s|6.04s|5.88s|5.77s|6.01s|6.20s|
-|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_NONE_R0.08.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F1B0_R0.08.png width=105px> |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B8_R0.12.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B12_R0.12.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B16_R0.2.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B20_R0.2.png width=105px>|
-
-<div align="center">
-  <p align="center">
-    DBCache, <b> L20x4 </b>, Steps: 20, case to show the texture recovery ability of DBCache
-  </p>
-</div>
-
-These case studies demonstrate that even with relatively high thresholds (such as 0.12, 0.15, 0.2, etc.) under the DBCache **F12B12** or **F8B16** configuration, the detailed texture of the kitten's fur, colored cloth, and the clarity of text can still be preserved. This suggests that users can leverage DBCache to effectively balance performance and precision in their workflows! 
-
-<div align="center">
-  <p align="center">
-    <h3>🔥DBPrune: Dynamic Block Prune with Residual Caching</h3>
-  </p>
-</div> 
-
-**DBPrune**: We have further implemented a new **Dynamic Block Prune** algorithm based on **Residual Caching** for Diffusion Transformers, referred to as DBPrune. DBPrune caches each block's hidden states and residuals, then **dynamically prunes** blocks during inference by computing the L1 distance between previous hidden states. When a block is pruned, its output is approximated using the cached residuals.
-
-<div align="center">
-  <p align="center">
-    DBPrune, <b> L20x1 </b>, Steps: 28, "A cat holding a sign that says hello world with complex background"
-  </p>
-</div>
-
-|Baseline(L20x1)|Pruned(24%)|Pruned(35%)|Pruned(38%)|Pruned(45%)|Pruned(60%)|
-|:---:|:---:|:---:|:---:|:---:|:---:|
-|24.85s|19.43s|16.82s|15.95s|14.24s|10.66s|
-|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBPRUNE_F1B0_R0.03_P24.0_T19.43s.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBPRUNE_F1B0_R0.04_P34.6_T16.82s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBPRUNE_F1B0_R0.05_P38.3_T15.95s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBPRUNE_F1B0_R0.06_P45.2_T14.24s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBPRUNE_F1B0_R0.2_P59.5_T10.66s.png width=105px>|
-
-<div align="center">
-  <p align="center">
-    <h3>🔥Context Parallelism and Torch Compile</h3>
-  </p>
-</div> 
-
-Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand with [ParaAttention](https://github.com/chengzeyi/ParaAttention). Users can easily tap into its **Context Parallelism** features for distributed inference. CacheDiT is designed to work compatibly with **torch.compile.** You can easily use CacheDiT with torch.compile to further achieve a better performance.
-
-<div align="center">
-  <p align="center">
-  DBPrune + <b>torch.compile + context parallelism</b> <br>Steps: 28, "A cat holding a sign that says hello world with complex background"
-  </p>
-</div>
-
-|Baseline|Pruned(24%)|Pruned(35%)|Pruned(38%)|Pruned(45%)|Pruned(60%)|
-|:---:|:---:|:---:|:---:|:---:|:---:|
-|+compile:20.43s|16.25s|14.12s|13.41s|12.00s|8.86s|
-|+L20x4:7.75s|6.62s|6.03s|5.81s|5.24s|3.93s|
-|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_NONE_R0.08_S0_T20.43s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.03_P24.0_T16.25s.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.04_P34.6_T14.12s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.045_P38.2_T13.41s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.055_P45.1_T12.00s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.2_P59.5_T8.86s.png width=105px>|
-
-## ©️Citations
-
-```BibTeX
-@misc{CacheDiT@2025,
-  title={CacheDiT: A Training-free and Easy-to-use cache acceleration Toolbox for Diffusion Transformers},
-  url={https://github.com/vipshop/cache-dit.git},
-  note={Open-source software available at https://github.com/vipshop/cache-dit.git},
-  author={vipshop.com},
-  year={2025}
-}
-```
-
-## 👋Reference
-
-<div id="reference"></div>
-
-The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache). Special thanks to their excellent work! 
-
 ## 📖Contents 
 
 <div id="contents"></div>  
@@ -171,6 +83,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 - [⚙️Metrics CLI](#metrics)
 - [👋Contribute](#contribute)
 - [©️License](#license)
+- [©️Citations](#citations)
 
 ## ⚙️Installation  
 
@@ -207,6 +120,32 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 
 ![](https://github.com/vipshop/cache-dit/raw/main/assets/dbcache-v1.png)
 
+
+**DBCache**: **Dual Block Caching** for Diffusion Transformers. We have enhanced `FBCache` into a more general and customizable cache algorithm, namely `DBCache`, enabling it to achieve fully `UNet-style` cache acceleration for DiT models. Different configurations of compute blocks (**F8B12**, etc.) can be customized in DBCache. Moreover, it can be entirely **training**-**free**. DBCache can strike a perfect **balance** between performance and precision!
+
+<div align="center">
+  <p align="center">
+    DBCache, <b> L20x1 </b>, Steps: 28, "A cat holding a sign that says hello world with complex background"
+  </p>
+</div>
+
+|Baseline(L20x1)|F1B0 (0.08)|F1B0 (0.20)|F8B8 (0.15)|F12B12 (0.20)|F16B16 (0.20)|
+|:---:|:---:|:---:|:---:|:---:|:---:|
+|24.85s|15.59s|8.58s|15.41s|15.11s|17.74s|
+|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.08_S11.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.2_S19.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F8B8S1_R0.15_S15.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F12B12S4_R0.2_S16.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F16B16S4_R0.2_S13.png width=105px>|
+|**Baseline(L20x1)**|**F1B0 (0.08)**|**F8B8 (0.12)**|**F8B12 (0.12)**|**F8B16 (0.20)**|**F8B20 (0.20)**|
+|27.85s|6.04s|5.88s|5.77s|6.01s|6.20s|
+|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_NONE_R0.08.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F1B0_R0.08.png width=105px> |<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B8_R0.12.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B12_R0.12.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B16_R0.2.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/TEXTURE_DBCACHE_F8B20_R0.2.png width=105px>|
+
+<div align="center">
+  <p align="center">
+    DBCache, <b> L20x4 </b>, Steps: 20, case to show the texture recovery ability of DBCache
+  </p>
+</div>
+
+These case studies demonstrate that even with relatively high thresholds (such as 0.12, 0.15, 0.2, etc.) under the DBCache **F12B12** or **F8B16** configuration, the detailed texture of the kitten's fur, colored cloth, and the clarity of text can still be preserved. This suggests that users can leverage DBCache to effectively balance performance and precision in their workflows! 
+
+
 **DBCache** provides configurable parameters for custom optimization, enabling a balanced trade-off between performance and precision:
 
 - **Fn**: 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 accurate information to subsequent blocks.
@@ -258,17 +197,6 @@ cache_options = {
 }
 ```
 
-<div align="center">
-  <p align="center">
-    DBCache, <b> L20x1 </b>, Steps: 28, "A cat holding a sign that says hello world with complex background"
-  </p>
-</div>
-
-|Baseline(L20x1)|F1B0 (0.08)|F1B0 (0.20)|F8B8 (0.15)|F12B12 (0.20)|F16B16 (0.20)|
-|:---:|:---:|:---:|:---:|:---:|:---:|
-|24.85s|15.59s|8.58s|15.41s|15.11s|17.74s|
-|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/NONE_R0.08_S0.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.08_S11.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F1B0S1_R0.2_S19.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F8B8S1_R0.15_S15.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F12B12S4_R0.2_S16.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/DBCACHE_F16B16S4_R0.2_S13.png width=105px>|
-
 ## 🔥Hybrid TaylorSeer
 
 <div id="taylorseer"></div>
@@ -489,6 +417,18 @@ Then, run the python test script with `torchrun`:
 torchrun --nproc_per_node=4 parallel_cache.py
 ```
 
+<div align="center">
+  <p align="center">
+  DBPrune + <b>torch.compile + context parallelism</b> <br>Steps: 28, "A cat holding a sign that says hello world with complex background"
+  </p>
+</div>
+
+|Baseline|Pruned(24%)|Pruned(35%)|Pruned(38%)|Pruned(45%)|Pruned(60%)|
+|:---:|:---:|:---:|:---:|:---:|:---:|
+|+compile:20.43s|16.25s|14.12s|13.41s|12.00s|8.86s|
+|+L20x4:7.75s|6.62s|6.03s|5.81s|5.24s|3.93s|
+|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_NONE_R0.08_S0_T20.43s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.03_P24.0_T16.25s.png width=105px> | <img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.04_P34.6_T14.12s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.045_P38.2_T13.41s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.055_P45.1_T12.00s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBPRUNE_F1B0_R0.2_P59.5_T8.86s.png width=105px>|
+
 ## 🔥Torch Compile
 
 <div id="compile"></div>  
@@ -550,5 +490,18 @@ How to contribute? Star ⭐️ this repo to support us or check [CONTRIBUTE.md](
 
 <div id="license"></div>
 
+The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi/ParaAttention/tree/main/src/para_attn/first_block_cache). Special thanks to their excellent work! We have followed the original License from [FBCache](https://github.com/chengzeyi/ParaAttention), please check [LICENSE](https://github.com/vipshop/cache-dit/raw/main/LICENSE) for more details.
 
-We have followed the original License from [ParaAttention](https://github.com/chengzeyi/ParaAttention), please check [LICENSE](https://github.com/vipshop/cache-dit/raw/main/LICENSE) for more details.
+## ©️Citations
+
+<div id="citations"></div>
+
+```BibTeX
+@misc{CacheDiT@2025,
+  title={CacheDiT: A Training-free and Easy-to-use cache acceleration Toolbox for Diffusion Transformers},
+  url={https://github.com/vipshop/cache-dit.git},
+  note={Open-source software available at https://github.com/vipshop/cache-dit.git},
+  author={vipshop.com},
+  year={2025}
+}
+```

{cache_dit-0.2.12.dist-info → cache_dit-0.2.14.dist-info}/RECORD

@@ -1,13 +1,14 @@
-cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=7CFcHKqzy7OglwTX58ipGg1TD8MpDORqWvEBO3W1dHI,513
+cache_dit/__init__.py,sha256=0-B173-fLi3IA8nJXoS71zK0zD33Xplysd9skmLfEOY,171
+cache_dit/_version.py,sha256=ut2sCt69XoYh0A1_KAmfCg1IKkN6zwqhu2eMFWAhMbQ,513
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
+cache_dit/utils.py,sha256=4cFNh0asch6Zgsixq0bS1ElfwBu_6BG5ZSmaa1khjyg,144
 cache_dit/cache_factory/__init__.py,sha256=iYQwLwB_XLoYl0OB9unZGDbBtrYvZaLkOAmhGRwdW2E,191
 cache_dit/cache_factory/adapters.py,sha256=QMCaXnmqM7NT7sx4bCF1mMLn-QcXX9h1RmgLAypDedg,5256
 cache_dit/cache_factory/taylorseer.py,sha256=LKSNo2ode69EVo9xrxjxAMEjz0yDGiGADeDYnEqddA8,3987
 cache_dit/cache_factory/utils.py,sha256=V-Mb5Jn07geEUUWo4QAfh6pmSzkL-2OGDn0VAXbG6hQ,1799
 cache_dit/cache_factory/dual_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=itVEb6gT2eZuncAHUmP51ZS0r6v6cGtRvnPjyeXqKH8,71156
+cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=sJ9yxQlcrX4qkPln94FrL0WDe2WIn3_UD2-Mk8YtjSw,73301
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=krNAICf-aS3JLmSG8vOB9tpLa04uYRcABsC8PMbVUKY,1870
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=3xUjvDzor9AkBkDUc0N7kZqM86MIdajuigesnicNzXE,2260
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=cIsov6Pf0dRyddqkzTA2CU-jSDotof8LQr-HIoY9T9M,2615
@@ -15,7 +16,7 @@ cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/mochi.py,sha256=8W9m-WeEVE2ytYi9udKEA8Wtb0EnvP3eT2A1Tu-d29k,2252
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py,sha256=EREHM5E1wxnL-uRXRAEege4HXraRp1oD_r1Zx4CsiKk,2596
 cache_dit/cache_factory/dynamic_block_prune/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/cache_factory/dynamic_block_prune/prune_context.py,sha256=so1wGdb8W0ATwrjv7E5IEZLPcobybaY1HJa6hBYlOOQ,34698
+cache_dit/cache_factory/dynamic_block_prune/prune_context.py,sha256=1qarKAsEFiaaN2_ghko2dqGz_R7BTQSOyGtb_eQq38Y,35716
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py,sha256=hVBTXj9MMGFGVezT3j8MntFRBiphSaUL4YhSOd8JtuY,1870
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py,sha256=KP8NxtHAKzzBOoX0lhvlMgY_5dmP4Z3T5TOfwl4SSyg,2273
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py,sha256=kCB7lL4OIq8TZn-baMIF8D_PVPTFW60omCMVQCb8ebs,2628
@@ -23,7 +24,7 @@ cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/hunyuan_video.py,
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py,sha256=zXgoRDDjus3a2WSjtNh4ERtQp20ceb6nzohHMDlo2zY,2265
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py,sha256=PA7nuLgfAelnaI8usQx0Kxi8XATzMapyR1WndEdFoZA,2604
 cache_dit/cache_factory/first_block_cache/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/cache_factory/first_block_cache/cache_context.py,sha256=tTPwhPLEA7LqGupps1Zy2MycCtLzs22wsW0yUhiiF-U,23217
+cache_dit/cache_factory/first_block_cache/cache_context.py,sha256=qn4zWJ_eEMIPYzrxXoslunxbzK0WueuNtC54Pp5Q57k,23241
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py,sha256=-FFgA2MoudEo7uDacg4aWgm1KwfLZFsEDTVxatgbq9M,2146
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/cogvideox.py,sha256=qO5CWyurtwW30mvOe6cxeQPTSXLDlPJcezm72zEjDq8,2375
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/flux.py,sha256=Dcd4OzABCtyQCZNX2KNnUTdVoO1E1ApM7P8gcVYzcK0,2733
@@ -38,10 +39,11 @@ cache_dit/metrics/__init__.py,sha256=RaUhl5dieF40RqnizGzR30qoJJ9dyMUEADwgwMaMQrE
 cache_dit/metrics/config.py,sha256=ieOgD9ayz722RjVzk24bSIqS2D6o7TZjGk8KeXV-OLQ,551
 cache_dit/metrics/fid.py,sha256=9Ivtazl6mW0Bon2VXa-Ia5Xj2ewxRD3V1Qkd69zYM3Y,17066
 cache_dit/metrics/inception.py,sha256=pBVe2X6ylLPIXTG4-GWDM9DWnCviMJbJ45R3ulhktR0,12759
-cache_dit/metrics/metrics.py,sha256=1TTbfaj_-vdUfxopLnc5kVrXs5rMpAoSi8D0ItYdPu8,26439
-cache_dit-0.2.12.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.12.dist-info/METADATA,sha256=-AIWGVOFsY-nhMkDeFErUFcELTWmza96-0IUN3od88A,28219
-cache_dit-0.2.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.12.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-0.2.12.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.12.dist-info/RECORD,,
+cache_dit/metrics/lpips.py,sha256=I2qCNi6qJh5TRsaIsdxO0WoRX1DN7U_H3zS0oCSahYM,1032
+cache_dit/metrics/metrics.py,sha256=8jvM1sF-nDxUuwCRy44QEoo4dYVLCQVh1QyAMs4eaQY,27840
+cache_dit-0.2.14.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.14.dist-info/METADATA,sha256=EyZN75JcVcvTc5bopHXfl6w-nA-ro9Uit2Sjy5DU66A,25198
+cache_dit-0.2.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.14.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.14.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.14.dist-info/RECORD,,