Function _apply_mag_cache_block_hook

src/diffusers/hooks/mag_cache.py:455–468

(
    block: torch.nn.Module,
    state_manager: StateManager,
    config: MagCacheConfig,
    is_tail: bool = False,
)

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

    # Automatically remove existing hook to allow re-application
    if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None:
        registry.remove_hook(_MAG_CACHE_BLOCK_HOOK)

    hook = MagCacheBlockHook(state_manager, is_tail, config)
    registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK)