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

Code before:

    Attributes:
        warmup_steps (`int`, defaults to `3`):
            Number of *outer* diffusion steps to run with full computation
            before enabling caching. During warmup, the Taylor series factors
            are still updated, but no predictions are used.

        predict_steps (`int`, defaults to `5`):
            Number of prediction (cached) steps to take between two full
            computations. That is, once a module state is refreshed, it will
            be reused for `predict_steps` subsequent outer steps, then a new
            full forward will be computed on the next step.

        stop_predicts (`int`, *optional*, defaults to `None`):
            Outer diffusion step index at which caching is disabled.
            If provided, for `true_step >= stop_predicts` all modules are
            evaluated normally (no predictions, no state updates).

        max_order (`int`, defaults to `1`):
            Maximum order of Taylor series expansion to approximate the
            features. Higher order gives closer approximation but more compute.

        num_inner_loops (`int`, defaults to `1`):
            Number of inner loops per outer diffusion step. For example,
            with classifier-free guidance (CFG) you typically have 2 inner
            loops: unconditional and conditional branches.

        taylor_factors_dtype (`torch.dtype`, defaults to `torch.float32`):
            Data type for computing Taylor series expansion factors.

        architecture (`str`, *optional*, defaults to `None`):
            If provided, will look up default `cache` and `skip` regex
            patterns in `_CACHE_TEMPLATES[architecture]`. These can be
            overridden by `skip_identifiers` and `cache_identifiers`.

        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.

    Notes:
        - Patterns are applied with `re.fullmatch` on `module_name`.
        - If either `skip_identifiers` or `cache_identifiers` is provided
          (or inferred from `architecture`), only modules matching at least
          one of those patterns will be hooked.
        - If neither is provided, all attention-like modules will be hooked.

Review comment (performance):
The StateManager object is designed to handle conditional/unconditional branching
https://github.com/huggingface/diffusers/blob/d176f61fcf231c97025f016f6ad520bcc6276db6/src/diffusers/hooks/first_block_cache.py#L227

I think we can remove the inner/outer loop design by using it. 

Code after:
            The denoising step index after which caching is disabled. If set, for steps >= this value, all modules run
            full computations without predictions or state updates, ensuring accuracy in later stages if needed.

        max_order (`int`, defaults to `1`):
            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 (