sde_lib.py 🚧¶

Source: LION/models/diffusion/NCSNpp_helpers/sde_lib.py

Warning

This file has not yet received a complete narrative and docstring audit. Its public source-level API is listed automatically below.

Abstract SDE classes, Reverse SDE, and VE/VP SDEs.

class SDE(N): SDE abstract class. Functions are designed for a mini-batch of inputs.

SDE.T(self): End time of the SDE.

SDE.sde(self, x, t): No docstring is available.

SDE.marginal_prob(self, x, t): Parameters to determine the marginal distribution of the SDE, $p_t(x)$.

SDE.prior_sampling(self, shape): Generate one sample from the prior distribution, $p_T(x)$.

SDE.prior_logp(self, z)

Compute log-density of the prior distribution.

Useful for computing the log-likelihood via probability flow ODE.

Args: z: latent code Returns: log probability density

SDE.discretize(self, x, t)

Discretize the SDE in the form: x_{i+1} = x_i + f_i(x_i) + G_i z_i.

Useful for reverse diffusion sampling and probabiliy flow sampling. Defaults to Euler-Maruyama discretization.

Args: x: a torch tensor t: a torch float representing the time step (from 0 to self.T)

Returns: f, G

SDE.reverse(self, score_fn, probability_flow=False)

Create the reverse-time SDE/ODE.

Args: score_fn: A time-dependent score-based model that takes x and t and returns the score. probability_flow: If True, create the reverse-time ODE used for probability flow sampling.

class VPSDE(beta_min=0.1, beta_max=20, N=1000): No docstring is available.

VPSDE.T(self): No docstring is available.

VPSDE.sde(self, x, t): No docstring is available.

VPSDE.marginal_prob(self, x, t): No docstring is available.

VPSDE.prior_sampling(self, shape): No docstring is available.

VPSDE.prior_logp(self, z): No docstring is available.

VPSDE.discretize(self, x, t): DDPM discretization.

class subVPSDE(beta_min=0.1, beta_max=20, N=1000): No docstring is available.

subVPSDE.T(self): No docstring is available.

subVPSDE.sde(self, x, t): No docstring is available.

subVPSDE.marginal_prob(self, x, t): No docstring is available.

subVPSDE.prior_sampling(self, shape): No docstring is available.

subVPSDE.prior_logp(self, z): No docstring is available.

class VESDE(sigma_min=0.01, sigma_max=50, N=1000): No docstring is available.

VESDE.T(self): No docstring is available.

VESDE.sde(self, x, t): No docstring is available.

VESDE.marginal_prob(self, x, t): No docstring is available.

VESDE.prior_sampling(self, shape): No docstring is available.

VESDE.prior_logp(self, z): No docstring is available.

VESDE.discretize(self, x, t): SMLD(NCSN) discretization.