minor fixes

Author: AlanPonnachan

src/diffusers/models/cache_utils.py

@@ -109,7 +109,7 @@ def disable_cache(self) -> None:
         )
         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_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
 
@@ -126,6 +126,7 @@ def disable_cache(self) -> None:
             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, TaylorSeerCacheConfig):

tests/hooks/test_mag_cache.py

@@ -70,7 +70,6 @@ def forward(self, hidden_states, encoder_hidden_states=None):
         return hidden_states, encoder_hidden_states
 
 
-@unittest.skip("MagCache unit tests are skipped.")
 class MagCacheTests(unittest.TestCase):
     def setUp(self):
         # Register standard dummy block