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

Code before:
            return (prev_sample,)

        return SchedulerOutput(prev_sample=prev_sample)

    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
            sample (`torch.Tensor`):
                The input sample.

        Returns:
            `torch.Tensor`:
                A scaled input sample.
        """
        return sample

    def _get_prev_sample(
        self, sample: torch.Tensor, timestep: int, prev_timestep: int, model_output: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute the previous sample x_(t-δ) from the current sample x_t using formula (9) from the PNDM paper.

        See formula (9) of [PNDM paper](https://huggingface.co/papers/2202.09778)

        Notation (<variable name> -> <name in paper>):
        - alpha_prod_t -> α_t
        - alpha_prod_t_prev -> α_(t−δ)
        - beta_prod_t -> (1 - α_t)
        - beta_prod_t_prev -> (1 - α_(t−δ))
        - sample -> x_t
        - model_output -> e_θ(x_t, t)
        - prev_sample -> x_(t−δ)

        Args:
            sample (`torch.Tensor`):
                The current sample x_t.
            timestep (`int`):
                The current timestep t.
            prev_timestep (`int`):
                The previous timestep (t-δ).
            model_output (`torch.Tensor`):
                The model output e_θ(x_t, t).

        Returns:
            `torch.Tensor`:
                The previous sample x_(t-δ).
        """
        alpha_prod_t = self.alphas_cumprod[timestep]

Review comment (suggestion):
```suggestion
        Compute the previous sample x_(t-δ) from the current sample x_t using formula (9) from the [PNDM paper](https://huggingface.co/papers/2202.09778).

```

Code after:
            return (prev_sample,)

        return SchedulerOutput(prev_sample=prev_sample)

    def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
            sample (`torch.Tensor`):
                The input sample.

        Returns:
            `torch.Tensor`:
                A scaled input sample.
        """
        return sample

    def _get_prev_sample(
        self, sample: torch.Tensor, timestep: int, prev_timestep: int, model_output: torch.Tensor
    ) -> torch.Tensor:
        """
        Compute the previous sample x_(t-δ) from the current sample x_t using formula (9) from the [PNDM
        paper](https://huggingface.co/papers/2202.09778).

        Args:
            sample (`torch.Tensor`):
                The current sample x_t.
            timestep (`int`):
                The current timestep t.
            prev_timestep (`int`):
                The previous timestep (t-δ).
            model_output (`torch.Tensor`):
                The model output e_θ(x_t, t).

        Returns:
            `torch.Tensor`:
                The previous sample x_(t-δ).
        """
        alpha_prod_t = self.alphas_cumprod[timestep]
        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
        beta_prod_t = 1 - alpha_prod_t
        beta_prod_t_prev = 1 - alpha_prod_t_prev

        if self.config.prediction_type == "v_prediction":
            model_output = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
        elif self.config.prediction_type != "epsilon":
            raise ValueError(
                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon` or `v_prediction`"
            )