nvidia/corrdiff

Pull and run nvidia/corrdiff using Docker (this will download the full model and run it in your local environment)

Get API KeyCopy Code
$ docker login nvcr.io
Username: $oauthtoken
Password: <PASTE_API_KEY_HERE>

Pull and run the NVIDIA Earth-2 CorrDiff NIM with the command below.

Copy Code
docker pull nvcr.io/nim/nvidia/corrdiff:1.0.0

This will download the optimized model for your infrastructure.

Copy Code
export NGC_API_KEY=<NGC API Key>

docker run --rm --runtime=nvidia --gpus all --shm-size 4g \
    -p 8000:8000 \
    -e NGC_API_KEY \
    nvcr.io/nim/nvidia/corrdiff:1.0.0

Check the health of the NIM with the following curl command:

Copy Code
curl -X 'GET' \
    'http://localhost:8000/v1/health/ready' \
    -H 'accept: application/json'

Generate an input numpy array for the model using the following Python script with Earth2Studio. See the NVIDIA NIM Docs quick start guide for more details.

Copy Code
from datetime import datetime, timedelta

import numpy as np
import torch
from earth2studio.data import GEFS_FX, GEFS_FX_721x1440

GEFS_SELECT_VARIABLES = ["u10m","v10m","t2m","r2m","sp","msl","tcwv"]
GEFS_VARIABLES = ["u1000","u925","u850","u700","u500","u250","v1000","v925","v850",\
    "v700","v500","v250","z1000","z925","z850","z700","z500","z200","t1000","t925",\
    "t850","t700","t500","t100","r1000","r925","r850","r700","r500","r100"]

ds_gefs = GEFS_FX(cache=True)
ds_gefs_select = GEFS_FX_721x1440(cache=True, product="gec00")

def fetch_input_gefs(
    time: datetime, lead_time: timedelta, content_dtype: str = "float32"
):
    dtype = np.dtype(getattr(np, content_dtype))
    # Fetch high-res select GEFS input data
    select_data = ds_gefs_select(time, lead_time, GEFS_SELECT_VARIABLES)
    select_data = select_data.values
    # Crop to bounding box [225, 21, 300, 53]
    select_data = select_data[:, 0, :, 148:277, 900:1201].astype(dtype)
    assert select_data.shape == (1, len(GEFS_SELECT_VARIABLES), 129, 301)

    # Fetch GEFS input data
    pressure_data = ds_gefs(time, lead_time, GEFS_VARIABLES)
    # Interpolate to 0.25 grid
    pressure_data = torch.nn.functional.interpolate(
        torch.Tensor(pressure_data.values),
        (len(GEFS_VARIABLES), 721, 1440),
        mode="nearest",
    )
    pressure_data = pressure_data.numpy()
    # Crop to bounding box [225, 21, 300, 53]
    pressure_data = pressure_data[:, 0, :, 148:277, 900:1201].astype(dtype)
    assert pressure_data.shape == (1, len(GEFS_VARIABLES), 129, 301)

    # Create lead time field
    lead_hour = int(lead_time.total_seconds() // (3 * 60 * 60)) * np.ones(
        (1, 1, 129, 301)
    ).astype(dtype)

    input_data = np.concatenate([select_data, pressure_data, lead_hour], axis=1)[None]
    return input_data 


input_array = fetch_input_gefs(datetime(2023, 1, 1), timedelta(hours=0))
np.save("corrdiff_inputs.npy", input_array)

You can now make a local API call using this curl command:

Copy Code
curl -X POST \
    -F "input_array=@corrdiff_inputs.npy" \
    -F "samples=2" \
    -F "steps=12" \
    -o output.tar \
    http://localhost:8000/v1/infer

For more details on getting started with this NIM, visit the NVIDIA NIM Docs. For more details on the model and its input / output tensors see the US CorrDiff Model Card.