class_docstr: DPM-Solver (and the improved version DPM-Solver++) is a fast dedicated high-order solver for diffusion ODEs with
  the convergence order guarantee. Empirically, sampling by DPM-Solver with only 20 steps can generate high-quality
  samples, and it can generate quite good samples even in only 10 steps.
  
  For more details, see the original paper: https://arxiv.org/abs/2206.00927 and https://arxiv.org/abs/2211.01095
  
  Currently, we support the multistep DPM-Solver for both noise prediction models and data prediction models. We
  recommend to use `solver_order=2` for guided sampling, and `solver_order=3` for unconditional sampling.
  
  We also support the "dynamic thresholding" method in Imagen (https://arxiv.org/abs/2205.11487). For pixel-space
  diffusion models, you can set both `algorithm_type="dpmsolver++"` and `thresholding=True` to use the dynamic
  thresholding. Note that the thresholding method is unsuitable for latent-space diffusion models (such as
  stable-diffusion).
  
  [`~ConfigMixin`] takes care of storing all config attributes that are passed in the scheduler's `__init__`
  function, such as `num_train_timesteps`. They can be accessed via `scheduler.config.num_train_timesteps`.
  [`SchedulerMixin`] provides general loading and saving functionality via the [`SchedulerMixin.save_pretrained`] and
  [`~SchedulerMixin.from_pretrained`] functions.
  
  For more details, see the original paper: https://arxiv.org/abs/2206.00927 and https://arxiv.org/abs/2211.01095
  
  Args:
      num_train_timesteps (`int`): number of diffusion steps used to train the model.
      beta_start (`float`): the starting `beta` value of inference.
      beta_end (`float`): the final `beta` value.
      beta_schedule (`str`):
          the beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
          `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
      trained_betas (`np.ndarray`, optional):
          option to pass an array of betas directly to the constructor to bypass `beta_start`, `beta_end` etc.
      solver_order (`int`, default `2`):
          the order of DPM-Solver; can be `1` or `2` or `3`. We recommend to use `solver_order=2` for guided
          sampling, and `solver_order=3` for unconditional sampling.
      prediction_type (`str`, default `epsilon`):
          indicates whether the model predicts the noise (epsilon), or the data / `x0`. One of `epsilon`, `sample`,
          or `v-prediction`.
      thresholding (`bool`, default `False`):
          whether to use the "dynamic thresholding" method (introduced by Imagen, https://arxiv.org/abs/2205.11487).
          For pixel-space diffusion models, you can set both `algorithm_type=dpmsolver++` and `thresholding=True` to
          use the dynamic thresholding. Note that the thresholding method is unsuitable for latent-space diffusion
          models (such as stable-diffusion).
      dynamic_thresholding_ratio (`float`, default `0.995`):
          the ratio for the dynamic thresholding method. Default is `0.995`, the same as Imagen
          (https://arxiv.org/abs/2205.11487).
      sample_max_value (`float`, default `1.0`):
          the threshold value for dynamic thresholding. Valid only when `thresholding=True` and
          `algorithm_type="dpmsolver++`.
      algorithm_type (`str`, default `dpmsolver++`):
          the algorithm type for the solver. Either `dpmsolver` or `dpmsolver++`. The `dpmsolver` type implements the
          algorithms in https://arxiv.org/abs/2206.00927, and the `dpmsolver++` type implements the algorithms in
          https://arxiv.org/abs/2211.01095. We recommend to use `dpmsolver++` with `solver_order=2` for guided
          sampling (e.g. stable-diffusion).
      solver_type (`str`, default `midpoint`):
          the solver type for the second-order solver. Either `midpoint` or `heun`. The solver type slightly affects
          the sample quality, especially for small number of steps. We empirically find that `midpoint` solvers are
          slightly better, so we recommend to use the `midpoint` type.
      lower_order_final (`bool`, default `True`):
          whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. We empirically
          find this trick can stabilize the sampling of DPM-Solver for steps < 15, especially for steps <= 10.
      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.
      dtype (`jnp.dtype`, *optional*, defaults to `jnp.float32`):
          the `dtype` used for params and computation.
class_docstr_tok: ["DPM-Solver","(","and","the","improved","version","DPM-Solver++",")","is","a","fast","dedicated","high-order","solver","for","diffusion","ODEs","with","the","convergence","order","guarantee",".","Empirically",",","sampling","by","DPM-Solver","with","only","20","steps","can","generate","high-quality","samples",",","and","it","can","generate","quite","good","samples","even","in","only","10","steps",".","For","more","details",",","see","the","original","paper",":","https",":","\/\/arxiv.org\/abs\/2206.00927","and","https",":","\/\/arxiv.org\/abs\/2211.01095","Currently",",","we","support","the","multistep","DPM-Solver","for","both","noise","prediction","models","and","data","prediction","models",".","We","recommend","to","use","`","solver_order=2","`","for","guided","sampling",",","and","`","solver_order=3","`","for","unconditional","sampling",".","We","also","support","the","``","dynamic","thresholding","''","method","in","Imagen","(","https",":","\/\/arxiv.org\/abs\/2205.11487",")",".","For","pixel-space","diffusion","models",",","you","can","set","both","`","algorithm_type=","''","dpmsolver++","''","`","and","`","thresholding=True","`","to","use","the","dynamic","thresholding",".","Note","that","the","thresholding","method","is","unsuitable","for","latent-space","diffusion","models","(","such","as","stable-diffusion",")",".","[","`","~ConfigMixin","`","]","takes","care","of","storing","all","config","attributes","that","are","passed","in","the","scheduler","'s","`","__init__","`","function",",","such","as","`","num_train_timesteps","`",".","They","can","be","accessed","via","`","scheduler.config.num_train_timesteps","`",".","[","`","SchedulerMixin","`","]","provides","general","loading","and","saving","functionality","via","the","[","`","SchedulerMixin.save_pretrained","`","]","and","[","`","~SchedulerMixin.from_pretrained","`","]","functions",".","For","more","details",",","see","the","original","paper",":","https",":","\/\/arxiv.org\/abs\/2206.00927","and","https",":","\/\/arxiv.org\/abs\/2211.01095","Args",":","num_train_timesteps","(","`","int","`",")",":","number","of","diffusion","steps","used","to","train","the","model",".","beta_start","(","`","float","`",")",":","the","starting","`","beta","`","value","of","inference",".","beta_end","(","`","float","`",")",":","the","final","`","beta","`","value",".","beta_schedule","(","`","str","`",")",":","the","beta","schedule",",","a","mapping","from","a","beta","range","to","a","sequence","of","betas","for","stepping","the","model",".","Choose","from","`","linear","`",",","`","scaled_linear","`",",","or","`","squaredcos_cap_v2","`",".","trained_betas","(","`","np.ndarray","`",",","optional",")",":","option","to","pass","an","array","of","betas","directly","to","the","constructor","to","bypass","`","beta_start","`",",","`","beta_end","`","etc",".","solver_order","(","`","int","`",",","default","`","2","`",")",":","the","order","of","DPM-Solver",";","can","be","`","1","`","or","`","2","`","or","`","3","`",".","We","recommend","to","use","`","solver_order=2","`","for","guided","sampling",",","and","`","solver_order=3","`","for","unconditional","sampling",".","prediction_type","(","`","str","`",",","default","`","epsilon","`",")",":","indicates","whether","the","model","predicts","the","noise","(","epsilon",")",",","or","the","data","\/","`","x0","`",".","One","of","`","epsilon","`",",","`","sample","`",",","or","`","v-prediction","`",".","thresholding","(","`","bool","`",",","default","`","False","`",")",":","whether","to","use","the","``","dynamic","thresholding","''","method","(","introduced","by","Imagen",",","https",":","\/\/arxiv.org\/abs\/2205.11487",")",".","For","pixel-space","diffusion","models",",","you","can","set","both","`","algorithm_type=dpmsolver++","`","and","`","thresholding=True","`","to","use","the","dynamic","thresholding",".","Note","that","the","thresholding","method","is","unsuitable","for","latent-space","diffusion","models","(","such","as","stable-diffusion",")",".","dynamic_thresholding_ratio","(","`","float","`",",","default","`","0.995","`",")",":","the","ratio","for","the","dynamic","thresholding","method",".","Default","is","`","0.995","`",",","the","same","as","Imagen","(","https",":","\/\/arxiv.org\/abs\/2205.11487",")",".","sample_max_value","(","`","float","`",",","default","`","1.0","`",")",":","the","threshold","value","for","dynamic","thresholding",".","Valid","only","when","`","thresholding=True","`","and","`","algorithm_type=","''","dpmsolver++","`",".","algorithm_type","(","`","str","`",",","default","`","dpmsolver++","`",")",":","the","algorithm","type","for","the","solver",".","Either","`","dpmsolver","`","or","`","dpmsolver++","`",".","The","`","dpmsolver","`","type","implements","the","algorithms","in","https",":","\/\/arxiv.org\/abs\/2206.00927",",","and","the","`","dpmsolver++","`","type","implements","the","algorithms","in","https",":","\/\/arxiv.org\/abs\/2211.01095",".","We","recommend","to","use","`","dpmsolver++","`","with","`","solver_order=2","`","for","guided","sampling","(","e.g",".","stable-diffusion",")",".","solver_type","(","`","str","`",",","default","`","midpoint","`",")",":","the","solver","type","for","the","second-order","solver",".","Either","`","midpoint","`","or","`","heun","`",".","The","solver","type","slightly","affects","the","sample","quality",",","especially","for","small","number","of","steps",".","We","empirically","find","that","`","midpoint","`","solvers","are","slightly","better",",","so","we","recommend","to","use","the","`","midpoint","`","type",".","lower_order_final","(","`","bool","`",",","default","`","True","`",")",":","whether","to","use","lower-order","solvers","in","the","final","steps",".","Only","valid","for","<","15","inference","steps",".","We","empirically","find","this","trick","can","stabilize","the","sampling","of","DPM-Solver","for","steps","<","15",",","especially","for","steps","<","=","10.","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",".","dtype","(","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","`","jnp.float32","`",")",":","the","`","dtype","`","used","for","params","and","computation","."]
docstr: Convert the model output to the corresponding type that the algorithm (DPM-Solver / DPM-Solver++) needs.
  
  DPM-Solver is designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to
  discretize an integral of the data prediction model. So we need to first convert the model output to the
  corresponding type to match the algorithm.
  
  Note that the algorithm type and the model type is decoupled. That is to say, we can use either DPM-Solver or
  DPM-Solver++ for both noise prediction model and data prediction model.
  
  Args:
      model_output (`jnp.ndarray`): direct output from learned diffusion model.
      timestep (`int`): current discrete timestep in the diffusion chain.
      sample (`jnp.ndarray`):
          current instance of sample being created by diffusion process.
  
  Returns:
      `jnp.ndarray`: the converted model output.
docstr_tok: ["Convert","the","model","output","to","the","corresponding","type","that","the","algorithm","(","DPM-Solver","\/","DPM-Solver++",")","needs",".","DPM-Solver","is","designed","to","discretize","an","integral","of","the","noise","prediction","model",",","and","DPM-Solver++","is","designed","to","discretize","an","integral","of","the","data","prediction","model",".","So","we","need","to","first","convert","the","model","output","to","the","corresponding","type","to","match","the","algorithm",".","Note","that","the","algorithm","type","and","the","model","type","is","decoupled",".","That","is","to","say",",","we","can","use","either","DPM-Solver","or","DPM-Solver++","for","both","noise","prediction","model","and","data","prediction","model",".","Args",":","model_output","(","`","jnp.ndarray","`",")",":","direct","output","from","learned","diffusion","model",".","timestep","(","`","int","`",")",":","current","discrete","timestep","in","the","diffusion","chain",".","sample","(","`","jnp.ndarray","`",")",":","current","instance","of","sample","being","created","by","diffusion","process",".","Returns",":","`","jnp.ndarray","`",":","the","converted","model","output","."]
code: def convert_model_output(
      self,
      state: DPMSolverMultistepSchedulerState,
      model_output: jnp.ndarray,
      timestep: int,
      sample: jnp.ndarray,
  ) -> jnp.ndarray:
      """
      Convert the model output to the corresponding type that the algorithm (DPM-Solver / DPM-Solver++) needs.
  
      DPM-Solver is designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to
      discretize an integral of the data prediction model. So we need to first convert the model output to the
      corresponding type to match the algorithm.
  
      Note that the algorithm type and the model type is decoupled. That is to say, we can use either DPM-Solver or
      DPM-Solver++ for both noise prediction model and data prediction model.
  
      Args:
          model_output (`jnp.ndarray`): direct output from learned diffusion model.
          timestep (`int`): current discrete timestep in the diffusion chain.
          sample (`jnp.ndarray`):
              current instance of sample being created by diffusion process.
  
      Returns:
          `jnp.ndarray`: the converted model output.
      """
      # DPM-Solver++ needs to solve an integral of the data prediction model.
      if self.config.algorithm_type == "dpmsolver++":
          if self.config.prediction_type == "epsilon":
              alpha_t, sigma_t = (
                  state.alpha_t[timestep],
                  state.sigma_t[timestep],
              )
              x0_pred = (sample - sigma_t * model_output) / alpha_t
          elif self.config.prediction_type == "sample":
              x0_pred = model_output
          elif self.config.prediction_type == "v_prediction":
              alpha_t, sigma_t = (
                  state.alpha_t[timestep],
                  state.sigma_t[timestep],
              )
              x0_pred = alpha_t * sample - sigma_t * model_output
          else:
              raise ValueError(
                  f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, "
                  " or `v_prediction` for the FlaxDPMSolverMultistepScheduler."
              )
  
          if self.config.thresholding:
              # Dynamic thresholding in https://arxiv.org/abs/2205.11487
              dynamic_max_val = jnp.percentile(
                  jnp.abs(x0_pred),
                  self.config.dynamic_thresholding_ratio,
                  axis=tuple(range(1, x0_pred.ndim)),
              )
              dynamic_max_val = jnp.maximum(
                  dynamic_max_val,
                  self.config.sample_max_value
                  * jnp.ones_like(dynamic_max_val),
              )
              x0_pred = (
                  jnp.clip(x0_pred, -dynamic_max_val, dynamic_max_val)
                  / dynamic_max_val
              )
          return x0_pred
      # DPM-Solver needs to solve an integral of the noise prediction model.
      elif self.config.algorithm_type == "dpmsolver":
          if self.config.prediction_type == "epsilon":
              return model_output
          elif self.config.prediction_type == "sample":
              alpha_t, sigma_t = (
                  state.alpha_t[timestep],
                  state.sigma_t[timestep],
              )
              epsilon = (sample - alpha_t * model_output) / sigma_t
              return epsilon
          elif self.config.prediction_type == "v_prediction":
              alpha_t, sigma_t = (
                  state.alpha_t[timestep],
                  state.sigma_t[timestep],
              )
              epsilon = alpha_t * model_output + sigma_t * sample
              return epsilon
          else:
              raise ValueError(
                  f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, "
                  " or `v_prediction` for the FlaxDPMSolverMultistepScheduler."
              )
  
imports: from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  import flax
  import jax
  import jax.numpy
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import CommonSchedulerState, FlaxKarrasDiffusionSchedulers, FlaxSchedulerMixin, FlaxSchedulerOutput, add_noise_common
label_desc: Use image node_id 5 for calling the FlaxDPMSolverMultistepScheduler obj's underlying member method code with example usage: obj.convert_model_output(state, model_output, timestep, sample) and returns: x0_pred, model_output, epsilon, epsilon