Skip to content

Add inverse SNR and TrigFlow components - #2095

Closed
CharlelieLrt wants to merge 1 commit into
codex/diffusion-structured-iofrom
codex/diffusion-trigflow
Closed

CharlelieLrt wants to merge 1 commit into
codex/diffusion-structured-iofrom
codex/diffusion-trigflow

Conversation

@CharlelieLrt

@CharlelieLrt CharlelieLrt commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Description

This PR adds TrigFlowNoiseScheduler and TrigFlowPreconditioner to
physicsnemo.diffusion. Users can pair them with the same sigma_data value to train
a backbone that predicts normalized velocity while exposing clean-data
predictions through the usual diffusion model interface.

Built-in linear Gaussian schedules also gain snr_inv(q), which converts the
amplitude ratio q = alpha / sigma back to diffusion time.
Custom schedules can supply their own inverse. This gives the private FastGen
adapters a way to translate noise levels between schedules without relying on
inference grids.

Checklist

  • This PR follows the contribution guidelines.
  • The documentation describes the API and its role in the stack.
  • This PR remains a draft.

Dependencies

If FastGen maintainers decline FastGen #44, we would close PhysicsNeMo #2100.
The preceding PhysicsNeMo stack would remain available.

Review process

Reviewers can compare this PR with its immediate parent. After a parent merges,
we update the remaining branches and their PR bases to keep each diff focused.

@CharlelieLrt CharlelieLrt added 2 - In Progress Currently a work in progress 5 - DO NOT MERGE Hold off on merging; see PR for details 5 - Merge After Dependencies Depends on another PR: do not merge out of order labels Oct 14, 2026
@copy-pr-bot

copy-pr-bot Bot commented Oct 14, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions

Copy link
Copy Markdown
Contributor

CODEOWNERS review map

Current for commit bda6a9107929. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @CharlelieLrt — 24 file(s)
  • docs/api/diffusion/noise_schedulers.rst
  • docs/api/diffusion/preconditioners.rst
  • physicsnemo/diffusion/noise_schedulers/init.py
  • physicsnemo/diffusion/noise_schedulers/edm.py
  • physicsnemo/diffusion/noise_schedulers/iddpm.py
  • physicsnemo/diffusion/noise_schedulers/linear_gaussian.py
  • physicsnemo/diffusion/noise_schedulers/rectified_flow.py
  • physicsnemo/diffusion/noise_schedulers/trigflow.py
  • physicsnemo/diffusion/noise_schedulers/ve.py
  • physicsnemo/diffusion/noise_schedulers/vp.py
  • physicsnemo/diffusion/preconditioners/init.py
  • physicsnemo/diffusion/preconditioners/preconditioners.py
  • test/diffusion/data/trigflow_precond_conv.mdlus
  • test/diffusion/data/trigflow_precond_conv_coefficients.pth
  • test/diffusion/data/trigflow_precond_conv_forward.pth
  • test/diffusion/data/trigflow_precond_conv_sigma.pth
  • test/diffusion/data/trigflow_precond_linear.mdlus
  • test/diffusion/data/trigflow_precond_linear_coefficients.pth
  • test/diffusion/data/trigflow_precond_linear_forward.pth
  • test/diffusion/data/trigflow_precond_linear_sigma.pth
  • test/diffusion/test_model_io.py
  • test/diffusion/test_noise_schedulers.py
  • test/diffusion/test_preconditioners.py
  • test/diffusion/test_trigflow.py
⏳ @megnvidia — 2 file(s)
  • docs/api/diffusion/noise_schedulers.rst
  • docs/api/diffusion/preconditioners.rst

No CODEOWNER

  • CHANGELOG.md

Comment /codeowners-info to refresh.

@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-trigflow branch from bda6a91 to ab1841d Compare October 16, 2026 01:19
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-structured-io branch from fe02e10 to 455f611 Compare October 16, 2026 01:19
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-trigflow branch from ab1841d to fbfcd1b Compare October 16, 2026 02:52
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-structured-io branch from 455f611 to 58c9ff2 Compare October 16, 2026 02:52
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-trigflow branch from fbfcd1b to 901fe94 Compare October 16, 2026 20:15
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-structured-io branch 2 times, most recently from a11d86e to b0b607e Compare October 17, 2026 01:44
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-trigflow branch from 901fe94 to a9b6d79 Compare October 17, 2026 01:44
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-trigflow branch from a9b6d79 to ff99456 Compare October 28, 2026 03:42
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-structured-io branch from b0b607e to 5b5e3a9 Compare October 28, 2026 03:42
Signed-off-by: Charlelie Laurent <claurent@nvidia.com>
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-structured-io branch from 5b5e3a9 to ee0ab8d Compare October 29, 2026 02:23
@CharlelieLrt
CharlelieLrt force-pushed the codex/diffusion-trigflow branch from ff99456 to d4a40d7 Compare October 29, 2026 02:23
@CharlelieLrt

Copy link
Copy Markdown
Collaborator Author

Continued in PhysicsNeMo #2200, from branch diffusion-trigflow in CharlelieLrt/modulus. The replacement preserves this implementation and includes the requested documentation clarifications. Please continue the review there. This PR retains the earlier discussion.

@CharlelieLrt
CharlelieLrt deleted the codex/diffusion-trigflow branch October 29, 2026 04:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

2 - In Progress Currently a work in progress 5 - DO NOT MERGE Hold off on merging; see PR for details 5 - Merge After Dependencies Depends on another PR: do not merge out of order

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant