Variance Exploding Stochastic Differential Equation (VE-SDE) scheduler
Overview
Original paper can be found here.
ScoreSdeVeScheduler
class diffusers.ScoreSdeVeScheduler
< source >( 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
< source >(
sample: FloatTensor
timestep: typing.Optional[int] = None
)
→
torch.FloatTensor
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the current timestep.
set_sigmas
< source >( 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
< source >( num_inference_steps: int sampling_eps: float = None device: typing.Union[str, torch.device] = None )
Sets the continuous timesteps used for the diffusion chain. Supporting function to be run before inference.
step_correct
< source >(
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
< source >(
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).