CMStochasticIterativeScheduler
Consistency Models by Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever introduced a multistep and onestep scheduler (Algorithm 1) that is capable of generating good samples in one or a small number of steps.
The abstract from the paper is:
Diffusion models have made significant breakthroughs in image, audio, and video generation, but they depend on an iterative generation process that causes slow sampling speed and caps their potential for realtime applications. To overcome this limitation, we propose consistency models, a new family of generative models that achieve high sample quality without adversarial training. They support fast onestep generation by design, while still allowing for fewstep sampling to trade compute for sample quality. They also support zeroshot data editing, like image inpainting, colorization, and superresolution, without requiring explicit training on these tasks. Consistency models can be trained either as a way to distill pretrained diffusion models, or as standalone generative models. Through extensive experiments, we demonstrate that they outperform existing distillation techniques for diffusion models in one and fewstep generation. For example, we achieve the new stateoftheart FID of 3.55 on CIFAR10 and 6.20 on ImageNet 64x64 for onestep generation. When trained as standalone generative models, consistency models also outperform singlestep, nonadversarial generative models on standard benchmarks like CIFAR10, ImageNet 64x64 and LSUN 256x256.
The original codebase can be found at openai/consistency_models.
CMStochasticIterativeScheduler
class diffusers.CMStochasticIterativeScheduler
< source >( num_train_timesteps: int = 40 sigma_min: float = 0.002 sigma_max: float = 80.0 sigma_data: float = 0.5 s_noise: float = 1.0 rho: float = 7.0 clip_denoised: bool = True )
Parameters

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. 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. Defaults to 7.0 from the original implementation. 
clip_denoised (
bool
, defaults toTrue
) — Whether to clip the denoised outputs to(1, 1)
. 
timesteps (
List
ornp.ndarray
ortorch.Tensor
, optional) — An explicit timestep schedule that can be optionally specified. The timesteps are expected to be in increasing order.
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.
get_scalings_for_boundary_condition
< source >(
sigma
)
→
tuple
Gets the scalings used in the consistency model parameterization (from Appendix C of the paper) to enforce boundary condition.
epsilon
in the equations for c_skip
and c_out
is set to sigma_min
.
scale_model_input
< source >(
sample: FloatTensor
timestep: typing.Union[float, torch.FloatTensor]
)
→
torch.FloatTensor
Scales the consistency model input by (sigma**2 + sigma_data**2) ** 0.5
.
set_timesteps
< source >( num_inference_steps: typing.Optional[int] = None device: typing.Union[str, torch.device] = None timesteps: typing.Optional[typing.List[int]] = None )
Parameters

num_inference_steps (
int
) — The number of diffusion steps used when generating samples with a pretrained model. 
device (
str
ortorch.device
, optional) — The device to which the timesteps should be moved to. IfNone
, the timesteps are not moved. 
timesteps (
List[int]
, optional) — Custom timesteps used to support arbitrary spacing between timesteps. IfNone
, then the default timestep spacing strategy of equal spacing between timesteps is used. Iftimesteps
is passed,num_inference_steps
must beNone
.
Sets the timesteps used for the diffusion chain (to be run before inference).
sigma_to_t
< source >(
sigmas: typing.Union[float, numpy.ndarray]
)
→
float
or np.ndarray
Gets scaled timesteps from the Karras sigmas for input to the consistency model.
step
< source >(
model_output: FloatTensor
timestep: typing.Union[float, torch.FloatTensor]
sample: FloatTensor
generator: typing.Optional[torch._C.Generator] = None
return_dict: bool = True
)
→
CMStochasticIterativeSchedulerOutput or tuple
Parameters

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 toTrue
) — Whether or not to return a CMStochasticIterativeSchedulerOutput ortuple
.
Returns
CMStochasticIterativeSchedulerOutput or tuple
If return_dict is True
,
CMStochasticIterativeSchedulerOutput is returned,
otherwise a tuple is returned where the first element is the sample tensor.
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).
CMStochasticIterativeSchedulerOutput
class diffusers.schedulers.scheduling_consistency_models.CMStochasticIterativeSchedulerOutput
< source >( prev_sample: FloatTensor )
Output class for the scheduler’s step
function.