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

Code before:

    def reset_state(self, module: torch.nn.Module) -> None:
        """
        Reset state between sampling runs.
        """
        self.dummy_outputs = None
        self.current_step = -1
        self.state_manager.reset()

    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
        state: TaylorSeerState = self.state_manager.get_state()
        state.current_step += 1
        current_step = state.current_step
        is_warmup_phase = current_step < self.warmup_steps
        should_compute = (
            is_warmup_phase
            or ((current_step - self.warmup_steps - 1) % self.predict_steps == 0)
            or (self.stop_predicts is not None and current_step >= self.stop_predicts)
        )
        if should_compute:
            outputs = self.fn_ref.original_forward(*args, **kwargs)
            if not self.is_skip:
                state.update((outputs,) if isinstance(outputs, torch.Tensor) else outputs, current_step)
            else:
                self.dummy_outputs = outputs
            return outputs

        if self.is_skip:
            return self.dummy_outputs

        outputs = state.predict(current_step)
        return outputs[0] if len(outputs) == 1 else outputs


def _resolve_patterns(config: TaylorSeerCacheConfig) -> Tuple[List[str], List[str]]:
    """
    Resolve effective skip and cache pattern lists from config + templates.
    """

    skip_patterns = config.skip_identifiers if config.skip_identifiers is not None else None
    cache_patterns = config.cache_identifiers if config.cache_identifiers is not None else None

    return skip_patterns or [], cache_patterns or []


def apply_taylorseer_cache(module: torch.nn.Module, config: TaylorSeerCacheConfig):
    """
    Applies the TaylorSeer cache to a given pipeline (typically the transformer / UNet).

    Args:
        module (torch.nn.Module): The model subtree to apply the hooks to.

Review comment (suggestion):
Nit. Prefer to avoid `not` conditionals if possible.  
```suggestion
        if should_compute:
            outputs = self.fn_ref.original_forward(*args, **kwargs)
            if self.is_skip:
                self.dummy_outputs = outputs
            else:
                state.update((outputs,) if isinstance(outputs, torch.Tensor) else outputs)

            return outputs

```

Code after:
                    factor = taylor_factors[order]
                    output = output + factor.to(output_dtype) * coeff
                outputs.append(output)
        return outputs


class TaylorSeerCacheHook(ModelHook):
    _is_stateful = True

    def __init__(
        self,
        cache_interval: int,
        disable_cache_before_step: int,
        taylor_factors_dtype: torch.dtype,
        state_manager: StateManager,
        disable_cache_after_step: Optional[int] = None,
    ):
        super().__init__()
        self.cache_interval = cache_interval
        self.disable_cache_before_step = disable_cache_before_step
        self.disable_cache_after_step = disable_cache_after_step
        self.taylor_factors_dtype = taylor_factors_dtype
        self.state_manager = state_manager

    def initialize_hook(self, module: torch.nn.Module):
        return module

    def reset_state(self, module: torch.nn.Module) -> None:
        """
        Reset state between sampling runs.
        """
        self.state_manager.reset()

    @torch.compiler.disable
    def _measure_should_compute(self) -> bool:
        state: TaylorSeerState = self.state_manager.get_state()
        state.current_step += 1
        current_step = state.current_step
        is_warmup_phase = current_step < self.disable_cache_before_step
        is_compute_interval = (current_step - self.disable_cache_before_step - 1) % self.cache_interval == 0
        is_cooldown_phase = self.disable_cache_after_step is not None and current_step >= self.disable_cache_after_step
        should_compute = is_warmup_phase or is_compute_interval or is_cooldown_phase
        return should_compute, state

    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
        should_compute, state = self._measure_should_compute()
        if should_compute:
            outputs = self.fn_ref.original_forward(*args, **kwargs)
            wrapped_outputs = (outputs,) if isinstance(outputs, torch.Tensor) else outputs
            state.update(wrapped_outputs)
            return outputs