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: Sets the timesteps used for the diffusion chain (to be run before inference).
  
  Args:
      num_inference_steps (`int`):
          The number of diffusion steps used when generating samples with a pre-trained model.
      device (`str` or `torch.device`, *optional*):
          The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
      timesteps (`List[int]`, *optional*):
          Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
          timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,
          `num_inference_steps` must be `None`.
docstr_tok: ["Sets","the","timesteps","used","for","the","diffusion","chain","(","to","be","run","before","inference",")",".","Args",":","num_inference_steps","(","`","int","`",")",":","The","number","of","diffusion","steps","used","when","generating","samples","with","a","pre-trained","model",".","device","(","`","str","`","or","`","torch.device","`",",","*","optional","*",")",":","The","device","to","which","the","timesteps","should","be","moved","to",".","If","`","None","`",",","the","timesteps","are","not","moved",".","timesteps","(","`","List","[","int","]","`",",","*","optional","*",")",":","Custom","timesteps","used","to","support","arbitrary","spacing","between","timesteps",".","If","`","None","`",",","then","the","default","timestep","spacing","strategy","of","equal","spacing","between","timesteps","is","used",".","If","`","timesteps","`","is","passed",",","`","num_inference_steps","`","must","be","`","None","`","."]
code: def set_timesteps(
      self,
      num_inference_steps: Optional[int] = None,
      device: Union[str, torch.device] = None,
      timesteps: Optional[List[int]] = None,
  ):
      """
      Sets the timesteps used for the diffusion chain (to be run before inference).
  
      Args:
          num_inference_steps (`int`):
              The number of diffusion steps used when generating samples with a pre-trained model.
          device (`str` or `torch.device`, *optional*):
              The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
          timesteps (`List[int]`, *optional*):
              Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
              timestep spacing strategy of equal spacing between timesteps is used. If `timesteps` is passed,
              `num_inference_steps` must be `None`.
      """
      if num_inference_steps is None and timesteps is None:
          raise ValueError(
              "Exactly one of `num_inference_steps` or `timesteps` must be supplied."
          )
  
      if num_inference_steps is not None and timesteps is not None:
          raise ValueError(
              "Can only pass one of `num_inference_steps` or `timesteps`."
          )
  
      # Follow DDPMScheduler custom timesteps logic
      if timesteps is not None:
          for i in range(1, len(timesteps)):
              if timesteps[i] >= timesteps[i - 1]:
                  raise ValueError(
                      "`timesteps` must be in descending order."
                  )
  
          if timesteps[0] >= self.config.num_train_timesteps:
              raise ValueError(
                  f"`timesteps` must start before `self.config.train_timesteps`:"
                  f" {self.config.num_train_timesteps}."
              )
  
          timesteps = np.array(timesteps, dtype=np.int64)
          self.custom_timesteps = True
      else:
          if num_inference_steps > self.config.num_train_timesteps:
              raise ValueError(
                  f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
                  f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
                  f" maximal {self.config.num_train_timesteps} timesteps."
              )
  
          self.num_inference_steps = num_inference_steps
  
          step_ratio = (
              self.config.num_train_timesteps
              // self.num_inference_steps
          )
          timesteps = (
              (np.arange(0, num_inference_steps) * step_ratio)
              .round()[::-1]
              .copy()
              .astype(np.int64)
          )
          self.custom_timesteps = False
  
      # Map timesteps to Karras sigmas directly for multistep sampling
      # See https://github.com/openai/consistency_models/blob/main/cm/karras_diffusion.py#L675
      num_train_timesteps = self.config.num_train_timesteps
      ramp = timesteps[::-1].copy()
      ramp = ramp / (num_train_timesteps - 1)
      sigmas = self._convert_to_karras(ramp)
      timesteps = self.sigma_to_t(sigmas)
  
      sigmas = np.concatenate([sigmas, [self.sigma_min]]).astype(
          np.float32
      )
      self.sigmas = torch.from_numpy(sigmas).to(device=device)
  
      if str(device).startswith("mps"):
          # mps does not support float64
          self.timesteps = torch.from_numpy(timesteps).to(
              device, dtype=torch.float32
          )
      else:
          self.timesteps = torch.from_numpy(timesteps).to(device=device)
  
      self._step_index = None
      self.sigmas.to("cpu")
  
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 6 for calling the CMStochasticIterativeScheduler obj's underlying member method code with example usage: obj.set_timesteps(num_inference_steps, device, timesteps) without return types