Repository: huggingface/diffusers
File: src/diffusers/hooks/taylorseer_cache.py (Python)

Code before:
            approximation but more compute.

        taylor_factors_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type for computing Taylor series expansion factors. Use lower precision to reduce memory usage. Use
            higher precision to improve numerical stability.

        skip_identifiers (`List[str]`, *optional*, defaults to `None`):
            Regex patterns (fullmatch) for module names to be placed in "skip" mode, where the module is evaluated
            during warmup / refresh, but then replaced by a cheap dummy tensor during prediction steps.

        cache_identifiers (`List[str]`, *optional*, defaults to `None`):
            Regex patterns (fullmatch) for module names to be placed in Taylor-series caching mode.

        lite (`bool`, *optional*, defaults to `False`):
            Whether to use a TaylorSeer Lite variant that reduces memory usage. This option overrides any user-provided
            `skip_identifiers` or `cache_identifiers` patterns.
    Notes:
        - Patterns are applied with `re.fullmatch` on `module_name`.
        - If either `skip_identifiers` or `cache_identifiers` is provided, only modules matching at least one of those
          patterns will be hooked.
        - If neither is provided, all attention-like modules will be hooked.
    """

    warmup_steps: int = 3
    predict_steps: int = 5
    stop_predicts: Optional[int] = None
    max_order: int = 1
    taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16
    skip_identifiers: Optional[List[str]] = None
    cache_identifiers: Optional[List[str]] = None
    lite: bool = False

    def __repr__(self) -> str:
        return (
            "TaylorSeerCacheConfig("
            f"warmup_steps={self.warmup_steps}, "
            f"predict_steps={self.predict_steps}, "
            f"stop_predicts={self.stop_predicts}, "
            f"max_order={self.max_order}, "
            f"taylor_factors_dtype={self.taylor_factors_dtype}, "
            f"skip_identifiers={self.skip_identifiers}, "
            f"cache_identifiers={self.cache_identifiers}, "
            f"lite={self.lite})"
        )


class TaylorSeerState:
    def __init__(
        self,
        taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16,
        max_order: int = 1,

Review comment (performance):
I think the purpose of this parameter is a bit more clear if we call it something like `disable_cache_after_step` 

Code after:
            The highest order in the Taylor series expansion for approximating module outputs. Higher orders provide
            better approximations but increase computation and memory usage.

        taylor_factors_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type used for storing and computing Taylor series factors. Lower precision reduces memory but may
            affect stability; higher precision improves accuracy at the cost of more memory.

        skip_predict_identifiers (`List[str]`, *optional*, defaults to `None`):
            Regex patterns (using `re.fullmatch`) for module names to place as "skip" in "cache" mode. In this mode,
            the module computes fully during initial or refresh steps but returns a zero tensor (matching recorded
            shape) during prediction steps to skip computation cheaply.

        cache_identifiers (`List[str]`, *optional*, defaults to `None`):
            Regex patterns (using `re.fullmatch`) for module names to place in Taylor-series caching mode, where
            outputs are approximated and cached for reuse.

        use_lite_mode (`bool`, *optional*, defaults to `False`):
            Enables a lightweight TaylorSeer variant that minimizes memory usage by applying predefined patterns for
            skipping and caching (e.g., skipping blocks and caching projections). This overrides any custom
            `inactive_identifiers` or `active_identifiers`.

    Notes:
        - Patterns are matched using `re.fullmatch` on the module name.
        - If `skip_predict_identifiers` or `cache_identifiers` are provided, only matching modules are hooked.
        - If neither is provided, all attention-like modules are hooked by default.

    Example of inactive and active usage:

    ```py
    def forward(x):
        x = self.module1(x)  # inactive module: returns zeros tensor based on shape recorded during full compute
        x = self.module2(x)  # active module: caches output here, avoiding recomputation of prior steps
        return x
    ```
    """

    cache_interval: int = 5
    disable_cache_before_step: int = 3
    disable_cache_after_step: Optional[int] = None
    max_order: int = 1
    taylor_factors_dtype: Optional[torch.dtype] = torch.bfloat16
    skip_predict_identifiers: Optional[List[str]] = None
    cache_identifiers: Optional[List[str]] = None
    use_lite_mode: bool = False

    def __repr__(self) -> str:
        return (
            "TaylorSeerCacheConfig("
            f"cache_interval={self.cache_interval}, "
            f"disable_cache_before_step={self.disable_cache_before_step}, "
            f"disable_cache_after_step={self.disable_cache_after_step}, "