cache-dit 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

cache_dit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__version__ = version = '0.2.1'
+__version_tuple__ = version_tuple = (0, 2, 1)

cache_dit/cache_factory/dual_block_cache/cache_context.py

@@ -9,6 +9,7 @@ from typing import Any, DefaultDict, Dict, List, Optional, Union
 import torch
 
 import cache_dit.primitives as DP
+from cache_dit.cache_factory.taylorseer import TaylorSeer
 from cache_dit.logger import init_logger
 
 logger = init_logger(__name__)
@@ -60,6 +61,18 @@ class DBCacheContext:
     residual_diffs: DefaultDict[str, float] = dataclasses.field(
         default_factory=lambda: defaultdict(float),
     )
+    # TODO: Support TaylorSeers and SLG in Dual Block Cache
+    # TaylorSeers: 
+    # Title: From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers
+    # Url: https://arxiv.org/pdf/2503.06923
+    taylorseer: Optional[TaylorSeer] = None
+    alter_taylorseer: Optional[TaylorSeer] = None
+
+    # Skip Layer Guidance, SLG
+    # https://github.com/huggingface/candle/issues/2588
+    slg_layers: Optional[List[int]] = None
+    slg_start: float = 0.0
+    slg_end: float = 0.1
 
     def get_incremental_name(self, name=None):
         if name is None:
@@ -700,7 +713,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 encoder_hidden_states,
                 hidden_states_residual,
                 encoder_hidden_states_residual,
-            ) = self.call_MN2n_transformer_blocks(  # middle
+            ) = self.call_Mn_transformer_blocks(  # middle
                 hidden_states,
                 encoder_hidden_states,
                 *args,
@@ -772,32 +785,32 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         return selected_Fn_transformer_blocks
 
     @torch.compiler.disable
-    def _MN2n_single_transformer_blocks(self):  # middle
+    def _Mn_single_transformer_blocks(self):  # middle blocks
         # M(N-2n): transformer_blocks [n,...] + single_transformer_blocks [0,...,N-n]
-        selected_MN2n_single_transformer_blocks = []
+        selected_Mn_single_transformer_blocks = []
         if self.single_transformer_blocks is not None:
             if Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
-                selected_MN2n_single_transformer_blocks = (
+                selected_Mn_single_transformer_blocks = (
                     self.single_transformer_blocks
                 )
             else:
-                selected_MN2n_single_transformer_blocks = (
+                selected_Mn_single_transformer_blocks = (
                     self.single_transformer_blocks[: -Bn_compute_blocks()]
                 )
-        return selected_MN2n_single_transformer_blocks
+        return selected_Mn_single_transformer_blocks
 
     @torch.compiler.disable
-    def _MN2n_transformer_blocks(self):
+    def _Mn_transformer_blocks(self): # middle blocks
         # M(N-2n): only transformer_blocks [n,...,N-n], middle
         if Bn_compute_blocks() == 0:  # WARN: x[:-0] = []
-            selected_MN2n_transformer_blocks = self.transformer_blocks[
+            selected_Mn_transformer_blocks = self.transformer_blocks[
                 Fn_compute_blocks() :
             ]
         else:
-            selected_MN2n_transformer_blocks = self.transformer_blocks[
+            selected_Mn_transformer_blocks = self.transformer_blocks[
                 Fn_compute_blocks() : -Bn_compute_blocks()
             ]
-        return selected_MN2n_transformer_blocks
+        return selected_Mn_transformer_blocks
 
     @torch.compiler.disable
     def _Bn_single_transformer_blocks(self):
@@ -845,7 +858,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
         return hidden_states, encoder_hidden_states
 
-    def call_MN2n_transformer_blocks(
+    def call_Mn_transformer_blocks(
         self,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
@@ -873,7 +886,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
             hidden_states = torch.cat(
                 [encoder_hidden_states, hidden_states], dim=1
             )
-            for block in self._MN2n_single_transformer_blocks():
+            for block in self._Mn_single_transformer_blocks():
                 hidden_states = block(
                     hidden_states,
                     *args,
@@ -887,7 +900,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 dim=1,
             )
         else:
-            for block in self._MN2n_transformer_blocks():
+            for block in self._Mn_transformer_blocks():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -1016,7 +1029,9 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
     def _compute_and_cache_single_transformer_block(
         self,
-        i: int,  # Block index in the transformer blocks
+        # Block index in the transformer blocks
+        # Bn: 8, block_id should be in [0, 8)
+        block_id: int,
         # Helper inputs for hidden states split and reshape
         original_hidden_states: torch.Tensor,
         original_encoder_hidden_states: torch.Tensor,
@@ -1042,7 +1057,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
             )
             # Cache residuals for the non-compute Bn blocks for
             # subsequent cache steps.
-            if i not in Bn_compute_blocks_ids():
+            if block_id not in Bn_compute_blocks_ids():
                 Bn_i_hidden_states = hidden_states
                 (
                     Bn_i_hidden_states_residual,
@@ -1057,16 +1072,16 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 # Save original_hidden_states for diff calculation.
                 set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"Bn_{i}_single_original",
+                    prefix=f"Bn_{block_id}_single_original",
                 )
 
                 set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{i}_single_residual",
+                    prefix=f"Bn_{block_id}_single_residual",
                 )
                 set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"Bn_{i}_single_residual",
+                    prefix=f"Bn_{block_id}_single_residual",
                 )
                 del Bn_i_hidden_states
                 del Bn_i_hidden_states_residual
@@ -1077,7 +1092,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         else:
             # Cache steps: Reuse the cached residuals.
             # Check if the block is in the Bn_compute_blocks_ids.
-            if i in Bn_compute_blocks_ids():
+            if block_id in Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     *args,
@@ -1091,7 +1106,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     hidden_states,  # curr step
                     parallelized=self._is_parallelized(),
                     threshold=non_compute_blocks_diff_threshold(),
-                    prefix=f"Bn_{i}_single_original",  # prev step
+                    prefix=f"Bn_{block_id}_single_original",  # prev step
                 ):
                     Bn_i_original_hidden_states = hidden_states
                     (
@@ -1106,7 +1121,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                         apply_hidden_states_residual(
                             Bn_i_original_hidden_states,
                             Bn_i_original_encoder_hidden_states,
-                            prefix=f"Bn_{i}_single_residual",
+                            prefix=f"Bn_{block_id}_single_residual",
                         )
                     )
                     hidden_states = torch.cat(
@@ -1125,7 +1140,9 @@ class DBCachedTransformerBlocks(torch.nn.Module):
 
     def _compute_and_cache_transformer_block(
         self,
-        i: int,  # Block index in the transformer blocks
+        # Block index in the transformer blocks
+        # Bn: 8, block_id should be in [0, 8)
+        block_id: int,  
         # Below are the inputs to the block
         block,  # The transformer block to be executed
         hidden_states: torch.Tensor,
@@ -1158,7 +1175,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     )
             # Cache residuals for the non-compute Bn blocks for
             # subsequent cache steps.
-            if i not in Bn_compute_blocks_ids():
+            if block_id not in Bn_compute_blocks_ids():
                 Bn_i_hidden_states_residual = (
                     hidden_states - Bn_i_original_hidden_states
                 )
@@ -1169,16 +1186,16 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                 # Save original_hidden_states for diff calculation.
                 set_Bn_buffer(
                     Bn_i_original_hidden_states,
-                    prefix=f"Bn_{i}_original",
+                    prefix=f"Bn_{block_id}_original",
                 )
 
                 set_Bn_buffer(
                     Bn_i_hidden_states_residual,
-                    prefix=f"Bn_{i}_residual",
+                    prefix=f"Bn_{block_id}_residual",
                 )
                 set_Bn_encoder_buffer(
                     Bn_i_encoder_hidden_states_residual,
-                    prefix=f"Bn_{i}_residual",
+                    prefix=f"Bn_{block_id}_residual",
                 )
                 del Bn_i_hidden_states_residual
                 del Bn_i_encoder_hidden_states_residual
@@ -1189,7 +1206,7 @@ class DBCachedTransformerBlocks(torch.nn.Module):
         else:
             # Cache steps: Reuse the cached residuals.
             # Check if the block is in the Bn_compute_blocks_ids.
-            if i in Bn_compute_blocks_ids():
+            if block_id in Bn_compute_blocks_ids():
                 hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
@@ -1211,13 +1228,13 @@ class DBCachedTransformerBlocks(torch.nn.Module):
                     hidden_states,  # curr step
                     parallelized=self._is_parallelized(),
                     threshold=non_compute_blocks_diff_threshold(),
-                    prefix=f"Bn_{i}_original",  # prev step
+                    prefix=f"Bn_{block_id}_original",  # prev step
                 ):
                     hidden_states, encoder_hidden_states = (
                         apply_hidden_states_residual(
                             hidden_states,
                             encoder_hidden_states,
-                            prefix=f"Bn_{i}_residual",
+                            prefix=f"Bn_{block_id}_residual",
                         )
                     )
                 else:

cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py

@@ -56,7 +56,7 @@ def apply_cache_on_pipe(
     shallow_patch: bool = False,
     residual_diff_threshold=0.03,
     downsample_factor=1,
-    # SLG is not supported in WAN with DBCache yet
+    # SLG is not supported in WAN with DBPrune yet
     # slg_layers=None,
     # slg_start: float = 0.0,
     # slg_end: float = 0.1,

{cache_dit-0.2.0.dist-info → cache_dit-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.0
+Version: 0.2.1
 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
@@ -44,20 +44,11 @@ Dynamic: requires-python
       <img src=https://img.shields.io/badge/PyPI-pass-brightgreen.svg >
       <img src=https://static.pepy.tech/badge/cache-dit >
       <img src=https://img.shields.io/badge/Python-3.10|3.11|3.12-9cf.svg >
-      <img src=https://img.shields.io/badge/Release-v0.2.0-brightgreen.svg >
+      <img src=https://img.shields.io/badge/Release-v0.2.1-brightgreen.svg >
  </div>
   <p align="center">
     DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT <br>offers a set of training-free cache accelerators for DiT: 🔥DBCache, DBPrune, FBCache, etc🔥
   </p>
-  <p align="center">
-  <h4> 🔥Supported Models🔥</h4>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀FLUX.1</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀Mochi</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀CogVideoX</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀CogVideoX1.5</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀Wan2.1</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀HunyuanVideo</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
-  </p>
 </div>
 
 ## 👋 Highlight
@@ -160,6 +151,7 @@ Moreover, **CacheDiT** are **plug-and-play** solutions that works hand-in-hand w
 <div id="contents"></div>  
 
 - [⚙️Installation](#️installation)
+- [🔥Supported Models](#supported)
 - [⚡️Dual Block Cache](#dbcache)
 - [🎉First Block Cache](#fbcache)
 - [⚡️Dynamic Block Prune](#dbprune)
@@ -183,6 +175,30 @@ Or you can install the latest develop version from GitHub:
 pip3 install git+https://github.com/vipshop/cache-dit.git
 ```
 
+## 🔥Supported Models  
+
+<div id="supported"></div>
+
+- [🚀FLUX.1](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀Mochi](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀CogVideoX](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀CogVideoX1.5](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀Wan2.1](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
+
+
+<!--
+<p align="center">
+  <h4> 🔥Supported Models🔥</h4>
+  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀FLUX.1</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
+  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀Mochi</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
+  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀CogVideoX</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
+  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀CogVideoX1.5</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
+  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀Wan2.1</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
+  <a href=https://github.com/vipshop/cache-dit/raw/main/examples> <b>🚀HunyuanVideo</b>: ✔️DBCache, ✔️DBPrune, ✔️FBCache🔥</a> <br>
+</p>
+-->
+
 ## ⚡️DBCache: Dual Block Cache  
 
 <div id="dbcache"></div>
@@ -340,6 +356,9 @@ cache_options = {
 apply_cache_on_pipe(pipe, **cache_options)
 ```
 
+> [!Important]
+> Please note that for GPUs with lower VRAM, DBPrune may not be suitable for use on video DiTs, as it caches the hidden states and residuals of each block, leading to higher GPU memory requirements. In such cases, please use DBCache, which only caches the hidden states and residuals of 2 blocks.
+
 <div align="center">
   <p align="center">
     DBPrune, <b> L20x1 </b>, Steps: 28, "A cat holding a sign that says hello world with complex background"

{cache_dit-0.2.0.dist-info → cache_dit-0.2.1.dist-info}/RECORD

@@ -1,12 +1,12 @@
 cache_dit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/_version.py,sha256=iB5DfB5V6YB5Wo4JmvS-txT42QtmGaWcWp3udRT7zCI,511
+cache_dit/_version.py,sha256=UoNvMtd4wCG76RwoSpNCUtaFyTwakGcZolfjXzNVSMY,511
 cache_dit/logger.py,sha256=dKfNe_RRk9HJwfgHGeRR1f0LbskJpKdGmISCbL9roQs,3443
 cache_dit/primitives.py,sha256=A2iG9YLot3gOsZSPp-_gyjqjLgJvWQRx8aitD4JQ23Y,3877
 cache_dit/cache_factory/__init__.py,sha256=5RNuhWakvvqrOV4vkqrEBA7d-V1LwcNSsjtW14mkqK8,5255
 cache_dit/cache_factory/taylorseer.py,sha256=0W29ykJg3MnyLAB2KFicsl11Xe41cDYPgI60bquG_NY,2495
 cache_dit/cache_factory/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=EJ-uhA2-sWMW1jNDhcBtjHDqSn8lUzfKbYoPfZDQhZU,49665
+cache_dit/cache_factory/dual_block_cache/cache_context.py,sha256=huudVxz-SF3wAY_9vkViFCBhFKm5IzLvXR686u82pbM,50430
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py,sha256=ySmO_0IuSm5VrYdi9ccGVHoFVkgtPZMJGq_OMoyl0Q8,2003
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/cogvideox.py,sha256=1_n-RFMiL3v2SjhSfFrPH5Mn5Dq9z4BesVK8GN_nh2g,2404
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/flux.py,sha256=UbE6nIF-EtA92QxIZVMzIssdZKQSPAVX1hchF9R8drU,2754
@@ -20,7 +20,7 @@ cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/cogvideox.py,sha2
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/flux.py,sha256=KbEkLSsHtS6xwLWNh3jlOlXRyGRdrI2pWV1zyQxMTj4,2757
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/hunyuan_video.py,sha256=v3SgLJQPbEqSVy2sYVGMhptJqCe-XPXL7LIV7GEacZg,10105
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/mochi.py,sha256=rgeXfww-7WX6URSDg7mF1HuxSmYmoJVjMVoNGuxjwxc,2395
-cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py,sha256=JCRIq7kt59nbv7t10FdfmFZd1GawUZu2tYQ_QBB8zXQ,2713
+cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/wan.py,sha256=3-Bg9oPdLcIFZqqSpBGU3Ps1DJ9J8rslP5X7Ow1EHmc,2713
 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=DpDhtK095PlrvACf7sbjOt2-QpVkV1arr1qGEKJqgaQ,23502
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/__init__.py,sha256=-FFgA2MoudEo7uDacg4aWgm1KwfLZFsEDTVxatgbq9M,2146
@@ -29,8 +29,8 @@ cache_dit/cache_factory/first_block_cache/diffusers_adapters/flux.py,sha256=Dcd4
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/hunyuan_video.py,sha256=OL7W4ukYlZz0IDmBR1zVV6XT3Mgciglj9Hqzv1wUAkQ,10092
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/mochi.py,sha256=lQTClo52OwPbNEE4jiBZQhfC7hbtYqnYIABp_vbm_dk,2363
 cache_dit/cache_factory/first_block_cache/diffusers_adapters/wan.py,sha256=dBNzHBECAuTTA1a7kLdvZL20YzaKTAS3iciVLzKKEWA,2638
-cache_dit-0.2.0.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-0.2.0.dist-info/METADATA,sha256=WK3Fu8euIwLlm3TXjJws9VzwNjEfcMNvkCSJRt7jEdo,21845
-cache_dit-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-0.2.0.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-0.2.0.dist-info/RECORD,,
+cache_dit-0.2.1.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.1.dist-info/METADATA,sha256=YbS-gmVFpGfTmaKNTUbXfWlfQa-RCoz0TyRqCtHEGJc,22700
+cache_dit-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.1.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.1.dist-info/RECORD,,