Update to PhysicsNeMo v2

[pzharrington]

Earth2Studio Pull Request

Description

Bumps dependencies and updates imports where needed to be compatible with PNM v2.0. InterpModAFNO, Corrdiff, StormCast and StormScope need updates to the imports/package loading and are thus not compatible with v1.x with these changes; other models are cross-compatible. FCN3 is a special case since makani source code is still reliant on PNM v1.x, so that will not be updated until makani's PNM dependency gets bumped.

Tested all the package load unit tests and examples using PNM models, things look good.

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.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[greptile-apps]

Greptile Summary

This PR systematically migrates Earth2Studio to PhysicsNeMo v2.0 by reorganizing imports, updating dependency constraints, and modernizing checkpoint loading patterns.

• Reorganized all PNM imports from legacy v1 paths (e.g., physicsnemo.models.module, physicsnemo.utils.corrdiff) to new v2 API structure (e.g., physicsnemo, physicsnemo.diffusion.generate.legacy_generate)
• Updated checkpoint loading to use specific model classes (EDMPrecondSR, UNet, StormCastUNet, EDMPrecond) instead of generic Module.from_checkpoint()
• Removed <2.0 version caps for most models and updated to >=2.0 for models requiring v2 (interp-modafno, stormcast, stormscope, corrdiff)
• Added UV conflict rules to prevent mixing physicsnemo v1 and v2 models (fcn3/sfno use v1, others use v2)
• Removed pytz dependencies (no longer needed in PNM v2)
• All exception handlers properly updated to include new imports

Temporary workarounds (tracked with TODOs):

• stormcast.py uses strict=False in checkpoint loading until checkpoints are updated to new diffusion API
• Git dependency for nvidia-physicsnemo until pip wheel is released

Confidence Score: 4/5

• This PR is safe to merge with minimal risk - changes are systematic and well-structured for the v2 migration
• Score reflects thorough and correct migration work with proper version constraints and conflict management. Not 5/5 due to temporary workarounds (strict=False, git dependency) that need follow-up, but these are acknowledged with TODO comments and don't introduce bugs
• Pay attention to stormcast.py - the strict=False parameter should be removed once checkpoints are updated (TODO on line 294)

Important Files Changed

Filename	Overview
pyproject.toml	Updated dependencies to PhysicsNeMo v2, removed version caps and pytz dependencies, added conflict rules to prevent mixing v1/v2 models
earth2studio/models/dx/corrdiff.py	Reorganized imports to new PNM v2 API paths, updated checkpoint loading to use specific model classes (EDMPrecondSR, UNet)
earth2studio/models/dx/corrdiff_cmip6.py	Updated imports to new PNM v2 paths, switched to specific model classes for checkpoint loading
earth2studio/models/px/stormcast.py	Updated to specific model classes (StormCastUNet, EDMPrecond), added strict=False to checkpoint loading with TODO noting temporary workaround

_{Last reviewed commit: afee8bd}

[greptile-apps]

_{9 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[NickGeneva]

LGTM! Hopefully the sfno and fcn3 env can get initialized... I actually dont know if they will be able to install since its forcing a git install. Pipeline will test it so approving.

[pzharrington]

/blossom-ci

[NickGeneva]

/blossom-ci