cache-dit 0.2.14__py3-none-any.whl → 0.2.15__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.14'
-__version_tuple__ = version_tuple = (0, 2, 14)
+__version__ = version = '0.2.15'
+__version_tuple__ = version_tuple = (0, 2, 15)

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/__init__.py

@@ -15,6 +15,8 @@ def apply_db_cache_on_transformer(transformer, *args, **kwargs):
         adapter_name = "wan"
     elif transformer_cls_name.startswith("HunyuanVideo"):
         adapter_name = "hunyuan_video"
+    elif transformer_cls_name.startswith("QwenImage"):
+        adapter_name = "qwen_image"
     else:
         raise ValueError(
             f"Unknown transformer class name: {transformer_cls_name}"
@@ -41,6 +43,8 @@ def apply_db_cache_on_pipe(pipe: DiffusionPipeline, *args, **kwargs):
         adapter_name = "wan"
     elif pipe_cls_name.startswith("HunyuanVideo"):
         adapter_name = "hunyuan_video"
+    elif pipe_cls_name.startswith("QwenImage"):
+        adapter_name = "qwen_image"
     else:
         raise ValueError(f"Unknown pipeline class name: {pipe_cls_name}")
 

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/qwen_image.py

@@ -0,0 +1,88 @@
+import functools
+import unittest
+
+import torch
+from diffusers import QwenImagePipeline, QwenImageTransformer2DModel
+
+from cache_dit.cache_factory.dual_block_cache import cache_context
+
+
+def apply_db_cache_on_transformer(
+    transformer: QwenImageTransformer2DModel,
+):
+    if getattr(transformer, "_is_cached", False):
+        return transformer
+
+    transformer_blocks = torch.nn.ModuleList(
+        [
+            cache_context.DBCachedTransformerBlocks(
+                transformer.transformer_blocks,
+                transformer=transformer,
+                return_hidden_states_first=False,
+            )
+        ]
+    )
+
+    original_forward = transformer.forward
+
+    @functools.wraps(transformer.__class__.forward)
+    def new_forward(
+        self,
+        *args,
+        **kwargs,
+    ):
+        with unittest.mock.patch.object(
+            self,
+            "transformer_blocks",
+            transformer_blocks,
+        ):
+            return original_forward(
+                *args,
+                **kwargs,
+            )
+
+    transformer.forward = new_forward.__get__(transformer)
+
+    transformer._is_cached = True
+
+    return transformer
+
+
+def apply_db_cache_on_pipe(
+    pipe: QwenImagePipeline,
+    *,
+    shallow_patch: bool = False,
+    residual_diff_threshold=0.06,
+    downsample_factor=1,
+    warmup_steps=0,
+    max_cached_steps=-1,
+    **kwargs,
+):
+    cache_kwargs, kwargs = cache_context.collect_cache_kwargs(
+        default_attrs={
+            "residual_diff_threshold": residual_diff_threshold,
+            "downsample_factor": downsample_factor,
+            "warmup_steps": warmup_steps,
+            "max_cached_steps": max_cached_steps,
+        },
+        **kwargs,
+    )
+    if not getattr(pipe, "_is_cached", False):
+        original_call = pipe.__class__.__call__
+
+        @functools.wraps(original_call)
+        def new_call(self, *args, **kwargs):
+            with cache_context.cache_context(
+                cache_context.create_cache_context(
+                    **cache_kwargs,
+                )
+            ):
+                return original_call(self, *args, **kwargs)
+
+        pipe.__class__.__call__ = new_call
+        pipe.__class__._is_cached = True
+
+    if not shallow_patch:
+        apply_db_cache_on_transformer(pipe.transformer)
+
+    return pipe

cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py

@@ -2,7 +2,7 @@ import functools
 import unittest
 
 import torch
-from diffusers import DiffusionPipeline, WanTransformer3DModel
+from diffusers import WanPipeline, WanTransformer3DModel
 
 from cache_dit.cache_factory.dual_block_cache import cache_context
 
@@ -49,7 +49,7 @@ def apply_db_cache_on_transformer(
 
 
 def apply_db_cache_on_pipe(
-    pipe: DiffusionPipeline,
+    pipe: WanPipeline,
     *,
     shallow_patch: bool = False,
     residual_diff_threshold=0.03,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 0.2.14
+Version: 0.2.15
 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
@@ -52,20 +52,14 @@ Dynamic: requires-python
       <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-brightgreen.svg >
  </div>
-  <p align="center">
-    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>
+   <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>🔥
 </div>
 
-<div align="center">
-  <p align="center">
-    <b>♥️ Please consider to leave a ⭐️ Star to support us ~ ♥️</b>
-  </p>
-</div> 
-
-## 🔥News🔥
-- [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.
+## 🔥News
+- [2025-08-11] 🔥[Qwen-Image](./examples/run_qwen_image.py) is supported! Please check [run_qwen_image.py](./examples/run_qwen_image.py) as an example.
+- [2025-08-10] 🔥[FLUX.1-Kontext-dev](./examples/run_flux_kontext.py) is supported! Please check [run_flux_kontext.py](./examples/run_flux_kontext.py) as an example.
+- [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! **3.3x** speedup for FLUX.1 on NVIDIA L20 with `cache-dit`.
 
 ## 📖Contents 
 
@@ -104,8 +98,10 @@ pip3 install git+https://github.com/vipshop/cache-dit.git
 
 <div id="supported"></div>
 
+- [🚀Qwen-Image](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)  
+- [🚀FLUX.1-Kontext-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)

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

@@ -1,5 +1,5 @@
 cache_dit/__init__.py,sha256=0-B173-fLi3IA8nJXoS71zK0zD33Xplysd9skmLfEOY,171
-cache_dit/_version.py,sha256=ut2sCt69XoYh0A1_KAmfCg1IKkN6zwqhu2eMFWAhMbQ,513
+cache_dit/_version.py,sha256=TspyaLI34cMH7HNr-lrUKFnY4CF39yGfG_421vC3fhg,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
@@ -9,12 +9,13 @@ cache_dit/cache_factory/taylorseer.py,sha256=LKSNo2ode69EVo9xrxjxAMEjz0yDGiGADeD
 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=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/__init__.py,sha256=uSqF5aD2-feHB25vEbx1STBQVjWVAOn_wYTdAEmS4NU,2045
 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
 cache_dit/cache_factory/dual_block_cache/diffusers_adapters/hunyuan_video.py,sha256=SO4q39PQuQ5QVHy5Z-ubiKdstzvQPedONN2J5oiGUh0,9955
 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/dual_block_cache/diffusers_adapters/qwen_image.py,sha256=ZIjB4GFIL0_xY8FkXdOJJ9-Xcft54rnBCrz43VWZLi0,2296
+cache_dit/cache_factory/dual_block_cache/diffusers_adapters/wan.py,sha256=EH2PpYFJ76KQLb35za6T-nM8Q10owLNatv6cd480ydE,2584
 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=1qarKAsEFiaaN2_ghko2dqGz_R7BTQSOyGtb_eQq38Y,35716
 cache_dit/cache_factory/dynamic_block_prune/diffusers_adapters/__init__.py,sha256=hVBTXj9MMGFGVezT3j8MntFRBiphSaUL4YhSOd8JtuY,1870
@@ -41,9 +42,9 @@ cache_dit/metrics/fid.py,sha256=9Ivtazl6mW0Bon2VXa-Ia5Xj2ewxRD3V1Qkd69zYM3Y,1706
 cache_dit/metrics/inception.py,sha256=pBVe2X6ylLPIXTG4-GWDM9DWnCviMJbJ45R3ulhktR0,12759
 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,,
+cache_dit-0.2.15.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-0.2.15.dist-info/METADATA,sha256=7cxw9ZaYpCFbtA1iDt-CWhk07mDXuRPECxfuO-wB0IE,25153
+cache_dit-0.2.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-0.2.15.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-0.2.15.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-0.2.15.dist-info/RECORD,,