Method enable_cache

src/diffusers/models/cache_utils.py:39–104

r""" Enable caching techniques on the model. Args: config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig`): The configuration for applying the caching technique. Currently supported caching techniques are

(self, config)

        return self._cache_config is not None

    def enable_cache(self, config) -> None:
        r"""
        Enable caching techniques on the model.

        Args:
            config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig`):
                The configuration for applying the caching technique. Currently supported caching techniques are:
                    - [`~hooks.PyramidAttentionBroadcastConfig`]
                    - [`~hooks.FasterCacheConfig`]
                    - [`~hooks.FirstBlockCacheConfig`]
                    - [`~hooks.TextKVCacheConfig`]

        Example:

        ```python
        >>> import torch
        >>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig

        >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
        >>> pipe.to("cuda")

        >>> config = PyramidAttentionBroadcastConfig(
        ...     spatial_attention_block_skip_range=2,
        ...     spatial_attention_timestep_skip_range=(100, 800),
        ...     current_timestep_callback=lambda: pipe.current_timestep,
        ... )
        >>> pipe.transformer.enable_cache(config)
        ```
        """

        from ..hooks import (
            FasterCacheConfig,
            FirstBlockCacheConfig,
            MagCacheConfig,
            PyramidAttentionBroadcastConfig,
            TaylorSeerCacheConfig,
            TextKVCacheConfig,
            apply_faster_cache,
            apply_first_block_cache,
            apply_mag_cache,
            apply_pyramid_attention_broadcast,
            apply_taylorseer_cache,
            apply_text_kv_cache,
        )

        if self.is_cache_enabled:
            raise ValueError(
                f"Caching has already been enabled with {type(self._cache_config)}. To apply a new caching technique, please disable the existing one first."
            )

        if isinstance(config, FasterCacheConfig):
            apply_faster_cache(self, config)
        elif isinstance(config, FirstBlockCacheConfig):
            apply_first_block_cache(self, config)
        elif isinstance(config, MagCacheConfig):
            apply_mag_cache(self, config)
        elif isinstance(config, TextKVCacheConfig):
            apply_text_kv_cache(self, config)