Repository: huggingface/diffusers
File: src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py (Python)

Code before:
            begin_index (`int`, defaults to `0`):
                The begin index for the scheduler.
        """
        self._begin_index = begin_index

    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs
    def precondition_inputs(self, sample, sigma):
        """
        Precondition the input sample by scaling it according to the EDM formulation.

        Args:
            sample (`torch.Tensor`):
                The input sample tensor to precondition.
            sigma (`float` or `torch.Tensor`):
                The current sigma (noise level) value.

        Returns:
            `torch.Tensor`:
                The scaled input sample.
        """
        c_in = self._get_conditioning_c_in(sigma)
        scaled_sample = sample * c_in
        return scaled_sample

    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_noise
    def precondition_noise(self, sigma):
        """
        Precondition the noise level by applying a logarithmic transformation.

        Args:
            sigma (`float` or `torch.Tensor`):
                The sigma (noise level) value to precondition.

        Returns:
            `torch.Tensor`:
                The preconditioned noise value computed as `0.25 * log(sigma)`.
        """
        if not isinstance(sigma, torch.Tensor):
            sigma = torch.tensor([sigma])

        c_noise = 0.25 * torch.log(sigma)

        return c_noise

    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs
    def precondition_outputs(
        self,
        sample: torch.Tensor,
        model_output: torch.Tensor,
        sigma: Union[float, torch.Tensor],
    ) -> torch.Tensor:

Review comment (suggestion):
The method signature is missing type hints that are present in the source method (`scheduling_edm_euler.EDMEulerScheduler.precondition_noise`). The signature should be:

```python
def precondition_noise(self, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
```

This ensures consistency with the source file and maintains proper type annotations.
```suggestion
    def precondition_noise(self, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
```

Code after:
            begin_index (`int`, defaults to `0`):
                The begin index for the scheduler.
        """
        self._begin_index = begin_index

    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_inputs
    def precondition_inputs(self, sample: torch.Tensor, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
        """
        Precondition the input sample by scaling it according to the EDM formulation.

        Args:
            sample (`torch.Tensor`):
                The input sample tensor to precondition.
            sigma (`float` or `torch.Tensor`):
                The current sigma (noise level) value.

        Returns:
            `torch.Tensor`:
                The scaled input sample.
        """
        c_in = self._get_conditioning_c_in(sigma)
        scaled_sample = sample * c_in
        return scaled_sample

    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_noise
    def precondition_noise(self, sigma: Union[float, torch.Tensor]) -> torch.Tensor:
        """
        Precondition the noise level by applying a logarithmic transformation.

        Args:
            sigma (`float` or `torch.Tensor`):
                The sigma (noise level) value to precondition.

        Returns:
            `torch.Tensor`:
                The preconditioned noise value computed as `0.25 * log(sigma)`.
        """
        if not isinstance(sigma, torch.Tensor):
            sigma = torch.tensor([sigma])

        c_noise = 0.25 * torch.log(sigma)

        return c_noise

    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs
    def precondition_outputs(
        self,
        sample: torch.Tensor,
        model_output: torch.Tensor,
        sigma: Union[float, torch.Tensor],
    ) -> torch.Tensor: