class_docstr: Multistep and onestep sampling for consistency models.
  
  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 40):
          The number of diffusion steps to train the model.
      sigma_min (`float`, defaults to 0.002):
          Minimum noise magnitude in the sigma schedule. Defaults to 0.002 from the original implementation.
      sigma_max (`float`, defaults to 80.0):
          Maximum noise magnitude in the sigma schedule. Defaults to 80.0 from the original implementation.
      sigma_data (`float`, defaults to 0.5):
          The standard deviation of the data distribution from the EDM
          [paper](https://huggingface.co/papers/2206.00364). Defaults to 0.5 from the original implementation.
      s_noise (`float`, defaults to 1.0):
          The amount of additional noise to counteract loss of detail during sampling. A reasonable range is [1.000,
          1.011]. Defaults to 1.0 from the original implementation.
      rho (`float`, defaults to 7.0):
          The parameter for calculating the Karras sigma schedule from the EDM
          [paper](https://huggingface.co/papers/2206.00364). Defaults to 7.0 from the original implementation.
      clip_denoised (`bool`, defaults to `True`):
          Whether to clip the denoised outputs to `(-1, 1)`.
      timesteps (`List` or `np.ndarray` or `torch.Tensor`, *optional*):
          An explicit timestep schedule that can be optionally specified. The timesteps are expected to be in
          increasing order.
class_docstr_tok: ["Multistep","and","onestep","sampling","for","consistency","models",".","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","40",")",":","The","number","of","diffusion","steps","to","train","the","model",".","sigma_min","(","`","float","`",",","defaults","to","0.002",")",":","Minimum","noise","magnitude","in","the","sigma","schedule",".","Defaults","to","0.002","from","the","original","implementation",".","sigma_max","(","`","float","`",",","defaults","to","80.0",")",":","Maximum","noise","magnitude","in","the","sigma","schedule",".","Defaults","to","80.0","from","the","original","implementation",".","sigma_data","(","`","float","`",",","defaults","to","0.5",")",":","The","standard","deviation","of","the","data","distribution","from","the","EDM","[","paper","]","(","https",":","\/\/huggingface.co\/papers\/2206.00364",")",".","Defaults","to","0.5","from","the","original","implementation",".","s_noise","(","`","float","`",",","defaults","to","1.0",")",":","The","amount","of","additional","noise","to","counteract","loss","of","detail","during","sampling",".","A","reasonable","range","is","[","1.000",",","1.011","]",".","Defaults","to","1.0","from","the","original","implementation",".","rho","(","`","float","`",",","defaults","to","7.0",")",":","The","parameter","for","calculating","the","Karras","sigma","schedule","from","the","EDM","[","paper","]","(","https",":","\/\/huggingface.co\/papers\/2206.00364",")",".","Defaults","to","7.0","from","the","original","implementation",".","clip_denoised","(","`","bool","`",",","defaults","to","`","True","`",")",":","Whether","to","clip","the","denoised","outputs","to","`","(","-1",",","1",")","`",".","timesteps","(","`","List","`","or","`","np.ndarray","`","or","`","torch.Tensor","`",",","*","optional","*",")",":","An","explicit","timestep","schedule","that","can","be","optionally","specified",".","The","timesteps","are","expected","to","be","in","increasing","order","."]
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 the learned diffusion model.
      timestep (`float`):
          The current 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`, *optional*, defaults to `True`):
          Whether or not to return a
          [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`.
  
  Returns:
      [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`:
          If return_dict is `True`,
          [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] 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","the","learned","diffusion","model",".","timestep","(","`","float","`",")",":","The","current","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","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput","`","]","or","`","tuple","`",".","Returns",":","[","`","~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput","`","]","or","`","tuple","`",":","If","return_dict","is","`","True","`",",","[","`","~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput","`","]","is","returned",",","otherwise","a","tuple","is","returned","where","the","first","element","is","the","sample","tensor","."]
code: def step(
      self,
      model_output: torch.FloatTensor,
      timestep: Union[float, torch.FloatTensor],
      sample: torch.FloatTensor,
      generator: Optional[torch.Generator] = None,
      return_dict: bool = True,
  ) -> Union[CMStochasticIterativeSchedulerOutput, 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 the learned diffusion model.
          timestep (`float`):
              The current 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`, *optional*, defaults to `True`):
              Whether or not to return a
              [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`.
  
      Returns:
          [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] or `tuple`:
              If return_dict is `True`,
              [`~schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput`] is returned,
              otherwise a tuple is returned where the first element is the sample tensor.
      """
  
      if (
          isinstance(timestep, int)
          or isinstance(timestep, torch.IntTensor)
          or isinstance(timestep, torch.LongTensor)
      ):
          raise ValueError(
              (
                  "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
                  f" `{self.__class__}.step()` is not supported. Make sure to pass"
                  " one of the `scheduler.timesteps` as a timestep."
              ),
          )
  
      if not self.is_scale_input_called:
          logger.warning(
              "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
              "See `StableDiffusionPipeline` for a usage example."
          )
  
      sigma_min = self.config.sigma_min
      sigma_max = self.config.sigma_max
  
      if self.step_index is None:
          self._init_step_index(timestep)
  
      # sigma_next corresponds to next_t in original implementation
      sigma = self.sigmas[self.step_index]
      if self.step_index + 1 < self.config.num_train_timesteps:
          sigma_next = self.sigmas[self.step_index + 1]
      else:
          # Set sigma_next to sigma_min
          sigma_next = self.sigmas[-1]
  
      # Get scalings for boundary conditions
      c_skip, c_out = self.get_scalings_for_boundary_condition(sigma)
  
      # 1. Denoise model output using boundary conditions
      denoised = c_out * model_output + c_skip * sample
      if self.config.clip_denoised:
          denoised = denoised.clamp(-1, 1)
  
      # 2. Sample z ~ N(0, s_noise^2 * I)
      # Noise is not used for onestep sampling.
      if len(self.timesteps) > 1:
          noise = randn_tensor(
              model_output.shape,
              dtype=model_output.dtype,
              device=model_output.device,
              generator=generator,
          )
      else:
          noise = torch.zeros_like(model_output)
      z = noise * self.config.s_noise
  
      sigma_hat = sigma_next.clamp(min=sigma_min, max=sigma_max)
  
      # 3. Return noisy sample
      # tau = sigma_hat, eps = sigma_min
      prev_sample = (
          denoised + z * (sigma_hat**2 - sigma_min**2) ** 0.5
      )
  
      # upon completion increase step index by one
      self._step_index += 1
  
      if not return_dict:
          return (prev_sample,)
  
      return CMStochasticIterativeSchedulerOutput(
          prev_sample=prev_sample
      )
  
imports: from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  import numpy
  import torch
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils import BaseOutput, logging
  from ..utils.torch_utils import randn_tensor
  from .scheduling_utils import SchedulerMixin
label_desc: Use image node_id 11 for calling the CMStochasticIterativeScheduler obj's underlying member method code with example usage: obj.step(model_output, timestep, sample, generator, return_dict) and returns: CMStochasticIterativeSchedulerOutput, prev_sample