Method disable_cache

src/diffusers/models/cache_utils.py:106–147

(self)

        self._cache_config = config

    def disable_cache(self) -> None:
        from ..hooks import (
            FasterCacheConfig,
            FirstBlockCacheConfig,
            HookRegistry,
            MagCacheConfig,
            PyramidAttentionBroadcastConfig,
            TaylorSeerCacheConfig,
            TextKVCacheConfig,
        )
        from ..hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK
        from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK
        from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK
        from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK
        from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK
        from ..hooks.text_kv_cache import _TEXT_KV_CACHE_BLOCK_HOOK, _TEXT_KV_CACHE_TRANSFORMER_HOOK

        if self._cache_config is None:
            logger.warning("Caching techniques have not been enabled, so there's nothing to disable.")
            return

        registry = HookRegistry.check_if_exists_or_initialize(self)
        if isinstance(self._cache_config, FasterCacheConfig):
            registry.remove_hook(_FASTER_CACHE_DENOISER_HOOK, recurse=True)
            registry.remove_hook(_FASTER_CACHE_BLOCK_HOOK, recurse=True)
        elif isinstance(self._cache_config, FirstBlockCacheConfig):
            registry.remove_hook(_FBC_LEADER_BLOCK_HOOK, recurse=True)
            registry.remove_hook(_FBC_BLOCK_HOOK, recurse=True)
        elif isinstance(self._cache_config, MagCacheConfig):
            registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK, recurse=True)
            registry.remove_hook(_MAG_CACHE_BLOCK_HOOK, recurse=True)
        elif isinstance(self._cache_config, PyramidAttentionBroadcastConfig):
            registry.remove_hook(_PYRAMID_ATTENTION_BROADCAST_HOOK, recurse=True)
        elif isinstance(self._cache_config, TextKVCacheConfig):
            registry.remove_hook(_TEXT_KV_CACHE_TRANSFORMER_HOOK, recurse=True)
            registry.remove_hook(_TEXT_KV_CACHE_BLOCK_HOOK, recurse=True)
        elif isinstance(self._cache_config, TaylorSeerCacheConfig):
            registry.remove_hook(_TAYLORSEER_CACHE_HOOK, recurse=True)
        else:
            raise ValueError(f"Cache config {type(self._cache_config)} is not supported.")

        self._cache_config = None

    def _reset_stateful_cache(self, recurse: bool = True) -> None:
        from ..hooks import HookRegistry