cache-dit 1.0.3__py3-none-any.whl → 1.0.14__py3-none-any.whl

cache_dit/{cache_factory → caching}/cache_interface.py

@@ -1,12 +1,17 @@
+import torch
 from typing import Any, Tuple, List, Union, Optional
-from diffusers import DiffusionPipeline
-from cache_dit.cache_factory.cache_types import CacheType
-from cache_dit.cache_factory.block_adapters import BlockAdapter
-from cache_dit.cache_factory.block_adapters import BlockAdapterRegistry
-from cache_dit.cache_factory.cache_adapters import CachedAdapter
-from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig
-from cache_dit.cache_factory.params_modifier import ParamsModifier
+from diffusers import DiffusionPipeline, ModelMixin
+from cache_dit.caching.cache_types import CacheType
+from cache_dit.caching.block_adapters import BlockAdapter
+from cache_dit.caching.block_adapters import BlockAdapterRegistry
+from cache_dit.caching.cache_adapters import CachedAdapter
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts import DBCacheConfig
+from cache_dit.caching.cache_contexts import DBPruneConfig
+from cache_dit.caching.cache_contexts import CalibratorConfig
+from cache_dit.caching.params_modifier import ParamsModifier
+from cache_dit.parallelism import ParallelismConfig
+from cache_dit.parallelism import enable_parallelism
 
 from cache_dit.logger import init_logger
 
@@ -18,9 +23,18 @@ def enable_cache(
     pipe_or_adapter: Union[
         DiffusionPipeline,
         BlockAdapter,
+        # Transformer-only
+        torch.nn.Module,
+        ModelMixin,
     ],
-    # Basic DBCache config: BasicCacheConfig
-    cache_config: BasicCacheConfig = BasicCacheConfig(),
+    # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
+    cache_config: Optional[
+        Union[
+            BasicCacheConfig,
+            DBCacheConfig,
+            DBPruneConfig,
+        ]
+    ] = None,
     # Calibrator config: TaylorSeerCalibratorConfig, etc.
     calibrator_config: Optional[CalibratorConfig] = None,
     # Modify cache context params for specific blocks.
@@ -31,10 +45,15 @@ def enable_cache(
             List[List[ParamsModifier]],
         ]
     ] = None,
+    # Config for Parallelism
+    parallelism_config: Optional[ParallelismConfig] = None,
     # Other cache context kwargs: Deprecated cache kwargs
     **kwargs,
 ) -> Union[
     DiffusionPipeline,
+    # Transformer-only
+    torch.nn.Module,
+    ModelMixin,
     BlockAdapter,
 ]:
     r"""
@@ -64,10 +83,9 @@ def enable_cache(
     with minimal code changes.
 
     Args:
-        pipe_or_adapter (`DiffusionPipeline` or `BlockAdapter`, *required*):
+        pipe_or_adapter (`DiffusionPipeline`, `BlockAdapter` or `Transformer`, *required*):
             The standard Diffusion Pipeline or custom BlockAdapter (from cache-dit or user-defined).
-            For example: cache_dit.enable_cache(FluxPipeline(...)). Please check https://github.com/vipshop/cache-dit/blob/main/docs/BlockAdapter.md
-            for the usgae of BlockAdapter.
+            For example: cache_dit.enable_cache(FluxPipeline(...)).
 
         cache_config (`BasicCacheConfig`, *required*, defaults to BasicCacheConfig()):
             Basic DBCache config for cache context, defaults to BasicCacheConfig(). The configurable params listed belows:
@@ -107,6 +125,10 @@ def enable_cache(
                     Whether to compute separate difference values for CFG and non-CFG steps, default is True.
                     If False, we will use the computed difference from the current non-CFG transformer step
                     for the current CFG step.
+                num_inference_steps (`int`, *optional*, defaults to None):
+                    num_inference_steps for DiffusionPipeline, used to adjust some internal settings
+                    for better caching performance. For example, we will refresh the cache once the
+                    executed steps exceed num_inference_steps if num_inference_steps is provided.
 
         calibrator_config (`CalibratorConfig`, *optional*, defaults to None):
             Config for calibrator. If calibrator_config is not None, it means the user wants to use DBCache
@@ -121,8 +143,29 @@ def enable_cache(
                 **kwargs: (`dict`, *optional*, defaults to {}):
                     The same as 'kwargs' param in cache_dit.enable_cache() interface.
 
+        parallelism_config (`ParallelismConfig`, *optional*, defaults to None):
+            Config for Parallelism. If parallelism_config is not None, it means the user wants to enable
+            parallelism for cache-dit. Please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/parallelism/parallel_config.py
+            for more details of ParallelismConfig.
+                backend: (`ParallelismBackend`, *required*, defaults to "ParallelismBackend.NATIVE_DIFFUSER"):
+                    Parallelism backend, currently only NATIVE_DIFFUSER and NVTIVE_PYTORCH are supported.
+                    For context parallelism, only NATIVE_DIFFUSER backend is supported, for tensor parallelism,
+                    only NATIVE_PYTORCH backend is supported.
+                ulysses_size: (`int`, *optional*, defaults to None):
+                    The size of Ulysses cluster. If ulysses_size is not None, enable Ulysses style parallelism.
+                    This setting is only valid when backend is NATIVE_DIFFUSER.
+                ring_size: (`int`, *optional*, defaults to None):
+                    The size of ring for ring parallelism. If ring_size is not None, enable ring attention.
+                    This setting is only valid when backend is NATIVE_DIFFUSER.
+                tp_size: (`int`, *optional*, defaults to None):
+                    The size of tensor parallelism. If tp_size is not None, enable tensor parallelism.
+                    This setting is only valid when backend is NATIVE_PYTORCH.
+                parallel_kwargs: (`dict`, *optional*, defaults to {}):
+                    Additional kwargs for parallelism backends. For example, for NATIVE_DIFFUSER backend,
+                    it can include `cp_plan` and `attention_backend` arguments for `Context Parallelism`.
+
         kwargs (`dict`, *optional*, defaults to {})
-            Other cache context kwargs, please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/cache_factory/cache_contexts/cache_context.py
+            Other cache context kwargs, please check https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/caching/cache_contexts/cache_context.py
             for more details.
 
     Examples:
@@ -135,15 +178,29 @@ def enable_cache(
     >>> stats = cache_dit.summary(pipe) # Then, get the summary of cache acceleration stats.
     >>> cache_dit.disable_cache(pipe) # Disable cache and run original pipe.
     """
+    # Precheck for compatibility of different configurations
+    if cache_config is None:
+        if parallelism_config is None:
+            # Set default cache config only when parallelism is not enabled
+            logger.info("cache_config is None, using default DBCacheConfig")
+            cache_config = DBCacheConfig()
+        else:
+            # Allow empty cache_config when parallelism is enabled
+            logger.warning(
+                "Parallelism is enabled and cache_config is None. Please manually "
+                "set cache_config to avoid potential compatibility issues. "
+                "Otherwise, cache will not be enabled."
+            )
+
     # Collect cache context kwargs
-    cache_context_kwargs = {}
-    if (cache_type := cache_context_kwargs.pop("cache_type", None)) is not None:
+    context_kwargs = {}
+    if (cache_type := context_kwargs.get("cache_type", None)) is not None:
         if cache_type == CacheType.NONE:
             return pipe_or_adapter
 
-    # WARNING: Deprecated cache config params. These parameters are now retained
+    # NOTE: Deprecated cache config params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
-    deprecated_cache_kwargs = {
+    deprecated_kwargs = {
         "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
         "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
         "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -159,27 +216,27 @@ def enable_cache(
         ),
     }
 
-    deprecated_cache_kwargs = {
-        k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+    deprecated_kwargs = {
+        k: v for k, v in deprecated_kwargs.items() if v is not None
     }
 
-    if deprecated_cache_kwargs:
+    if deprecated_kwargs:
         logger.warning(
             "Manually settup DBCache context without BasicCacheConfig is "
             "deprecated and will be removed in the future, please use "
             "`cache_config` parameter instead!"
         )
         if cache_config is not None:
-            cache_config.update(**deprecated_cache_kwargs)
+            cache_config.update(**deprecated_kwargs)
         else:
-            cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+            cache_config = BasicCacheConfig(**deprecated_kwargs)
 
     if cache_config is not None:
-        cache_context_kwargs["cache_config"] = cache_config
+        context_kwargs["cache_config"] = cache_config
 
-    # WARNING: Deprecated taylorseer params. These parameters are now retained
+    # NOTE: Deprecated taylorseer params. These parameters are now retained
     # for backward compatibility but will be removed in the future.
-    if (
+    if cache_config is not None and (
         kwargs.get("enable_taylorseer", None) is not None
         or kwargs.get("enable_encoder_taylorseer", None) is not None
     ):
@@ -188,7 +245,7 @@ def enable_cache(
             "deprecated and will be removed in the future, please use "
             "`calibrator_config` parameter instead!"
         )
-        from cache_dit.cache_factory.cache_contexts.calibrators import (
+        from cache_dit.caching.cache_contexts.calibrators import (
             TaylorSeerCalibratorConfig,
         )
 
@@ -202,23 +259,79 @@ def enable_cache(
         )
 
     if calibrator_config is not None:
-        cache_context_kwargs["calibrator_config"] = calibrator_config
+        context_kwargs["calibrator_config"] = calibrator_config
 
     if params_modifiers is not None:
-        cache_context_kwargs["params_modifiers"] = params_modifiers
+        context_kwargs["params_modifiers"] = params_modifiers
 
-    if isinstance(pipe_or_adapter, (DiffusionPipeline, BlockAdapter)):
-        return CachedAdapter.apply(
+    if cache_config is not None:
+        if isinstance(
             pipe_or_adapter,
-            **cache_context_kwargs,
-        )
+            (DiffusionPipeline, BlockAdapter, torch.nn.Module, ModelMixin),
+        ):
+            pipe_or_adapter = CachedAdapter.apply(
+                pipe_or_adapter,
+                **context_kwargs,
+            )
+        else:
+            raise ValueError(
+                f"type: {type(pipe_or_adapter)} is not valid, "
+                "Please pass DiffusionPipeline or BlockAdapter"
+                "for the 1's position param: pipe_or_adapter"
+            )
     else:
-        raise ValueError(
-            f"type: {type(pipe_or_adapter)} is not valid, "
-            "Please pass DiffusionPipeline or BlockAdapter"
-            "for the 1's position param: pipe_or_adapter"
+        logger.warning(
+            "cache_config is None, skip enabling cache for "
+            f"{pipe_or_adapter.__class__.__name__}."
         )
 
+    # NOTE: Users should always enable parallelism after applying
+    # cache to avoid hooks conflict.
+    if parallelism_config is not None:
+        assert isinstance(
+            parallelism_config, ParallelismConfig
+        ), "parallelism_config should be of type ParallelismConfig."
+
+        transformers = []
+        if isinstance(pipe_or_adapter, DiffusionPipeline):
+            adapter = BlockAdapterRegistry.get_adapter(pipe_or_adapter)
+            if adapter is None:
+                assert hasattr(pipe_or_adapter, "transformer"), (
+                    "The given DiffusionPipeline does not have "
+                    "a 'transformer' attribute, cannot enable "
+                    "parallelism."
+                )
+                transformers = [pipe_or_adapter.transformer]
+            else:
+                adapter = BlockAdapter.normalize(adapter, unique=False)
+                transformers = BlockAdapter.flatten(adapter.transformer)
+        else:
+            if not BlockAdapter.is_normalized(pipe_or_adapter):
+                pipe_or_adapter = BlockAdapter.normalize(
+                    pipe_or_adapter, unique=False
+                )
+            transformers = BlockAdapter.flatten(pipe_or_adapter.transformer)
+
+        if len(transformers) == 0:
+            logger.warning(
+                "No transformer is detected in the "
+                "BlockAdapter, skip enabling parallelism."
+            )
+            return pipe_or_adapter
+
+        if len(transformers) > 1:
+            logger.warning(
+                "Multiple transformers are detected in the "
+                "BlockAdapter, all transfomers will be "
+                "enabled for parallelism."
+            )
+        for i, transformer in enumerate(transformers):
+            # Enable parallelism for the transformer inplace
+            transformers[i] = enable_parallelism(
+                transformer, parallelism_config
+            )
+    return pipe_or_adapter
+
 
 def disable_cache(
     pipe_or_adapter: Union[

cache_dit/{cache_factory → caching}/cache_types.py

@@ -6,7 +6,8 @@ logger = init_logger(__name__)
 
 class CacheType(Enum):
     NONE = "NONE"
-    DBCache = "Dual_Block_Cache"
+    DBCache = "DBCache"  # "Dual_Block_Cache"
+    DBPrune = "DBPrune"  # "Dynamic_Block_Prune"
 
     @staticmethod
     def type(type_hint: "CacheType | str") -> "CacheType":
@@ -14,6 +15,9 @@ class CacheType(Enum):
             return type_hint
         return cache_type(type_hint)
 
+    def __str__(self) -> str:
+        return self.value
+
 
 def cache_type(type_hint: "CacheType | str") -> "CacheType":
     if type_hint is None:
@@ -21,7 +25,6 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
 
     if isinstance(type_hint, CacheType):
         return type_hint
-
     elif type_hint.upper() in (
         "DUAL_BLOCK_CACHE",
         "DB_CACHE",
@@ -29,6 +32,20 @@ def cache_type(type_hint: "CacheType | str") -> "CacheType":
         "DB",
     ):
         return CacheType.DBCache
+    elif type_hint.upper() in (
+        "DYNAMIC_BLOCK_PRUNE",
+        "DB_PRUNE",
+        "DBPRUNE",
+        "DBP",
+    ):
+        return CacheType.DBPrune
+    elif type_hint.upper() in (
+        "NONE",
+        "NO_CACHE",
+        "NOCACHE",
+        "NC",
+    ):
+        return CacheType.NONE
     return CacheType.NONE
 
 

cache_dit/{cache_factory → caching}/forward_pattern.py

@@ -20,33 +20,33 @@ class ForwardPattern(Enum):
 
     Pattern_0 = (
         True,  # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_1 = (
-        False, # Return_H_First
-        False, # Return_H_Only
-        False, # Forward_H_only
+        False,  # Return_H_First
+        False,  # Return_H_Only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_2 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
-        False, # Forward_H_only
+        False,  # Forward_H_only
         ("hidden_states", "encoder_hidden_states"),  # In
-        ("hidden_states",),                          # Out
+        ("hidden_states",),  # Out
         True,  # Supported
     )
 
     Pattern_3 = (
-        False, # Return_H_First
+        False,  # Return_H_First
         True,  # Return_H_Only
         True,  # Forward_H_only
         ("hidden_states",),  # In
@@ -56,18 +56,18 @@ class ForwardPattern(Enum):
 
     Pattern_4 = (
         True,  # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("hidden_states", "encoder_hidden_states"),  # Out
         True,  # Supported
     )
 
     Pattern_5 = (
-        False, # Return_H_First
-        False, # Return_H_Only
+        False,  # Return_H_First
+        False,  # Return_H_Only
         True,  # Forward_H_only
-        ("hidden_states",),                          # In
+        ("hidden_states",),  # In
         ("encoder_hidden_states", "hidden_states"),  # Out
         True,  # Supported
     )

cache_dit/{cache_factory → caching}/params_modifier.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
-from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
-from cache_dit.cache_factory.cache_contexts import CalibratorConfig
+from cache_dit.caching.cache_contexts import BasicCacheConfig
+from cache_dit.caching.cache_contexts import CalibratorConfig
 
 from cache_dit.logger import init_logger
 
@@ -11,7 +11,7 @@ logger = init_logger(__name__)
 class ParamsModifier:
     def __init__(
         self,
-        # Basic DBCache config: BasicCacheConfig
+        # BasicCacheConfig, DBCacheConfig, DBPruneConfig, etc.
         cache_config: BasicCacheConfig = None,
         # Calibrator config: TaylorSeerCalibratorConfig, etc.
         calibrator_config: Optional[CalibratorConfig] = None,
@@ -22,7 +22,7 @@ class ParamsModifier:
 
         # WARNING: Deprecated cache config params. These parameters are now retained
         # for backward compatibility but will be removed in the future.
-        deprecated_cache_kwargs = {
+        deprecated_kwargs = {
             "Fn_compute_blocks": kwargs.get("Fn_compute_blocks", None),
             "Bn_compute_blocks": kwargs.get("Bn_compute_blocks", None),
             "max_warmup_steps": kwargs.get("max_warmup_steps", None),
@@ -40,20 +40,20 @@ class ParamsModifier:
             ),
         }
 
-        deprecated_cache_kwargs = {
-            k: v for k, v in deprecated_cache_kwargs.items() if v is not None
+        deprecated_kwargs = {
+            k: v for k, v in deprecated_kwargs.items() if v is not None
         }
 
-        if deprecated_cache_kwargs:
+        if deprecated_kwargs:
             logger.warning(
                 "Manually settup DBCache context without BasicCacheConfig is "
                 "deprecated and will be removed in the future, please use "
                 "`cache_config` parameter instead!"
             )
             if cache_config is not None:
-                cache_config.update(**deprecated_cache_kwargs)
+                cache_config.update(**deprecated_kwargs)
             else:
-                cache_config = BasicCacheConfig(**deprecated_cache_kwargs)
+                cache_config = BasicCacheConfig(**deprecated_kwargs)
 
         if cache_config is not None:
             self._context_kwargs["cache_config"] = cache_config
@@ -68,7 +68,7 @@ class ParamsModifier:
                 "deprecated and will be removed in the future, please use "
                 "`calibrator_config` parameter instead!"
             )
-            from cache_dit.cache_factory.cache_contexts.calibrators import (
+            from cache_dit.caching.cache_contexts.calibrators import (
                 TaylorSeerCalibratorConfig,
             )
 

cache_dit/caching/patch_functors/__init__.py

@@ -0,0 +1,15 @@
+from cache_dit.caching.patch_functors.functor_base import PatchFunctor
+from cache_dit.caching.patch_functors.functor_dit import DiTPatchFunctor
+from cache_dit.caching.patch_functors.functor_flux import FluxPatchFunctor
+from cache_dit.caching.patch_functors.functor_chroma import (
+    ChromaPatchFunctor,
+)
+from cache_dit.caching.patch_functors.functor_hidream import (
+    HiDreamPatchFunctor,
+)
+from cache_dit.caching.patch_functors.functor_hunyuan_dit import (
+    HunyuanDiTPatchFunctor,
+)
+from cache_dit.caching.patch_functors.functor_qwen_image_controlnet import (
+    QwenImageControlNetPatchFunctor,
+)

cache_dit/{cache_factory → caching}/patch_functors/functor_chroma.py

@@ -13,7 +13,7 @@ from diffusers.utils import (
     unscale_lora_layers,
 )
 
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_dit.py

@@ -6,7 +6,7 @@ from diffusers.models.transformers.dit_transformer_2d import (
     DiTTransformer2DModel,
     Transformer2DModelOutput,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_flux.py

@@ -14,7 +14,7 @@ from diffusers.utils import (
     unscale_lora_layers,
 )
 
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_hidream.py

@@ -13,7 +13,7 @@ from diffusers.utils import (
     scale_lora_layers,
     unscale_lora_layers,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_hunyuan_dit.py

@@ -5,7 +5,7 @@ from diffusers.models.transformers.hunyuan_transformer_2d import (
     HunyuanDiTBlock,
     Transformer2DModelOutput,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/patch_functors/functor_qwen_image_controlnet.py

@@ -11,7 +11,7 @@ from diffusers.utils import (
     scale_lora_layers,
     unscale_lora_layers,
 )
-from cache_dit.cache_factory.patch_functors.functor_base import (
+from cache_dit.caching.patch_functors.functor_base import (
     PatchFunctor,
 )
 from cache_dit.logger import init_logger

cache_dit/{cache_factory → caching}/utils.py

@@ -7,10 +7,6 @@ def load_cache_options_from_yaml(yaml_file_path):
             kwargs: dict = yaml.safe_load(f)
 
         required_keys = [
-            "max_warmup_steps",
-            "max_cached_steps",
-            "Fn_compute_blocks",
-            "Bn_compute_blocks",
             "residual_diff_threshold",
         ]
         for key in required_keys:
@@ -21,7 +17,7 @@ def load_cache_options_from_yaml(yaml_file_path):
 
         cache_context_kwargs = {}
         if kwargs.get("enable_taylorseer", False):
-            from cache_dit.cache_factory.cache_contexts.calibrators import (
+            from cache_dit.caching.cache_contexts.calibrators import (
                 TaylorSeerCalibratorConfig,
             )
 
@@ -38,10 +34,25 @@ def load_cache_options_from_yaml(yaml_file_path):
                 )
             )
 
-        from cache_dit.cache_factory.cache_contexts import BasicCacheConfig
+        if "cache_type" not in kwargs:
+            from cache_dit.caching.cache_contexts import BasicCacheConfig
 
-        cache_context_kwargs["cache_config"] = BasicCacheConfig()
-        cache_context_kwargs["cache_config"].update(**kwargs)
+            cache_context_kwargs["cache_config"] = BasicCacheConfig()
+            cache_context_kwargs["cache_config"].update(**kwargs)
+        else:
+            cache_type = kwargs.pop("cache_type")
+            if cache_type == "DBCache":
+                from cache_dit.caching.cache_contexts import DBCacheConfig
+
+                cache_context_kwargs["cache_config"] = DBCacheConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            elif cache_type == "DBPrune":
+                from cache_dit.caching.cache_contexts import DBPruneConfig
+
+                cache_context_kwargs["cache_config"] = DBPruneConfig()
+                cache_context_kwargs["cache_config"].update(**kwargs)
+            else:
+                raise ValueError(f"Unsupported cache_type: {cache_type}.")
 
         return cache_context_kwargs
 

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/__init__.py

@@ -0,0 +1,3 @@
+from cache_dit.parallelism.parallel_backend import ParallelismBackend
+from cache_dit.parallelism.parallel_config import ParallelismConfig
+from cache_dit.parallelism.parallel_interface import enable_parallelism

cache_dit/parallelism/backends/native_diffusers/__init__.py

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