Add exponential Euler, stochastic exponential Euler, DPM-Solver++(2M), and UniC-2 diffusion samplers - #1932
Add exponential Euler, stochastic exponential Euler, DPM-Solver++(2M), and UniC-2 diffusion samplers#1932CharlelieLrt wants to merge 25 commits into
Conversation
…on solvers. Configurations reproduce DDIM, the re-noising sampler of distilled models, and DPM-Solver++(2M). Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
… ExponentialAB2Solver with ExponentialEulerSolver, EDMStochasticExponentialEulerSolver, and a dedicated DPMPlusPlus2M. Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
|
Too many files changed for review (133 files, 100 file limit). Bypass the limit by tagging |
CODEOWNERS review mapCurrent for commit ⏳ @CharlelieLrt — 133 file(s)
⏳ @megnvidia — 1 file(s)
Comment |
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
|
Thanks for integrating the DPM++2M work into this PR. I tested the public
The default path. The class docstring says that omitting On the analytic EDM problem in the attached reproduction, the public path has 
Standard DPM++2M is still available, which I want to be clear about. If I This recipe is not documented, and its time steps do not follow the documented The scheduler-supplied This is also the configuration used by the solver test matrix Test coverage. As a check, I replaced the DPM++ multistep coefficient Could you clarify:
A short CPU reproduction pinned to |
…ring the existing get_denoiser coverage. Also check in the not-yet-committed golden files for the get_linear_denoiser and DPMPlusPlus2M test suites. Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
…the solver and sampler suites. Also unify the DPMPlusPlus2M step into a single masked update path, so a compiled solver needs no caller-side priming. Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
|
I reintroduced the extrapolation variable The time steps grid does not need to be in lambda coordinates. The solver always operates in (x, t) coordinates. The lambda conversion is purely internal (only used to approximate the moments/integrals). |
|
@megnvidia Could I get your review on this PR for docs modifications? Thanks! |
|
Thanks @CharlelieLrt, current solver multiplier: current solver output: To be clear the kernel can express canonical DPM2++2M but what I believe we can not reproduce is DPM++2M using the documented |
|
@gertln this AI slop is really not helpful. Can you please actually look at the code and let me know where you think there's an error? |
laserkelvin
left a comment
laserkelvin
left a comment
There was a problem hiding this comment.
I think more generally, adding regression tests that check the solver value at t=0, t=0.5, and t=1 against simple cases would be good for these three solvers (ideally all solvers). If you can resolve that and make sure all tests pass, it'll give me some more time to review as well
…evaluate the nonlinear weight as an exponential-kernel integral via a private Gauss-Legendre quadrature, extrapolate the predictor-like term in DPM-Solver++(2M), and add a snr method to the linear-Gaussian scheduler. Regenerate the solver and sampler golden files for the new numerics. Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
CharlelieLrt
changed the title
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
laserkelvin
left a comment
laserkelvin
left a comment
There was a problem hiding this comment.
Ooga! Me look long at this PR. Short version: me like. Math check out — derived by hand twice (independent reviewers): DDIM reduction, DPM-Solver++(2M) update, UniC-2 corrector coefficients, EDM churn — all correct. VP sigma_to_t rewrite is a real fix (old code return slightly negative t at sigma=0). Golden files clean (334 new, all referenced, zero orphans). Compile design with masked updates + stable cache storage is smart.
Me leave comments inline. Two big asks, rest small.
BIG ASK 1 — string keys not do what name say. sample(..., solver="dpmpp_2m") with no solver_options build solver with empty pockets (no scheduler callbacks) → quietly run classical AB2, not DPM-Solver++(2M). Same story for the other three new keys. Defaults are a fine feature (me see the dpmpp_2m_ab2 test use them on purpose) — but a user who not read the class docstrings get a different method than the name promise, and nothing tell them. Ask: warn when the new keys are used without callbacks + per-key honesty in the sample() docstring and docs page. Details inline at samplers.py.
BIG ASK 2 — the new consistency test cannot see the fancy parts. TestStepConsistency use slope=0, which make the nonlinear term N≡0 — so the J quadrature, the DPM extrapolation, and the UniC-2 corrector all get multiplied by zero, and the single-step design never reach UniC-2 steady state. The "second-order" / "third-order" claims today rest only on self-generated goldens, which catch changes, not wrong-from-the-start. Ask: one empirical order-of-accuracy test — golden-free mock-up inline at test_solvers.py, ~50 lines, CPU.
Small asks inline: dedup _nonlinear_weight ×4 into _utils.py; validate denoising_type; CHANGELOG fixes; strengthen the gradient-flow assertion; one SDE end-to-end config.
One question inline (dpmpp_2m_unic2.py): please confirm grads-through-time are out of scope — the copy_ cache updates would crash backward loudly in that config; if out of scope, worth a one-line invariant comment.
And noting here because the file is not in this diff: DomainParallelNoiseScheduler forwards get_denoiser but not the new get_linear_denoiser/snr — domain-parallel users must find the .inner_scheduler trick themselves. Two small pass-through methods would keep parity.
Nice work overall — the numerics are solid; the asks above are about what users and tests can see.
…ver. Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
…ian.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
…o make them differentiable wrt timesteps Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
|
/ok to test f64647e |
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
|
/ok to test eb84621 |
This branch has not been deployed
4 participants
PhysicsNeMo Pull Request
Description
This PR adds four composable diffusion solvers to
physicsnemo.diffusion.samplers. Each solver can be instantiated directly or selected by passing its string key tosample().Solvers
ExponentialEulerSolver("exponential_euler") is a first-order solver for semi-linear ODEs. It supports DDIM-like sampling and distilled few-step models.EDMStochasticExponentialEulerSolver("edm_stochastic_exponential_euler") adds EDM-style churn and controllable re-noising to exponential Euler. It supports stochastic, ancestral, and full re-noising configurations for distilled and consistency models.DPMPlusPlus2M("dpmpp_2m") is a configurable second-order multistep solver with one denoiser evaluation per step. Its defaults recover classical two-step Adams-Bashforth; configuring it with the schedule's semi-linear structure recovers the published DPM-Solver++(2M) formulation.DPMPlusPlus2MUniC2("dpmpp_2m_unic2") adds a UniC-2 corrector stage on top ofDPMPlusPlus2M, raising the update to third order while keeping one denoiser evaluation per step, at the cost of a slightly larger memory footprint.API integration
LinearGaussianNoiseScheduler.get_linear_denoiser()to expose the semi-linear structure of the ODE/SDE right-hand side that the new exponential and multistep solvers rely on. It complementsget_denoiser(), which returns the full right-hand side.LinearGaussianNoiseScheduler.snr()to expose the schedule's signal-to-noise ratio, used as the multistep extrapolation coordinate byDPMPlusPlus2MandDPMPlusPlus2MUniC2.sample()with string-key selection and solver-specific options for all four solvers.Relationship to #1896
This PR supersedes #1896. It builds on the DPM-Solver++(2M) implementation from that PR through a leaner and more compositional API. See #1896 for quantitative evidence of the sampler's benefits.
Checklist
Dependencies
None.
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.
AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.