ls_mlkit.diffuser.sde.sde_lib module

class ls_mlkit.diffuser.sde.sde_lib.SubVPSDE(beta_min: float = 0.1, beta_max: float = 20, n_discretization_steps: int = 1000, ndim_micro_shape: int = 2)[source]

Bases: SDE

property T: float

End time of the SDE.

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

Get the drift and diffusion of the SDE.

Parameters:

  • x (Tensor) – the sample.

  • t (Tensor) – the time step.

  • mask (Tensor, optional) – the mask of the sample. Defaults to None.

Returns:

the drift and diffusion of the SDE.

Return type:

Tuple[Tensor, Tensor]

marginal_prob(x, t, mask=None)[source]
prior_logp(z)[source]
prior_sampling(shape)[source]
class ls_mlkit.diffuser.sde.sde_lib.VESDE(sigma_min=0.01, sigma_max=50, n_discretization_steps=1000, ndim_micro_shape=2, drop_first_step=False)[source]

Bases: SDE

property T: float

End time of the SDE.

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

SMLD(NCSN) discretization. .. math:

x_t &= x_0 + g \epsilon

x_t &\sim \mathcal{N}(x_0, \sigma_t^2)

\sigma_t^2 &= \sigma_{t-1}^2 + g^2

g &= \sqrt{\sigma_t^2 - \sigma_{t-1}^2}
get_drift_and_diffusion(x: Tensor, t: Tensor, mask=None) → Tuple[Tensor, Tensor][source]
\[ \begin{align}\begin{aligned}dx = 0 dt + \sigma_{min} \left(\frac{\sigma_{max}}{\sigma_{min}}\right)^t \sqrt{2 \log(\frac{\sigma_{max}}{\sigma_{min}})} dw \sigma_t = \sigma_{min} \left(\frac{\sigma_{max}}{\sigma_{min}}\right)^t\\diffusion = \sigma_t * \sqrt{2 \log(\frac{\sigma_{max}}{\sigma_{min}})}\end{aligned}\end{align} \]
marginal_prob(x, t, mask=None)[source]
prior_logp(z)[source]
prior_sampling(shape)[source]
class ls_mlkit.diffuser.sde.sde_lib.VPSDE(beta_min: float = 0.1, beta_max: float = 20, n_discretization_steps: int = 1000, ndim_micro_shape: int = 2)[source]

Bases: SDE

property T: float

End time of the SDE.

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

DDPM discretization.

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

continuous DDPM SDE

\[dx &= -\frac{1}{2}\beta_t x dt + \sqrt{\beta_t} dw\]
Parameters:

  • x

  • t – (macro_shape)

  • mask

Returns:

shape = x.shape diffusion: shape=x.macro_shape

Return type:

drift

get_target_score(x_0: Tensor, x_t: Tensor, t: Tensor, mask: Tensor, continuous: bool = False) → Tensor[source]
\[p_{0t} (x_t|x_0) = \nabla_{x_t} \ln p_{0t} (x_t|x_0)\]
marginal_prob(x_0: Tensor, t: Tensor, mask: Tensor = None) → Tuple[Tensor, Tensor][source]
\[p_{0t} (x_t|x_0)\]
\[ \begin{align}\begin{aligned}\gamma = -\frac{1}{4}t^2 (\beta_1 - \beta_0) - \frac{1}{2} t \beta_0\\mean = e^{\gamma} * x\\std = \sqrt{1 - e^{2 \gamma }}\end{aligned}\end{align} \]
prior_logp(z: Tensor) → Tensor[source]
\[(2\pi)^{-k/2} \det(\Sigma)^{-1/2} \exp\left( -\frac{1}{2} (\mathbf{x} - \boldsymbol{\mu})^\mathrm{T} \Sigma^{-1} (\mathbf{x} - \boldsymbol{\mu}) \right)\]

where \(\Sigma = I\) and \(\mathbf{\mu} = 0\)

prior_sampling(shape: Tuple) → Tensor[source]
\[\epsilon \sim \mathbfcal{N}(0,1)\]