1 / 1

Karras Euler Sampler

Diffusion ModelsSamplingODE Solvers

Hard

Build the Karras noise schedule and integrate the probability flow ODE from sigma_max down to 0 with Euler or Heun steps.

With $N$ = num_steps and $\rho$ = rho, the noise levels for $i = 0, \dots, N - 1$ are

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

followed by an appended $\sigma_N = 0$ . The data is Gaussian with mean $\mu$ (data_mean) and standard deviation $s$ (data_std) in every coordinate, so the exact denoiser is $D(x, \sigma) = (s^2 x + \sigma^2 \mu) / (s^2 + \sigma^2)$ .

Start from $x = \sigma_{\max}$ times initial_noise. For each $i$ from 0 to $N - 1$ , compute the slope $d = (x - D(x, \sigma_i)) / \sigma_i$ and the Euler proposal $x' = x + (\sigma_{i+1} - \sigma_i)\, d$ . For method "euler" the new $x$ is $x'$ . For method "heun", when $\sigma_{i+1} > 0$ , compute $d' = (x' - D(x', \sigma_{i+1})) / \sigma_{i+1}$ and set the new $x$ to $x + (\sigma_{i+1} - \sigma_i)(d + d') / 2$ . When $\sigma_{i+1} = 0$ , both methods set the new $x$ to the Euler proposal.

Return a dict with "sigmas", a PyTorch tensor of shape (num_steps + 1,), and "samples", a PyTorch tensor with the shape of initial_noise, both with the dtype and device of initial_noise.

Examples

Input: initial_noise = [[1.2,-0.5,0.3],[-0.8,0.1,2.0]], num_steps = 5, sigma_min = 0.02, sigma_max = 10.0, rho = 7.0, data_mean = 0.5, data_std = 0.5, method = "euler"

Output: {"sigmas":[10.0,3.282849,0.872338,0.169858,0.02,0.0],"samples":[[0.882868,0.316889,0.583232],[0.217011,0.516646,1.149211]]}

Explanation: The five Karras noise levels from 10 down to 0.02 are about 10, 3.283, 0.872, 0.170 and 0.02, followed by the appended 0. The samples start at 10 times the noise, and the last Euler step to sigma 0 returns the denoiser output at sigma 0.02, so the first noise value 1.2 ends at about 0.883.

Input: initial_noise = [[1.2,-0.5,0.3],[-0.8,0.1,2.0]], num_steps = 5, sigma_min = 0.02, sigma_max = 10.0, rho = 7.0, data_mean = 0.5, data_std = 0.5, method = "heun"

Output: {"sigmas":[10.0,3.282849,0.872338,0.169858,0.02,0.0],"samples":[[1.277256,0.128269,0.668969],[-0.074493,0.533794,1.817955]]}

Input: initial_noise = [[[[0.5,-1.0],[1.5,0.0]]]], num_steps = 2, sigma_min = 1.0, sigma_max = 4.0, rho = 1.0, data_mean = -1.0, data_std = 2.0, method = "heun"

Output: {"sigmas":[4.0,1.0,0.0],"samples":[[[[0.392,-2.392],[2.248,-0.536]]]]}

Hint 1

torch.linspace from 0 to 1 gives the ramp i / (N - 1) for every noise level at once.

Hint 2

torch.cat([sigmas, sigmas.new_zeros(1)]) appends an exact 0 with the same dtype.

Hint 3

A small inner function D(x, sigma) lets you call the denoiser twice in the Heun branch.

Requirements

Use PyTorch tensor operations.
Order the noise levels from sigma_max down to sigma_min.
Return a dict with no keys other than "sigmas" and "samples".
Leave initial_noise unchanged.

Constraints

initial_noise is a floating-point PyTorch tensor of shape (B, ...) with B at least 1 and at least one non-batch dimension.
num_steps is a Python int between 2 and 25.
sigma_min and sigma_max are Python floats with 0 < sigma_min < sigma_max <= 80.
rho is a Python float between 1 and 7.
data_mean is a Python float, and data_std is a Python float above 0.
method is exactly "euler" or "heun".
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

DDIM StepHard DDPM Ancestral StepMedium Classifier-Free GuidanceMedium Toy DDPM SamplerHard Annealed Langevin DynamicsMedium Tiny U-Net Forward PassHard DiT Patchify and adaLN-Zero BlockHard Toy Diffusion Training StepHard Gaussian Mixture ScoreMedium Cosine Noise ScheduleMedium Rescale to Zero Terminal SNRMedium Time-Conditioned Residual BlockMedium

Karras Euler Sampler

Diffusion ModelsSamplingODE Solvers

Hard

Build the Karras noise schedule and integrate the probability flow ODE from sigma_max down to 0 with Euler or Heun steps.

With $N$ = num_steps and $\rho$ = rho, the noise levels for $i = 0, \dots, N - 1$ are

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

Return a dict with "sigmas", a PyTorch tensor of shape (num_steps + 1,), and "samples", a PyTorch tensor with the shape of initial_noise, both with the dtype and device of initial_noise.

Examples

Input: initial_noise = [[1.2,-0.5,0.3],[-0.8,0.1,2.0]], num_steps = 5, sigma_min = 0.02, sigma_max = 10.0, rho = 7.0, data_mean = 0.5, data_std = 0.5, method = "euler"

Output: {"sigmas":[10.0,3.282849,0.872338,0.169858,0.02,0.0],"samples":[[0.882868,0.316889,0.583232],[0.217011,0.516646,1.149211]]}

Input: initial_noise = [[1.2,-0.5,0.3],[-0.8,0.1,2.0]], num_steps = 5, sigma_min = 0.02, sigma_max = 10.0, rho = 7.0, data_mean = 0.5, data_std = 0.5, method = "heun"

Output: {"sigmas":[10.0,3.282849,0.872338,0.169858,0.02,0.0],"samples":[[1.277256,0.128269,0.668969],[-0.074493,0.533794,1.817955]]}

Input: initial_noise = [[[[0.5,-1.0],[1.5,0.0]]]], num_steps = 2, sigma_min = 1.0, sigma_max = 4.0, rho = 1.0, data_mean = -1.0, data_std = 2.0, method = "heun"

Output: {"sigmas":[4.0,1.0,0.0],"samples":[[[[0.392,-2.392],[2.248,-0.536]]]]}

Hint 1

torch.linspace from 0 to 1 gives the ramp i / (N - 1) for every noise level at once.

Hint 2

torch.cat([sigmas, sigmas.new_zeros(1)]) appends an exact 0 with the same dtype.

Hint 3

A small inner function D(x, sigma) lets you call the denoiser twice in the Heun branch.

Requirements

Use PyTorch tensor operations.
Order the noise levels from sigma_max down to sigma_min.
Return a dict with no keys other than "sigmas" and "samples".
Leave initial_noise unchanged.

Constraints

initial_noise is a floating-point PyTorch tensor of shape (B, ...) with B at least 1 and at least one non-batch dimension.
num_steps is a Python int between 2 and 25.
sigma_min and sigma_max are Python floats with 0 < sigma_min < sigma_max <= 80.
rho is a Python float between 1 and 7.
data_mean is a Python float, and data_std is a Python float above 0.
method is exactly "euler" or "heun".
Supported dtypes are float32 and float64. All values are finite.

Try Similar Problems

initial_noise =

Accepts: array

num_steps =

Accepts: integer

sigma_min =

Accepts: number

sigma_max =

Accepts: number

rho =

Accepts: number

data_mean =

Accepts: number

data_std =

Accepts: number

method =

Accepts: string