Function apply_mag_cache

src/diffusers/hooks/mag_cache.py:397–441

Applies MagCache to a given module (typically a Transformer). Args: module (`torch.nn.Module`): The module to apply MagCache to. config (`MagCacheConfig`): The configuration for MagCache.

(module: torch.nn.Module, config: MagCacheConfig)

def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None:
    """
    Applies MagCache to a given module (typically a Transformer).

    Args:
        module (`torch.nn.Module`):
            The module to apply MagCache to.
        config (`MagCacheConfig`):
            The configuration for MagCache.
    """
    # Initialize registry on the root module so the Pipeline can set context.
    HookRegistry.check_if_exists_or_initialize(module)

    state_manager = StateManager(MagCacheState, (), {})
    remaining_blocks = []

    for name, submodule in module.named_children():
        if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList):
            continue
        for index, block in enumerate(submodule):
            remaining_blocks.append((f"{name}.{index}", block))

    if not remaining_blocks:
        logger.warning("MagCache: No transformer blocks found to apply hooks.")
        return

    # Handle single-block models
    if len(remaining_blocks) == 1:
        name, block = remaining_blocks[0]
        logger.info(f"MagCache: Applying Head+Tail Hooks to single block '{name}'")
        _apply_mag_cache_block_hook(block, state_manager, config, is_tail=True)
        _apply_mag_cache_head_hook(block, state_manager, config)
        return

    head_block_name, head_block = remaining_blocks.pop(0)
    tail_block_name, tail_block = remaining_blocks.pop(-1)

    logger.info(f"MagCache: Applying Head Hook to {head_block_name}")
    _apply_mag_cache_head_hook(head_block, state_manager, config)

    for name, block in remaining_blocks:
        _apply_mag_cache_block_hook(block, state_manager, config)

    logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}")
    _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True)


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