Fixes 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 fixes several device-compatibility and DDP/torch.compile compatibility issues in the multi-diffusion training path. Key changes: reset_patch_indices and _CompiledPatchX now operate on the unwrapped MultiDiffusionModel2D (peeling off DDP/compiled wrappers via the new _unwrap_multi_diffusion helper); _compute_global_index propagates device from patch_indices to avoid cross-device errors; random.randint is replaced by torch.randint for GPU-native patch sampling with optional generator support; and RandomPatching2D.forward is rewritten from unfold+advanced-indexing to a torch.gather-based approach that is more compatible with torch.compile. All remaining findings are P2 suggestions.

Important Files Changed

Filename	Overview
physicsnemo/diffusion/multi_diffusion/losses.py	Adds _unwrap_multi_diffusion helper to peel DDP/compiled wrappers, and uses the unwrapped model for reset_patch_indices and _CompiledPatchX; constructor type annotations still declare MultiDiffusionModel2D but the logic now accepts wrapped variants.
physicsnemo/diffusion/multi_diffusion/patching.py	Replaces random.randint with torch.randint (GPU-compatible, supports PRNG generator), propagates device from patch_indices in _compute_global_index, defers _global_index recomputation lazily via a plain Python flag, and rewrites forward from unfold-based to torch.gather-based indexing.
physicsnemo/diffusion/utils/model_wrappers.py	Trivial: adds a blank line after the license header. No functional change.

Comments Outside Diff (2)

1. physicsnemo/diffusion/multi_diffusion/losses.py, line 228-229 (link)

   Type annotation doesn't match new behaviour

   The model parameter is still typed as MultiDiffusionModel2D, but _unwrap_multi_diffusion was introduced precisely because this constructor should now also accept DistributedDataParallel and torch.compile-wrapped models. Passing a DDP-wrapped model currently satisfies the runtime logic but violates the declared type, which will mislead type checkers and users reading the docstring.

   The same applies to the docstring Parameters section (model : MultiDiffusionModel2D) — updating it to model : torch.nn.Module or model : MultiDiffusionModel2D | torch.nn.Module would keep the contract accurate.

2. physicsnemo/diffusion/multi_diffusion/losses.py, line 451-452 (link)

   Type annotation doesn't match new behaviour

   Same annotation mismatch as in MultiDiffusionMSEDSMLoss: the parameter is typed MultiDiffusionModel2D but the class now unwraps DDP/compiled wrappers. Consider updating to torch.nn.Module.

_{Reviews (1): Last reviewed commit: "Fixes for multi-diffusion" | Re-trigger Greptile}

[pzharrington]

LGTM!

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blosson-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci