Method new_forward

src/diffusers/hooks/mag_cache.py:301–361

(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")
        state: MagCacheState = self.state_manager.get_state()

        if not state.should_compute:
            arg_name = self._metadata.hidden_states_argument_name
            hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)

            if self.is_tail:
                # Still need to advance step index even if we skip
                self._advance_step(state)

            if self._metadata.return_encoder_hidden_states_index is not None:
                encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
                    "encoder_hidden_states", args, kwargs
                )
                max_idx = max(
                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
                )
                ret_list = [None] * (max_idx + 1)
                ret_list[self._metadata.return_hidden_states_index] = hidden_states
                ret_list[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states
                return tuple(ret_list)

            return hidden_states

        output = self.fn_ref.original_forward(*args, **kwargs)

        if self.is_tail:
            # Calculate residual for next steps
            if isinstance(output, tuple):
                out_hidden = output[self._metadata.return_hidden_states_index]
            else:
                out_hidden = output

            in_hidden = state.head_block_input

            if in_hidden is None:
                return output

            # Determine residual
            if out_hidden.shape == in_hidden.shape:
                residual = out_hidden - in_hidden
            elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]:
                diff = in_hidden.shape[1] - out_hidden.shape[1]
                if diff == 0:
                    residual = out_hidden - in_hidden
                else:
                    residual = out_hidden - in_hidden  # Fallback to matching tail
            else:
                # Fallback for completely mismatched shapes
                residual = out_hidden

            if self.config.calibrate:
                self._perform_calibration_step(state, residual)

            state.previous_residual = residual