Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DPMSolver. Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. timestep (`int`)
(
self,
model_output: torch.Tensor,
timestep: Union[int, torch.Tensor],
sample: torch.Tensor,
generator=None,
variance_noise: Optional[torch.Tensor] = None,
return_dict: bool = True,
)
| 704 | |
| 705 | # Modified from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.step |
| 706 | def step( |
| 707 | self, |
| 708 | model_output: torch.Tensor, |
| 709 | timestep: Union[int, torch.Tensor], |
| 710 | sample: torch.Tensor, |
| 711 | generator=None, |
| 712 | variance_noise: Optional[torch.Tensor] = None, |
| 713 | return_dict: bool = True, |
| 714 | ) -> Union[SchedulerOutput, Tuple]: |
| 715 | """ |
| 716 | Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with |
| 717 | the multistep DPMSolver. |
| 718 | Args: |
| 719 | model_output (`torch.Tensor`): |
| 720 | The direct output from learned diffusion model. |
| 721 | timestep (`int`): |
| 722 | The current discrete timestep in the diffusion chain. |
| 723 | sample (`torch.Tensor`): |
| 724 | A current instance of a sample created by the diffusion process. |
| 725 | generator (`torch.Generator`, *optional*): |
| 726 | A random number generator. |
| 727 | variance_noise (`torch.Tensor`): |
| 728 | Alternative to generating noise with `generator` by directly providing the noise for the variance |
| 729 | itself. Useful for methods such as [`LEdits++`]. |
| 730 | return_dict (`bool`): |
| 731 | Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. |
| 732 | Returns: |
| 733 | [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: |
| 734 | If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a |
| 735 | tuple is returned where the first element is the sample tensor. |
| 736 | """ |
| 737 | if self.num_inference_steps is None: |
| 738 | raise ValueError( |
| 739 | "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" |
| 740 | ) |
| 741 | |
| 742 | if self.step_index is None: |
| 743 | self._init_step_index(timestep) |
| 744 | |
| 745 | # Improve numerical stability for small number of steps |
| 746 | lower_order_final = (self.step_index == len(self.timesteps) - 1) and ( |
| 747 | self.config.euler_at_final or |
| 748 | (self.config.lower_order_final and len(self.timesteps) < 15) or |
| 749 | self.config.final_sigmas_type == "zero") |
| 750 | lower_order_second = ((self.step_index == len(self.timesteps) - 2) and |
| 751 | self.config.lower_order_final and |
| 752 | len(self.timesteps) < 15) |
| 753 | |
| 754 | model_output = self.convert_model_output(model_output, sample=sample) |
| 755 | for i in range(self.config.solver_order - 1): |
| 756 | self.model_outputs[i] = self.model_outputs[i + 1] |
| 757 | self.model_outputs[-1] = model_output |
| 758 | |
| 759 | # Upcast to avoid precision issues when computing prev_sample |
| 760 | sample = sample.to(torch.float32) |
| 761 | if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++" |
| 762 | ] and variance_noise is None: |
| 763 | noise = randn_tensor( |
no test coverage detected