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

Code before:
        if not return_dict:
            return (
                pred_prev_sample,
                pred_original_sample,
            )

        return DDPMParallelSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)

    def batch_step_no_noise(
        self,
        model_output: torch.Tensor,
        timesteps: torch.Tensor,
        sample: torch.Tensor,
    ) -> torch.Tensor:
        """
        Batched version of the `step` function, to be able to reverse the SDE for multiple samples/timesteps at once.
        Also, does not add any noise to the predicted sample, which is necessary for parallel sampling where the noise
        is pre-sampled by the pipeline.

        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            model_output (`torch.Tensor`): direct output from learned diffusion model.
            timesteps (`torch.Tensor`):
                current discrete timesteps in the diffusion chain. This is a tensor of integers.
            sample (`torch.Tensor`):
                current instance of sample being created by diffusion process.

        Returns:
            `torch.Tensor`: sample tensor at previous timestep.
        """
        t = timesteps
        num_inference_steps = self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
        prev_t = t - self.config.num_train_timesteps // num_inference_steps

        t = t.view(-1, *([1] * (model_output.ndim - 1)))
        prev_t = prev_t.view(-1, *([1] * (model_output.ndim - 1)))

        if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [
            "learned",
            "learned_range",
        ]:
            model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
        else:
            pass

        # 1. compute alphas, betas
        self.alphas_cumprod = self.alphas_cumprod.to(model_output.device)
        alpha_prod_t = self.alphas_cumprod[t]
        alpha_prod_t_prev = self.alphas_cumprod[torch.clip(prev_t, min=0)]

Review comment (suggestion):
```suggestion
                Current discrete timesteps in the diffusion chain. This is a tensor of integers.
```

Code after:
        if not return_dict:
            return (
                pred_prev_sample,
                pred_original_sample,
            )

        return DDPMParallelSchedulerOutput(prev_sample=pred_prev_sample, pred_original_sample=pred_original_sample)

    def batch_step_no_noise(
        self,
        model_output: torch.Tensor,
        timesteps: torch.Tensor,
        sample: torch.Tensor,
    ) -> torch.Tensor:
        """
        Batched version of the `step` function, to be able to reverse the SDE for multiple samples/timesteps at once.
        Also, does not add any noise to the predicted sample, which is necessary for parallel sampling where the noise
        is pre-sampled by the pipeline.

        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
        process from the learned model outputs (most often the predicted noise).

        Args:
            model_output (`torch.Tensor`): direct output from learned diffusion model.
            timesteps (`torch.Tensor`):
                Current discrete timesteps in the diffusion chain. This is a tensor of integers.
            sample (`torch.Tensor`):
                current instance of sample being created by diffusion process.

        Returns:
            `torch.Tensor`: sample tensor at previous timestep.
        """
        t = timesteps
        num_inference_steps = self.num_inference_steps if self.num_inference_steps else self.config.num_train_timesteps
        prev_t = t - self.config.num_train_timesteps // num_inference_steps

        t = t.view(-1, *([1] * (model_output.ndim - 1)))
        prev_t = prev_t.view(-1, *([1] * (model_output.ndim - 1)))

        if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in [
            "learned",
            "learned_range",
        ]:
            model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
        else:
            pass

        # 1. compute alphas, betas
        self.alphas_cumprod = self.alphas_cumprod.to(model_output.device)
        alpha_prod_t = self.alphas_cumprod[t]
        alpha_prod_t_prev = self.alphas_cumprod[torch.clip(prev_t, min=0)]