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

Code before:
            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,
    ):
        self.taylor_factors_dtype = taylor_factors_dtype
        self.max_order = max_order

        self.module_dtypes: Tuple[torch.dtype, ...] = ()
        self.last_update_step: Optional[int] = None
        self.taylor_factors: Dict[int, Dict[int, torch.Tensor]] = {}

        # For skip-mode modules
        self.device: Optional[torch.device] = None
        self.dummy_tensors: Optional[Tuple[torch.Tensor, ...]] = None

        self.current_step = -1

    def reset(self) -> None:
        self.last_update_step = None
        self.taylor_factors = {}
        self.device = None
        self.dummy_tensors = None
        self.current_step = -1

    def update(
        self,
        outputs: Tuple[torch.Tensor, ...],
        current_step: int,
    ) -> None:
        self.module_dtypes = tuple(output.dtype for output in outputs)
        for i in range(len(outputs)):
            features = outputs[i].to(self.taylor_factors_dtype)
            new_factors: Dict[int, torch.Tensor] = {0: features}
            is_first_update = self.last_update_step is None
            if not is_first_update:
                delta_step = current_step - self.last_update_step
                if delta_step == 0:
                    raise ValueError("Delta step cannot be zero for TaylorSeer update.")


Review comment (performance):
Used to store and return cached outputs. Would rename

```suggestion
        self.cached_output: Optional[Tuple[torch.Tensor, ...]] = None
```

Code after:
    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}, "
            f"max_order={self.max_order}, "
            f"taylor_factors_dtype={self.taylor_factors_dtype}, "
            f"skip_predict_identifiers={self.skip_predict_identifiers}, "
            f"cache_identifiers={self.cache_identifiers}, "
            f"use_lite_mode={self.use_lite_mode})"
        )


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

        self.module_dtypes: Tuple[torch.dtype, ...] = ()
        self.last_update_step: Optional[int] = None
        self.taylor_factors: Dict[int, Dict[int, torch.Tensor]] = {}
        self.inactive_shapes: Optional[Tuple[Tuple[int, ...], ...]] = None
        self.device: Optional[torch.device] = None
        self.current_step: int = -1

    def reset(self) -> None:
        self.current_step = -1
        self.last_update_step = None
        self.taylor_factors = {}
        self.inactive_shapes = None
        self.device = None

    def update(
        self,
        outputs: Tuple[torch.Tensor, ...],