Class MagCacheHeadHook

src/diffusers/hooks/mag_cache.py:171–284

class MagCacheHeadHook(ModelHook):
    _is_stateful = True

    def __init__(self, state_manager: StateManager, config: MagCacheConfig):
        self.state_manager = state_manager
        self.config = config
        self._metadata = None

    def initialize_hook(self, module):
        unwrapped_module = unwrap_module(module)
        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
        return module

    @torch.compiler.disable
    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
        if self.state_manager._current_context is None:
            self.state_manager.set_context("inference")

        arg_name = self._metadata.hidden_states_argument_name
        hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)

        state: MagCacheState = self.state_manager.get_state()
        state.head_block_input = hidden_states

        should_compute = True

        if self.config.calibrate:
            # Never skip during calibration
            should_compute = True
        else:
            # MagCache Logic
            current_step = state.step_index
            if current_step >= len(self.config.mag_ratios):
                current_scale = 1.0
            else:
                current_scale = self.config.mag_ratios[current_step]

            retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5)

            if current_step >= retention_step:
                state.accumulated_ratio *= current_scale
                state.accumulated_steps += 1
                state.accumulated_err += abs(1.0 - state.accumulated_ratio)

                if (
                    state.previous_residual is not None
                    and state.accumulated_err <= self.config.threshold
                    and state.accumulated_steps <= self.config.max_skip_steps
                ):
                    should_compute = False
                else:
                    state.accumulated_ratio = 1.0
                    state.accumulated_steps = 0
                    state.accumulated_err = 0.0

        state.should_compute = should_compute

        if not should_compute: