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

Code before:
import math
import re
from dataclasses import dataclass
from typing import Optional, List, Dict, Tuple

import torch

from .hooks import ModelHook
from ..models.attention import Attention
from ..models.attention import AttentionModuleMixin
from ._common import _ATTENTION_CLASSES
from ..hooks import HookRegistry
from ..utils import logging


logger = logging.get_logger(__name__)

_TAYLORSEER_ATTENTION_CACHE_HOOK = "taylorseer_attention_cache"

# Predefined cache templates for optimized architectures
_CACHE_TEMPLATES: Dict[str, Dict[str, List[str]]] = {
    "flux": {
        "cache": [
            r"transformer_blocks\.\d+\.attn",
            r"transformer_blocks\.\d+\.ff",
            r"transformer_blocks\.\d+\.ff_context",
            r"single_transformer_blocks\.\d+\.proj_out",
        ],
        "skip": [
            r"single_transformer_blocks\.\d+\.attn",
            r"single_transformer_blocks\.\d+\.proj_mlp",
            r"single_transformer_blocks\.\d+\.act_mlp",
        ],
    },
}


@dataclass
class TaylorSeerCacheConfig:
    """
    Configuration for TaylorSeer cache.
    See: https://huggingface.co/papers/2503.06923

    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

Review comment (performance):
Let's avoid using model specific templates. We can apply the cache to global templates that should work across all models. Similar to how it's done here
https://github.com/huggingface/diffusers/blob/d176f61fcf231c97025f016f6ad520bcc6276db6/src/diffusers/hooks/faster_cache.py#L33-L46
 
For more optimized cacheing patterns users can pass them in directly through `cache_identifiers` and `skip_identifiers`

Code after:
import math
import re
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn

from ..utils import logging
from .hooks import HookRegistry, ModelHook, StateManager


logger = logging.get_logger(__name__)
_TAYLORSEER_CACHE_HOOK = "taylorseer_cache"
_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = (
    "^blocks.*attn",
    "^transformer_blocks.*attn",
    "^single_transformer_blocks.*attn",
)
_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("^temporal_transformer_blocks.*attn",)
_TRANSFORMER_BLOCK_IDENTIFIERS = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS + _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
_BLOCK_IDENTIFIERS = ("^[^.]*block[^.]*\\.[^.]+$",)
_PROJ_OUT_IDENTIFIERS = ("^proj_out$",)


@dataclass
class TaylorSeerCacheConfig:
    """
    Configuration for TaylorSeer cache. See: https://huggingface.co/papers/2503.06923

    Attributes:
        cache_interval (`int`, defaults to `5`):
            The interval between full computation steps. After a full computation, the cached (predicted) outputs are
            reused for this many subsequent denoising steps before refreshing with a new full forward pass.

        disable_cache_before_step (`int`, defaults to `3`):
            The denoising step index before which caching is disabled, meaning full computation is performed for the
            initial steps (0 to disable_cache_before_step - 1) to gather data for Taylor series approximations. During
            these steps, Taylor factors are updated, but caching/predictions are not applied. Caching begins at this
            step.

        disable_cache_after_step (`int`, *optional*, defaults to `None`):
            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