Refactored diffusion sampler

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

Overview

This PR introduces a new modular framework for diffusion model sampling, designed to decouple the core components of the sampling process for maximum flexibility and reusability.

Key Contributions

Decoupled Architecture

• Model-agnostic: The sampling function now works with any model implementing the DiffusionDenoiser protocol—shape-agnostic, dimension-agnostic, and supports arbitrary conditioning/guidance
• Solver ↔ Noise Schedule decoupling: Numerical solvers and noise schedules are now independent, allowing any combination
• Flexible time-steps: Time-step discretization can optionally be overridden independently of both solver and noise schedule

New Abstractions

Solver Protocol (physicsnemo.diffusion.samplers.solvers)

Defines a minimal interface for numerical integration methods; 4 concrete solver implementations: Euler, Heun, and their respective EDM-stochastic versions.

NoiseScheduler Protocol (physicsnemo.diffusion.noise_schedulers)

Defines minimal interface for training (add_noise, sample_time) and sampling (timesteps, init_latents, get_denoiser). A base class (LinearGaussianNoiseScheduler) implements the protocol for the family of linear-Gaussian schedules; 6 concrete implementation of such noise schedulers: LinearGaussianNoiseScheduler, EDMNoiseScheduler, VENoiseScheduler, VPNoiseScheduler, IDDPMNoiseScheduler, StudentTEDMNoiseScheduler.

DiffusionDenoiser Protocol (physicsnemo.diffusion.base)

Minimal interface for denoisers.

New sample Function (physicsnemo.diffusion.samplers.sample)

Unified entry point for sampling with:

• Configurable solver (string key or custom instance)
• Configurable noise scheduler
• Optional custom time-steps
• Support for intermediate sample collection

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

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.

[greptile-apps]

Additional Comments (1)

physicsnemo/diffusion/noise_schedulers/noise_schedulers.py
1725 lines of noise scheduler implementations without tests. MOD-008a, MOD-008b require comprehensive tests for all 6 scheduler classes.

Add tests verifying:

• Constructor/attributes with default and custom parameters
• Non-regression tests for timesteps(), sample_time(), add_noise(), init_latents(), get_denoiser() against reference data
• Abstract method implementations across configurations

[pzharrington]

Nice job with this, really only minor comments/clarifications. While I know you have the CI tests as a remaining to-do, would be a good sanity check to make sure piping together the new machinery can reproduce the result from one of the previous more rigid/standalone samplers.

[NickGeneva]

Approving to unblock, pending small usability improvements for the get_denoiser + adding (or renaming) a DiffusionOperator call to help clear up type confusion.

[CharlelieLrt]

/blossom-ci