Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,155 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Downscaling
+- Diffusion Model
+- ERA5
+- CorrDiff
+- arxiv:2309.15214
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">CorrDiff</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+CorrDiff (Residual Corrective Diffusion Model) is a two-stage generative model for kilometer-scale regional weather downscaling. It downscales approximately 25 km ERA5 fields to approximately 2 km regional weather fields over Taiwan and synthesizes radar reflectivity that is not present in the input.
+
+Paper: Residual Corrective Diffusion Modeling for Km-scale Atmospheric Downscaling
+
+https://arxiv.org/abs/2309.15214
+
+# Model Description
+
+CorrDiff was proposed by NVIDIA and its collaborators for kilometer-scale regional weather downscaling and was trained with coarse-resolution ERA5 reanalysis and high-resolution WRF regional model data from Taiwan's Central Weather Administration.
+The model is suitable for converting coarse-resolution global weather fields into high-resolution regional weather fields and producing probabilistic weather predictions.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather forecast training | Train CorrDiff with ERA5 HDF5 data. |
+| Local quick validation | Use synthetic data to check data loading, model training, inference, and result visualization. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch multi-process training with `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/CorrDiff --local-dir ./CorrDiff
+cd CorrDiff
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The OneScience community provides an ERA5 data slice that can be downloaded as follows:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+When real data is unavailable, generate synthetic data for pipeline validation:
+
+```bash
+python scripts/fake_data.py --output data/era5_corrdiff.npz
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training outputs:
+
+```text
+data/checkpoints/model_bak.pth
+data/checkpoints/trloss.npy
+data/checkpoints/valoss.npy
+```
+
+### Training Weights
+
+This repository provides weights trained on 39 years of ERA5 reanalysis data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are written to `result/output/` by default.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- This repository is a reproduction of the original CorrDiff paper.

conf/config.yaml

@@ -0,0 +1,21 @@
+# Lightweight OneScience configuration for the official CorrDiff example.
+model:
+  regression_config: conf/config_training_taiwan_regression.yaml
+  diffusion_config: conf/config_training_taiwan_diffusion.yaml
+  generation_config: conf/config_generate_taiwan.yaml
+  input_channels: 12
+  output_channels: 4
+  input_shape: [36, 36]
+  output_shape: [448, 448]
+  ensemble_size: 32
+
+data:
+  type: cwb_zarr
+  path: ./data/era5_corrdiff.npz
+  input_variables: [tcwv, t2m, u10m, v10m, t500, z500, u500, v500, t850, z850, u850, v850]
+  output_variables: [t2m, u10m, v10m, maximum_radar_reflectivity]
+
+runtime:
+  device: auto
+  output_dir: ./result/output
+  seed: 42

config.json

@@ -0,0 +1,53 @@
+{
+  "model_name": "CorrDiff",
+  "model_type": "corrdiff",
+  "architectures": [
+    "CorrDiff"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "regional-weather-downscaling",
+  "implementation": {
+    "entry_point": "model/corrdiff.py",
+    "scope": "compact two-stage reproduction of the CorrDiff regression and residual diffusion architecture"
+  },
+  "architecture": {
+    "family": "residual corrective diffusion model",
+    "input_grid_shape": [
+      36,
+      36
+    ],
+    "output_grid_shape": [
+      448,
+      448
+    ],
+    "input_channels": 12,
+    "output_channels": 4,
+    "base_channels": 32,
+    "mean_model": "RegressionUNet",
+    "residual_model": "ResidualDiffusionUNet",
+    "conditioning": [
+      "bilinearly upsampled coarse input",
+      "high-resolution regression mean",
+      "noise level sigma"
+    ],
+    "normalization": "group_norm",
+    "activation": "silu",
+    "sampling_steps": 4,
+    "sigma_max": 1.0,
+    "sigma_min": 0.01
+  },
+  "data": {
+    "dataset": "ERA5 and CWA-WRF",
+    "input_spatial_resolution_km": 25,
+    "output_spatial_resolution_km": 2,
+    "time_step_hours": 1,
+    "input_steps": 1,
+    "output_steps": 1,
+    "protocol": "synthetic_era5_corrdiff"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/corrdiff.py"
+  ]
+}

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"weather_downscaling"}

model/corrdiff.py

@@ -0,0 +1,102 @@
+"""Compact CorrDiff implementation for the OneScience model package."""
+
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1),
+            nn.GroupNorm(8, out_channels),
+            nn.SiLU(),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1),
+            nn.GroupNorm(8, out_channels),
+            nn.SiLU(),
+        )
+        self.skip = nn.Conv2d(in_channels, out_channels, 1)
+
+    def forward(self, x):
+        return self.block(x) + self.skip(x)
+
+
+class RegressionUNet(nn.Module):
+    """Mean predictor: 12x36x36 coarse fields to 4x448x448 fields."""
+
+    def __init__(self, in_channels=12, out_channels=4, base_channels=32):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            ConvBlock(in_channels, base_channels),
+            ConvBlock(base_channels, base_channels * 2),
+            ConvBlock(base_channels * 2, base_channels * 4),
+        )
+        self.head = nn.Sequential(
+            ConvBlock(base_channels * 4, base_channels * 2),
+            nn.Conv2d(base_channels * 2, out_channels, 1),
+        )
+
+    def forward(self, x):
+        x = self.encoder[0](x)
+        x = F.avg_pool2d(x, 2)
+        x = self.encoder[1](x)
+        x = F.avg_pool2d(x, 2)
+        x = self.encoder[2](x)
+        x = self.head[0](x)
+        x = F.interpolate(x, size=(448, 448), mode="bilinear", align_corners=False)
+        return self.head[1](x)
+
+
+class ResidualDiffusionUNet(nn.Module):
+    """Conditional residual denoiser used by the compact EDM-style sampler."""
+
+    def __init__(self, condition_channels=16, out_channels=4, base_channels=32):
+        super().__init__()
+        self.time = nn.Sequential(
+            nn.Linear(1, base_channels), nn.SiLU(), nn.Linear(base_channels, base_channels)
+        )
+        self.net = nn.Sequential(
+            ConvBlock(condition_channels, base_channels),
+            ConvBlock(base_channels, base_channels * 2),
+            ConvBlock(base_channels * 2, base_channels),
+            nn.Conv2d(base_channels, out_channels, 1),
+        )
+
+    def forward(self, noisy_residual, condition, sigma):
+        x = torch.cat((noisy_residual, condition), dim=1)
+        embedding = self.time(torch.log(sigma).reshape(-1, 1).float())
+        x = self.net[0](x)
+        x = x + embedding[:, :, None, None]
+        for layer in self.net[1:]:
+            x = layer(x)
+        return x
+
+
+class CorrDiff(nn.Module):
+    """Two-stage mean plus stochastic residual atmospheric downscaler."""
+
+    def __init__(self, in_channels=12, out_channels=4, base_channels=32):
+        super().__init__()
+        self.regression = RegressionUNet(in_channels, out_channels, base_channels)
+        self.diffusion = ResidualDiffusionUNet(in_channels + 2 * out_channels, out_channels, base_channels)
+
+    def mean(self, coarse):
+        return self.regression(coarse)
+
+    def denoise(self, noisy_residual, coarse, mean, sigma):
+        coarse = F.interpolate(coarse, size=mean.shape[-2:], mode="bilinear", align_corners=False)
+        return self.diffusion(noisy_residual, torch.cat((coarse, mean), dim=1), sigma)
+
+    def forward(self, coarse, steps=4):
+        mean = self.mean(coarse)
+        residual = torch.randn_like(mean)
+        schedule = torch.linspace(1.0, 0.01, steps, device=coarse.device)
+        for sigma in schedule:
+            sigma_batch = sigma.expand(coarse.shape[0])
+            residual = residual - 0.2 * self.denoise(residual, coarse, mean, sigma_batch)
+        return mean + residual
+
+
+__all__ = ["CorrDiff", "RegressionUNet", "ResidualDiffusionUNet"]

scripts/fake_data.py

@@ -0,0 +1,23 @@
+from pathlib import Path
+import argparse
+import numpy as np
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output", default="data/fake_corrdiff.npz")
+    parser.add_argument("--samples", type=int, default=1)
+    args = parser.parse_args()
+    path = Path(args.output)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    rng = np.random.default_rng(42)
+    inputs = rng.standard_normal((args.samples, 12, 36, 36)).astype("float32")
+    targets = rng.standard_normal((args.samples, 4, 448, 448)).astype("float32")
+    np.savez_compressed(path, input=inputs, target=targets)
+    print(f"input: {inputs.shape}")
+    print(f"target: {targets.shape}")
+    print(f"saved: {path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,37 @@
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.corrdiff import CorrDiff
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run CorrDiff inference")
+    parser.add_argument("--data", default="data/era5_corrdiff.npz")
+    parser.add_argument("--checkpoint", default="data/checkpoints/model_bak.pth")
+    parser.add_argument("--ensemble-size", type=int, default=1)
+    parser.add_argument("--output", default="result/output/predictions.npz")
+    args = parser.parse_args()
+    data = np.load(args.data)
+    coarse = torch.from_numpy(data["input"])
+    model = CorrDiff()
+    checkpoint = Path(args.checkpoint)
+    if checkpoint.exists():
+        model.load_state_dict(torch.load(checkpoint, map_location="cpu")["model"])
+    model.eval()
+    samples = []
+    with torch.no_grad():
+        for _ in range(args.ensemble_size):
+            samples.append(model(coarse).numpy().astype("float32"))
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, prediction=np.stack(samples), input=data["input"])
+    print(f"prediction: {np.stack(samples).shape}")
+    print(f"saved: {output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,21 @@
+import argparse
+from pathlib import Path
+import numpy as np
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate CorrDiff NPZ predictions")
+    parser.add_argument("--prediction", default="result/output/predictions.npz")
+    parser.add_argument("--target", default="data/era5_corrdiff.npz")
+    args = parser.parse_args()
+    prediction = np.load(args.prediction)["prediction"]
+    target = np.load(args.target)["target"]
+    mae = np.abs(prediction.mean(axis=0) - target).mean()
+    output = Path(args.prediction).with_name("metrics.npz")
+    np.savez(output, mae=np.array(mae, dtype="float32"))
+    print(f"ensemble={prediction.shape[0]} mae={mae:.6f}")
+    print(f"saved: {output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,44 @@
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.corrdiff import CorrDiff
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train compact CorrDiff on NPZ data")
+    parser.add_argument("--data", default="data/era5_corrdiff.npz")
+    parser.add_argument("--steps", type=int, default=2)
+    parser.add_argument("--output", default="data/checkpoints/model_bak.pth")
+    args = parser.parse_args()
+    data = np.load(args.data)
+    coarse = torch.from_numpy(data["input"])
+    target = torch.from_numpy(data["target"])
+    model = CorrDiff()
+    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
+    model.train()
+    for step in range(args.steps):
+        mean = model.mean(coarse)
+        residual = target - mean
+        sigma = torch.rand(coarse.shape[0]).clamp_min(0.01)
+        noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual)
+        predicted = model.denoise(noisy, coarse, mean.detach(), sigma)
+        loss = F.mse_loss(mean, target) + F.mse_loss(predicted, residual)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        print(f"step={step + 1} loss={loss.item():.6f}")
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({"model": model.state_dict(), "format": "corrdiff-compact-v1"}, output)
+    np.save(output.parent / "trloss.npy", np.asarray([loss.item()], dtype="float32"))
+    np.save(output.parent / "valoss.npy", np.asarray([loss.item()], dtype="float32"))
+    print(f"saved: {output}")
+
+
+if __name__ == "__main__":
+    main()