mudhil/era5-conus404-diffusion-downscaling

ERA5→CONUS404 Diffusion Downscaling — Current Checkpoints

A three-stage "Latent CorrDiff" pipeline for downscaling ERA5 (0.25°, ~27km) to CONUS404 (4km) over the contiguous United States:

1. DRN (Deterministic Regression Network) — predicts the conditional mean E[CONUS404 | ERA5] via a 4-level UNet (49.6M params).
2. VAE — compresses the DRN residual (CONUS404 − DRN prediction) from 256×256 to a 64×64×8 latent space (39.7M params).
3. Diffusion UNet — an EDM-style (Karras et al. 2022) conditional diffusion model that generates the latent residual, decoded back through the VAE and added to the DRN mean for the final prediction (142.4M params).

Trained on 6 variables: t2m→T2, d2m→TD2, u10→U10, v10→V10, sp→PSFC, tp→PREC_ACC_NC (2m temperature, 2m dewpoint, 10m winds, surface pressure, precipitation), plus 6 static fields (terrain, orographic variance, lat, lon, LAI, land-sea mask). Trained on CONUS404 years 1980–2014, validated 2015–2017, tested 2018–2020.

Files

File	Size	Description
diffusion_best.pt	2.3 GB	Diffusion UNet, best validation checkpoint
diffusion_latest.pt	2.3 GB	Diffusion UNet, latest checkpoint
drn_best.pt	596 MB	DRN, best validation checkpoint
vae_best.pt	476 MB	VAE, best validation checkpoint

Usage

git clone https://github.com/milhud/diffusion_downscaling_model
cd diffusion_downscaling_model
hf download mudhil/era5-conus404-diffusion-downscaling --local-dir checkpoints/
python -m src.inference.sample_nc --input era5_input.nc --output downscaled.nc

See the repo's docs/ARCHITECTURE.md and docs/PUBLISHABLE_POINTS.md (How to Run section) for training/inference details.