fix kandinsky errors

Author: AlanPonnachan

src/diffusers/hooks/_helpers.py

@@ -26,6 +26,7 @@ class AttentionProcessorMetadata:
 class TransformerBlockMetadata:
     return_hidden_states_index: int = None
     return_encoder_hidden_states_index: int = None
+    hidden_states_argument_name: str = "hidden_states"
 
     _cls: Type = None
     _cached_parameter_indices: Dict[str, int] = None
@@ -346,6 +347,7 @@ def _register_transformer_blocks_metadata():
         metadata=TransformerBlockMetadata(
             return_hidden_states_index=0,
             return_encoder_hidden_states_index=None,
+            hidden_states_argument_name="visual_embed",
         ),
     )
 

src/diffusers/hooks/mag_cache.py

@@ -182,8 +182,9 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         if self.state_manager._current_context is None:
             self.state_manager.set_context("inference")
 
-        # Capture input hidden_states
-        hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs)
+        arg_name = self._metadata.hidden_states_argument_name
+        hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
 
         state: MagCacheState = self.state_manager.get_state()
         state.head_block_input = hidden_states
@@ -297,7 +298,9 @@ def new_forward(self, module: torch.nn.Module, *args, **kwargs):
         state: MagCacheState = self.state_manager.get_state()
 
         if not state.should_compute:
-            hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs)
+            arg_name = self._metadata.hidden_states_argument_name
+            hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
             if self.is_tail:
                 # Still need to advance step index even if we skip
                 self._advance_step(state)

src/diffusers/models/cache_utils.py

@@ -66,10 +66,12 @@ def enable_cache(self, config) -> None:
         from ..hooks import (
             FasterCacheConfig,
             FirstBlockCacheConfig,
+            MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
             apply_faster_cache,
             apply_first_block_cache,
+            apply_mag_cache,
             apply_pyramid_attention_broadcast,
             apply_taylorseer_cache,
         )
@@ -83,6 +85,8 @@ def enable_cache(self, config) -> None:
             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, PyramidAttentionBroadcastConfig):
             apply_pyramid_attention_broadcast(self, config)
         elif isinstance(config, TaylorSeerCacheConfig):
@@ -97,11 +101,13 @@ def disable_cache(self) -> None:
             FasterCacheConfig,
             FirstBlockCacheConfig,
             HookRegistry,
+            MagCacheConfig,
             PyramidAttentionBroadcastConfig,
             TaylorSeerCacheConfig,
         )
         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.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK
         from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK
 
@@ -116,6 +122,8 @@ def disable_cache(self) -> None:
         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)
         elif isinstance(self._cache_config, PyramidAttentionBroadcastConfig):
             registry.remove_hook(_PYRAMID_ATTENTION_BROADCAST_HOOK, recurse=True)
         elif isinstance(self._cache_config, TaylorSeerCacheConfig):