Implements Predictor specialization for 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 introduces MultiDiffusionPredictor, a Predictor-protocol-compatible wrapper for test-time patch-based diffusion sampling using MultiDiffusionModel2D. It pre-patches the condition and positional embedding once at construction time to avoid redundant work per step, refactors the duplicate PE injection logic into a shared _inject_patched_pos_embd helper, and generalises the DDP/compile-unwrapping helper into a reusable _unwrap_module utility. A thorough test suite covering constructor validation, non-regression, gradient flow, torch.compile, and checkpoint round-trips is included.

• P1 (predictor.py line 406): _skip_positional_embedding_injection = True is permanently written to the underlying MultiDiffusionModel2D and never restored. If the same model object is used for loss computation after the predictor is created (e.g., continued fine-tuning), PE injection is silently skipped, producing wrong training outputs without any error.

Important Files Changed

Filename	Overview
physicsnemo/diffusion/multi_diffusion/predictor.py	New file implementing MultiDiffusionPredictor — a Predictor-protocol wrapper for patch-based sampling. Correctly pre-patches condition and PE at construction, but permanently mutates the underlying model's _skip_positional_embedding_injection flag with no restore mechanism (P1), and the dead is_compiling() guard in init is misleading.
physicsnemo/diffusion/multi_diffusion/models.py	Adds _skip_positional_embedding_injection flag and refactors PE injection into a shared _inject_patched_pos_embd helper. The refactoring is clean and the logic is preserved correctly.
physicsnemo/diffusion/utils/utils.py	Extracts the DDP/torch.compile unwrapping logic into a reusable generic _unwrap_module function. Clean generalization with correct TypeVar usage.
physicsnemo/diffusion/multi_diffusion/losses.py	Replaces the local _unwrap_multi_diffusion helper with the new generic _unwrap_module utility. Equivalent behavior, no logic changes.
physicsnemo/diffusion/guidance/dps_guidance.py	One-line change: DPSScorePredictor now explicitly inherits from Predictor, correctly aligning it with the protocol.
physicsnemo/diffusion/multi_diffusion/init.py	Adds MultiDiffusionPredictor to the public module exports.
test/diffusion/test_multi_diffusion_predictor.py	Comprehensive test suite for MultiDiffusionPredictor covering constructor validation, non-regression outputs, gradient flow, torch.compile compatibility, and checkpoint round-trips.
test/diffusion/test_multi_diffusion_sampling.py	End-to-end sampling tests for MultiDiffusionPredictor integrated with all three noise schedulers (EDM, VE, VP), two solvers, and compiled predictor paths.

Comments Outside Diff (1)

1. physicsnemo/diffusion/multi_diffusion/predictor.py, line 405-406 (link)

   Persistent flag side-effect on underlying model

   _skip_positional_embedding_injection = True is written to _md_model and never restored. If the same MultiDiffusionModel2D is used for loss computation after the predictor is created — for example, after sampling a checkpoint to validate a mid-training result or during continued fine-tuning — the model's forward will silently skip PE injection in both the "no patching" and "with patching" paths, producing wrong outputs from the loss without any error or warning.

   A minimal safeguard would be to save and restore the original flag, or at minimum include a runtime warning that the model is permanently modified:

   # Save original value so callers can restore it
   self._prev_skip_pe = self._md_model._skip_positional_embedding_injection
   self._md_model._skip_positional_embedding_injection = True

   Or add a close() / __del__ to restore the flag, or document that the model must not be used for training after predictor construction.

_{Reviews (1): Last reviewed commit: "Implements Predictor specialization for ..." | Re-trigger Greptile}

[CharlelieLrt]

/blossom-ci