cache-dit 0.2.3__tar.gz → 0.2.5__tar.gz

{cache_dit-0.2.3 → cache_dit-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.3
+Version: 0.2.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.2.2-brightgreen.svg >
+      <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🔥
@@ -154,6 +154,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 - [🔥Supported Models](#supported)
 - [⚡️Dual Block Cache](#dbcache)
 - [🔥Hybrid TaylorSeer](#taylorseer)
+- [⚡️Hybrid Cache CFG](#cfg)
 - [🎉First Block Cache](#fbcache)
 - [⚡️Dynamic Block Prune](#dbprune)
 - [🎉Context Parallelism](#context-parallelism)  
@@ -168,7 +169,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 You can install the stable release of `cache-dit` from PyPI:
 
 ```bash
-pip3 install cache-dit
+pip3 install -U cache-dit
 ```
 Or you can install the latest develop version from GitHub:
 
@@ -180,11 +181,13 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 
 <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)
+- [🚀FLUX.1-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀FLUX.1-Fill-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀mochi-1-preview](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)
+- [🚀Wan2.1-T2V](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀Wan2.1-FLF2V](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
 
 
@@ -280,7 +283,7 @@ cache_options = {
     "taylorseer_kwargs": {
         "n_derivatives": 2, # default is 2.
     },
-    "warmup_steps": 3, # n_derivatives + 1
+    "warmup_steps": 3, # prefer: >= n_derivatives + 1
     "residual_diff_threshold": 0.12,
 }
 ``` 
@@ -299,6 +302,30 @@ cache_options = {
 |24.85s|12.85s|12.86s|10.27s|10.28s|8.48s|
 |<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/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.12_S14_T12.85s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.12_S14_T12.86s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.15_S17_T10.27s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T10.28s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T8.48s.png width=105px>|
 
+## ⚡️Hybrid Cache CFG
+
+<div id="cfg"></div>
+
+CacheDiT supports caching for **CFG (classifier-free guidance)**. For models that fuse CFG and non-CFG into a single forward step, or models that do not include CFG (classifier-free guidance) in the forward step, please set `do_separate_classifier_free_guidance` param to **False (default)**. Otherwise, set it to True. For examples:
+
+```python
+cache_options = {
+    # CFG: classifier free guidance or not
+    # For model that fused CFG and non-CFG into single forward step,
+    # should set do_separate_classifier_free_guidance as False.
+    # For example, set it as True for Wan 2.1 and set it as False 
+    # for FLUX.1, HunyuanVideo, CogVideoX, Mochi.
+    "do_separate_classifier_free_guidance": True,  # Wan 2.1
+    # Compute cfg forward first or not, default False, namely, 
+    # 0, 2, 4, ..., -> non-CFG step; 1, 3, 5, ... -> CFG step.
+    "cfg_compute_first": False,
+    # Compute spearate diff values for CFG and non-CFG step, 
+    # default True. If False, we will use the computed diff from 
+    # current non-CFG transformer step for current CFG step.
+    "cfg_diff_compute_separate": True,
+}
+```
+
 ## 🎉FBCache: First Block Cache  
 
 <div id="fbcache"></div>

{cache_dit-0.2.3 → cache_dit-0.2.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.2.2-brightgreen.svg >
+      <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🔥
@@ -119,6 +119,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 - [🔥Supported Models](#supported)
 - [⚡️Dual Block Cache](#dbcache)
 - [🔥Hybrid TaylorSeer](#taylorseer)
+- [⚡️Hybrid Cache CFG](#cfg)
 - [🎉First Block Cache](#fbcache)
 - [⚡️Dynamic Block Prune](#dbprune)
 - [🎉Context Parallelism](#context-parallelism)  
@@ -133,7 +134,7 @@ The **CacheDiT** codebase is adapted from [FBCache](https://github.com/chengzeyi
 You can install the stable release of `cache-dit` from PyPI:
 
 ```bash
-pip3 install cache-dit
+pip3 install -U cache-dit
 ```
 Or you can install the latest develop version from GitHub:
 
@@ -145,11 +146,13 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 
 <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)
+- [🚀FLUX.1-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀FLUX.1-Fill-dev](https://github.com/vipshop/cache-dit/raw/main/examples)  
+- [🚀mochi-1-preview](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)
+- [🚀Wan2.1-T2V](https://github.com/vipshop/cache-dit/raw/main/examples)
+- [🚀Wan2.1-FLF2V](https://github.com/vipshop/cache-dit/raw/main/examples)
 - [🚀HunyuanVideo](https://github.com/vipshop/cache-dit/raw/main/examples)
 
 
@@ -245,7 +248,7 @@ cache_options = {
     "taylorseer_kwargs": {
         "n_derivatives": 2, # default is 2.
     },
-    "warmup_steps": 3, # n_derivatives + 1
+    "warmup_steps": 3, # prefer: >= n_derivatives + 1
     "residual_diff_threshold": 0.12,
 }
 ``` 
@@ -264,6 +267,30 @@ cache_options = {
 |24.85s|12.85s|12.86s|10.27s|10.28s|8.48s|
 |<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/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.12_S14_T12.85s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.12_S14_T12.86s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T0ET0_R0.15_S17_T10.27s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C0_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T10.28s.png width=105px>|<img src=https://github.com/vipshop/cache-dit/raw/main/assets/U0_C1_DBCACHE_F1B0S1W0T1ET1_R0.15_S17_T8.48s.png width=105px>|
 
+## ⚡️Hybrid Cache CFG
+
+<div id="cfg"></div>
+
+CacheDiT supports caching for **CFG (classifier-free guidance)**. For models that fuse CFG and non-CFG into a single forward step, or models that do not include CFG (classifier-free guidance) in the forward step, please set `do_separate_classifier_free_guidance` param to **False (default)**. Otherwise, set it to True. For examples:
+
+```python
+cache_options = {
+    # CFG: classifier free guidance or not
+    # For model that fused CFG and non-CFG into single forward step,
+    # should set do_separate_classifier_free_guidance as False.
+    # For example, set it as True for Wan 2.1 and set it as False 
+    # for FLUX.1, HunyuanVideo, CogVideoX, Mochi.
+    "do_separate_classifier_free_guidance": True,  # Wan 2.1
+    # Compute cfg forward first or not, default False, namely, 
+    # 0, 2, 4, ..., -> non-CFG step; 1, 3, 5, ... -> CFG step.
+    "cfg_compute_first": False,
+    # Compute spearate diff values for CFG and non-CFG step, 
+    # default True. If False, we will use the computed diff from 
+    # current non-CFG transformer step for current CFG step.
+    "cfg_diff_compute_separate": True,
+}
+```
+
 ## 🎉FBCache: First Block Cache  
 
 <div id="fbcache"></div>

{cache_dit-0.2.3 → cache_dit-0.2.5}/examples/.gitignore

@@ -166,3 +166,4 @@ report*.html
 .DS_Store
 *.png
 *.mp4
+tmp*

cache_dit-0.2.5/examples/README.md

@@ -0,0 +1,65 @@
+# Examples for CacheDiT  
+
+## Install requirements  
+
+```bash
+pip3 install -r requirements.txt
+```
+
+## Run examples  
+
+- FLUX.1-dev 
+
+```bash
+python3 run_flux.py # baseline
+python3 run_flux.py --cache --Fn 8 --Bn 8
+python3 run_flux.py --cache --Fn 8 --Bn 0 --taylorseer
+```
+
+- FLUX.1-Fill-dev 
+
+```bash
+python3 run_flux_fill.py # baseline
+python3 run_flux_fill.py --cache --Fn 8 --Bn 8
+python3 run_flux_fill.py --cache --Fn 8 --Bn 0 --taylorseer
+```
+
+- CogVideoX 
+
+```bash
+python3 run_cogvideox.py # baseline
+python3 run_cogvideox.py --cache --Fn 8 --Bn 8
+python3 run_cogvideox.py --cache --Fn 8 --Bn 0 --taylorseer
+```
+
+- Wan2.1 T2V
+
+```bash
+python3 run_wan.py # baseline
+python3 run_wan.py --cache --Fn 8 --Bn 8
+python3 run_wan.py --cache --Fn 8 --Bn 0 --taylorseer
+```
+
+- Wan2.1 FLF2V
+
+```bash
+python3 run_wan_flf2v.py # baseline
+python3 run_wan_flf2v.py --cache --Fn 8 --Bn 8
+python3 run_wan_flf2v.py --cache --Fn 8 --Bn 0 --taylorseer
+```
+
+- mochi-1-preview
+
+```bash
+python3 run_mochi.py # baseline
+python3 run_mochi.py --cache --Fn 8 --Bn 8
+python3 run_mochi.py --cache --Fn 8 --Bn 0 --taylorseer
+```
+
+- HunyuanVideo
+
+```bash
+python3 run_hunyuan_video.py # baseline
+python3 run_hunyuan_video.py --cache --Fn 8 --Bn 8
+python3 run_hunyuan_video.py --cache --Fn 8 --Bn 0 --taylorseer
+```

{cache_dit-0.2.3 → cache_dit-0.2.5}/examples/requirements.txt

@@ -1,4 +1,4 @@
 imageio-ffmpeg
 # wan currently requires installing from source
-diffusers @ git+https://github.com/huggingface/diffusers
+diffusers>=0.34.0
 ftfy

cache_dit-0.2.5/examples/run_cogvideox.py

@@ -0,0 +1,142 @@
+import os
+import time
+import torch
+import argparse
+from diffusers.utils import export_to_video
+from diffusers import CogVideoXPipeline, AutoencoderKLCogVideoX
+from cache_dit.cache_factory import apply_cache_on_pipe, CacheType
+
+
+def get_args() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    # General arguments
+    parser.add_argument("--cache", action="store_true", default=False)
+    parser.add_argument("--taylorseer", action="store_true", default=False)
+    parser.add_argument("--taylorseer-order", "--order", type=int, default=2)
+    parser.add_argument("--Fn-compute-blocks", "--Fn", type=int, default=1)
+    parser.add_argument("--Bn-compute-blocks", "--Bn", type=int, default=0)
+    parser.add_argument("--rdt", type=float, default=0.08)
+    parser.add_argument("--warmup-steps", type=int, default=0)
+    return parser.parse_args()
+
+
+args = get_args()
+print(args)
+
+
+model_id = os.environ.get("COGVIDEOX_DIR", "THUDM/CogVideoX-5b")
+
+
+def is_cogvideox_1_5():
+    return "CogVideoX1.5" in model_id or "THUDM/CogVideoX1.5" in model_id
+
+
+def get_gpu_memory_in_gib():
+    if not torch.cuda.is_available():
+        return 0
+
+    try:
+        total_memory_bytes = torch.cuda.get_device_properties(
+            torch.cuda.current_device(),
+        ).total_memory
+        total_memory_gib = total_memory_bytes / (1024**3)
+        return int(total_memory_gib)
+    except Exception:
+        return 0
+
+
+pipe = CogVideoXPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+
+if args.cache:
+    cache_options = {
+        "cache_type": CacheType.DBCache,
+        "warmup_steps": args.warmup_steps,
+        "max_cached_steps": -1,  # -1 means no limit
+        # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
+        "Fn_compute_blocks": args.Fn_compute_blocks,  # Fn, F8, etc.
+        "Bn_compute_blocks": args.Bn_compute_blocks,  # Bn, B16, etc.
+        "residual_diff_threshold": args.rdt,
+        # releative token diff threshold, default is 0.0
+        "important_condition_threshold": 0.05,
+        # CFG: classifier free guidance or not
+        # CogVideoX fused CFG and non-CFG into single forward step
+        # so, we set do_separate_classifier_free_guidance as False.
+        "do_separate_classifier_free_guidance": False,
+        "cfg_compute_first": False,
+        "enable_taylorseer": args.taylorseer,
+        "enable_encoder_taylorseer": args.taylorseer,
+        # Taylorseer cache type cache be hidden_states or residual
+        "taylorseer_cache_type": "hidden_states",
+        "taylorseer_kwargs": {
+            "n_derivatives": args.taylorseer_order,
+        },
+    }
+    cache_type_str = "DBCACHE"
+    cache_type_str = (
+        f"{cache_type_str}_F{args.Fn_compute_blocks}"
+        f"B{args.Bn_compute_blocks}W{args.warmup_steps}"
+        f"T{int(args.taylorseer)}O{args.taylorseer_order}"
+    )
+    print(f"cache options:\n{cache_options}")
+
+    apply_cache_on_pipe(pipe, **cache_options)
+else:
+    cache_type_str = "NONE"
+
+
+pipe.enable_model_cpu_offload()
+assert isinstance(pipe.vae, AutoencoderKLCogVideoX)  # enable type check for IDE
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
+
+start = time.time()
+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=(
+        # Avoid OOM for CogVideoX1.5 model on 48GB GPU
+        16
+        if (is_cogvideox_1_5() and get_gpu_memory_in_gib() < 48)
+        else 49
+    ),
+    guidance_scale=6,
+    generator=torch.Generator("cpu").manual_seed(0),
+).frames[0]
+end = time.time()
+
+if hasattr(pipe.transformer, "_cached_steps"):
+    cached_steps = pipe.transformer._cached_steps
+    residual_diffs = pipe.transformer._residual_diffs
+    print(f"Cache Steps: {len(cached_steps)}, {cached_steps}")
+    print(f"Residual Diffs: {len(residual_diffs)}, {residual_diffs}")
+if hasattr(pipe.transformer, "_cfg_cached_steps"):
+    cfg_cached_steps = pipe.transformer._cfg_cached_steps
+    cfg_residual_diffs = pipe.transformer._cfg_residual_diffs
+    print(f"CFG Cache Steps: {len(cfg_cached_steps)}, {cfg_cached_steps} ")
+    print(
+        f"CFG Residual Diffs: {len(cfg_residual_diffs)}, {cfg_residual_diffs}"
+    )
+
+time_cost = end - start
+save_path = f"cogvideox.{cache_type_str}.mp4"
+print(f"Time cost: {time_cost:.2f}s")
+print(f"Saving video to {save_path}")
+export_to_video(video, save_path, fps=8)

cache_dit-0.2.5/examples/run_flux.py

@@ -0,0 +1,96 @@
+import os
+import time
+import torch
+import argparse
+from diffusers import FluxPipeline
+from cache_dit.cache_factory import apply_cache_on_pipe, CacheType
+
+
+def get_args() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    # General arguments
+    parser.add_argument("--cache", action="store_true", default=False)
+    parser.add_argument("--taylorseer", action="store_true", default=False)
+    parser.add_argument("--taylorseer-order", "--order", type=int, default=2)
+    parser.add_argument("--Fn-compute-blocks", "--Fn", type=int, default=1)
+    parser.add_argument("--Bn-compute-blocks", "--Bn", type=int, default=0)
+    parser.add_argument("--rdt", type=float, default=0.08)
+    parser.add_argument("--warmup-steps", type=int, default=0)
+    return parser.parse_args()
+
+
+args = get_args()
+print(args)
+
+
+pipe = FluxPipeline.from_pretrained(
+    os.environ.get(
+        "FLUX_DIR",
+        "black-forest-labs/FLUX.1-dev",
+    ),
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+
+if args.cache:
+    cache_options = {
+        "cache_type": CacheType.DBCache,
+        "warmup_steps": args.warmup_steps,
+        "max_cached_steps": -1,  # -1 means no limit
+        # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
+        "Fn_compute_blocks": args.Fn_compute_blocks,  # Fn, F8, etc.
+        "Bn_compute_blocks": args.Bn_compute_blocks,  # Bn, B16, etc.
+        "residual_diff_threshold": args.rdt,
+        # CFG: classifier free guidance or not
+        # FLUX.1 dev don not have CFG, so, we set
+        # do_separate_classifier_free_guidance as False.
+        "do_separate_classifier_free_guidance": False,
+        "cfg_compute_first": False,
+        "enable_taylorseer": args.taylorseer,
+        "enable_encoder_taylorseer": args.taylorseer,
+        # Taylorseer cache type cache be hidden_states or residual
+        "taylorseer_cache_type": "residual",
+        "taylorseer_kwargs": {
+            "n_derivatives": args.taylorseer_order,
+        },
+    }
+    cache_type_str = "DBCACHE"
+    cache_type_str = (
+        f"{cache_type_str}_F{args.Fn_compute_blocks}"
+        f"B{args.Bn_compute_blocks}W{args.warmup_steps}"
+        f"T{int(args.taylorseer)}O{args.taylorseer_order}"
+    )
+    print(f"cache options:\n{cache_options}")
+
+    apply_cache_on_pipe(pipe, **cache_options)
+else:
+    cache_type_str = "NONE"
+
+
+start = time.time()
+image = pipe(
+    "A cat holding a sign that says hello world",
+    num_inference_steps=28,
+    generator=torch.Generator("cpu").manual_seed(0),
+).images[0]
+
+end = time.time()
+
+if hasattr(pipe.transformer, "_cached_steps"):
+    cached_steps = pipe.transformer._cached_steps
+    residual_diffs = pipe.transformer._residual_diffs
+    print(f"Cache Steps: {len(cached_steps)}, {cached_steps}")
+    print(f"Residual Diffs: {len(residual_diffs)}, {residual_diffs}")
+if hasattr(pipe.transformer, "_cfg_cached_steps"):
+    cfg_cached_steps = pipe.transformer._cfg_cached_steps
+    cfg_residual_diffs = pipe.transformer._cfg_residual_diffs
+    print(f"CFG Cache Steps: {len(cfg_cached_steps)}, {cfg_cached_steps} ")
+    print(
+        f"CFG Residual Diffs: {len(cfg_residual_diffs)}, {cfg_residual_diffs}"
+    )
+
+time_cost = end - start
+save_path = f"flux.{cache_type_str}.png"
+print(f"Time cost: {time_cost:.2f}s")
+print(f"Saving image to {save_path}")
+image.save(save_path)

cache_dit-0.2.5/examples/run_flux_fill.py

@@ -0,0 +1,100 @@
+import os
+import time
+import torch
+import argparse
+from diffusers import FluxFillPipeline
+from diffusers.utils import load_image
+from cache_dit.cache_factory import apply_cache_on_pipe, CacheType
+
+
+def get_args() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    # General arguments
+    parser.add_argument("--cache", action="store_true", default=False)
+    parser.add_argument("--taylorseer", action="store_true", default=False)
+    parser.add_argument("--taylorseer-order", "--order", type=int, default=2)
+    parser.add_argument("--Fn-compute-blocks", "--Fn", type=int, default=1)
+    parser.add_argument("--Bn-compute-blocks", "--Bn", type=int, default=0)
+    parser.add_argument("--rdt", type=float, default=0.08)
+    parser.add_argument("--warmup-steps", type=int, default=0)
+    return parser.parse_args()
+
+
+args = get_args()
+print(args)
+
+
+pipe = FluxFillPipeline.from_pretrained(
+    os.environ.get(
+        "FLUX_FILL_DIR",
+        "black-forest-labs/FLUX.1-Fill-dev",
+    ),
+    torch_dtype=torch.bfloat16,
+).to("cuda")
+
+
+if args.cache:
+    cache_options = {
+        "cache_type": CacheType.DBCache,
+        "warmup_steps": args.warmup_steps,
+        "max_cached_steps": -1,  # -1 means no limit
+        # Fn=1, Bn=0, means FB Cache, otherwise, Dual Block Cache
+        "Fn_compute_blocks": args.Fn_compute_blocks,  # Fn, F8, etc.
+        "Bn_compute_blocks": args.Bn_compute_blocks,  # Bn, B16, etc.
+        "residual_diff_threshold": args.rdt,
+        # CFG: classifier free guidance or not
+        # FLUX.1 dev don not have CFG, so, we set
+        # do_separate_classifier_free_guidance as False.
+        "do_separate_classifier_free_guidance": False,
+        "cfg_compute_first": False,
+        "enable_taylorseer": args.taylorseer,
+        "enable_encoder_taylorseer": args.taylorseer,
+        # Taylorseer cache type cache be hidden_states or residual
+        "taylorseer_cache_type": "residual",
+        "taylorseer_kwargs": {
+            "n_derivatives": args.taylorseer_order,
+        },
+    }
+    cache_type_str = "DBCACHE"
+    cache_type_str = (
+        f"{cache_type_str}_F{args.Fn_compute_blocks}"
+        f"B{args.Bn_compute_blocks}W{args.warmup_steps}"
+        f"T{int(args.taylorseer)}O{args.taylorseer_order}"
+    )
+    print(f"cache options:\n{cache_options}")
+
+    apply_cache_on_pipe(pipe, **cache_options)
+else:
+    cache_type_str = "NONE"
+
+start = time.time()
+image = pipe(
+    prompt="a white paper cup",
+    image=load_image("data/cup.png"),
+    mask_image=load_image("data/cup_mask.png"),
+    guidance_scale=30,
+    num_inference_steps=28,
+    max_sequence_length=512,
+    generator=torch.Generator("cpu").manual_seed(0),
+).images[0]
+
+end = time.time()
+
+if hasattr(pipe.transformer, "_cached_steps"):
+    cached_steps = pipe.transformer._cached_steps
+    residual_diffs = pipe.transformer._residual_diffs
+    print(f"Cache Steps: {len(cached_steps)}, {cached_steps}")
+    print(f"Residual Diffs: {len(residual_diffs)}, {residual_diffs}")
+if hasattr(pipe.transformer, "_cfg_cached_steps"):
+    cfg_cached_steps = pipe.transformer._cfg_cached_steps
+    cfg_residual_diffs = pipe.transformer._cfg_residual_diffs
+    print(f"CFG Cache Steps: {len(cfg_cached_steps)}, {cfg_cached_steps} ")
+    print(
+        f"CFG Residual Diffs: {len(cfg_residual_diffs)}, {cfg_residual_diffs}"
+    )
+
+time_cost = end - start
+save_path = f"flux-fill.{cache_type_str}.png"
+print(f"Time cost: {time_cost:.2f}s")
+print(f"Saving image to {save_path}")
+image.save(save_path)