Restructure diffusion subpackage

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

Restructure Diffusion Subpackage

This PR establishes the foundational structure for the physicsnemo.diffusion subpackage, organizing diffusion-related functionality into a clean, layered architecture.

Summary

• Created physicsnemo.diffusion package with dedicated submodules:

  • generate/ - Higher-level generation utilities and pipelines
  • samplers/ - Samplers based on ODE and SDE solvers
  • metrics/ - Loss functions and diffusion-specific metrics (FID, etc.)
  • preconditioners/ - Preconditioning model wrappers
  • multi_diffusion/ - Patch-based diffusion utilities
  • denoisers/ - Placeholder for future denoiser implementations; to be used in samplers
  • noise_schedulers/ - Placeholder for future noise scheduler implementations
  • utils/ - Shared utilities

• Added warning system for incremental migration:

  • FutureFeatureWarning for placeholder modules not yet implemented
  • LegacyFeatureWarning for existing code that will be deprecated in future releases

• Reorganized UNet models:

  • Moved SongUNet models and application-specific UNet wrappers to physicsnemo.models.diffusion_unets

• Extracted reusable layers into physicsnemo.nn:
  This includes for example:

  • group_norm.py - GroupNorm variants with optional FP32 conversion
  • unet_layers.py - UNet encoder/decoder blocks, attention blocks
  • embedding_layers.py - Positional and Fourier embeddings
    Etc...

• Cleaned up legacy code:

  • Deleted unused utility functions from old physicsnemo.models.diffusion/
  • Moved example-specific utilities directly into examples/ directory

• Added import-linter contracts for the diffusion package:

  • Enforces layered dependency structure: generate → samplers/metrics → preconditioners/multi_diffusion/denoisers/noise_schedulers → utils
  • Prevents circular imports and maintains clean architecture

• Updated all examples and tests to use the new package structure

• 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.

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.

[megnvidia]

I had small changes to the two rst files. If those get merged well, I approve.

[pzharrington]

Nice, LGTM!

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[ktangsali]

Reviewed the changes in the docs/ folder and the pyproject.toml file and all of them look good to me. Thanks!

[CharlelieLrt]

/blossom-ci