samplers

samplers

Backward-compatibility shim: samplers moved to deltaflow.solvers.

EulerSolver

EulerSolver(model, time_scale=1.0)

Bases: BaseSolver

Explicit (forward) Euler integrator for the flow-matching ODE.

Integrates \(\mathrm{d}x/\mathrm{d}t = v_\theta(x, t)\) with the first-order update

\[ x_{n+1} = x_n + \Delta t\; v_\theta(x_n, t_n), \]

using a single velocity evaluation per step. It is cheap but incurs \(\mathcal{O}(\Delta t^2)\) local truncation error (\(\mathcal{O}(\Delta t)\) global), so prefer HeunSolver when accuracy at low step counts matters.

Source code in deltaflow/core/base_solver.py

def __init__(self, model: Callable, time_scale: float = 1.0):
    self.model = model
    self.time_scale = time_scale

sample

sample(x, n_steps=50, x_cond=None, t_start=0.0, t_end=1.0, show_progress=True, **cond)

Generate samples by integrating the velocity field.

Parameters:

Name	Type	Description	Default
x	Tensor	initial state (shape/device/dtype template if x_cond is set).	required
n_steps	int	number of Euler integration steps.	50
x_cond	Optional[Tensor]	optional partially-noised starting point, if given the initial state becomes (1 - t_start) * noise + t_start * x_cond.	None
t_start	float	starting time, used together with x_cond.	0.0
t_end	float	end time, defaults to 1.0.	1.0
**cond	Any	extra keyword arguments forwarded to the velocity model.	{}

Source code in deltaflow/solvers/euler.py

@torch.no_grad()
def sample(
    self,
    x: torch.Tensor,
    n_steps: int = 50,
    x_cond: Optional[torch.Tensor] = None,
    t_start: float = 0.0,
    t_end: float = 1.0,
    show_progress: bool = True,
    **cond: Any,
) -> torch.Tensor:
    """Generate samples by integrating the velocity field.

    Args:
        x: initial state (shape/device/dtype template if ``x_cond`` is set).
        n_steps: number of Euler integration steps.
        x_cond: optional partially-noised starting point, if given the
            initial state becomes ``(1 - t_start) * noise + t_start * x_cond``.
        t_start: starting time, used together with ``x_cond``.
        t_end: end time, defaults to ``1.0``.
        **cond: extra keyword arguments forwarded to the velocity model.
    """
    if x_cond is not None:
        x0 = torch.randn_like(x_cond)
        x = (1 - t_start) * x0 + t_start * x_cond
    return super().sample(
        x,
        n_steps=n_steps,
        t_start=t_start,
        t_end=t_end,
        show_progress=show_progress,
        progress_desc="EulerSolver",
        **cond,
    )

HeunSolver

HeunSolver(model, time_scale=1.0)

Bases: BaseSolver

Heun (improved-Euler) second-order predictor-corrector integrator.

Each step takes an Euler predictor and averages the velocity at the current and predicted states,

\[ \begin{aligned} k_1 &= v_\theta(x_n, t_n), \\ k_2 &= v_\theta\bigl(x_n + \Delta t\,k_1,\; t_n + \Delta t\bigr), \\ x_{n+1} &= x_n + \tfrac{\Delta t}{2}\,(k_1 + k_2). \end{aligned} \]

This costs two velocity evaluations per step but has \(\mathcal{O}(\Delta t^3)\) local truncation error (\(\mathcal{O}(\Delta t^2)\) global), so it typically matches Euler's quality at half the number of steps. It is the trapezoidal-rule integrator widely used in EDM-style samplers.

Source code in deltaflow/core/base_solver.py

def __init__(self, model: Callable, time_scale: float = 1.0):
    self.model = model
    self.time_scale = time_scale

PosteriorSolver

PosteriorSolver(base_solver, likelihood, tweedie=None, guidance_scale=1.0, grad_normalize=False)

Bases: BaseSolver

Base-solver wrapper that adds a per-step measurement-likelihood gradient.

For an inverse problem with measurement \(y = A(x) + n\), this samples from the posterior \(p(x \mid y)\) by nudging an unconditional flow-matching solver toward the data-consistency term at every step. Given the current state \(x_t\) and velocity \(v_\theta(x_t, t)\), one step is

\[ \hat{x}_1 = \mathcal{T}(x_t, v_\theta(x_t, t), t), \qquad x_{t+\Delta t} = \underbrace{\text{Base}(x_t, t, \Delta t)}_{\text{unconditional step}} \;-\; \eta\,\nabla_{x_t}\bigl[-\log p(y \mid \hat{x}_1)\bigr], \]

where \(\mathcal{T}\) is the flow-matching Tweedie decomposition that maps \((x_t, v_t, t)\) to the clean-signal estimate \(\hat{x}_1\) (see deltaflow.inverse.tweedie), and \(\eta\) is the guidance_scale. The likelihood gradient is obtained by autograd, so if the velocity field operates on VAE latents while \(A\) is defined on pixels, passing a decoder to the likelihood object pulls the gradient back into latent space automatically. Only the sampling-time ODE is modified. The pretrained velocity field is untouched.

References

Kim et al., "FlowDPS: Flow-Driven Posterior Sampling for Inverse Problems" (2025), https://arxiv.org/abs/2503.08136. "Flower: A Flow-Matching Solver for Inverse Problems" (2025), https://arxiv.org/abs/2509.26287.

Parameters:

Name	Type	Description	Default
base_solver	BaseSolver	any BaseSolver that already integrates the unconditional flow (Euler, Heun, ...).	required
likelihood	Likelihood	object with a .neg_log_prob(x_clean) method that returns a per-sample (or reducible) scalar tensor with requires_grad=True support. See deltaflow.inverse.likelihood.	required
tweedie	Optional[BaseTweedie]	flow-matching Tweedie decomposition \(\mathcal{T}\) to derive \(\hat{x}_1\) from \((x_t, v_t, t)\). Defaults to LinearTweedie, matching a LinearInterpolant training path.	None
guidance_scale	float	step size \(\eta\) on the likelihood gradient. Larger values snap harder to the measurement but risk over-shooting.	1.0
grad_normalize	bool	if True, the injected gradient is rescaled to match the norm of the base step. This is a stability trick used in some DPS variants when likelihood magnitudes vary wildly.	False

Source code in deltaflow/solvers/posterior_solver.py

def __init__(
    self,
    base_solver: BaseSolver,
    likelihood: Likelihood,
    tweedie: Optional[BaseTweedie] = None,
    guidance_scale: float = 1.0,
    grad_normalize: bool = False,
):
    # Re-use the base solver's model reference and time_scale.
    super().__init__(model=base_solver.model, time_scale=base_solver.time_scale)
    self.base = base_solver
    self.likelihood = likelihood
    self.tweedie = tweedie or LinearTweedie()
    self.guidance_scale = guidance_scale
    self.grad_normalize = grad_normalize