Add epsilon predictor support to diffusion framework

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

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]

Greptile Summary

This PR adds epsilon-predictor support to the diffusion framework by introducing four new conversion methods (epsilon_to_score, score_to_epsilon, epsilon_to_x0, x0_to_epsilon) on LinearGaussianNoiseScheduler, an epsilon_predictor option in get_denoiser, and "epsilon" as a valid prediction_type in both MSEDSMLoss and WeightedMSEDSMLoss. The math is correct, the implementation is clean, and test coverage is comprehensive across schedulers, spatial configs, and compile checks. The only gaps are documentation: both loss classes have outdated docstrings that still list prediction_type as Literal["x0", "score"] and omit the new epsilon_to_x0_fn parameter from their Parameters sections.

Important Files Changed

Filename	Overview
physicsnemo/diffusion/metrics/losses.py	Adds epsilon_to_x0_fn parameter and "epsilon" prediction type to both MSEDSMLoss and WeightedMSEDSMLoss; logic is correct but docstrings in both classes were not updated to reflect the new parameter or the extended type literal.
physicsnemo/diffusion/noise_schedulers/noise_schedulers.py	Adds four new conversion methods (epsilon_to_score, score_to_epsilon, epsilon_to_x0, x0_to_epsilon) to LinearGaussianNoiseScheduler, and adds epsilon_predictor support to get_denoiser; math and implementation are correct.
test/diffusion/test_losses.py	Extends all existing loss test parametrizations to include "epsilon" prediction type; coverage is thorough across schedulers, spatial configs, gradient flow, and compile tests.
test/diffusion/test_noise_schedulers.py	Adds roundtrip tests for epsilon_to_score/score_to_epsilon, epsilon_to_x0/x0_to_epsilon, cross-consistency between paths, and epsilon_predictor in get_denoiser; well-structured and comprehensive.

Comments Outside Diff (2)

1. physicsnemo/diffusion/metrics/losses.py, line 109-128 (link)

   Docstring not updated for epsilon support in MSEDSMLoss

   The Parameters and Raises sections still only mention "x0" and "score" — "epsilon" and the new epsilon_to_x0_fn parameter are undocumented. A user reading the docstring would not discover the epsilon prediction type or know that epsilon_to_x0_fn is required when using it.

2. physicsnemo/diffusion/metrics/losses.py, line 380-385 (link)

   Docstring not updated for epsilon support in WeightedMSEDSMLoss

   Same issue as in MSEDSMLoss: prediction_type still lists only Literal["x0", "score"] and epsilon_to_x0_fn is not documented anywhere in this class's Parameters section.

_{Reviews (1): Last reviewed commit: "Add epsilon predictor support to diffusi..." | Re-trigger Greptile}

[laserkelvin]

Left some design considerations - let me know what you think.

[laserkelvin]

Personal style, but I think if there are three cases it's probably better to do

match prediction_type:
    case "x0": ...
    case "score": ...
    case "epsilon": ...
    case _: raise Exception

[laserkelvin]

Maybe consider aliasing prediction_type, so that you can re-use it in the Loss __init__?

[laserkelvin]

This is pretty hard to read IMO

[CharlelieLrt]

What would you suggest to improve readability?

[laserkelvin]

expected_shape = (-1,) + (1,) * (epsilon.ndim - 1)
t_bc = t.reshape(expected_shape)

[laserkelvin]

Understandably subjective :)

[laserkelvin]

Same readability comment

[laserkelvin]

Approving