cache-dit 0.1.2__tar.gz → 0.1.5__tar.gz

{cache_dit-0.1.2 → cache_dit-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.1.2
+Version: 0.1.5
 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
@@ -37,13 +37,14 @@ Dynamic: requires-python
   <p align="center">
     <h3>🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration <br>Toolbox for Diffusion Transformers</h3>
   </p>
-  <img src=https://github.com/vipshop/cache-dit/raw/dev/assets/cache-dit.png >
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cache-dit.png >
   <div align='center'>
-        <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
-        <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
-        <img src=https://img.shields.io/badge/PyPI-pass-brightgreen.svg >
-        <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.1.2-brightgreen.svg >
+      <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
+      <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
+      <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.1.5-brightgreen.svg >
  </div>
   <p align="center">
     DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT provides <br>a series of training-free, UNet-style cache accelerators for DiT: DBCache, DBPrune, FBCache, etc.
@@ -166,7 +167,7 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 - **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.
 - **Bn**: Further fuses approximate information in the **last n** Transformer blocks to enhance prediction accuracy. These blocks act as an auto-scaler for approximate hidden states that use residual cache.
 - **warmup_steps**: (default: 0) DBCache does not apply the caching strategy when the number of running steps is less than or equal to this value, ensuring the model sufficiently learns basic features during warmup.
-- **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the running steps exceed this value to prevent precision degradation.
+- **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the previous cached steps exceed this value to prevent precision degradation.
 - **residual_diff_threshold**: The value of residual diff threshold, a higher value leads to faster performance at the cost of lower precision.
 
 For a good balance between performance and precision, DBCache is configured by default with **F8B8**, 8 warmup steps, and unlimited cached steps.
@@ -202,8 +203,9 @@ Moreover, users configuring higher **Bn** values (e.g., **F8B16**) while aiming
 cache_options = {
     # 0, 2, 4, ..., 14, 15, etc. [0,16)
     "Bn_compute_blocks_ids": CacheType.range(0, 16, 2),
-    # Skip Bn blocks (1, 3, 5 ,..., etc.) only if the L1 diff 
-    # lower than this value, otherwise, compute it.
+    # If the L1 difference is below this threshold, skip Bn blocks 
+    # not in `Bn_compute_blocks_ids`(1, 3,..., etc), Otherwise, 
+    # compute these blocks.
     "non_compute_blocks_diff_threshold": 0.08,
 }
 ```
@@ -259,12 +261,47 @@ pipe = FluxPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
 ).to("cuda")
 
-# Using DBPrune
+# Using DBPrune with default options
 cache_options = CacheType.default_options(CacheType.DBPrune)
 
 apply_cache_on_pipe(pipe, **cache_options)
 ```
 
+We have also brought the designs from DBCache to DBPrune to make it a more general and customizable block prune algorithm. You can specify the values of **Fn** and **Bn** for higher precision, or set up the non-prune blocks list **non_prune_blocks_ids** to avoid aggressive pruning. For example:
+
+```python
+# Custom options for DBPrune
+cache_options = {
+    "cache_type": CacheType.DBPrune,
+    "residual_diff_threshold": 0.05,
+    # Never prune the first `Fn` and last `Bn` blocks.
+    "Fn_compute_blocks": 8,  # default 1
+    "Bn_compute_blocks": 8,  # default 0
+    "warmup_steps": 8,  # default -1
+    # Disables the pruning strategy when the previous 
+    # pruned steps greater than this value.
+    "max_pruned_steps": 12,  # default, -1 means no limit
+    # Enable dynamic prune threshold within step, higher 
+    # `max_dynamic_prune_threshold` value may introduce a more 
+    # ageressive pruning strategy.
+    "enable_dynamic_prune_threshold": True,
+    "max_dynamic_prune_threshold": 2 * 0.05,
+    # (New thresh) = mean(previous_block_diffs_within_step) * 1.25
+    # (New thresh) = ((New thresh) if (New thresh) <
+    # max_dynamic_prune_threshold else residual_diff_threshold)
+    "dynamic_prune_threshold_relax_ratio": 1.25,
+    # The step interval to update residual cache. For example, 
+    # 2: means the update steps will be [0, 2, 4, ...].
+    "residual_cache_update_interval": 1,
+    # You can set non-prune blocks to avoid ageressive pruning. 
+    # For example, FLUX.1 has 19 + 38 blocks, so we can set it 
+    # to 0, 2, 4, ..., 56, etc.
+    "non_prune_blocks_ids": [],
+}
+
+apply_cache_on_pipe(pipe, **cache_options)
+```
+
 <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.1.2 → cache_dit-0.1.5}/README.md

@@ -2,13 +2,14 @@
   <p align="center">
     <h3>🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration <br>Toolbox for Diffusion Transformers</h3>
   </p>
-  <img src=https://github.com/vipshop/cache-dit/raw/dev/assets/cache-dit.png >
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cache-dit.png >
   <div align='center'>
-        <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
-        <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
-        <img src=https://img.shields.io/badge/PyPI-pass-brightgreen.svg >
-        <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.1.2-brightgreen.svg >
+      <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
+      <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
+      <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.1.5-brightgreen.svg >
  </div>
   <p align="center">
     DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT provides <br>a series of training-free, UNet-style cache accelerators for DiT: DBCache, DBPrune, FBCache, etc.
@@ -131,7 +132,7 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 - **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.
 - **Bn**: Further fuses approximate information in the **last n** Transformer blocks to enhance prediction accuracy. These blocks act as an auto-scaler for approximate hidden states that use residual cache.
 - **warmup_steps**: (default: 0) DBCache does not apply the caching strategy when the number of running steps is less than or equal to this value, ensuring the model sufficiently learns basic features during warmup.
-- **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the running steps exceed this value to prevent precision degradation.
+- **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the previous cached steps exceed this value to prevent precision degradation.
 - **residual_diff_threshold**: The value of residual diff threshold, a higher value leads to faster performance at the cost of lower precision.
 
 For a good balance between performance and precision, DBCache is configured by default with **F8B8**, 8 warmup steps, and unlimited cached steps.
@@ -167,8 +168,9 @@ Moreover, users configuring higher **Bn** values (e.g., **F8B16**) while aiming
 cache_options = {
     # 0, 2, 4, ..., 14, 15, etc. [0,16)
     "Bn_compute_blocks_ids": CacheType.range(0, 16, 2),
-    # Skip Bn blocks (1, 3, 5 ,..., etc.) only if the L1 diff 
-    # lower than this value, otherwise, compute it.
+    # If the L1 difference is below this threshold, skip Bn blocks 
+    # not in `Bn_compute_blocks_ids`(1, 3,..., etc), Otherwise, 
+    # compute these blocks.
     "non_compute_blocks_diff_threshold": 0.08,
 }
 ```
@@ -224,12 +226,47 @@ pipe = FluxPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
 ).to("cuda")
 
-# Using DBPrune
+# Using DBPrune with default options
 cache_options = CacheType.default_options(CacheType.DBPrune)
 
 apply_cache_on_pipe(pipe, **cache_options)
 ```
 
+We have also brought the designs from DBCache to DBPrune to make it a more general and customizable block prune algorithm. You can specify the values of **Fn** and **Bn** for higher precision, or set up the non-prune blocks list **non_prune_blocks_ids** to avoid aggressive pruning. For example:
+
+```python
+# Custom options for DBPrune
+cache_options = {
+    "cache_type": CacheType.DBPrune,
+    "residual_diff_threshold": 0.05,
+    # Never prune the first `Fn` and last `Bn` blocks.
+    "Fn_compute_blocks": 8,  # default 1
+    "Bn_compute_blocks": 8,  # default 0
+    "warmup_steps": 8,  # default -1
+    # Disables the pruning strategy when the previous 
+    # pruned steps greater than this value.
+    "max_pruned_steps": 12,  # default, -1 means no limit
+    # Enable dynamic prune threshold within step, higher 
+    # `max_dynamic_prune_threshold` value may introduce a more 
+    # ageressive pruning strategy.
+    "enable_dynamic_prune_threshold": True,
+    "max_dynamic_prune_threshold": 2 * 0.05,
+    # (New thresh) = mean(previous_block_diffs_within_step) * 1.25
+    # (New thresh) = ((New thresh) if (New thresh) <
+    # max_dynamic_prune_threshold else residual_diff_threshold)
+    "dynamic_prune_threshold_relax_ratio": 1.25,
+    # The step interval to update residual cache. For example, 
+    # 2: means the update steps will be [0, 2, 4, ...].
+    "residual_cache_update_interval": 1,
+    # You can set non-prune blocks to avoid ageressive pruning. 
+    # For example, FLUX.1 has 19 + 38 blocks, so we can set it 
+    # to 0, 2, 4, ..., 56, etc.
+    "non_prune_blocks_ids": [],
+}
+
+apply_cache_on_pipe(pipe, **cache_options)
+```
+
 <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.1.2 → cache_dit-0.1.5}/bench/bench.py

@@ -25,6 +25,7 @@ def get_args() -> argparse.ArgumentParser:
     parser.add_argument("--Bn-steps", "--BnS", type=int, default=1)
     parser.add_argument("--warmup-steps", type=int, default=0)
     parser.add_argument("--max-cached-steps", type=int, default=-1)
+    parser.add_argument("--max-pruned-steps", type=int, default=-1)
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument(
         "--compile",
@@ -79,7 +80,7 @@ def get_cache_options(cache_type: CacheType, args: argparse.Namespace):
             "Fn_compute_blocks": args.Fn_compute_blocks,
             "Bn_compute_blocks": args.Bn_compute_blocks,
             "warmup_steps": args.warmup_steps,
-            "max_pruned_steps": args.max_cached_steps,  # -1 means no limit
+            "max_pruned_steps": args.max_pruned_steps,  # -1 means no limit
             # releative token diff threshold, default is 0.0
             "important_condition_threshold": 0.00,
             "enable_dynamic_prune_threshold": (

cache_dit-0.1.5/examples/run_cogvideox.py

@@ -0,0 +1,30 @@
+import torch
+from diffusers import CogVideoXPipeline
+from diffusers.utils import export_to_video
+from cache_dit.cache_factory import apply_cache_on_pipe, CacheType
+
+pipe = CogVideoXPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b",
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+# Default options, F8B8, good balance between performance and precision
+cache_options = CacheType.default_options(CacheType.DBCache)
+
+apply_cache_on_pipe(pipe, **cache_options)
+
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
+
+prompt = "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance."
+video = pipe(
+    prompt=prompt,
+    num_videos_per_prompt=1,
+    num_inference_steps=50,
+    num_frames=49,
+    guidance_scale=6,
+    generator=torch.Generator("cuda").manual_seed(0),
+).frames[0]
+
+print("Saving video to cogvideox.mp4")
+export_to_video(video, "cogvideox.mp4", fps=8)

cache_dit-0.1.5/examples/run_mochi.py

@@ -0,0 +1,25 @@
+import torch
+from diffusers import MochiPipeline
+from diffusers.utils import export_to_video
+from cache_dit.cache_factory import apply_cache_on_pipe, CacheType
+
+pipe = MochiPipeline.from_pretrained(
+    "genmo/mochi-1-preview",
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+# Default options, F8B8, good balance between performance and precision
+cache_options = CacheType.default_options(CacheType.DBCache)
+
+apply_cache_on_pipe(pipe, **cache_options)
+
+pipe.enable_vae_tiling()
+
+prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k."
+video = pipe(
+    prompt,
+    num_frames=84,
+).frames[0]
+
+print("Saving video to mochi.mp4")
+export_to_video(video, "mochi.mp4", fps=30)

{cache_dit-0.1.2 → cache_dit-0.1.5}/src/cache_dit/_version.py

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

{cache_dit-0.1.2 → cache_dit-0.1.5}/src/cache_dit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.1.2
+Version: 0.1.5
 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
@@ -37,13 +37,14 @@ Dynamic: requires-python
   <p align="center">
     <h3>🤗 CacheDiT: A Training-free and Easy-to-use Cache Acceleration <br>Toolbox for Diffusion Transformers</h3>
   </p>
-  <img src=https://github.com/vipshop/cache-dit/raw/dev/assets/cache-dit.png >
+  <img src=https://github.com/vipshop/cache-dit/raw/main/assets/cache-dit.png >
   <div align='center'>
-        <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
-        <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
-        <img src=https://img.shields.io/badge/PyPI-pass-brightgreen.svg >
-        <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.1.2-brightgreen.svg >
+      <img src=https://img.shields.io/badge/Language-Python-brightgreen.svg >
+      <img src=https://img.shields.io/badge/PRs-welcome-9cf.svg >
+      <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.1.5-brightgreen.svg >
  </div>
   <p align="center">
     DeepCache is for UNet not DiT. Most DiT cache speedups are complex and not training-free. CacheDiT provides <br>a series of training-free, UNet-style cache accelerators for DiT: DBCache, DBPrune, FBCache, etc.
@@ -166,7 +167,7 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 - **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.
 - **Bn**: Further fuses approximate information in the **last n** Transformer blocks to enhance prediction accuracy. These blocks act as an auto-scaler for approximate hidden states that use residual cache.
 - **warmup_steps**: (default: 0) DBCache does not apply the caching strategy when the number of running steps is less than or equal to this value, ensuring the model sufficiently learns basic features during warmup.
-- **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the running steps exceed this value to prevent precision degradation.
+- **max_cached_steps**:  (default: -1) DBCache disables the caching strategy when the previous cached steps exceed this value to prevent precision degradation.
 - **residual_diff_threshold**: The value of residual diff threshold, a higher value leads to faster performance at the cost of lower precision.
 
 For a good balance between performance and precision, DBCache is configured by default with **F8B8**, 8 warmup steps, and unlimited cached steps.
@@ -202,8 +203,9 @@ Moreover, users configuring higher **Bn** values (e.g., **F8B16**) while aiming
 cache_options = {
     # 0, 2, 4, ..., 14, 15, etc. [0,16)
     "Bn_compute_blocks_ids": CacheType.range(0, 16, 2),
-    # Skip Bn blocks (1, 3, 5 ,..., etc.) only if the L1 diff 
-    # lower than this value, otherwise, compute it.
+    # If the L1 difference is below this threshold, skip Bn blocks 
+    # not in `Bn_compute_blocks_ids`(1, 3,..., etc), Otherwise, 
+    # compute these blocks.
     "non_compute_blocks_diff_threshold": 0.08,
 }
 ```
@@ -259,12 +261,47 @@ pipe = FluxPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
 ).to("cuda")
 
-# Using DBPrune
+# Using DBPrune with default options
 cache_options = CacheType.default_options(CacheType.DBPrune)
 
 apply_cache_on_pipe(pipe, **cache_options)
 ```
 
+We have also brought the designs from DBCache to DBPrune to make it a more general and customizable block prune algorithm. You can specify the values of **Fn** and **Bn** for higher precision, or set up the non-prune blocks list **non_prune_blocks_ids** to avoid aggressive pruning. For example:
+
+```python
+# Custom options for DBPrune
+cache_options = {
+    "cache_type": CacheType.DBPrune,
+    "residual_diff_threshold": 0.05,
+    # Never prune the first `Fn` and last `Bn` blocks.
+    "Fn_compute_blocks": 8,  # default 1
+    "Bn_compute_blocks": 8,  # default 0
+    "warmup_steps": 8,  # default -1
+    # Disables the pruning strategy when the previous 
+    # pruned steps greater than this value.
+    "max_pruned_steps": 12,  # default, -1 means no limit
+    # Enable dynamic prune threshold within step, higher 
+    # `max_dynamic_prune_threshold` value may introduce a more 
+    # ageressive pruning strategy.
+    "enable_dynamic_prune_threshold": True,
+    "max_dynamic_prune_threshold": 2 * 0.05,
+    # (New thresh) = mean(previous_block_diffs_within_step) * 1.25
+    # (New thresh) = ((New thresh) if (New thresh) <
+    # max_dynamic_prune_threshold else residual_diff_threshold)
+    "dynamic_prune_threshold_relax_ratio": 1.25,
+    # The step interval to update residual cache. For example, 
+    # 2: means the update steps will be [0, 2, 4, ...].
+    "residual_cache_update_interval": 1,
+    # You can set non-prune blocks to avoid ageressive pruning. 
+    # For example, FLUX.1 has 19 + 38 blocks, so we can set it 
+    # to 0, 2, 4, ..., 56, etc.
+    "non_prune_blocks_ids": [],
+}
+
+apply_cache_on_pipe(pipe, **cache_options)
+```
+
 <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.1.2 → cache_dit-0.1.5}/src/cache_dit.egg-info/SOURCES.txt

@@ -40,7 +40,9 @@ bench/.gitignore
 bench/bench.py
 docs/.gitignore
 examples/.gitignore
+examples/run_cogvideox.py
 examples/run_flux.py
+examples/run_mochi.py
 src/cache_dit/__init__.py
 src/cache_dit/_version.py
 src/cache_dit/logger.py