cache-dit 1.0.4__py3-none-any.whl → 1.0.5__py3-none-any.whl

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0.4'
-__version_tuple__ = version_tuple = (1, 0, 4)
+__version__ = version = '1.0.5'
+__version_tuple__ = version_tuple = (1, 0, 5)
 
 __commit_id__ = commit_id = None

cache_dit/cache_factory/block_adapters/__init__.py

@@ -12,7 +12,10 @@ def flux_adapter(pipe, **kwargs) -> BlockAdapter:
     from cache_dit.utils import is_diffusers_at_least_0_3_5
 
     assert isinstance(pipe.transformer, FluxTransformer2DModel)
-    if is_diffusers_at_least_0_3_5():
+    transformer_cls_name: str = pipe.transformer.__class__.__name__
+    if is_diffusers_at_least_0_3_5() and not transformer_cls_name.startswith(
+        "Nunchaku"
+    ):
         return BlockAdapter(
             pipe=pipe,
             transformer=pipe.transformer,

cache_dit/cache_factory/cache_adapters/cache_adapter.py

@@ -1,3 +1,4 @@
+import copy
 import torch
 import unittest
 import functools
@@ -197,7 +198,6 @@ class CachedAdapter:
         flatten_contexts, contexts_kwargs = cls.modify_context_params(
             block_adapter, **context_kwargs
         )
-
         original_call = block_adapter.pipe.__class__.__call__
 
         @functools.wraps(original_call)
@@ -238,7 +238,7 @@ class CachedAdapter:
             block_adapter.unique_blocks_name
         )
         contexts_kwargs = [
-            context_kwargs.copy()
+            copy.deepcopy(context_kwargs)  # must deep copy
             for _ in range(
                 len(flatten_contexts),
             )
@@ -259,9 +259,41 @@ class CachedAdapter:
         for i in range(
             min(len(contexts_kwargs), len(flatten_modifiers)),
         ):
-            contexts_kwargs[i].update(
-                flatten_modifiers[i]._context_kwargs,
-            )
+            if "cache_config" in flatten_modifiers[i]._context_kwargs:
+                modifier_cache_config = flatten_modifiers[
+                    i
+                ]._context_kwargs.get("cache_config", None)
+                modifier_calibrator_config = flatten_modifiers[
+                    i
+                ]._context_kwargs.get("calibrator_config", None)
+                if modifier_cache_config is not None:
+                    assert isinstance(
+                        modifier_cache_config, BasicCacheConfig
+                    ), (
+                        f"cache_config must be BasicCacheConfig, but got "
+                        f"{type(modifier_cache_config)}."
+                    )
+                    contexts_kwargs[i]["cache_config"].update(
+                        **modifier_cache_config.as_dict()
+                    )
+                if modifier_calibrator_config is not None:
+                    assert isinstance(
+                        modifier_calibrator_config, CalibratorConfig
+                    ), (
+                        f"calibrator_config must be CalibratorConfig, but got "
+                        f"{type(modifier_calibrator_config)}."
+                    )
+                    if (
+                        contexts_kwargs[i].get("calibrator_config", None)
+                        is None
+                    ):
+                        contexts_kwargs[i][
+                            "calibrator_config"
+                        ] = modifier_calibrator_config
+                    else:
+                        contexts_kwargs[i]["calibrator_config"].update(
+                            **modifier_calibrator_config.as_dict()
+                        )
             cls._config_messages(**contexts_kwargs[i])
 
         return flatten_contexts, contexts_kwargs

cache_dit/cache_factory/cache_contexts/cache_config.py

@@ -60,9 +60,25 @@ class BasicCacheConfig:
     def update(self, **kwargs) -> "BasicCacheConfig":
         for key, value in kwargs.items():
             if hasattr(self, key):
-                setattr(self, key, value)
+                if value is not None:
+                    setattr(self, key, value)
         return self
 
+    def empty(self, **kwargs) -> "BasicCacheConfig":
+        # Set all fields to None
+        for field in dataclasses.fields(self):
+            if hasattr(self, field.name):
+                setattr(self, field.name, None)
+        if kwargs:
+            self.update(**kwargs)
+        return self
+
+    def reset(self, **kwargs) -> "BasicCacheConfig":
+        return self.empty(**kwargs)
+
+    def as_dict(self) -> dict:
+        return dataclasses.asdict(self)
+
     def strify(self) -> str:
         return (
             f"{self.cache_type}_"

cache_dit/cache_factory/cache_contexts/calibrators/__init__.py

@@ -45,6 +45,28 @@ class CalibratorConfig:
     def to_kwargs(self) -> Dict:
         return self.calibrator_kwargs.copy()
 
+    def as_dict(self) -> dict:
+        return dataclasses.asdict(self)
+
+    def update(self, **kwargs) -> "CalibratorConfig":
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                if value is not None:
+                    setattr(self, key, value)
+        return self
+
+    def empty(self, **kwargs) -> "CalibratorConfig":
+        # Set all fields to None
+        for field in dataclasses.fields(self):
+            if hasattr(self, field.name):
+                setattr(self, field.name, None)
+        if kwargs:
+            self.update(**kwargs)
+        return self
+
+    def reset(self, **kwargs) -> "CalibratorConfig":
+        return self.empty(**kwargs)
+
 
 @dataclasses.dataclass
 class TaylorSeerCalibratorConfig(CalibratorConfig):

cache_dit/cache_factory/cache_contexts/prune_config.py

@@ -50,12 +50,6 @@ class DBPruneConfig(BasicCacheConfig):
     #     to at least 2 to reduce the VRAM usage of the calibrator.
     force_reduce_calibrator_vram: bool = False
 
-    def update(self, **kwargs) -> "DBPruneConfig":
-        for key, value in kwargs.items():
-            if hasattr(self, key):
-                setattr(self, key, value)
-        return self
-
     def strify(self) -> str:
         return (
             f"{self.cache_type}_"

cache_dit/quantize/quantize_ao.py

@@ -83,11 +83,18 @@ def quantize_ao(
     def _quantization_fn():
         try:
             if quant_type == "fp8_w8a8_dq":
-                from torchao.quantization import (
-                    float8_dynamic_activation_float8_weight,
-                    PerTensor,
-                    PerRow,
-                )
+                try:
+                    from torchao.quantization import (
+                        float8_dynamic_activation_float8_weight,
+                        PerTensor,
+                        PerRow,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Float8DynamicActivationFloat8WeightConfig as float8_dynamic_activation_float8_weight,
+                        PerTensor,
+                        PerRow,
+                    )
 
                 if per_row:  # Ensure bfloat16
                     module.to(torch.bfloat16)
@@ -109,7 +116,12 @@ def quantize_ao(
                 )
 
             elif quant_type == "fp8_w8a16_wo":
-                from torchao.quantization import float8_weight_only
+                try:
+                    from torchao.quantization import float8_weight_only
+                except ImportError:
+                    from torchao.quantization import (
+                        Float8WeightOnlyConfig as float8_weight_only,
+                    )
 
                 quantization_fn = float8_weight_only(
                     weight_dtype=kwargs.get(
@@ -119,14 +131,25 @@ def quantize_ao(
                 )
 
             elif quant_type == "int8_w8a8_dq":
-                from torchao.quantization import (
-                    int8_dynamic_activation_int8_weight,
-                )
+                try:
+                    from torchao.quantization import (
+                        int8_dynamic_activation_int8_weight,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Int8DynamicActivationInt8WeightConfig as int8_dynamic_activation_int8_weight,
+                    )
 
                 quantization_fn = int8_dynamic_activation_int8_weight()
 
             elif quant_type == "int8_w8a16_wo":
-                from torchao.quantization import int8_weight_only
+
+                try:
+                    from torchao.quantization import int8_weight_only
+                except ImportError:
+                    from torchao.quantization import (
+                        Int8WeightOnlyConfig as int8_weight_only,
+                    )
 
                 quantization_fn = int8_weight_only(
                     # group_size is None -> per_channel, else per group
@@ -134,23 +157,41 @@ def quantize_ao(
                 )
 
             elif quant_type == "int4_w4a8_dq":
-                from torchao.quantization import (
-                    int8_dynamic_activation_int4_weight,
-                )
+
+                try:
+                    from torchao.quantization import (
+                        int8_dynamic_activation_int4_weight,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Int8DynamicActivationInt4WeightConfig as int8_dynamic_activation_int4_weight,
+                    )
 
                 quantization_fn = int8_dynamic_activation_int4_weight(
                     group_size=kwargs.get("group_size", 32),
                 )
 
             elif quant_type == "int4_w4a4_dq":
-                from torchao.quantization import (
-                    int4_dynamic_activation_int4_weight,
-                )
+
+                try:
+                    from torchao.quantization import (
+                        int4_dynamic_activation_int4_weight,
+                    )
+                except ImportError:
+                    from torchao.quantization import (
+                        Int4DynamicActivationInt4WeightConfig as int4_dynamic_activation_int4_weight,
+                    )
 
                 quantization_fn = int4_dynamic_activation_int4_weight()
 
             elif quant_type == "int4_w4a16_wo":
-                from torchao.quantization import int4_weight_only
+
+                try:
+                    from torchao.quantization import int4_weight_only
+                except ImportError:
+                    from torchao.quantization import (
+                        Int4WeightOnlyConfig as int4_weight_only,
+                    )
 
                 quantization_fn = int4_weight_only(
                     group_size=kwargs.get("group_size", 32),

{cache_dit-1.0.4.dist-info → cache_dit-1.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 1.0.4
+Version: 1.0.5
 Summary: A Unified, Flexible and Training-free Cache Acceleration Framework for 🤗Diffusers.
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
@@ -194,7 +194,7 @@ You can install the stable release of cache-dit from PyPI, or the latest develop
 - **[🎉Easy New Model Integration](./docs/User_Guide.md#automatic-block-adapter)**: Features like **Unified Cache APIs**, **Forward Pattern Matching**, **Automatic Block Adapter**, **Hybrid Forward Pattern**, and **Patch Functor** make it highly functional and flexible. For example, we achieved 🎉 Day 1 support for [HunyuanImage-2.1](https://github.com/Tencent-Hunyuan/HunyuanImage-2.1) with 1.7x speedup w/o precision loss—even before it was available in the Diffusers library.  
 - **[🎉State-of-the-Art Performance](./bench/)**: Compared with algorithms including Δ-DiT, Chipmunk, FORA, DuCa, TaylorSeer and FoCa, cache-dit achieved the **SOTA** performance w/ **7.4x↑🎉** speedup on ClipScore!
 - **[🎉Support for 4/8-Steps Distilled Models](./bench/)**: Surprisingly, cache-dit's **DBCache** works for extremely few-step distilled models—something many other methods fail to do.  
-- **[🎉Compatibility with Other Optimizations](./docs/User_Guide.md#️torch-compile)**: Designed to work seamlessly with torch.compile, model CPU offload, sequential CPU offload, group offloading, etc.  
+- **[🎉Compatibility with Other Optimizations](./docs/User_Guide.md#️torch-compile)**: Designed to work seamlessly with torch.compile, model CPU offload, sequential CPU offload, group offloading, Quantization(**[torchao](./examples/quantize/)**, **[🔥nunchaku](./examples/quantize/)**), etc.  
 - **[🎉Hybrid Cache Acceleration](./docs/User_Guide.md#taylorseer-calibrator)**: Now supports hybrid **Block-wise Cache + Calibrator** schemes (e.g., DBCache or DBPrune + TaylorSeerCalibrator). DBCache or DBPrune acts as the **Indicator** to decide *when* to cache, while the Calibrator decides *how* to cache. More mainstream cache acceleration algorithms (e.g., FoCa) will be supported in the future, along with additional benchmarks—stay tuned for updates!  
 - **[🤗Diffusers Ecosystem Integration](https://huggingface.co/docs/diffusers/main/en/optimization/cache_dit)**: 🔥**cache-dit** has joined the Diffusers community ecosystem as the **first** DiT-specific cache acceleration framework! Check out the documentation here: <a href="https://huggingface.co/docs/diffusers/main/en/optimization/cache_dit"><img src=https://img.shields.io/badge/🤗Diffusers-ecosystem-yellow.svg ></a>
 
@@ -202,6 +202,7 @@ You can install the stable release of cache-dit from PyPI, or the latest develop
 
 ## 🔥Important News
 
+- 2025.10.15: 🎉cache-dit now supported [**🔥nunchaku**](https://github.com/nunchaku-tech/nunchaku): Qwen-Image/FLUX.1 [4-bits examples](./examples/quantize/)
 - 2025.10.13: 🎉cache-dit achieved the **SOTA** performance w/ **7.4x↑🎉** speedup on ClipScore!
 - 2025.10.10: 🔥[**Qwen-Image-ControlNet-Inpainting**](https://huggingface.co/InstantX/Qwen-Image-ControlNet-Inpainting) **2.3x↑🎉** speedup! Check the [example](https://github.com/vipshop/cache-dit/blob/main/examples/pipeline/run_qwen_image_controlnet_inpaint.py).
 - 2025.09.26: 🔥[**Qwen-Image-Edit-Plus(2509)**](https://github.com/QwenLM/Qwen-Image) **2.1x↑🎉** speedup! Please check the [example](https://github.com/vipshop/cache-dit/blob/main/examples/pipeline/run_qwen_image_edit_plus.py).

{cache_dit-1.0.4.dist-info → cache_dit-1.0.5.dist-info}/RECORD

@@ -1,5 +1,5 @@
 cache_dit/__init__.py,sha256=JQLxwr5aqoMFp-BNR58J0i6NutbRmNXKsaRJKCZQDCg,1638
-cache_dit/_version.py,sha256=jp1Oow7okdi1HqeKIp8SmyysmUf-oq2X9syICfrgATI,704
+cache_dit/_version.py,sha256=uYAmg6HpJuEIdeeqRyfS9YeZSA-4pXIPoW7JKBe0Uvw,704
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/utils.py,sha256=0YNFr84pxYoHOCZvnONKKXYN3PZY4kao9Tq2yEfHHR8,16986
 cache_dit/cache_factory/.gitignore,sha256=5Cb-qT9wsTUoMJ7vACDF7ZcLpAXhi5v-xdcWSRit988,23
@@ -9,11 +9,11 @@ cache_dit/cache_factory/cache_types.py,sha256=QnWfaS52UOXQtnoCUOwwz4ziY0dyBta6vQ
 cache_dit/cache_factory/forward_pattern.py,sha256=FumlCuZ-TSmSYH0hGBHctSJ-oGLCftdZjLygqhsmdR4,2258
 cache_dit/cache_factory/params_modifier.py,sha256=2T98IbepAolWW6GwQsqUDsRzu0k65vo7BOrN3V8mKog,3606
 cache_dit/cache_factory/utils.py,sha256=S3SD6Zhexzhkqnmfo830v6oNLm8stZe32nF4VdxD_bA,2497
-cache_dit/cache_factory/block_adapters/__init__.py,sha256=vM3aDMzPY79Tw4L0hlV2PdA3MFYomnf0eo0BGBo9P78,18087
+cache_dit/cache_factory/block_adapters/__init__.py,sha256=zs-cYacRL_hWlhUXmKc0TZNDAKzzWuznvHeuDpAmuwc,18221
 cache_dit/cache_factory/block_adapters/block_adapters.py,sha256=2TVK_KqiYXC7AKZ2s07fzdOzUoeUBc9P1SzQtLVzhf4,22249
 cache_dit/cache_factory/block_adapters/block_registers.py,sha256=2L7QeM4ygnaKQpC9PoJod0QRYyxidUKU2AYpysDCUwE,2572
 cache_dit/cache_factory/cache_adapters/__init__.py,sha256=py71WGD3JztQ1uk6qdLVbzYcQ1rvqFidNNaQYo7tqTo,79
-cache_dit/cache_factory/cache_adapters/cache_adapter.py,sha256=Za9HixVkEKldYzyDA57xvF91fm9dao2S-Fz5QBIT02M,22123
+cache_dit/cache_factory/cache_adapters/cache_adapter.py,sha256=-KTKukcNaVYk92np-QnGJ-EJbTdRQplqIHtqEORPiAo,23745
 cache_dit/cache_factory/cache_blocks/__init__.py,sha256=cpxzmDcUhbXcReHqaKSnWyEEbIg1H91Pz5hE3z9Xj3k,9984
 cache_dit/cache_factory/cache_blocks/offload_utils.py,sha256=wusgcqaCrwEjvv7Guy-6VXhNOgPPUrBV2sSVuRmGuvo,3513
 cache_dit/cache_factory/cache_blocks/pattern_0_1_2.py,sha256=j4bTafqU5DLQhzP_X5XwOk-QUVLWkGrX-Q6JZvBGHh0,666
@@ -21,14 +21,14 @@ cache_dit/cache_factory/cache_blocks/pattern_3_4_5.py,sha256=2qPnXVZwpQIm2oJ-Yrn
 cache_dit/cache_factory/cache_blocks/pattern_base.py,sha256=9H87qBRpa6UWRkUKXLVO0_9NJgxCVKkFSzaQxM9YPw8,25487
 cache_dit/cache_factory/cache_blocks/pattern_utils.py,sha256=qOxoVTlYPQzPMrR06-7_Ce_lwNg6n5pt1KQrvxzAJhE,3124
 cache_dit/cache_factory/cache_contexts/__init__.py,sha256=7uY8fX9uhpC71VNm1HH4aDIicYn-dD3kRpPQhvc9-EI,853
-cache_dit/cache_factory/cache_contexts/cache_config.py,sha256=WBHU2XVuYSFUSkrrJk8c4952LTeqvgetdkdtch_uSmg,5238
+cache_dit/cache_factory/cache_contexts/cache_config.py,sha256=G0PVWgckDqeyARc72Ne_0lRtO_LftsOeMERRhbh2gCA,5739
 cache_dit/cache_factory/cache_contexts/cache_context.py,sha256=fjZMEHaT1DZvUKnzY41GP0Ep8tmPEZTOsCSvG-5it5k,11269
 cache_dit/cache_factory/cache_contexts/cache_manager.py,sha256=tKtP35GDwZDoxGrQ_Okg_enlh3L-t-iqpytx8TFO_fw,30519
 cache_dit/cache_factory/cache_contexts/context_manager.py,sha256=j5zP_kwZAKla3EXbfr6JKI1vIxZuUEbZVhAPrtC4COw,853
-cache_dit/cache_factory/cache_contexts/prune_config.py,sha256=efFO_tu6AFJxIDp0OxExWKPzOFj95-NSrLGXggimBMA,3407
+cache_dit/cache_factory/cache_contexts/prune_config.py,sha256=WMTh6zb480a0oJiYMlgI0cwCsDSVvs6UjyeJLiXbjP8,3216
 cache_dit/cache_factory/cache_contexts/prune_context.py,sha256=ywiT9P0w_GjIFLowzUDa6jhTohNsSGfTbanZcs9wMic,6359
 cache_dit/cache_factory/cache_contexts/prune_manager.py,sha256=rZG7HD9ATqgH4VZdMq1XtP_h2pokaotFOVx1svB3J7E,5478
-cache_dit/cache_factory/cache_contexts/calibrators/__init__.py,sha256=mzYXO8tbytGpJJ9rpPu20kMoj1Iu_7Ym9tjfzV8rA98,5574
+cache_dit/cache_factory/cache_contexts/calibrators/__init__.py,sha256=QTbyT8xcFEjfIp9xjbnsnlnVCNvMjUc20NjB0W-s95k,6269
 cache_dit/cache_factory/cache_contexts/calibrators/base.py,sha256=mn6ZBkChGpGwN5csrHTUGMoX6BBPvqHXSLbIExiW-EU,748
 cache_dit/cache_factory/cache_contexts/calibrators/foca.py,sha256=nhHGs_hxwW1M942BQDMJb9-9IuHdnOxp774Jrna1bJI,891
 cache_dit/cache_factory/cache_contexts/calibrators/taylorseer.py,sha256=l1QSNaBwtGtpZZFAgCE7Hu8Nf1oL4QAcYu7lShpFGyw,5850
@@ -53,11 +53,11 @@ cache_dit/metrics/inception.py,sha256=pBVe2X6ylLPIXTG4-GWDM9DWnCviMJbJ45R3ulhktR
 cache_dit/metrics/lpips.py,sha256=hrHrmdM-f2B4TKDs0xLqJO5JFaYcCjq2qNIR8oCrVkc,811
 cache_dit/metrics/metrics.py,sha256=AZbQyoavE-djvyRUZ_EfCIrWSQbiWQFo7n2dhn7XptE,40466
 cache_dit/quantize/__init__.py,sha256=kWYoMAyZgBXu9BJlZjTQ0dRffW9GqeeY9_iTkXrb70A,59
-cache_dit/quantize/quantize_ao.py,sha256=Pr3u3Qr6qLvFkd8k-_rfcz4Mkjlg36U9BHG2t6Bl-6M,6301
+cache_dit/quantize/quantize_ao.py,sha256=LlkKh2uAtLihNRXHWIVggYR7ls12ak_VVLCY8trp2LY,7996
 cache_dit/quantize/quantize_interface.py,sha256=2s_R7xPSKuJeFpEGeLwRxnq_CqJcBG3a3lzyW5wh-UM,1241
-cache_dit-1.0.4.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-1.0.4.dist-info/METADATA,sha256=f04uCgApjgfHTC7Ll9aPejXCFFXbFTpTy-rjd5I_iwM,28376
-cache_dit-1.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-1.0.4.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-1.0.4.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-1.0.4.dist-info/RECORD,,
+cache_dit-1.0.5.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-1.0.5.dist-info/METADATA,sha256=qalYkFx9Y0cSmoBqHx-ZmI2RdVxcFhuNkWZVs7laExI,28632
+cache_dit-1.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-1.0.5.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-1.0.5.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-1.0.5.dist-info/RECORD,,