ls_mlkit.diffuser.sde.predictor module

class ls_mlkit.diffuser.sde.predictor.NonePredictor(sde, score_fn, use_probability_flow=False)

Bases: Predictor

update_fn(x, t, mask=None)

One update of the predictor.

Parameters:

• x – A PyTorch tensor representing the current state

• t – A Pytorch tensor representing the current time step.

Returns:

A PyTorch tensor of the next state. x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.

Return type:

x

class ls_mlkit.diffuser.sde.predictor.Predictor(sde: SDE, score_fn: object, use_probability_flow=False)

Bases: ABC

abstractmethod update_fn(x: Tensor, t: Tensor, mask=None) → Tuple[Tensor, Tensor]

One update of the predictor.

Parameters:

• x – A PyTorch tensor representing the current state

• t – A Pytorch tensor representing the current time step.

Returns:

A PyTorch tensor of the next state. x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.

Return type:

x

class ls_mlkit.diffuser.sde.predictor.ReverseDiffusionPredictor(sde: SDE, score_fn, use_probability_flow=False, n_dim: int = 3)

Bases: Predictor

update_fn(x: Tensor, t: Tensor, mask=None) → Tuple[Tensor, Tensor]

\[ \begin{align}\begin{aligned}x_{t+\Delta t} &= x_t + f(x_t, t)(\Delta t) + g(x_t, t) \epsilon, \epsilon \sim \mathcal{N}(0,\sqrt{\Delta t}))\\f &= f(x_t, t)|\Delta t|\\g &= g(x_t, t)\sqrt{|\Delta t|}\end{aligned}\end{align} \]