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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.1.3
+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
@@ -44,7 +44,7 @@ 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.1.3-brightgreen.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.
@@ -167,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.
@@ -261,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.3 → cache_dit-0.1.5}/README.md

@@ -9,7 +9,7 @@
       <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.3-brightgreen.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.
@@ -132,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.
@@ -226,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.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.3 → 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.3'
-__version_tuple__ = version_tuple = (0, 1, 3)
+__version__ = version = '0.1.5'
+__version_tuple__ = version_tuple = (0, 1, 5)

{cache_dit-0.1.3 → 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.3
+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
@@ -44,7 +44,7 @@ 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.1.3-brightgreen.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.
@@ -167,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.
@@ -261,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.3 → 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