UniPC Predictor-Corrector Step

Problem

Implement one step of the UniPC (Unified Predictor-Corrector) sampler for diffusion models. Given the current noisy sample, a model prediction, previous model outputs, the current and previous timesteps, perform one predictor-corrector update step.

Solution

import math

def unipc_step(
    x_t: list[float], model_output: list[float],
    prev_outputs: list[list[float]], timesteps: list[float],
    alpha_t: float, sigma_t: float, alpha_s: float, sigma_s: float
) -> list[float]:
    h = math.log(alpha_t / sigma_t) - math.log(alpha_s / sigma_s)
    ratio = alpha_t / alpha_s
    exp_h = math.exp(h) - 1.0
    x_pred = [ratio * x - sigma_t * exp_h * m for x, m in zip(x_t, model_output)]
    if prev_outputs:
        prev = prev_outputs[-1]
        x_pred = [xp + -0.5 * sigma_t * exp_h * (m - p) for xp, m, p in zip(x_pred, model_output, prev)]
    return x_pred

Explanation

1. Log-SNR computation: Compute the log signal-to-noise ratio lambda at the current and previous timesteps, and the step size h between them.
2. Predictor step: Apply the first-order ODE solver (analogous to DDIM) to predict the sample at the next timestep using the scaling ratio alpha_t/alpha_s and the exponential integrator term.
3. Corrector step: If a previous model output is available, apply a second-order correction using the difference between the current and previous model outputs, which approximates the derivative of the model output trajectory.
4. This predictor-corrector scheme achieves higher-order accuracy with minimal extra computation.

Complexity

• Time: O(d) where d is the dimensionality of the sample
• Space: O(d) for the predicted output