Method new_forward

src/diffusers/hooks/mag_cache.py:185–280

(self, module: torch.nn.Module, *args, **kwargs)

    @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:
            logger.debug(f"MagCache: Skipping step {state.step_index}")
            # Apply MagCache: Output = Input + Previous Residual

            output = hidden_states
            res = state.previous_residual

            if res.device != output.device:
                res = res.to(output.device)

            # Attempt to apply residual handling shape mismatches (e.g., text+image vs image only)
            if res.shape == output.shape:
                output = output + res
            elif (
                output.ndim == 3