class_docstr: `RePaintScheduler` is a scheduler for DDPM inpainting inside a given mask.
  
  This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
  methods the library implements for all schedulers such as loading and saving.
  
  Args:
      num_train_timesteps (`int`, defaults to 1000):
          The number of diffusion steps to train the model.
      beta_start (`float`, defaults to 0.0001):
          The starting `beta` value of inference.
      beta_end (`float`, defaults to 0.02):
          The final `beta` value.
      beta_schedule (`str`, defaults to `"linear"`):
          The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
          `linear`, `scaled_linear`, `squaredcos_cap_v2`, or `sigmoid`.
      eta (`float`):
          The weight of noise for added noise in diffusion step. If its value is between 0.0 and 1.0 it corresponds
          to the DDIM scheduler, and if its value is between -0.0 and 1.0 it corresponds to the DDPM scheduler.
      trained_betas (`np.ndarray`, *optional*):
          Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
      clip_sample (`bool`, defaults to `True`):
          Clip the predicted sample between -1 and 1 for numerical stability.
docstr: 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.FloatTensor`):
          The direct output from learned diffusion model.
      timestep (`int`):
          The current discrete timestep in the diffusion chain.
      sample (`torch.FloatTensor`):
          A current instance of a sample created by the diffusion process.
      original_image (`torch.FloatTensor`):
          The original image to inpaint on.
      mask (`torch.FloatTensor`):
          The mask where a value of 0.0 indicates which part of the original image to inpaint.
      generator (`torch.Generator`, *optional*):
          A random number generator.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`.
  
  Returns:
      [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`:
          If return_dict is `True`, [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] is returned,
          otherwise a tuple is returned where the first element is the sample tensor.
docstr_tok: ["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.FloatTensor","`",")",":","The","direct","output","from","learned","diffusion","model",".","timestep","(","`","int","`",")",":","The","current","discrete","timestep","in","the","diffusion","chain",".","sample","(","`","torch.FloatTensor","`",")",":","A","current","instance","of","a","sample","created","by","the","diffusion","process",".","original_image","(","`","torch.FloatTensor","`",")",":","The","original","image","to","inpaint","on",".","mask","(","`","torch.FloatTensor","`",")",":","The","mask","where","a","value","of","0.0","indicates","which","part","of","the","original","image","to","inpaint",".","generator","(","`","torch.Generator","`",",","*","optional","*",")",":","A","random","number","generator",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~schedulers.scheduling_repaint.RePaintSchedulerOutput","`","]","or","`","tuple","`",".","Returns",":","[","`","~schedulers.scheduling_repaint.RePaintSchedulerOutput","`","]","or","`","tuple","`",":","If","return_dict","is","`","True","`",",","[","`","~schedulers.scheduling_repaint.RePaintSchedulerOutput","`","]","is","returned",",","otherwise","a","tuple","is","returned","where","the","first","element","is","the","sample","tensor","."]
code: def step(
      self,
      model_output: torch.FloatTensor,
      timestep: int,
      sample: torch.FloatTensor,
      original_image: torch.FloatTensor,
      mask: torch.FloatTensor,
      generator: Optional[torch.Generator] = None,
      return_dict: bool = True,
  ) -> Union[RePaintSchedulerOutput, Tuple]:
      """
      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.FloatTensor`):
              The direct output from learned diffusion model.
          timestep (`int`):
              The current discrete timestep in the diffusion chain.
          sample (`torch.FloatTensor`):
              A current instance of a sample created by the diffusion process.
          original_image (`torch.FloatTensor`):
              The original image to inpaint on.
          mask (`torch.FloatTensor`):
              The mask where a value of 0.0 indicates which part of the original image to inpaint.
          generator (`torch.Generator`, *optional*):
              A random number generator.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`.
  
      Returns:
          [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] or `tuple`:
              If return_dict is `True`, [`~schedulers.scheduling_repaint.RePaintSchedulerOutput`] is returned,
              otherwise a tuple is returned where the first element is the sample tensor.
  
      """
      t = timestep
      prev_timestep = (
          timestep
          - self.config.num_train_timesteps // self.num_inference_steps
      )
  
      # 1. compute alphas, betas
      alpha_prod_t = self.alphas_cumprod[t]
      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
  
      # 2. compute predicted original sample from predicted noise also called
      # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
      pred_original_sample = (
          sample - beta_prod_t**0.5 * model_output
      ) / alpha_prod_t**0.5
  
      # 3. Clip "predicted x_0"
      if self.config.clip_sample:
          pred_original_sample = torch.clamp(
              pred_original_sample, -1, 1
          )
  
      # We choose to follow RePaint Algorithm 1 to get x_{t-1}, however we
      # substitute formula (7) in the algorithm coming from DDPM paper
      # (formula (4) Algorithm 2 - Sampling) with formula (12) from DDIM paper.
      # DDIM schedule gives the same results as DDPM with eta = 1.0
      # Noise is being reused in 7. and 8., but no impact on quality has
      # been observed.
  
      # 5. Add noise
      device = model_output.device
      noise = randn_tensor(
          model_output.shape,
          generator=generator,
          device=device,
          dtype=model_output.dtype,
      )
      std_dev_t = self.eta * self._get_variance(timestep) ** 0.5
  
      variance = 0
      if t > 0 and self.eta > 0:
          variance = std_dev_t * noise
  
      # 6. compute "direction pointing to x_t" of formula (12)
      # from https://arxiv.org/pdf/2010.02502.pdf
      pred_sample_direction = (
          1 - alpha_prod_t_prev - std_dev_t**2
      ) ** 0.5 * model_output
  
      # 7. compute x_{t-1} of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
      prev_unknown_part = (
          alpha_prod_t_prev**0.5 * pred_original_sample
          + pred_sample_direction
          + variance
      )
  
      # 8. Algorithm 1 Line 5 https://arxiv.org/pdf/2201.09865.pdf
      prev_known_part = (alpha_prod_t_prev**0.5) * original_image + (
          (1 - alpha_prod_t_prev) ** 0.5
      ) * noise
  
      # 9. Algorithm 1 Line 8 https://arxiv.org/pdf/2201.09865.pdf
      pred_prev_sample = (
          mask * prev_known_part + (1.0 - mask) * prev_unknown_part
      )
  
      if not return_dict:
          return (
              pred_prev_sample,
              pred_original_sample,
          )
  
      return RePaintSchedulerOutput(
          prev_sample=pred_prev_sample,
          pred_original_sample=pred_original_sample,
      )
  
imports: import math
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import numpy
  import torch
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils import BaseOutput
  from ..utils.torch_utils import randn_tensor
  from .scheduling_utils import SchedulerMixin
label_desc: Use image node_id 5 for calling the RePaintScheduler obj's underlying member method code with example usage: obj.step(model_output, timestep, sample, original_image, mask, generator, return_dict) and returns: RePaintSchedulerOutput, pred_prev_sample,  pred_original_sample