class_docstr: KDPM2DiscreteScheduler with ancestral sampling is inspired by the DPMSolver2 and Algorithm 2 from the [Elucidating
  the Design Space of Diffusion-Based Generative Models](https://huggingface.co/papers/2206.00364) paper.
  
  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.00085):
          The starting `beta` value of inference.
      beta_end (`float`, defaults to 0.012):
          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` or `scaled_linear`.
      trained_betas (`np.ndarray`, *optional*):
          Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
      use_karras_sigmas (`bool`, *optional*, defaults to `False`):
          Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
          the sigmas are determined according to a sequence of noise levels {σi}.
      prediction_type (`str`, defaults to `epsilon`, *optional*):
          Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
          `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
          Video](https://imagen.research.google/video/paper.pdf) paper).
      timestep_spacing (`str`, defaults to `"linspace"`):
          The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
          Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
      steps_offset (`int`, defaults to 0):
          An offset added to the inference steps. You can use a combination of `offset=1` and
          `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
          Diffusion.
class_docstr_tok: ["KDPM2DiscreteScheduler","with","ancestral","sampling","is","inspired","by","the","DPMSolver2","and","Algorithm","2","from","the","[","Elucidating","the","Design","Space","of","Diffusion-Based","Generative","Models","]","(","https",":","\/\/huggingface.co\/papers\/2206.00364",")","paper",".","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.00085",")",":","The","starting","`","beta","`","value","of","inference",".","beta_end","(","`","float","`",",","defaults","to","0.012",")",":","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","`","or","`","scaled_linear","`",".","trained_betas","(","`","np.ndarray","`",",","*","optional","*",")",":","Pass","an","array","of","betas","directly","to","the","constructor","to","bypass","`","beta_start","`","and","`","beta_end","`",".","use_karras_sigmas","(","`","bool","`",",","*","optional","*",",","defaults","to","`","False","`",")",":","Whether","to","use","Karras","sigmas","for","step","sizes","in","the","noise","schedule","during","the","sampling","process",".","If","`","True","`",",","the","sigmas","are","determined","according","to","a","sequence","of","noise","levels","{","\u03c3i","}",".","prediction_type","(","`","str","`",",","defaults","to","`","epsilon","`",",","*","optional","*",")",":","Prediction","type","of","the","scheduler","function",";","can","be","`","epsilon","`","(","predicts","the","noise","of","the","diffusion","process",")",",","`","sample","`","(","directly","predicts","the","noisy","sample","`",")","or","`","v_prediction","`","(","see","section","2.4","of","[","Imagen","Video","]","(","https",":","\/\/imagen.research.google\/video\/paper.pdf",")","paper",")",".","timestep_spacing","(","`","str","`",",","defaults","to","`","``","linspace","''","`",")",":","The","way","the","timesteps","should","be","scaled",".","Refer","to","Table","2","of","the","[","Common","Diffusion","Noise","Schedules","and","Sample","Steps","are","Flawed","]","(","https",":","\/\/huggingface.co\/papers\/2305.08891",")","for","more","information",".","steps_offset","(","`","int","`",",","defaults","to","0",")",":","An","offset","added","to","the","inference","steps",".","You","can","use","a","combination","of","`","offset=1","`","and","`","set_alpha_to_one=False","`","to","make","the","last","step","use","step","0","for","the","previous","alpha","product","like","in","Stable","Diffusion","."]
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 (`float`):
          The current discrete timestep in the diffusion chain.
      sample (`torch.FloatTensor`):
          A current instance of a sample created by the diffusion process.
      generator (`torch.Generator`, *optional*):
          A random number generator.
      return_dict (`bool`):
          Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or tuple.
  
  Returns:
      [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
          If return_dict is `True`, [`~schedulers.scheduling_ddim.SchedulerOutput`] 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","(","`","float","`",")",":","The","current","discrete","timestep","in","the","diffusion","chain",".","sample","(","`","torch.FloatTensor","`",")",":","A","current","instance","of","a","sample","created","by","the","diffusion","process",".","generator","(","`","torch.Generator","`",",","*","optional","*",")",":","A","random","number","generator",".","return_dict","(","`","bool","`",")",":","Whether","or","not","to","return","a","[","`","~schedulers.scheduling_utils.SchedulerOutput","`","]","or","tuple",".","Returns",":","[","`","~schedulers.scheduling_utils.SchedulerOutput","`","]","or","`","tuple","`",":","If","return_dict","is","`","True","`",",","[","`","~schedulers.scheduling_ddim.SchedulerOutput","`","]","is","returned",",","otherwise","a","tuple","is","returned","where","the","first","element","is","the","sample","tensor","."]
code: def step(
      self,
      model_output: Union[torch.FloatTensor, np.ndarray],
      timestep: Union[float, torch.FloatTensor],
      sample: Union[torch.FloatTensor, np.ndarray],
      generator: Optional[torch.Generator] = None,
      return_dict: bool = True,
  ) -> Union[SchedulerOutput, 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 (`float`):
              The current discrete timestep in the diffusion chain.
          sample (`torch.FloatTensor`):
              A current instance of a sample created by the diffusion process.
          generator (`torch.Generator`, *optional*):
              A random number generator.
          return_dict (`bool`):
              Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or tuple.
  
      Returns:
          [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
              If return_dict is `True`, [`~schedulers.scheduling_ddim.SchedulerOutput`] is returned, otherwise a
              tuple is returned where the first element is the sample tensor.
      """
      if self.step_index is None:
          self._init_step_index(timestep)
  
      # advance index counter by 1
      timestep_int = (
          timestep.cpu().item()
          if torch.is_tensor(timestep)
          else timestep
      )
      self._index_counter[timestep_int] += 1
  
      if self.state_in_first_order:
          sigma = self.sigmas[self.step_index]
          sigma_interpol = self.sigmas_interpol[self.step_index]
          sigma_up = self.sigmas_up[self.step_index]
          sigma_down = self.sigmas_down[self.step_index - 1]
      else:
          # 2nd order / KPDM2's method
          sigma = self.sigmas[self.step_index - 1]
          sigma_interpol = self.sigmas_interpol[self.step_index - 1]
          sigma_up = self.sigmas_up[self.step_index - 1]
          sigma_down = self.sigmas_down[self.step_index - 1]
  
      # currently only gamma=0 is supported. This usually works best anyways.
      # We can support gamma in the future but then need to scale the timestep before
      # passing it to the model which requires a change in API
      gamma = 0
      sigma_hat = sigma * (
          gamma + 1
      )  # Note: sigma_hat == sigma for now
  
      device = model_output.device
      noise = randn_tensor(
          model_output.shape,
          dtype=model_output.dtype,
          device=device,
          generator=generator,
      )
  
      # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
      if self.config.prediction_type == "epsilon":
          sigma_input = (
              sigma_hat if self.state_in_first_order else sigma_interpol
          )
          pred_original_sample = sample - sigma_input * model_output
      elif self.config.prediction_type == "v_prediction":
          sigma_input = (
              sigma_hat if self.state_in_first_order else sigma_interpol
          )
          pred_original_sample = model_output * (
              -sigma_input / (sigma_input**2 + 1) ** 0.5
          ) + (sample / (sigma_input**2 + 1))
      elif self.config.prediction_type == "sample":
          raise NotImplementedError(
              "prediction_type not implemented yet: sample"
          )
      else:
          raise ValueError(
              f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
          )
  
      if self.state_in_first_order:
          # 2. Convert to an ODE derivative for 1st order
          derivative = (sample - pred_original_sample) / sigma_hat
          # 3. delta timestep
          dt = sigma_interpol - sigma_hat
  
          # store for 2nd order step
          self.sample = sample
          self.dt = dt
          prev_sample = sample + derivative * dt
      else:
          # DPM-Solver-2
          # 2. Convert to an ODE derivative for 2nd order
          derivative = (sample - pred_original_sample) / sigma_interpol
          # 3. delta timestep
          dt = sigma_down - sigma_hat
  
          sample = self.sample
          self.sample = None
  
          prev_sample = sample + derivative * dt
          prev_sample = prev_sample + noise * sigma_up
  
      # upon completion increase step index by one
      self._step_index += 1
  
      if not return_dict:
          return (prev_sample,)
  
      return SchedulerOutput(prev_sample=prev_sample)
  
imports: import math
  from collections import defaultdict
  from typing import List, Optional, Tuple, Union
  import numpy
  import torch
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils.torch_utils import randn_tensor
  from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
label_desc: Use image node_id 11 for calling the KDPM2AncestralDiscreteScheduler obj's underlying member method code with example usage: obj.step(model_output, timestep, sample, generator, return_dict) and returns: SchedulerOutput, prev_sample