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

Code before:
    def step_index(self):
        """
        The index counter for current timestep. It will increase 1 after each scheduler step.
        """
        return self._step_index

    @property
    def begin_index(self):
        """
        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
        """
        return self._begin_index

    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
    def set_begin_index(self, begin_index: int = 0):
        """
        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

        Args:
            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.

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

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

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

Code after:
    def step_index(self):
        """
        The index counter for current timestep. It will increase 1 after each scheduler step.
        """
        return self._step_index

    @property
    def begin_index(self):
        """
        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
        """
        return self._begin_index

    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
    def set_begin_index(self, begin_index: int = 0):
        """
        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

        Args:
            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.