CorrDiff Inference Package

This package provides the pre-trained CorrDiff checkpoints for Taiwan, as well as a small dataset sample to easily explore the inference capabilities. This checkpoint was trained on ERA5 → WRF data that spans 2018-2021 at one hour time resolution.

Getting started

Follow these steps to set up all needed to run the CorrDiff inference:

1. Clone the NVIDIA Modulus package, check out a specific hash, and install it:

git clone https://github.com/NVIDIA/modulus
cd modulus
git checkout ee6b3ac0d4c32bca1c81f401d2bfd461887eb9a3
pip install .

2. Get this package using direct download or via wget:

wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/modulus/corrdiff_inference_package/versions/1/zip -O corrdiff_inference_package_1.zip

3. Unzip the contents in the same directory:

unzip corrdiff_inference_package_1.zip
unzip corrdiff_inference_package.zip

4. Assuming you have both the Modulus and the unzipped CorrDiff inference packages in the <root_dir> directory, go to the directory that includes the CorrDiff inference script:

cd <root_dir>/modulus/examples/generative/corrdiff

5. Run the inference:

python generate.py dataset.data_path=<root_dir>/corrdiff_inference_package/dataset/2023-01-24-cwb-4years_5times.zarr res_ckpt_filename=<root_dir>/corrdiff_inference_package/checkpoints/diffusion.mdlus reg_ckpt_filename=<root_dir>/corrdiff_inference_package/checkpoints/regression.mdlus seed_batch_size=5 use_torch_compile=false

License

The checkpoints are distributed under the Apache 2.0 license, and the sample dataset is distributed under the CC BY-NC-ND 4.0 license.