Noise Schedulers#

Noise schedulers are the central abstraction in the PhysicsNeMo diffusion framework. A noise scheduler defines the forward diffusion process (how noise is added to data) and provides the ingredients needed for both training and inference. It is the only component that participates in every stage of the pipeline: training, loss computation, denoiser construction, and sampling.

If you are familiar with the Scheduler API in HuggingFace Diffusers, PhysicsNeMo noise schedulers serve a broadly similar role. A key design difference is that the PhysicsNeMo NoiseScheduler also owns the get_denoiser() factory method, which converts a trained Predictor into a Denoiser suitable for the solver. This keeps the coupling between noise schedule and reverse process in one place and avoids the need for schedule-specific sampler code.

Role in Training#

During training, the noise scheduler is responsible for three operations:

Sampling diffusion times via sample_time().
Adding noise via add_noise().
Computing loss weights via loss_weight().

These three methods are consumed by the MSEDSMLoss training loss.

from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler
from physicsnemo.diffusion.metrics.losses import MSEDSMLoss

scheduler = EDMNoiseScheduler()
loss_fn = MSEDSMLoss(model, scheduler)

# Inside the training loop:
loss = loss_fn(x0)   # Internally samples t, adds noise, computes weighted MSE
loss.backward()

Role in Inference (Sampling)#

During sampling, the noise scheduler provides:

Time-step schedule via timesteps().
Initial latent state via init_latents().
Denoiser factory via get_denoiser().

from functools import partial
from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler
from physicsnemo.diffusion.samplers import sample

scheduler = EDMNoiseScheduler()
num_steps = 50

t_steps = scheduler.timesteps(num_steps)
tN = t_steps[0].expand(4)                      # Batch of 4 samples
xN = scheduler.init_latents((3, 64, 64), tN)    # 3-channel, 64x64 images

x0_predictor = partial(trained_model, condition=condition)
denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)
samples = sample(denoiser, xN, scheduler, num_steps=num_steps)

Three Levels of Customization#

Following the framework’s design philosophy, noise schedulers are available at three levels:

Protocol (NoiseScheduler): The minimal interface. Any object that implements the six required methods can be used as a noise scheduler. This is the right choice for fully custom forward processes (for example, non-Gaussian, non-linear, discrete).
Abstract base class (LinearGaussianNoiseScheduler): For the common family of linear-Gaussian forward processes of the form \(\mathbf{x}(t) = \alpha(t)\,\mathbf{x}_0 + \sigma(t)\,\boldsymbol{\epsilon}\). This base class implements noise injection, score conversion, and denoiser construction. Subclasses only need to define the schedule-specific quantities:
- the signal coefficient \(\alpha(t)\)
- the noise level \(\sigma(t)\)
- their time derivatives \(\dot{\alpha}(t)\) and \(\dot{\sigma}(t)\)
- the inverse mapping \(\sigma^{-1}(\sigma) = t\) from noise level back to time
- the discretization of the diffusion time grid
Ready-to-use schedules: Five concrete implementations that work out of the box:
- EDMNoiseScheduler — \(\alpha(t)=1\), \(\sigma(t)=t\). The recommended default for most applications.
- VENoiseScheduler — Variance Exploding schedule.
- VPNoiseScheduler — Variance Preserving schedule.
- IDDPMNoiseScheduler — Improved DDPM schedule.
- StudentTEDMNoiseScheduler — EDM variant with Student-t noise for heavy-tailed data.

API Reference#

`NoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.NoiseScheduler(*args, **kwargs)[source]#

Protocol defining the minimal interface for noise schedulers.

A noise scheduler defines methods for training (adding noise, sampling diffusion time) and for sampling (generating diffusion time-steps, initializing latent state, obtaining a denoiser). This interface is generic and does not assume any specific form of noise schedule.

Any object that implements this interface can be used with the diffusion training and sampling utilities.

Training methods:

sample_time(): Sample diffusion time values for training
add_noise(): Add noise to clean data at given diffusion time
loss_weight(): Compute per-sample loss weight for training

Sampling methods:

timesteps(): Generate discrete time-steps for sampling
init_latents(): Initialize noisy latent state \(\mathbf{x}_N\)
get_denoiser(): Convert a predictor (e.g. model that predicts
clean, data, score, etc.) to a sampling-compatible denoiser

`LinearGaussianNoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.LinearGaussianNoiseScheduler(*args, **kwargs)[source]#

Bases: ABC, NoiseScheduler

Abstract base class for linear-Gaussian noise schedules.

It implements the NoiseScheduler interface and it can be subclassed to define custom linear-Gaussian noise schedules of the form:

\[\mathbf{x}(t) = \alpha(t) \mathbf{x}_0 + \sigma(t) \boldsymbol{\epsilon}\]

where \(\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\) is standard Gaussian noise, \(\alpha(t)\) is the signal coefficient, and \(\sigma(t)\) is the noise level.

Training:

The add_noise() method implements the forward diffusion process using the formula above. The sample_time() method samples diffusion times.

Sampling:

For ODE-based sampling, the reverse process follows the probability flow ODE:

\[\frac{d\mathbf{x}}{dt} = f(\mathbf{x}, t) - \frac{1}{2} g^2(\mathbf{x}, t) \nabla_{\mathbf{x}} \log p(\mathbf{x})\]

For SDE-based sampling:

\[d\mathbf{x} = \left[ f(\mathbf{x}, t) - g^2(\mathbf{x}, t) \nabla_{\mathbf{x}} \log p(\mathbf{x}) \right] dt + g(\mathbf{x}, t) d\mathbf{W}\]

The get_denoiser() factory converts a predictor (either a score-predictor or an x0-predictor) into the appropriate ODE/SDE right-hand side.

Abstract methods (must be implemented by subclasses):

sigma(): Map time to noise level \(\sigma(t)\)
sigma_inv(): Map noise level back to time
sigma_dot(): Time derivative \(\dot{\sigma}(t)\)
alpha(): Compute the signal coefficient \(\alpha(t)\)
alpha_dot(): Time derivative \(\dot{\alpha}(t)\)
timesteps(): Generate discrete time-steps for sampling
sample_time(): Sample diffusion times for training
loss_weight(): Compute loss weight for training

Concrete methods (have default implementations, but can be overridden for custom behavior):

drift(): Drift term \(f(\mathbf{x}, t)\) for ODE/SDE
diffusion(): Squared diffusion term \(g^2(\mathbf{x}, t)\)
x0_to_score(): Convert x0-prediction to score
score_to_x0(): Convert score to x0-prediction
add_noise(): Add noise to clean data (training)
init_latents(): Initialize latent state (sampling)
get_denoiser(): Get ODE/SDE RHS (sampling)

Examples

Example 1: A minimal EDM-like noise schedule. Only the abstract methods need to be implemented since defaults work for EDM:

>>> import torch
>>> from physicsnemo.diffusion.noise_schedulers import (
...     LinearGaussianNoiseScheduler,
... )
>>>
>>> class SimpleEDMScheduler(LinearGaussianNoiseScheduler):
...     def __init__(self, sigma_min=0.002, sigma_max=80.0, rho=7.0):
...         self.sigma_min = sigma_min
...         self.sigma_max = sigma_max
...         self.rho = rho
...
...     def sigma(self, t): return t
...     def sigma_inv(self, sigma): return sigma
...     def sigma_dot(self, t): return torch.ones_like(t)
...     def alpha(self, t): return torch.ones_like(t)
...     def alpha_dot(self, t): return torch.zeros_like(t)
...
...     def timesteps(self, num_steps, *, device=None, dtype=None):
...         i = torch.arange(num_steps, device=device, dtype=dtype)
...         smax_rho = self.sigma_max**(1/self.rho)
...         smin_rho = self.sigma_min**(1/self.rho)
...         frac = i/(num_steps-1)
...         t = (smax_rho + frac * (smin_rho - smax_rho))**self.rho
...         return torch.cat([t, torch.zeros(1, device=device)])
...
...     def sample_time(self, N, *, device=None, dtype=None):
...         u = torch.rand(N, device=device, dtype=dtype)
...         return self.sigma_min * (self.sigma_max/self.sigma_min)**u
...     def loss_weight(self, t):
...         return 1 / t**2
...
>>> scheduler = SimpleEDMScheduler()
>>> t_steps = scheduler.timesteps(10)
>>> t_steps.shape
torch.Size([11])

Example 2: Customizing behavior by overriding concrete methods. This shows how to override the drift term for a custom diffusion process:

>>> class CustomDriftScheduler(SimpleEDMScheduler):
...     def drift(self, x, t):
...         # Custom drift: f(x, t) = -0.5 * x (Ornstein-Uhlenbeck style)
...         return -0.5 * x
...
>>> custom = CustomDriftScheduler()
>>>
>>> # The custom drift is used internally by get_denoiser
>>> score_pred = lambda x, t: -x / (1 + t.view(-1, 1)**2)  # Toy score predictor
>>> denoiser = custom.get_denoiser(score_predictor=score_pred)
>>> x = torch.randn(2, 4)
>>> t = torch.tensor([1.0, 1.0])
>>> out = denoiser(x, t)  # Uses custom drift in ODE RHS computation
>>> out.shape
torch.Size([2, 4])

add_noise( x0: Float[Tensor, 'B *dims'], time: Float[Tensor, 'B'], ) → Float[Tensor, 'B *dims'][source]#

Add noise to clean data at the given diffusion times.

Used in training to create noisy samples from clean data. Implements:

\[\mathbf{x}(t) = \alpha(t) \mathbf{x}_0 + \sigma(t) \boldsymbol{\epsilon}\]

Usually does not need to be overridden in subclasses: overriding the alpha() and sigma() methods is sufficient for most use cases.

Parameters:

x0 (Tensor) – Clean latent state of shape \((B, *)\).
time (Tensor) – Diffusion time values of shape \((B,)\).

Returns:

Noisy latent state of shape \((B, *)\).

Return type:

Tensor

abstractmethod alpha( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#

Compute the signal coefficient \(\alpha(t)\).

Used in both training and sampling.

Parameters:: t (Tensor) – Diffusion time tensor of any shape.
Returns:: Signal coefficient \(\alpha(t)\) with same shape as t.
Return type:: Tensor

abstractmethod alpha_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#

Compute time derivative of signal coefficient \(\dot{\alpha}(t)\).

Used in sampling.

Parameters:: t (Tensor) – Diffusion time tensor of any shape.
Returns:: Time derivative \(\dot{\alpha}(t)\) with same shape as t.
Return type:: Tensor

diffusion( x: Float[Tensor, 'B *dims'], t: Float[Tensor, 'B'], ) → Float[Tensor, 'B *_'][source]#

Compute squared diffusion term \(g^2(\mathbf{x}, t)\).

Used by get_denoiser() to build the ODE/SDE right-hand side.

By default: \(g^2 = 2 \dot{\sigma} \sigma - 2 \frac{\dot{\alpha}} {\alpha} \sigma^2\). This method can be overridden to implement different diffusion terms.

Parameters:

x (Tensor) – Latent state of shape \((B, *)\).
t (Tensor) – Diffusion time of shape \((B,)\).

Returns:

Squared diffusion term, broadcastable to shape of x.

Return type:

Tensor

drift( x: Float[Tensor, 'B *dims'], t: Float[Tensor, 'B'], ) → Float[Tensor, 'B *dims'][source]#

Compute drift term \(f(\mathbf{x}, t)\) for ODE/SDE sampling.

Used by get_denoiser() to build the ODE/SDE right-hand side.

By default: \(f(\mathbf{x}, t) = \frac{\dot{\alpha}(t)}{\alpha(t)} \mathbf{x}\).

This method can be overridden to implement different drift terms.

Parameters:

x (Tensor) – Latent state of shape \((B, *)\).
t (Tensor) – Diffusion time of shape \((B,)\).

Returns:

Drift term with same shape as x.

Return type:

Tensor

get_denoiser(

*,

score_predictor: Predictor | None = None,

x0_predictor: Predictor | None = None,

denoising_type: Literal['ode', 'sde'] = 'ode',

**kwargs: Any,

) → Denoiser[source]#

Factory that converts a predictor to a denoiser for sampling.

Accepts either a score-predictor or an x0-predictor (exactly one must be provided). The returned denoiser computes the right-hand side of the reverse ODE or SDE.

For ODE (denoising_type="ode"):

\[\frac{d\mathbf{x}}{dt} = f(\mathbf{x}, t) - \frac{1}{2} g^2(t) s(\mathbf{x}, t)\]

For SDE (denoising_type="sde"):

\[d\mathbf{x} = \left[ f(\mathbf{x}, t) - g^2(t) s(\mathbf{x}, t) \right] dt + g(t) d\mathbf{W}\]

where \(s(\mathbf{x}, t)\) is the score. When an x0-predictor is provided, the score is computed internally via x0_to_score(). When a score-predictor is provided, it is used directly. Note: As usually done in SDE integration, the stochastic term \(g(t) d\mathbf{W}\) is handled by the solver, not returned by the denoiser itself.

Parameters:

score_predictor (Predictor, optional) – A score-predictor that takes (x_t, t) and returns a score (e.g. \(\nabla_{\mathbf{x}} \log p(\mathbf{x}_t)\)). Can be unconditional, conditional, guidance-augmented, etc. Mutually exclusive with x0_predictor.
x0_predictor (Predictor, optional) – An x0-predictor that takes (x_t, t) and returns an estimate of clean data \(\hat{\mathbf{x}}_0\). The score is computed internally via x0_to_score(). Mutually exclusive with score_predictor.
denoising_type ({"ode", "sde"}, default="ode") – Type of reverse process. Use "ode" for deterministic sampling, "sde" for stochastic sampling.
**kwargs (Any) – Ignored.

Returns:

A denoiser computing the RHS of the reverse ODE/SDE. Implements the Denoiser interface.

Return type:

Denoiser

Raises:

ValueError – If both or neither score_predictor and x0_predictor are provided.

Examples

Generate ODE RHS from a score-predictor:

>>> import torch
>>> scheduler = EDMNoiseScheduler()
>>> score_pred = lambda x, t: -x / t.view(-1, 1, 1, 1)**2  # Toy score-predictor
>>> denoiser = scheduler.get_denoiser(
...     score_predictor=score_pred, denoising_type="ode")
>>> x = torch.randn(2, 3, 8, 8)
>>> t = torch.tensor([1.0, 1.0])
>>> dx_dt = denoiser(x, t)  # Returns ODE RHS for sampling
>>> dx_dt.shape
torch.Size([2, 3, 8, 8])

Generate ODE RHS from an x0-predictor (score conversion is done internally):

>>> x0_pred = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2)  # Toy x0-predictor
>>> denoiser = scheduler.get_denoiser(
...     x0_predictor=x0_pred, denoising_type="ode")
>>> dx_dt = denoiser(x, t)  # Returns ODE RHS for sampling
>>> dx_dt.shape
torch.Size([2, 3, 8, 8])

init_latents( spatial_shape: Tuple[int, ...], tN: Float[Tensor, 'B'], *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'B *spatial_shape'][source]#

Initialize the noisy latent state \(\mathbf{x}_N\) for sampling.

Generates:

\[\mathbf{x}_N = \sigma(t_N) \cdot \boldsymbol{\epsilon}\]

where \(\boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\).

Parameters:

spatial_shape (Tuple[int, ...]) – Spatial shape of the latent state, e.g., (C, H, W).
tN (Tensor) – Initial diffusion time of shape \((B,)\).
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Initial noisy latent of shape \((B, *spatial\_shape)\).

Return type:

Tensor

abstractmethod loss_weight( t: Float[Tensor, 'N'], ) → Float[Tensor, 'N'][source]#

Compute loss weight for denoising score matching training.

Used in training to weight the per-sample loss in MSEDSMLoss. The loss weight is designed for training an x0-predictor (clean data predictor). For training a score-predictor, additionally provide a score_to_x0_fn callback to MSEDSMLoss.

Parameters:: t (Tensor) – Diffusion time values of shape \((N,)\).
Returns:: Loss weight of shape \((N,)\), same device and dtype as t.
Return type:: Tensor

abstractmethod sample_time( N: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N'][source]#

Sample N diffusion time values for training.

Used in training to sample random diffusion times for the denoising score matching loss.

Parameters:

N (int) – Number of time values to sample.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Sampled diffusion times of shape \((N,)\).

Return type:

Tensor

score_to_x0( score: Float[Tensor, 'B *dims'], x_t: Float[Tensor, 'B *dims'], t: Float[Tensor, 'B'], ) → Float[Tensor, 'B *dims'][source]#

Convert score to x0-prediction.

This is the inverse of x0_to_score(). Given a score prediction \(s(\mathbf{x}_t, t)\) and the noisy state \(\mathbf{x}_t\), recover the corresponding \(\hat{\mathbf{x}}_0\) estimate:

\[\hat{\mathbf{x}}_0 = \frac{\mathbf{x}_t + \sigma^2(t) \nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)} {\alpha(t)}\]

A common use case is with MSEDSMLoss to train a score-predictor instead of an x0-predictor: pass this method as the score_to_x0_fn argument with prediction_type="score".

This is a helper method that usually does not need to be overridden in subclasses.

Parameters:

score (Tensor) – Predicted score \(s(\mathbf{x}_t, t)\) of shape \((B, *)\).
x_t (Tensor) – Current noisy state \(\mathbf{x}_t\) with same shape as score.
t (Tensor) – Diffusion time with shape \((B,)\).

Returns:

Estimated clean data \(\hat{\mathbf{x}}_0\) with same shape as score.

Return type:

Tensor

Examples

>>> scheduler = EDMNoiseScheduler()
>>> # If you have a score-predictor, convert to x0 for DSM loss:
>>> def score_predictor(x, t):
...     return -x / (1 + t.view(-1, *([1] * (x.ndim - 1)))**2)
>>> x_t = torch.randn(2, 4)
>>> t = torch.tensor([1.0, 1.0])
>>> score = score_predictor(x_t, t)
>>> x0_est = scheduler.score_to_x0(score, x_t, t)
>>> x0_est.shape
torch.Size([2, 4])

abstractmethod sigma( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#

Map diffusion time to noise level \(\sigma(t)\).

Used in both training and sampling.

Parameters:: t (Tensor) – Diffusion time tensor of any shape.
Returns:: Noise coefficient \(\sigma(t)\) with same shape as t.
Return type:: Tensor

abstractmethod sigma_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#

Compute time derivative of noise level \(\dot{\sigma}(t)\).

Used in sampling.

Parameters:: t (Tensor) – Diffusion time tensor of any shape.
Returns:: Time derivative \(\dot{\sigma}(t)\) with same shape as t.
Return type:: Tensor

abstractmethod sigma_inv( sigma: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#

Map noise level back to diffusion time.

Used in both training and sampling.

Parameters:: sigma (Tensor) – Noise level tensor of any shape.
Returns:: Diffusion time with same shape as sigma.
Return type:: Tensor

abstractmethod timesteps( num_steps: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N+1'][source]#

Generate discrete time-steps for sampling.

Used in sampling to produce the sequence of diffusion times. Returns a tensor of shape \((N + 1,)\) in decreasing order, with the last element being 0.

Parameters:

num_steps (int) – Number of sampling steps.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Time-steps tensor of shape \((N + 1,)\).

Return type:

Tensor

x0_to_score( x0: Float[Tensor, 'B *dims'], x_t: Float[Tensor, 'B *dims'], t: Float[Tensor, 'B'], ) → Float[Tensor, 'B *dims'][source]#

Convert x0-predictor output to score.

This conversion is done automatically by get_denoiser() when x0_predictor is provided, but can also be called manually.

The score is: \(\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t) = \frac{\alpha(t) \hat{\mathbf{x}}_0 - \mathbf{x}_t}{\sigma^2(t)}\).

This is a helper method that usually does not need to be overridden in subclasses.

Parameters:

x0 (Tensor) – Predicted clean data \(\hat{\mathbf{x}}_0\) of shape \((B, *)\).
x_t (Tensor) – Current noisy state \(\mathbf{x}_t\) of shape \((B, *)\).
t (Tensor) – Diffusion time of shape \((B,)\).

Returns:

Score with same shape as x0.

Return type:

Tensor

Examples

>>> scheduler = EDMNoiseScheduler()
>>> # If you have an x0-predictor, wrap it for manual conversion
>>> # (done automatically by get_denoiser):
>>> def x0_predictor(x, t):
...     t_bc = t.view(-1, *([1] * (x.ndim - 1)))
...     return x / (1 + t_bc**2)
>>> def score_predictor(x, t):
...     x0_pred = x0_predictor(x, t)
...     return scheduler.x0_to_score(x0_pred, x, t)
>>> # Or simply: scheduler.get_denoiser(x0_predictor=x0_predictor)

`EDMNoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.EDMNoiseScheduler( sigma_min: float = 0.002, sigma_max: float = 80.0, rho: float = 7.0, sigma_data: float = 0.5, P_mean: float = -1.2, P_std: float = 1.2, )[source]#

Bases: LinearGaussianNoiseScheduler

EDM noise scheduler.

The EDM formulation uses \(\alpha(t) = 1\) (no signal attenuation) and \(\sigma(t) = t\) (identity mapping between time and noise level).

Sampling time-steps are computed with polynomial spacing:

\[t_i = \left(\sigma_{\max}^{1/\rho} + \frac{i}{N-1} \left(\sigma_{\min}^{1/\rho} - \sigma_{\max}^{1/\rho}\right) \right)^{\rho}\]

Training times are sampled from a log-normal distribution with mean \(P_{\text{mean}}\) and standard deviation \(P_{\text{std}}\).

Parameters:

sigma_min (float, optional) – Minimum noise level for sampling time-steps, by default 0.002.
sigma_max (float, optional) – Maximum noise level for sampling time-steps, by default 80.
rho (float, optional) – Exponent controlling time-step spacing. Larger values concentrate more steps at lower noise levels (better for fine details). By default 7.
sigma_data (float, optional) – Expected standard deviation of the training data, by default 0.5. Used by loss_weight() to compute the per-sample loss weight.
P_mean (float, optional) – Mean of the log-normal distribution used to sample training times, by default -1.2.
P_std (float, optional) – Standard deviation of the log-normal distribution used to sample training times, by default 1.2.

Note

Reference: Elucidating the Design Space of Diffusion-Based Generative Models

Examples

Basic training and sampling workflow using the EDM noise scheduler:

>>> import torch
>>> from physicsnemo.diffusion.noise_schedulers import EDMNoiseScheduler
>>>
>>> scheduler = EDMNoiseScheduler(sigma_min=0.002, sigma_max=80.0, rho=7)
>>>
>>> # Training: sample times and add noise
>>> x0 = torch.randn(4, 3, 8, 8)  # Clean data
>>> t = scheduler.sample_time(4)    # Sample diffusion times
>>> x_t = scheduler.add_noise(x0, t)  # Create noisy samples
>>> x_t.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Sampling: generate timesteps and initial latents
>>> t_steps = scheduler.timesteps(10)
>>> tN = t_steps[0].expand(4)  # Initial time for batch of 4
>>> xN = scheduler.init_latents((3, 8, 8), tN)  # Initial noise
>>> xN.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Convert x0-predictor to denoiser for sampling
>>> x0_predictor = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2)  # Toy x0-predictor
>>> denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)
>>> denoiser(xN, tN).shape  # ODE RHS for sampling
torch.Size([4, 3, 8, 8])

alpha( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant signal coefficient: \(\alpha(t) = 1\).

alpha_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Zero derivative: \(\dot{\alpha}(t) = 0\).

loss_weight( t: Float[Tensor, 'N'], ) → Float[Tensor, 'N'][source]#

Compute EDM loss weight.

\[w(t) = \frac{\sigma(t)^2 + \sigma_{\text{data}}^2} {\left(\sigma(t) \cdot \sigma_{\text{data}}\right)^2}\]

Important

This loss weight is designed for training an x0-predictor (clean data predictor) wrapped with EDMPreconditioner. It is not suitable for training a score-predictor, or a model without a pre-conditioner.

Parameters:: t (Tensor) – Diffusion time values of shape \((N,)\).
Returns:: Loss weight of shape \((N,)\).
Return type:: Tensor

sample_time( N: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N'][source]#

Sample N diffusion times from a log-normal distribution: \(\ln(\sigma) \sim \mathcal{N}(P_{\text{mean}}, P_{\text{std}}^2)\).

Parameters:

N (int) – Number of time values to sample.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Sampled diffusion times of shape \((N,)\).

Return type:

Tensor

sigma( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Identity mapping: \(\sigma(t) = t\).

sigma_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant derivative: \(\dot{\sigma}(t) = 1\).

sigma_inv( sigma: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Identity mapping: \(t = \sigma\).

timesteps( num_steps: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N+1'][source]#

Generate EDM time-steps with polynomial spacing.

Parameters:

num_steps (int) – Number of sampling steps.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Time-steps tensor of shape \((N + 1,)\) where \(N\) is num_steps.

Return type:

torch.Tensor

`VENoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.VENoiseScheduler(sigma_min: float = 0.02, sigma_max: float = 100.0)[source]#

Bases: LinearGaussianNoiseScheduler

Variance Exploding (VE) noise scheduler.

Implements the VE formulation with \(\sigma(t) = \sqrt{t}\) and \(\alpha(t) = 1\) (no signal attenuation).

Sampling time-steps use geometric spacing in \(\sigma^2\) space:

\[\sigma_i^2 = \sigma_{\max}^2 \cdot \left(\frac{\sigma_{\min}^2}{\sigma_{\max}^2}\right)^{i/(N-1)}\]

Training times are sampled log-uniformly between sigma_min and sigma_max, then mapped to time via \(t = \sigma^2\).

Parameters:

sigma_min (float, optional) – Minimum noise level, by default 0.02.
sigma_max (float, optional) – Maximum noise level, by default 100.

Note

Reference: Score-Based Generative Modeling through Stochastic Differential Equations

Examples

Basic training and sampling workflow using the VE noise scheduler:

>>> import torch
>>> from physicsnemo.diffusion.noise_schedulers import VENoiseScheduler
>>>
>>> scheduler = VENoiseScheduler(sigma_min=0.02, sigma_max=100.0)
>>>
>>> # Training: sample times and add noise
>>> x0 = torch.randn(4, 3, 8, 8)  # Clean data
>>> t = scheduler.sample_time(4)    # Sample diffusion times
>>> x_t = scheduler.add_noise(x0, t)  # Create noisy samples
>>> x_t.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Sampling: generate timesteps and initial latents
>>> t_steps = scheduler.timesteps(10)
>>> tN = t_steps[0].expand(4)  # Initial time for batch of 4
>>> xN = scheduler.init_latents((3, 8, 8), tN)  # Initial noise
>>> xN.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Convert x0-predictor to denoiser for sampling
>>> x0_predictor = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2)  # Toy x0-predictor
>>> denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)
>>> denoiser(xN, tN).shape  # ODE RHS for sampling
torch.Size([4, 3, 8, 8])

alpha( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant signal coefficient: \(\alpha(t) = 1\).

alpha_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Zero derivative: \(\dot{\alpha}(t) = 0\).

loss_weight( t: Float[Tensor, 'N'], ) → Float[Tensor, 'N'][source]#

Compute VE loss weight: \(w(t) = 1 / \sigma(t)^2\).

Important

This loss weight is designed for training an x0-predictor (clean data predictor) wrapped with VEPreconditioner. It is not suitable for training a score-predictor, or a model without a pre-conditioner.

Parameters:: t (Tensor) – Diffusion time values of shape \((N,)\).
Returns:: Loss weight of shape \((N,)\).
Return type:: Tensor

sample_time( N: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N'][source]#

Sample N diffusion times log-uniformly in sigma space, mapped to time.

Parameters:

N (int) – Number of time values to sample.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Sampled diffusion times of shape \((N,)\).

Return type:

Tensor

sigma( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: VE noise coefficient: \(\sigma(t) = \sqrt{t}\).

sigma_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Time derivative: \(\dot{\sigma}(t) = 1/(2\sqrt{t})\).

sigma_inv( sigma: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Inverse VE mapping: \(t = \sigma^2\).

timesteps( num_steps: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N+1'][source]#

Generate VE time-steps with geometric spacing in \(\sigma^2\).

Parameters:

num_steps (int) – Number of sampling steps.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Time-steps tensor of shape \((N + 1,)\).

Return type:

torch.Tensor

`VPNoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.VPNoiseScheduler( beta_min: float = 0.1, beta_d: float = 19.1, epsilon_s: float = 0.001, t_max: float = 1.0, )[source]#

Bases: LinearGaussianNoiseScheduler

Variance Preserving (VP) noise scheduler.

Implements the VP formulation where the total variance is preserved: \(\alpha(t)^2 + \sigma(t)^2 = 1\). This is based on a linear beta schedule: \(\beta(t) = \beta_{\min} + t \cdot \beta_d\).

The noise and signal coefficients are:

\[\alpha(t) = \exp\left(-\frac{1}{2} \left(\frac{\beta_d}{2} t^2 + \beta_{\min} t\right)\right)\]

\[\sigma(t) = \sqrt{1 - \alpha(t)^2} = \sqrt{1 - \exp\left(-\frac{\beta_d}{2} t^2 - \beta_{\min} t\right)}\]

Sampling time-steps are linearly spaced from t_max (usually 1) to epsilon_s (small positive value to avoid singularities).

Training times are sampled uniformly between epsilon_s and t_max.

Parameters:

beta_min (float, optional) – Minimum beta value for the linear schedule, by default 0.1.
beta_d (float, optional) – Beta slope (delta) for the linear schedule, by default 19.1.
epsilon_s (float, optional) – Small positive value for minimum time, by default 1e-3.
t_max (float, optional) – Maximum diffusion time, by default 1.0.

Note

Reference: Score-Based Generative Modeling through Stochastic Differential Equations

Examples

Basic training and sampling workflow using the VP noise scheduler:

>>> import torch
>>> from physicsnemo.diffusion.noise_schedulers import VPNoiseScheduler
>>>
>>> scheduler = VPNoiseScheduler(beta_min=0.1, beta_d=19.1)
>>>
>>> # Training: sample times and add noise
>>> x0 = torch.randn(4, 3, 8, 8)  # Clean data
>>> t = scheduler.sample_time(4)    # Sample diffusion times
>>> x_t = scheduler.add_noise(x0, t)  # Create noisy samples
>>> x_t.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Sampling: generate timesteps and initial latents
>>> t_steps = scheduler.timesteps(10)
>>> tN = t_steps[0].expand(4)  # Initial time for batch of 4
>>> xN = scheduler.init_latents((3, 8, 8), tN)  # Initial noise
>>> xN.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Convert x0-predictor to denoiser for sampling
>>> x0_predictor = lambda x, t: x * 0.9  # Toy x0-predictor
>>> denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)
>>> denoiser(xN, tN).shape  # ODE RHS for sampling
torch.Size([4, 3, 8, 8])

alpha( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Signal coefficient: \(\alpha(t) = \exp(-a(t)/2)\).

alpha_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Derivative: \(\dot{\alpha}(t) = -\frac{\beta(t)}{2} \alpha(t)\).

loss_weight( t: Float[Tensor, 'N'], ) → Float[Tensor, 'N'][source]#

Compute VP loss weight: \(w(t) = \alpha(t)^2 / \sigma(t)^2\).

Important

This loss weight is designed for training an x0-predictor (clean data predictor) wrapped with VPPreconditioner. It is not suitable for training a score-predictor, or a model without a pre-conditioner.

Parameters:: t (Tensor) – Diffusion time values of shape \((N,)\).
Returns:: Loss weight of shape \((N,)\).
Return type:: Tensor

sample_time( N: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N'][source]#

Sample N diffusion times uniformly in \([\epsilon_s, t_{max}]\).

Parameters:

N (int) – Number of time values to sample.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Sampled diffusion times of shape \((N,)\).

Return type:

Tensor

sigma( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Noise level: \(\sigma(t) = \sqrt{1 - \alpha(t)^2}\).

sigma_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Derivative: \(\dot{\sigma}(t) = -\alpha(t) \dot{\alpha}(t) / \sigma(t)\).

sigma_inv( sigma: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#

Inverse mapping from sigma to time.

Solves: \(\sigma^2 = 1 - \exp(-a(t))\) for \(t\).

timesteps( num_steps: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N+1'][source]#

Generate VP time-steps with linear spacing.

Parameters:

num_steps (int) – Number of sampling steps.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Time-steps tensor of shape \((N + 1,)\).

Return type:

torch.Tensor

`IDDPMNoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.IDDPMNoiseScheduler( sigma_min: float = 0.002, sigma_max: float = 81.0, C_1: float = 0.001, C_2: float = 0.008, M: int = 1000, )[source]#

Bases: LinearGaussianNoiseScheduler

Improved DDPM (iDDPM) noise scheduler.

Uses identity mappings \(\sigma(t) = t\) and \(\alpha(t) = 1\). The key feature is a precomputed noise level schedule derived from a cosine schedule, providing improved sample quality in comparison to original DDPM.

Sampling time-steps are selected from a precomputed schedule of \(M\) discrete noise levels, subsampled to num_steps.

Training times are sampled uniformly from the precomputed schedule.

Parameters:

sigma_min (float, optional) – Minimum noise level for filtering, by default 0.002.
sigma_max (float, optional) – Maximum noise level for filtering, by default 81.
C_1 (float, optional) – Clipping threshold for alpha ratio, by default 0.001.
C_2 (float, optional) – Cosine schedule parameter, by default 0.008.
M (int, optional) – Number of precomputed discretization steps, by default 1000.

Note

Reference: Improved Denoising Diffusion Probabilistic Models

Examples

Basic training and sampling workflow using the iDDPM noise scheduler:

>>> import torch
>>> from physicsnemo.diffusion.noise_schedulers import IDDPMNoiseScheduler
>>>
>>> scheduler = IDDPMNoiseScheduler(C_1=0.001, C_2=0.008, M=1000)
>>>
>>> # Training: sample times and add noise
>>> x0 = torch.randn(4, 3, 8, 8)  # Clean data
>>> t = scheduler.sample_time(4)    # Sample diffusion times
>>> x_t = scheduler.add_noise(x0, t)  # Create noisy samples
>>> x_t.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Sampling: generate timesteps and initial latents
>>> t_steps = scheduler.timesteps(10)
>>> tN = t_steps[0].expand(4)  # Initial time for batch of 4
>>> xN = scheduler.init_latents((3, 8, 8), tN)  # Initial noise
>>> xN.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Convert x0-predictor to denoiser for sampling
>>> x0_predictor = lambda x, t: x / (1 + t.view(-1, 1, 1, 1)**2)  # Toy x0-predictor
>>> denoiser = scheduler.get_denoiser(x0_predictor=x0_predictor)
>>> denoiser(xN, tN).shape  # ODE RHS for sampling
torch.Size([4, 3, 8, 8])

alpha( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant signal coefficient: \(\alpha(t) = 1\).

alpha_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Zero derivative: \(\dot{\alpha}(t) = 0\).

loss_weight( t: Float[Tensor, 'N'], ) → Float[Tensor, 'N'][source]#

Compute iDDPM loss weight: \(w(t) = 1 / \sigma(t)^2\).

Important

This loss weight is designed for training an x0-predictor (clean data predictor) wrapped with IDDPMPreconditioner. It is not suitable for training a score-predictor, or a model without a pre-conditioner.

Parameters:: t (Tensor) – Diffusion time values of shape \((N,)\).
Returns:: Loss weight of shape \((N,)\).
Return type:: Tensor

sample_time( N: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N'][source]#

Sample N diffusion times uniformly from precomputed schedule.

Parameters:

N (int) – Number of time values to sample.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Sampled diffusion times of shape \((N,)\).

Return type:

Tensor

sigma( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: For iDDPM, \(\sigma(t) = t\) (identity mapping).

sigma_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant derivative: \(\dot{\sigma}(t) = 1\).

sigma_inv( sigma: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: For iDDPM, \(t = \sigma\) (identity mapping).

timesteps( num_steps: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N+1'][source]#

Generate iDDPM time-steps from precomputed schedule.

Subsamples num_steps values from the precomputed schedule of \(M\) noise levels.

Parameters:

num_steps (int) – Number of sampling steps.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Time-steps tensor of shape \((N + 1,)\).

Return type:

torch.Tensor

`StudentTEDMNoiseScheduler`#

class physicsnemo.diffusion.noise_schedulers.StudentTEDMNoiseScheduler( sigma_min: float = 0.002, sigma_max: float = 80.0, rho: float = 7.0, nu: int = 10, sigma_data: float = 0.5, P_mean: float = -1.2, P_std: float = 1.2, )[source]#

Bases: LinearGaussianNoiseScheduler

Student-t EDM noise scheduler for heavy-tailed diffusion models.

This scheduler is a variant of EDMNoiseScheduler that uses Student-t noise instead of Gaussian noise. It is useful for modeling heavy-tailed distributions and can improve sample quality for certain data types.

Important

Despite inheriting from LinearGaussianNoiseScheduler, this scheduler is not truly Gaussian. It uses the same linear structure (identity mappings \(\sigma(t) = t\) and \(\alpha(t) = 1\)) but replaces Gaussian noise with Student-t noise. The “Linear” part of LinearGaussianNoiseScheduler still applies, but the “Gaussian” part does not.

This scheduler uses a non-gaussian forward process:

\[\mathbf{x}(t) = \mathbf{x}_0 + \sigma(t) \mathbf{n}, \quad \mathbf{n} \sim \text{Student-}t(\nu)\]

The marginal distribution \(p(\mathbf{x}_t | \mathbf{x}_0)\) is therefore a scaled Student-t distribution, not Gaussian.

Comparison with EDMNoiseScheduler:

This scheduler shares the same time-to-noise mappings as EDMNoiseScheduler. The only differences are in add_noise() and init_latents(), which use Student-t noise instead of Gaussian noise.

Parameters:

sigma_min (float, optional) – Minimum noise level for sampling time-steps, by default 0.002.
sigma_max (float, optional) – Maximum noise level for sampling time-steps, by default 80.
rho (float, optional) – Exponent controlling time-step spacing. Larger values concentrate more steps at lower noise levels (better for fine details). By default 7.
nu (int, optional) – Degrees of freedom for Student-t distribution. Must be > 2. As nu increases, the distribution approaches Gaussian. Lower values produce heavier tails. By default 10.
sigma_data (float, optional) – Expected standard deviation of the training data, by default 0.5. Used by loss_weight() to compute the per-sample loss weight.
P_mean (float, optional) – Mean of the log-normal distribution used to sample training times, by default -1.2.
P_std (float, optional) – Standard deviation of the log-normal distribution used to sample training times, by default 1.2.

Note

Reference: Heavy-Tailed Diffusion Models

Examples

Basic training and sampling workflow with Student-t noise:

>>> import torch
>>> from physicsnemo.diffusion.noise_schedulers import (
...     StudentTEDMNoiseScheduler,
... )
>>>
>>> scheduler = StudentTEDMNoiseScheduler(nu=10)
>>>
>>> # Training: sample times and add Student-t noise
>>> x0 = torch.randn(4, 3, 8, 8)  # Clean data
>>> t = scheduler.sample_time(4)    # Sample diffusion times
>>> x_t = scheduler.add_noise(x0, t)  # Adds Student-t noise
>>> x_t.shape
torch.Size([4, 3, 8, 8])
>>>
>>> # Sampling: generate timesteps and Student-t initial latents
>>> t_steps = scheduler.timesteps(10)
>>> tN = t_steps[0].expand(4)
>>> xN = scheduler.init_latents((3, 8, 8), tN)  # Student-t latents
>>> xN.shape
torch.Size([4, 3, 8, 8])

add_noise( x0: Float[Tensor, 'B *dims'], time: Float[Tensor, 'B'], ) → Float[Tensor, 'B *dims'][source]#

Add Student-t noise to clean data at the given diffusion times.

Unlike the Gaussian case in LinearGaussianNoiseScheduler, this method uses Student-t noise:

\[\mathbf{x}(t) = \mathbf{x}_0 + \sigma(t) \mathbf{n}, \quad \mathbf{n} \sim \text{Student-}t(\nu)\]

Parameters:

x0 (Tensor) – Clean latent state of shape \((B, *)\).
time (Tensor) – Diffusion time values of shape \((B,)\).

Returns:

Noisy latent state of shape \((B, *)\).

Return type:

Tensor

alpha( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant signal coefficient: \(\alpha(t) = 1\).

alpha_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Zero derivative: \(\dot{\alpha}(t) = 0\).

init_latents( spatial_shape: Tuple[int, ...], tN: Float[Tensor, 'B'], *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'B *spatial_shape'][source]#

Initialize noisy latent state with Student-t noise.

Unlike the Gaussian case in LinearGaussianNoiseScheduler, this method uses Student-t noise:

\[\mathbf{x}_N = \sigma(t_N) \cdot \mathbf{n}, \quad \mathbf{n} \sim \text{Student-}t(\nu)\]

Parameters:

spatial_shape (Tuple[int, ...]) – Spatial shape of the latent state, e.g., (C, H, W).
tN (Tensor) – Initial diffusion time of shape \((B,)\).
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Initial noisy latent of shape \((B, *spatial\_shape)\).

Return type:

Tensor

loss_weight( t: Float[Tensor, 'N'], ) → Float[Tensor, 'N'][source]#

Compute Student-t EDM loss weight: \(w(t) = \frac{\tilde{\sigma}(t)^2 + \sigma_{\text{data}}^2} {\left(\tilde{\sigma}(t) \cdot \sigma_{\text{data}}\right)^2}\)

where \(\tilde{\sigma}(t) = \sigma(t) \cdot \sqrt{\frac{\nu}{\nu - 2}}\) is the scaled noise level.

Important

This loss weight is designed for training an x0-predictor (clean data predictor) wrapped with EDMPreconditioner. It is not suitable for training a score-predictor, or a model without a pre-conditioner.

Parameters:: t (Tensor) – Diffusion time values of shape \((N,)\).
Returns:: Loss weight of shape \((N,)\).
Return type:: Tensor

sample_time( N: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N'][source]#

Sample N diffusion times from a log-normal distribution: \(\ln(\sigma) \sim \mathcal{N}(P_{\text{mean}}, P_{\text{std}}^2)\).

Parameters:

N (int) – Number of time values to sample.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Sampled diffusion times of shape \((N,)\).

Return type:

Tensor

sigma( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Identity mapping: \(\sigma(t) = t\).

sigma_dot( t: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Constant derivative: \(\dot{\sigma}(t) = 1\).

sigma_inv( sigma: Float[Tensor, '*shape'], ) → Float[Tensor, '*shape'][source]#: Identity mapping: \(t = \sigma\).

timesteps( num_steps: int, *, device: device | None = None, dtype: dtype | None = None, ) → Float[Tensor, 'N+1'][source]#

Generate EDM time-steps with polynomial spacing.

Parameters:

num_steps (int) – Number of sampling steps.
device (torch.device, optional) – Device to place the tensor on.
dtype (torch.dtype, optional) – Data type of the tensor.

Returns:

Time-steps tensor of shape \((N + 1,)\) where \(N\) is num_steps.

Return type:

torch.Tensor

Noise Schedulers#

Role in Training#

Role in Inference (Sampling)#

Three Levels of Customization#

API Reference#

NoiseScheduler#

LinearGaussianNoiseScheduler#

EDMNoiseScheduler#

VENoiseScheduler#

VPNoiseScheduler#

IDDPMNoiseScheduler#

StudentTEDMNoiseScheduler#

`NoiseScheduler`#

`LinearGaussianNoiseScheduler`#

`EDMNoiseScheduler`#

`VENoiseScheduler`#

`VPNoiseScheduler`#

`IDDPMNoiseScheduler`#

`StudentTEDMNoiseScheduler`#