Function apply_faster_cache

src/diffusers/hooks/faster_cache.py:486–574

r""" Applies [FasterCache](https://huggingface.co/papers/2410.19355) to a given pipeline. Args: module (`torch.nn.Module`): The pytorch module to apply FasterCache to. Typically, this should be a transformer architecture supported in Diffusers, such as `CogVi

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

def apply_faster_cache(module: torch.nn.Module, config: FasterCacheConfig) -> None:
    r"""
    Applies [FasterCache](https://huggingface.co/papers/2410.19355) to a given pipeline.

    Args:
        module (`torch.nn.Module`):
            The pytorch module to apply FasterCache to. Typically, this should be a transformer architecture supported
            in Diffusers, such as `CogVideoXTransformer3DModel`, but external implementations may also work.
        config (`FasterCacheConfig`):
            The configuration to use for FasterCache.

    Example:
    ```python
    >>> import torch
    >>> from diffusers import CogVideoXPipeline, FasterCacheConfig, apply_faster_cache

    >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
    >>> pipe.to("cuda")

    >>> config = FasterCacheConfig(
    ...     spatial_attention_block_skip_range=2,
    ...     spatial_attention_timestep_skip_range=(-1, 681),
    ...     low_frequency_weight_update_timestep_range=(99, 641),
    ...     high_frequency_weight_update_timestep_range=(-1, 301),
    ...     spatial_attention_block_identifiers=["transformer_blocks"],
    ...     attention_weight_callback=lambda _: 0.3,
    ...     tensor_format="BFCHW",
    ... )
    >>> apply_faster_cache(pipe.transformer, config)
    ```
    """

    logger.warning(
        "FasterCache is a purely experimental feature and may not work as expected. Not all models support FasterCache. "
        "The API is subject to change in future releases, with no guarantee of backward compatibility. Please report any issues at "
        "https://github.com/huggingface/diffusers/issues."
    )

    if config.attention_weight_callback is None:
        # If the user has not provided a weight callback, we default to 0.5 for all timesteps.
        # In the paper, they recommend using a gradually increasing weight from 0 to 1 as the inference progresses, but
        # this depends from model-to-model. It is required by the user to provide a weight callback if they want to
        # use a different weight function. Defaulting to 0.5 works well in practice for most cases.
        logger.warning(
            "No `attention_weight_callback` provided when enabling FasterCache. Defaulting to using a weight of 0.5 for all timesteps."
        )
        config.attention_weight_callback = lambda _: 0.5

    if config.low_frequency_weight_callback is None:
        logger.debug(
            "Low frequency weight callback not provided when enabling FasterCache. Defaulting to behaviour described in the paper."
        )

        def low_frequency_weight_callback(module: torch.nn.Module) -> float:
            is_within_range = (
                config.low_frequency_weight_update_timestep_range[0]
                < config.current_timestep_callback()
                < config.low_frequency_weight_update_timestep_range[1]