Function _apply_mag_cache_head_hook

src/diffusers/hooks/mag_cache.py:444–452

(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig)

def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None:
    registry = HookRegistry.check_if_exists_or_initialize(block)

    # Automatically remove existing hook to allow re-application (e.g. switching modes)
    if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None:
        registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)

    hook = MagCacheHeadHook(state_manager, config)
    registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK)


def _apply_mag_cache_block_hook(