autocast.processors.noise

Noise schedules for diffusion processors.

class NoiseSchedule

Bases: Schedule

Noise Schedule Module.

Implements Azula’s Schedule interface for local schedules.

forward(t)

Get alpha and sigma for given time steps t.

Parameters:

t (Tensor)

Return type:

tuple[Tensor, Tensor]

class LogLinearSchedule(sigma_min=0.002, sigma_max=80.0)

Bases: NoiseSchedule

Log-Linear Noise Schedule.

Implements a log-linear schedule for alpha and sigma.

Parameters:

• sigma_min (float)

• sigma_max (float)

forward(t)

Get alpha and sigma for given time steps t.

Parameters:

t (Tensor)

Return type:

tuple[Tensor, Tensor]

class LogLogitSchedule(sigma_min=0.001, sigma_max=1000.0, scale=1.0, shift=0.0)

Bases: NoiseSchedule

LoLA log-logit noise schedule.

Matches LogLogitSchedule in francois-rozet/lola.

Parameters:

• sigma_min (float)

• sigma_max (float)

• scale (float)

• shift (float)

forward(t)

Get alpha and sigma for given time steps t.

Parameters:

t (Tensor)

Return type:

tuple[Tensor, Tensor]