Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise). Args: model_output (`torch.Tensor`): The direct output from learne
(
self,
model_output: torch.Tensor,
timestep: int,
sample: torch.Tensor,
generator: Optional[torch.Generator] = None,
return_dict: bool = True,
)
| 496 | return c_skip, c_out |
| 497 | |
| 498 | def step( |
| 499 | self, |
| 500 | model_output: torch.Tensor, |
| 501 | timestep: int, |
| 502 | sample: torch.Tensor, |
| 503 | generator: Optional[torch.Generator] = None, |
| 504 | return_dict: bool = True, |
| 505 | ) -> Union[LCMSchedulerOutput, Tuple]: |
| 506 | """ |
| 507 | Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion |
| 508 | process from the learned model outputs (most often the predicted noise). |
| 509 | |
| 510 | Args: |
| 511 | model_output (`torch.Tensor`): |
| 512 | The direct output from learned diffusion model. |
| 513 | timestep (`float`): |
| 514 | The current discrete timestep in the diffusion chain. |
| 515 | sample (`torch.Tensor`): |
| 516 | A current instance of a sample created by the diffusion process. |
| 517 | generator (`torch.Generator`, *optional*): |
| 518 | A random number generator. |
| 519 | return_dict (`bool`, *optional*, defaults to `True`): |
| 520 | Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. |
| 521 | Returns: |
| 522 | [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: |
| 523 | If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a |
| 524 | tuple is returned where the first element is the sample tensor. |
| 525 | """ |
| 526 | if self.num_inference_steps is None: |
| 527 | raise ValueError( |
| 528 | "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" |
| 529 | ) |
| 530 | |
| 531 | if self.step_index is None: |
| 532 | self._init_step_index(timestep) |
| 533 | |
| 534 | # 1. get previous step value |
| 535 | prev_step_index = self.step_index + 1 |
| 536 | if prev_step_index < len(self.timesteps): |
| 537 | prev_timestep = self.timesteps[prev_step_index] |
| 538 | else: |
| 539 | prev_timestep = timestep |
| 540 | |
| 541 | # 2. compute alphas, betas |
| 542 | alpha_prod_t = self.alphas_cumprod[timestep] |
| 543 | alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod |
| 544 | |
| 545 | beta_prod_t = 1 - alpha_prod_t |
| 546 | beta_prod_t_prev = 1 - alpha_prod_t_prev |
| 547 | |
| 548 | # 3. Get scalings for boundary conditions |
| 549 | c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) |
| 550 | |
| 551 | # 4. Compute the predicted original sample x_0 based on the model parameterization |
| 552 | if self.config.prediction_type == "epsilon": # noise-prediction |
| 553 | predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() |
| 554 | elif self.config.prediction_type == "sample": # x-prediction |
| 555 | predicted_original_sample = model_output |
no test coverage detected