Predict the sample at the previous timestep by reversing the ODE. 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. I
(
self,
model_output: torch.Tensor,
timestep: int,
x_alpha: torch.Tensor,
)
| 16 | """ |
| 17 | |
| 18 | def step( |
| 19 | self, |
| 20 | model_output: torch.Tensor, |
| 21 | timestep: int, |
| 22 | x_alpha: torch.Tensor, |
| 23 | ) -> torch.Tensor: |
| 24 | """ |
| 25 | Predict the sample at the previous timestep by reversing the ODE. Core function to propagate the diffusion |
| 26 | process from the learned model outputs (most often the predicted noise). |
| 27 | |
| 28 | Args: |
| 29 | model_output (`torch.Tensor`): direct output from learned diffusion model. It is the direction from x0 to x1. |
| 30 | timestep (`float`): current timestep in the diffusion chain. |
| 31 | x_alpha (`torch.Tensor`): x_alpha sample for the current timestep |
| 32 | |
| 33 | Returns: |
| 34 | `torch.Tensor`: the sample at the previous timestep |
| 35 | |
| 36 | """ |
| 37 | if self.num_inference_steps is None: |
| 38 | raise ValueError( |
| 39 | "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" |
| 40 | ) |
| 41 | |
| 42 | alpha = timestep / self.num_inference_steps |
| 43 | alpha_next = (timestep + 1) / self.num_inference_steps |
| 44 | |
| 45 | d = model_output |
| 46 | |
| 47 | x_alpha = x_alpha + (alpha_next - alpha) * d |
| 48 | |
| 49 | return x_alpha |
| 50 | |
| 51 | def set_timesteps(self, num_inference_steps: int): |
| 52 | self.num_inference_steps = num_inference_steps |