cache-dit 1.0.7__py3-none-any.whl → 1.0.8__py3-none-any.whl

cache_dit/__init__.py

@@ -26,12 +26,24 @@ from cache_dit.cache_factory import FoCaCalibratorConfig
 from cache_dit.cache_factory import supported_pipelines
 from cache_dit.cache_factory import get_adapter
 from cache_dit.compile import set_compile_configs
-from cache_dit.quantize import quantize
 from cache_dit.parallelism import ParallelismBackend
 from cache_dit.parallelism import ParallelismConfig
 from cache_dit.utils import summary
 from cache_dit.utils import strify
 
+try:
+    from cache_dit.quantize import quantize
+except ImportError as e:  # noqa: F841
+    err_msg = str(e)
+
+    def quantize(*args, **kwargs):
+        raise ImportError(
+            "Quantization requires additional dependencies. "
+            "Please install cache-dit[quantization] or cache-dit[all] "
+            f"to use this feature. Error message: {err_msg}"
+        )
+
+
 NONE = CacheType.NONE
 DBCache = CacheType.DBCache
 DBPrune = CacheType.DBPrune

cache_dit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.0.7'
-__version_tuple__ = version_tuple = (1, 0, 7)
+__version__ = version = '1.0.8'
+__version_tuple__ = version_tuple = (1, 0, 8)
 
 __commit_id__ = commit_id = None

cache_dit/metrics/__init__.py

@@ -1,3 +1,14 @@
+try:
+    import ImageReward
+    import lpips
+    import skimage
+    import scipy
+except ImportError:
+    raise ImportError(
+        "Metrics functionality requires the 'metrics' extra dependencies. "
+        "Install with:\npip install cache-dit[metrics]"
+    )
+
 from cache_dit.metrics.metrics import compute_psnr
 from cache_dit.metrics.metrics import compute_ssim
 from cache_dit.metrics.metrics import compute_mse

cache_dit/parallelism/backends/native_diffusers/__init__.py

@@ -0,0 +1,6 @@
+from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
+    maybe_enable_parallelism,
+)
+from cache_dit.parallelism.backends.native_diffusers.parallel_difffusers import (
+    native_diffusers_parallelism_available,
+)

cache_dit/parallelism/backends/{parallel_difffusers.py → native_diffusers/parallel_difffusers.py}

@@ -63,8 +63,16 @@ def maybe_enable_parallelism(
                         "Set attention backend to _native_cudnn for parallelism because of "
                         "the issue: https://github.com/huggingface/diffusers/pull/12443"
                     )
-
-                transformer.enable_parallelism(config=cp_config)
+                cp_plan = parallelism_config.parallel_kwargs.get(
+                    "cp_plan", None
+                )
+                if cp_plan is not None:
+                    logger.info(
+                        f"Using custom context parallelism plan: {cp_plan}"
+                    )
+                transformer.enable_parallelism(
+                    config=cp_config, cp_plan=cp_plan
+                )
             else:
                 raise ValueError(
                     f"{transformer.__class__.__name__} does not support context parallelism."

cache_dit/parallelism/parallel_config.py

@@ -1,4 +1,5 @@
 import dataclasses
+from typing import Optional, Dict, Any
 from cache_dit.parallelism.parallel_backend import ParallelismBackend
 from cache_dit.logger import init_logger
 
@@ -20,6 +21,13 @@ class ParallelismConfig:
     # tp_size (`int`, *optional*):
     #     The degree of tensor parallelism.
     tp_size: int = None
+    # parallel_kwargs (`dict`, *optional*):
+    #     Additional kwargs for parallelism backends. For example, for
+    #     NATIVE_DIFFUSER backend, it can include `cp_plan` and other
+    #     arguments for `Context Parallelism`.
+    parallel_kwargs: Optional[Dict[str, Any]] = dataclasses.field(
+        default_factory=dict
+    )
 
     def __post_init__(self):
         assert ParallelismBackend.is_supported(self.backend), (

cache_dit/parallelism/parallel_interface.py

@@ -22,7 +22,7 @@ def enable_parallelism(
         return transformer
 
     if parallelism_config.backend == ParallelismBackend.NATIVE_DIFFUSER:
-        from cache_dit.parallelism.backends.parallel_difffusers import (
+        from cache_dit.parallelism.backends.native_diffusers import (
             maybe_enable_parallelism,
             native_diffusers_parallelism_available,
         )

cache_dit/quantize/__init__.py

@@ -1 +1,8 @@
+try:
+    import torchao
+except ImportError:
+    raise ImportError(
+        "Quantization functionality requires the 'quantization' extra dependencies. "
+        "Install with: pip install cache-dit[quantization]"
+    )
 from cache_dit.quantize.quantize_interface import quantize

{cache_dit-1.0.7.dist-info → cache_dit-1.0.8.dist-info}/METADATA

@@ -1,37 +1,33 @@
 Metadata-Version: 2.4
 Name: cache_dit
-Version: 1.0.7
+Version: 1.0.8
 Summary: A Unified, Flexible and Training-free Cache Acceleration Framework for 🤗Diffusers.
 Author: DefTruth, vipshop.com, etc.
 Maintainer: DefTruth, vipshop.com, etc
-Project-URL: Repository, https://github.com/vipshop/cache-dit.git
-Project-URL: Homepage, https://github.com/vipshop/cache-dit.git
+Project-URL: Repository, https://github.com/vipshop/cache-dit
+Project-URL: Homepage, https://github.com/vipshop/cache-dit
+Project-URL: GitHub, https://github.com/vipshop/cache-dit
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: packaging
 Requires-Dist: pyyaml
 Requires-Dist: torch>=2.7.1
-Requires-Dist: transformers>=4.55.2
 Requires-Dist: diffusers>=0.35.1
-Requires-Dist: scikit-image
-Requires-Dist: scipy
-Requires-Dist: lpips==0.1.4
-Requires-Dist: torchao>=0.12.0
-Requires-Dist: image-reward
-Provides-Extra: all
+Requires-Dist: transformers>=4.55.2
+Provides-Extra: quantization
+Requires-Dist: torchao>=0.12.0; extra == "quantization"
 Provides-Extra: metrics
+Requires-Dist: scipy; extra == "metrics"
+Requires-Dist: scikit-image; extra == "metrics"
 Requires-Dist: image-reward; extra == "metrics"
-Requires-Dist: pytorch-fid; extra == "metrics"
 Requires-Dist: lpips==0.1.4; extra == "metrics"
 Provides-Extra: dev
+Requires-Dist: packaging; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: pytest<8.0.0,>=7.0.0; extra == "dev"
 Requires-Dist: pytest-html; extra == "dev"
 Requires-Dist: expecttest; extra == "dev"
 Requires-Dist: hypothesis; extra == "dev"
-Requires-Dist: transformers; extra == "dev"
-Requires-Dist: diffusers; extra == "dev"
 Requires-Dist: accelerate; extra == "dev"
 Requires-Dist: peft; extra == "dev"
 Requires-Dist: protobuf; extra == "dev"
@@ -39,10 +35,10 @@ Requires-Dist: sentencepiece; extra == "dev"
 Requires-Dist: opencv-python-headless; extra == "dev"
 Requires-Dist: ftfy; extra == "dev"
 Requires-Dist: scikit-image; extra == "dev"
-Requires-Dist: pytorch-fid; extra == "dev"
+Provides-Extra: all
+Requires-Dist: cache-dit[quantization]; extra == "all"
+Requires-Dist: cache-dit[metrics]; extra == "all"
 Dynamic: license-file
-Dynamic: provides-extra
-Dynamic: requires-dist
 Dynamic: requires-python
 
 📚English | <a href="./README_CN.md">📚中文阅读 </a> 

{cache_dit-1.0.7.dist-info → cache_dit-1.0.8.dist-info}/RECORD

@@ -1,5 +1,5 @@
-cache_dit/__init__.py,sha256=HZb04M7AHCfk9DaEAGApGJ2lCM-rsP6pbsNQxsQudi0,1743
-cache_dit/_version.py,sha256=xUX1oSOk6hTPREy9SfhUBjaOBMJucMgoQViQ3e2Ce9A,704
+cache_dit/__init__.py,sha256=Azqj-3QMQK4HZDTGgyUtAfatUwuU-YQ4w8erJSyrsbE,2082
+cache_dit/_version.py,sha256=09KGe_qkXH8vhvX180khtkldtbrAX-u8refqVsC-Ky4,704
 cache_dit/logger.py,sha256=0zsu42hN-3-rgGC_C29ms1IvVpV4_b4_SwJCKSenxBE,4304
 cache_dit/utils.py,sha256=3NcEb324fNY0NYnrBTjsLURKQuckKeFe3V9Dfc_g4sc,17851
 cache_dit/cache_factory/.gitignore,sha256=5Cb-qT9wsTUoMJ7vACDF7ZcLpAXhi5v-xdcWSRit988,23
@@ -44,7 +44,7 @@ cache_dit/compile/__init__.py,sha256=FcTVzCeyypl-mxlc59_ehHL3lBNiDAFsXuRoJ-5Cfi0
 cache_dit/compile/utils.py,sha256=nN2OIrSdwRR5zGxJinKDqb07pXpvTNTF3g_OgLkeeBU,3858
 cache_dit/custom_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cache_dit/custom_ops/triton_taylorseer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-cache_dit/metrics/__init__.py,sha256=UjPJ69DyyjZDfERTpKAjZKOxOTx58aWnkze7VfH3en8,673
+cache_dit/metrics/__init__.py,sha256=Y_JrBr9XE6NKXwyXc7d_-PaX9c_rk5FKms-IYgCyHmY,936
 cache_dit/metrics/clip_score.py,sha256=ERNCFQFJKzJdbIX9OAg-1LiSPuXUVHLOFxbf2gcENpc,3938
 cache_dit/metrics/config.py,sha256=ieOgD9ayz722RjVzk24bSIqS2D6o7TZjGk8KeXV-OLQ,551
 cache_dit/metrics/fid.py,sha256=ZM_FM0XERtpnkMUfphmw2aOdljrh1uba-pnYItu0q6M,18219
@@ -54,15 +54,17 @@ cache_dit/metrics/lpips.py,sha256=hrHrmdM-f2B4TKDs0xLqJO5JFaYcCjq2qNIR8oCrVkc,81
 cache_dit/metrics/metrics.py,sha256=AZbQyoavE-djvyRUZ_EfCIrWSQbiWQFo7n2dhn7XptE,40466
 cache_dit/parallelism/__init__.py,sha256=dheBG5_TZCuwctviMslpAEgB-B3N8F816bE51qsw_fU,210
 cache_dit/parallelism/parallel_backend.py,sha256=js1soTMenLeAyPMsBgdI3gWcdXoqjWgBD-PuFEywMr0,508
-cache_dit/parallelism/parallel_config.py,sha256=bu24sRSzJMmH7FZqzUPTcT6tAzQ20-FAqAEvGV3Q1Fw,1733
-cache_dit/parallelism/parallel_interface.py,sha256=tsiIdHosTmRbeRg0z9q0eMQlx-7vefmSIlc56OWnuMg,2205
-cache_dit/parallelism/backends/parallel_difffusers.py,sha256=YQkCJ1yq1OomZLyRLtGMaPSNWbDeAWGx9XuObVJ_85I,2499
-cache_dit/quantize/__init__.py,sha256=kWYoMAyZgBXu9BJlZjTQ0dRffW9GqeeY9_iTkXrb70A,59
+cache_dit/parallelism/parallel_config.py,sha256=cOAXaniGf4CDPG5sbVktZy2rTZe49jSsnUIW2IBerGM,2106
+cache_dit/parallelism/parallel_interface.py,sha256=WPPYYaodo0PCYrn6-Haz8GcNZ2RK3EG7q6P_cH51Qj0,2202
+cache_dit/parallelism/backends/native_diffusers/__init__.py,sha256=T_6GeBA7TRiVbvtqGLLH2flkRiK0o7JBREt2xhS_-YE,242
+cache_dit/parallelism/backends/native_diffusers/parallel_difffusers.py,sha256=wHRjxRWK5E92cdSwDkZJpKQCQGZfxY53woW47rMFH2I,2844
+cache_dit/parallelism/backends/native_pytorch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+cache_dit/quantize/__init__.py,sha256=rUu0V9VRjOgwXuIUHHAI-osivNjAdUsi-jpkDbFp6Gk,278
 cache_dit/quantize/quantize_ao.py,sha256=bbEUwsrMp3bMuRw8qJZREIvCHaJRQoZyfMjlu4ImRMI,6315
 cache_dit/quantize/quantize_interface.py,sha256=2s_R7xPSKuJeFpEGeLwRxnq_CqJcBG3a3lzyW5wh-UM,1241
-cache_dit-1.0.7.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
-cache_dit-1.0.7.dist-info/METADATA,sha256=I0Vb-ZqUHblKOWwXyCyZVfcllq1lLm7ML2X7U6TJs4s,29475
-cache_dit-1.0.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-cache_dit-1.0.7.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
-cache_dit-1.0.7.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
-cache_dit-1.0.7.dist-info/RECORD,,
+cache_dit-1.0.8.dist-info/licenses/LICENSE,sha256=Dqb07Ik2dV41s9nIdMUbiRWEfDqo7-dQeRiY7kPO8PE,3769
+cache_dit-1.0.8.dist-info/METADATA,sha256=cs6CuqsbXQ8w69_kg2e7yALeczPAZjd5hz8odNrsvZA,29462
+cache_dit-1.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+cache_dit-1.0.8.dist-info/entry_points.txt,sha256=FX2gysXaZx6NeK1iCLMcIdP8Q4_qikkIHtEmi3oWn8o,65
+cache_dit-1.0.8.dist-info/top_level.txt,sha256=ZJDydonLEhujzz0FOkVbO-BqfzO9d_VqRHmZU-3MOZo,10
+cache_dit-1.0.8.dist-info/RECORD,,