Diffusers documentation

Variance Exploding Stochastic Differential Equation (VE-SDE) scheduler

You are viewing v0.16.0 version. A newer version v0.31.0 is available.
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Variance Exploding Stochastic Differential Equation (VE-SDE) scheduler

Overview

Original paper can be found here.

ScoreSdeVeScheduler

class diffusers.ScoreSdeVeScheduler

< >

( num_train_timesteps: int = 2000 snr: float = 0.15 sigma_min: float = 0.01 sigma_max: float = 1348.0 sampling_eps: float = 1e-05 correct_steps: int = 1 )

Parameters

  • num_train_timesteps (int) — number of diffusion steps used to train the model.
  • snr (float) — coefficient weighting the step from the model_output sample (from the network) to the random noise.
  • sigma_min (float) — initial noise scale for sigma sequence in sampling procedure. The minimum sigma should mirror the distribution of the data.
  • sigma_max (float) — maximum value used for the range of continuous timesteps passed into the model.
  • sampling_eps (float) — the end value of sampling, where timesteps decrease progressively from 1 to epsilon. —
  • correct_steps (int) — number of correction steps performed on a produced sample.

The variance exploding stochastic differential equation (SDE) scheduler.

For more information, see the original paper: https://arxiv.org/abs/2011.13456

~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 from_pretrained() functions.

scale_model_input

< >

( sample: FloatTensor timestep: typing.Optional[int] = None ) torch.FloatTensor

Parameters

  • sample (torch.FloatTensor) — input sample
  • timestep (int, optional) — current timestep

Returns

torch.FloatTensor

scaled input sample

Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep.

set_sigmas

< >

( num_inference_steps: int sigma_min: float = None sigma_max: float = None sampling_eps: float = None )

Parameters

  • num_inference_steps (int) — the number of diffusion steps used when generating samples with a pre-trained model.
  • sigma_min (float, optional) — initial noise scale value (overrides value given at Scheduler instantiation).
  • sigma_max (float, optional) — final noise scale value (overrides value given at Scheduler instantiation).
  • sampling_eps (float, optional) — final timestep value (overrides value given at Scheduler instantiation).

Sets the noise scales used for the diffusion chain. Supporting function to be run before inference.

The sigmas control the weight of the drift and diffusion components of sample update.

set_timesteps

< >

( num_inference_steps: int sampling_eps: float = None device: typing.Union[str, torch.device] = None )

Parameters

  • num_inference_steps (int) — the number of diffusion steps used when generating samples with a pre-trained model.
  • sampling_eps (float, optional) — final timestep value (overrides value given at Scheduler instantiation).

Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference.

step_correct

< >

( model_output: FloatTensor sample: FloatTensor generator: typing.Optional[torch._C.Generator] = None return_dict: bool = True ) SdeVeOutput or tuple

Parameters

  • model_output (torch.FloatTensor) — direct output from learned diffusion model.
  • sample (torch.FloatTensor) — current instance of sample being created by diffusion process. generator — random number generator.
  • return_dict (bool) — option for returning tuple rather than SchedulerOutput class

Returns

SdeVeOutput or tuple

SdeVeOutput if return_dict is True, otherwise a tuple. When returning a tuple, the first element is the sample tensor.

Correct the predicted sample based on the output model_output of the network. This is often run repeatedly after making the prediction for the previous timestep.

step_pred

< >

( model_output: FloatTensor timestep: int sample: FloatTensor generator: typing.Optional[torch._C.Generator] = None return_dict: bool = True ) SdeVeOutput or tuple

Parameters

  • model_output (torch.FloatTensor) — direct output from learned diffusion model.
  • timestep (int) — current discrete timestep in the diffusion chain.
  • sample (torch.FloatTensor) — current instance of sample being created by diffusion process. generator — random number generator.
  • return_dict (bool) — option for returning tuple rather than SchedulerOutput class

Returns

SdeVeOutput or tuple

SdeVeOutput if return_dict is True, otherwise a tuple. When returning a tuple, the first element is the sample tensor.

Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process from the learned model outputs (most often the predicted noise).