Cordiff usability and performance enhancements for custom dataset training

[CharlelieLrt]

Modulus Pull Request

Description

This PR provides enhancements for usability of CorrDiff training on custom datasets. It includes:

• Documentation improvements (readme, docstrings)
• A new general purpose patching API with enhanced usability and performance
• Refactor CorrDiff samplers and losses to delegate all patching logic to the dedicated API
• Refactor CorrDiff UNets positional embeddings for performance
• Config files refactoring and simplification
• Legacy code elimination
• Support for non-square images
• Experimental support for non-square patches
• Other bugfixes

Closes #726.
Closes #447.
Closes #417.
Closes #448.

Note

This PR introduces breaking changes on the way corrdiff config files are handled.
For training, older version of the config files is still supported.
For generation, a new parameter generation.patching: true/false is required.

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.

Dependencies

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[pzharrington]

Overall this is looking nice, thanks for doing a lot of cleanup and docstring/readme improvements! I checked for compatibility/overlap with other diffusion weather models (StormCast, ReGen) and it looks fine in that regard. I had a couple minor comments about the unit tests for patching utilities but other than that no major qualms. It would be good to address Akshay's comment about performance if we can, it does seem like avoiding allocating the full meshgrid is worth doing since the whole goal is to avoid too much memory allocation and ops at the full image resolution

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

@pzharrington thanks for your comments! I made the patching tests more robust and addressed @akshaysubr suggestions for performance improvements.

[CharlelieLrt]

/blossom-ci