# coding=utf-8 # Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import gc import pytest import torch from diffusers.hooks import ( FasterCacheConfig, FirstBlockCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, TaylorSeerCacheConfig, ) from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK from diffusers.hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK from diffusers.hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from diffusers.hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK from diffusers.hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK from diffusers.models.cache_utils import CacheMixin from ...testing_utils import assert_tensors_close, backend_empty_cache, is_cache, torch_device def require_cache_mixin(func): """Decorator to skip tests if model doesn't use CacheMixin.""" def wrapper(self, *args, **kwargs): if not issubclass(self.model_class, CacheMixin): pytest.skip(f"{self.model_class.__name__} does not use CacheMixin.") return func(self, *args, **kwargs) return wrapper class CacheTesterMixin: """ Base mixin class providing common test implementations for cache testing. Cache-specific mixins should: 1. Inherit from their respective config mixin (e.g., PyramidAttentionBroadcastConfigMixin) 2. Inherit from this mixin 3. Define the cache config to use for tests Expected class attributes: - model_class: The model class to test (must use CacheMixin) Expected methods in test classes: - get_init_dict(): Returns dict of arguments to initialize the model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Optional overrides: - cache_input_key: Property returning the input tensor key to vary between passes (default: "hidden_states") """ @property def cache_input_key(self): return "hidden_states" def setup_method(self): gc.collect() backend_empty_cache(torch_device) def teardown_method(self): gc.collect() backend_empty_cache(torch_device) def _get_cache_config(self): """ Get the cache config for testing. Should be implemented by subclasses. """ raise NotImplementedError("Subclass must implement _get_cache_config") def _get_hook_names(self): """ Get the hook names to check for this cache type. Should be implemented by subclasses. Returns a list of hook name strings. """ raise NotImplementedError("Subclass must implement _get_hook_names") def _test_cache_enable_disable_state(self): """Test that cache enable/disable updates the is_cache_enabled state correctly.""" init_dict = self.get_init_dict() model = self.model_class(**init_dict).to(torch_device) # Initially cache should not be enabled assert not model.is_cache_enabled, "Cache should not be enabled initially." config = self._get_cache_config() # Enable cache model.enable_cache(config) assert model.is_cache_enabled, "Cache should be enabled after enable_cache()." # Disable cache model.disable_cache() assert not model.is_cache_enabled, "Cache should not be enabled after disable_cache()." def _test_cache_double_enable_raises_error(self): """Test that enabling cache twice raises an error.""" init_dict = self.get_init_dict() model = self.model_class(**init_dict).to(torch_device) config = self._get_cache_config() model.enable_cache(config) # Trying to enable again should raise ValueError with pytest.raises(ValueError, match="Caching has already been enabled"): model.enable_cache(config) # Cleanup model.disable_cache() def _test_cache_hooks_registered(self): """Test that cache hooks are properly registered and removed.""" init_dict = self.get_init_dict() model = self.model_class(**init_dict).to(torch_device) config = self._get_cache_config() hook_names = self._get_hook_names() model.enable_cache(config) # Check that at least one hook was registered hook_count = 0 for module in model.modules(): if hasattr(module, "_diffusers_hook"): for hook_name in hook_names: hook = module._diffusers_hook.get_hook(hook_name) if hook is not None: hook_count += 1 assert hook_count > 0, f"At least one cache hook should be registered. Hook names: {hook_names}" # Disable and verify hooks are removed model.disable_cache() hook_count_after = 0 for module in model.modules(): if hasattr(module, "_diffusers_hook"): for hook_name in hook_names: hook = module._diffusers_hook.get_hook(hook_name) if hook is not None: hook_count_after += 1 assert hook_count_after == 0, "Cache hooks should be removed after disable_cache()." @torch.no_grad() def _test_cache_inference(self): """Test that model can run inference with cache enabled.""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) # First pass populates the cache _ = model(**inputs_dict, return_dict=False)[0] # Create modified inputs for second pass (vary input tensor to simulate denoising) inputs_dict_step2 = inputs_dict.copy() if self.cache_input_key in inputs_dict_step2: inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( inputs_dict_step2[self.cache_input_key] ) # Second pass uses cached attention with different inputs (produces approximated output) output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] assert output_with_cache is not None, "Model output should not be None with cache enabled." assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." # Run same inputs without cache to compare model.disable_cache() output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] # Cached output should be different from non-cached output (due to approximation) assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( "Cached output should be different from non-cached output due to cache approximation." ) @torch.no_grad() def _test_cache_context_manager(self, atol=1e-5, rtol=0): """Test the cache_context context manager properly isolates cache state.""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) # Run inference in first context with model.cache_context("context_1"): output_ctx1 = model(**inputs_dict, return_dict=False)[0] # Run same inference in second context (cache should be reset) with model.cache_context("context_2"): output_ctx2 = model(**inputs_dict, return_dict=False)[0] # Both contexts should produce the same output (first pass in each) assert_tensors_close( output_ctx1, output_ctx2, atol=atol, rtol=rtol, msg="First pass in different cache contexts should produce the same output.", ) model.disable_cache() @torch.no_grad() def _test_reset_stateful_cache(self): """Test that _reset_stateful_cache resets the cache state.""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) _ = model(**inputs_dict, return_dict=False)[0] model._reset_stateful_cache() model.disable_cache() @is_cache class PyramidAttentionBroadcastConfigMixin: """ Base mixin providing PyramidAttentionBroadcast cache config. Expected class attributes: - model_class: The model class to test (must use CacheMixin) """ # Default PAB config - can be overridden by subclasses PAB_CONFIG = { "spatial_attention_block_skip_range": 2, } # Store timestep for callback (must be within default range (100, 800) for skipping to trigger) _current_timestep = 500 def _get_cache_config(self): config_kwargs = self.PAB_CONFIG.copy() config_kwargs["current_timestep_callback"] = lambda: self._current_timestep return PyramidAttentionBroadcastConfig(**config_kwargs) def _get_hook_names(self): return [_PYRAMID_ATTENTION_BROADCAST_HOOK] @is_cache class PyramidAttentionBroadcastTesterMixin(PyramidAttentionBroadcastConfigMixin, CacheTesterMixin): """ Mixin class for testing PyramidAttentionBroadcast caching on models. Expected class attributes: - model_class: The model class to test (must use CacheMixin) Expected methods to be implemented by subclasses: - get_init_dict(): Returns dict of arguments to initialize the model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: cache Use `pytest -m "not cache"` to skip these tests """ @require_cache_mixin def test_pab_cache_enable_disable_state(self): self._test_cache_enable_disable_state() @require_cache_mixin def test_pab_cache_double_enable_raises_error(self): self._test_cache_double_enable_raises_error() @require_cache_mixin def test_pab_cache_hooks_registered(self): self._test_cache_hooks_registered() @require_cache_mixin def test_pab_cache_inference(self): self._test_cache_inference() @require_cache_mixin def test_pab_cache_context_manager(self): self._test_cache_context_manager() @require_cache_mixin def test_pab_reset_stateful_cache(self): self._test_reset_stateful_cache() @is_cache class FirstBlockCacheConfigMixin: """ Base mixin providing FirstBlockCache config. Expected class attributes: - model_class: The model class to test (must use CacheMixin) """ # Default FBC config - can be overridden by subclasses # Higher threshold makes FBC more aggressive about caching (skips more often) FBC_CONFIG = { "threshold": 1.0, } def _get_cache_config(self): return FirstBlockCacheConfig(**self.FBC_CONFIG) def _get_hook_names(self): return [_FBC_LEADER_BLOCK_HOOK, _FBC_BLOCK_HOOK] @is_cache class FirstBlockCacheTesterMixin(FirstBlockCacheConfigMixin, CacheTesterMixin): """ Mixin class for testing FirstBlockCache on models. Expected class attributes: - model_class: The model class to test (must use CacheMixin) Expected methods to be implemented by subclasses: - get_init_dict(): Returns dict of arguments to initialize the model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: cache Use `pytest -m "not cache"` to skip these tests """ @torch.no_grad() def _test_cache_inference(self): """Test that model can run inference with FBC cache enabled (requires cache_context).""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) # FBC requires cache_context to be set for inference with model.cache_context("fbc_test"): # First pass populates the cache _ = model(**inputs_dict, return_dict=False)[0] # Create modified inputs for second pass inputs_dict_step2 = inputs_dict.copy() if self.cache_input_key in inputs_dict_step2: inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( inputs_dict_step2[self.cache_input_key] ) # Second pass - FBC should skip remaining blocks and use cached residuals output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] assert output_with_cache is not None, "Model output should not be None with cache enabled." assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." # Run same inputs without cache to compare model.disable_cache() output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] # Cached output should be different from non-cached output (due to approximation) assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( "Cached output should be different from non-cached output due to cache approximation." ) @torch.no_grad() def _test_reset_stateful_cache(self): """Test that _reset_stateful_cache resets the FBC cache state (requires cache_context).""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) with model.cache_context("fbc_test"): _ = model(**inputs_dict, return_dict=False)[0] model._reset_stateful_cache() model.disable_cache() @require_cache_mixin def test_fbc_cache_enable_disable_state(self): self._test_cache_enable_disable_state() @require_cache_mixin def test_fbc_cache_double_enable_raises_error(self): self._test_cache_double_enable_raises_error() @require_cache_mixin def test_fbc_cache_hooks_registered(self): self._test_cache_hooks_registered() @require_cache_mixin def test_fbc_cache_inference(self): self._test_cache_inference() @require_cache_mixin def test_fbc_cache_context_manager(self): self._test_cache_context_manager() @require_cache_mixin def test_fbc_reset_stateful_cache(self): self._test_reset_stateful_cache() @is_cache class FasterCacheConfigMixin: """ Base mixin providing FasterCache config. Expected class attributes: - model_class: The model class to test (must use CacheMixin) """ # Default FasterCache config - can be overridden by subclasses FASTER_CACHE_CONFIG = { "spatial_attention_block_skip_range": 2, "spatial_attention_timestep_skip_range": (-1, 901), "tensor_format": "BCHW", } def _get_cache_config(self, current_timestep_callback=None): config_kwargs = self.FASTER_CACHE_CONFIG.copy() if current_timestep_callback is None: current_timestep_callback = lambda: 1000 # noqa: E731 config_kwargs["current_timestep_callback"] = current_timestep_callback return FasterCacheConfig(**config_kwargs) def _get_hook_names(self): return [_FASTER_CACHE_DENOISER_HOOK, _FASTER_CACHE_BLOCK_HOOK] @is_cache class FasterCacheTesterMixin(FasterCacheConfigMixin, CacheTesterMixin): """ Mixin class for testing FasterCache on models. Note: FasterCache is designed for pipeline-level inference with proper CFG batch handling and timestep management. Inference tests are skipped at model level - FasterCache should be tested via pipeline tests (e.g., FluxPipeline, HunyuanVideoPipeline). Expected class attributes: - model_class: The model class to test (must use CacheMixin) Expected methods to be implemented by subclasses: - get_init_dict(): Returns dict of arguments to initialize the model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: cache Use `pytest -m "not cache"` to skip these tests """ @torch.no_grad() def _test_cache_inference(self): """Test that model can run inference with FasterCache enabled.""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() current_timestep = [1000] config = self._get_cache_config(current_timestep_callback=lambda: current_timestep[0]) model.enable_cache(config) # First pass with timestep outside skip range - computes and populates cache current_timestep[0] = 1000 _ = model(**inputs_dict, return_dict=False)[0] # Move timestep inside skip range so subsequent passes use cache current_timestep[0] = 500 # Create modified inputs for second pass inputs_dict_step2 = inputs_dict.copy() if self.cache_input_key in inputs_dict_step2: inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( inputs_dict_step2[self.cache_input_key] ) # Second pass uses cached attention with different inputs output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] assert output_with_cache is not None, "Model output should not be None with cache enabled." assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." # Run same inputs without cache to compare model.disable_cache() output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] # Cached output should be different from non-cached output (due to approximation) assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( "Cached output should be different from non-cached output due to cache approximation." ) @torch.no_grad() def _test_reset_stateful_cache(self): """Test that _reset_stateful_cache resets the FasterCache state.""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) _ = model(**inputs_dict, return_dict=False)[0] model._reset_stateful_cache() model.disable_cache() @require_cache_mixin def test_faster_cache_enable_disable_state(self): self._test_cache_enable_disable_state() @require_cache_mixin def test_faster_cache_double_enable_raises_error(self): self._test_cache_double_enable_raises_error() @require_cache_mixin def test_faster_cache_hooks_registered(self): self._test_cache_hooks_registered() @require_cache_mixin def test_faster_cache_inference(self): self._test_cache_inference() @require_cache_mixin def test_faster_cache_context_manager(self): self._test_cache_context_manager() @require_cache_mixin def test_faster_cache_reset_stateful_cache(self): self._test_reset_stateful_cache() @is_cache class MagCacheConfigMixin: """ Base mixin providing MagCache config. Expected class attributes: - model_class: The model class to test (must use CacheMixin) """ # Default MagCache config - can be overridden by subclasses. # Uses neutral ratios [1.0, 1.0] and a high threshold so the second # inference step is always skipped, which is required by _test_cache_inference. MAG_CACHE_CONFIG = { "num_inference_steps": 2, "retention_ratio": 0.0, "threshold": 100.0, "mag_ratios": [1.0, 1.0], } def _get_cache_config(self): return MagCacheConfig(**self.MAG_CACHE_CONFIG) def _get_hook_names(self): return [_MAG_CACHE_LEADER_BLOCK_HOOK, _MAG_CACHE_BLOCK_HOOK] @is_cache class MagCacheTesterMixin(MagCacheConfigMixin, CacheTesterMixin): """ Mixin class for testing MagCache on models. Expected class attributes: - model_class: The model class to test (must use CacheMixin) Expected methods to be implemented by subclasses: - get_init_dict(): Returns dict of arguments to initialize the model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: cache Use `pytest -m "not cache"` to skip these tests """ @require_cache_mixin def test_mag_cache_enable_disable_state(self): self._test_cache_enable_disable_state() @require_cache_mixin def test_mag_cache_double_enable_raises_error(self): self._test_cache_double_enable_raises_error() @require_cache_mixin def test_mag_cache_hooks_registered(self): self._test_cache_hooks_registered() @require_cache_mixin def test_mag_cache_inference(self): self._test_cache_inference() @require_cache_mixin def test_mag_cache_context_manager(self): self._test_cache_context_manager() @require_cache_mixin def test_mag_cache_reset_stateful_cache(self): self._test_reset_stateful_cache() @is_cache class TaylorSeerCacheConfigMixin: """ Base mixin providing TaylorSeerCache config. Expected class attributes: - model_class: The model class to test (must use CacheMixin) """ # Default TaylorSeerCache config - can be overridden by subclasses. # Uses a low cache_interval and disable_cache_before_step=0 so the second # inference step is always predicted, which is required by _test_cache_inference. TAYLORSEER_CACHE_CONFIG = { "cache_interval": 3, "disable_cache_before_step": 1, "max_order": 1, } def _get_cache_config(self): return TaylorSeerCacheConfig(**self.TAYLORSEER_CACHE_CONFIG) def _get_hook_names(self): return [_TAYLORSEER_CACHE_HOOK] @is_cache class TaylorSeerCacheTesterMixin(TaylorSeerCacheConfigMixin, CacheTesterMixin): """ Mixin class for testing TaylorSeerCache on models. Expected class attributes: - model_class: The model class to test (must use CacheMixin) Expected methods to be implemented by subclasses: - get_init_dict(): Returns dict of arguments to initialize the model - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass Pytest mark: cache Use `pytest -m "not cache"` to skip these tests """ @torch.no_grad() def _test_cache_inference(self): """Test that model can run inference with TaylorSeer cache enabled (requires cache_context).""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) # TaylorSeer requires cache_context to be set for inference with model.cache_context("taylorseer_test"): # First pass populates the cache _ = model(**inputs_dict, return_dict=False)[0] # Create modified inputs for second pass inputs_dict_step2 = inputs_dict.copy() if self.cache_input_key in inputs_dict_step2: inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like( inputs_dict_step2[self.cache_input_key] ) # Second pass - TaylorSeer should use cached Taylor series predictions output_with_cache = model(**inputs_dict_step2, return_dict=False)[0] assert output_with_cache is not None, "Model output should not be None with cache enabled." assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled." # Run same inputs without cache to compare model.disable_cache() output_without_cache = model(**inputs_dict_step2, return_dict=False)[0] # Cached output should be different from non-cached output (due to approximation) assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), ( "Cached output should be different from non-cached output due to cache approximation." ) @torch.no_grad() def _test_reset_stateful_cache(self): """Test that _reset_stateful_cache resets the TaylorSeer cache state (requires cache_context).""" init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) model.eval() config = self._get_cache_config() model.enable_cache(config) with model.cache_context("taylorseer_test"): _ = model(**inputs_dict, return_dict=False)[0] model._reset_stateful_cache() model.disable_cache() @require_cache_mixin def test_taylorseer_cache_enable_disable_state(self): self._test_cache_enable_disable_state() @require_cache_mixin def test_taylorseer_cache_double_enable_raises_error(self): self._test_cache_double_enable_raises_error() @require_cache_mixin def test_taylorseer_cache_hooks_registered(self): self._test_cache_hooks_registered() @require_cache_mixin def test_taylorseer_cache_inference(self): self._test_cache_inference() @require_cache_mixin def test_taylorseer_cache_context_manager(self): self._test_cache_context_manager() @require_cache_mixin def test_taylorseer_cache_reset_stateful_cache(self): self._test_reset_stateful_cache()