Multi diffusion

[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 implements the multi-diffusion functionality for patch-based diffusion model training and inference, replacing the previous placeholder implementation. The changes introduce comprehensive support for decomposing 2D images into smaller patches, processing them through diffusion models, and fusing them back together.

Key changes:

• Implemented MultiDiffusionModel2D wrapper that handles patching strategies (random for training, grid for inference), condition preprocessing, and optional positional embeddings
• Added two loss classes (MultiDiffusionMSEDSMLoss and MultiDiffusionWeightedMSEDSMLoss) for patch-based denoising score matching training
• Refactored patching.py from a placeholder warning to a full implementation with BasePatching2D, RandomPatching2D, and GridPatching2D classes
• Added comprehensive test suite with reference data for non-regression testing
• Improved test organization by moving shared fixtures to conftest.py
• Updated import linter configuration to include the new multi-diffusion module

Code quality:

• Excellent documentation with detailed docstrings following NumPy style
• Comprehensive examples in docstrings demonstrating various use cases
• Proper use of jaxtyping for type hints
• Good test coverage with reference data for validation

Minor issues found:

• Two instances of typo: "postional" should be "positional" in comment lines

Important Files Changed

Filename	Overview
physicsnemo/diffusion/multi_diffusion/losses.py	New file implementing multi-diffusion denoising score matching losses with comprehensive docstrings and examples. Contains minor typo in comments.
physicsnemo/diffusion/multi_diffusion/models.py	New file implementing MultiDiffusionModel2D wrapper for patch-based diffusion with excellent documentation and examples.
physicsnemo/diffusion/multi_diffusion/patching.py	Refactored from placeholder to full implementation of 2D image patching utilities (random and grid-based) with comprehensive documentation.
physicsnemo/diffusion/multi_diffusion/init.py	Exports new multi-diffusion components: loss classes, model wrapper, and patching utilities.
test/diffusion/test_patching.py	New comprehensive test suite for patching utilities with reference data for non-regression testing.

_{Last reviewed commit: 2e6b052}

[greptile-apps]

_{30 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[pzharrington]

Nice, just minor comments!

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci