CompVis/weather-rf

Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

We propose FREUD, a FRame-wise Encoder, United Decoder rectified flow-based first stage for precipitation nowcasting. Weather forecasting requires probabilistic prediction. Our generative decoder allows uncertainty-aware compression. Our FREUD design enables variable-length inputs, robustness to frame drops, and preserves temporal consistency. We enable simple training - no loss weight tuning, only a simple, stable rectified flow objective. A rectified-flow model in FREUD latent space achieves state-of-the-art distributional and perceptual forecasting quality.

Our generative decoder can quantify uncertainty about compression and covers the true precipitation in heavy-rain scenarios, while deterministic decoding collapses to incorrect modes.

Forecasts remain realistic over time and ensemble members capture different plausible outcomes.

Paper and Abstract

The FREUD model was presented in the paper Probabilistic Precipitation Nowcasting with Rectified Flow Transformers, accepted at CVPR 2026.

Abstract:

Accurate weather forecasts are essential across various domains and are safety-critical in extreme weather conditions. Compared to simulation-based forecasting, data-driven approaches show greater efficiency, enabling short-term, high-resolution nowcasting. In particular, diffusion models proved effective in weather nowcasting due to their strong probabilistic foundation. However, existing methods rely on deterministic compression to reduce the complexity of high-dimensional weather data, limiting their ability to capture uncertainty in the decoding process. In this work, we introduce, a Frame-wise Encoder and United Decoder model based on rectified flow transformers for efficient compression of spatio-temporal weather data. Frame-wise encoding enables continuous forecast updates, while the unified video decoder ensures temporal consistency. Our uncertainty-preserving first stage allows us to capture aleatoric uncertainty through ensembling, which is particularly beneficial for extreme weather events with high decoding variability. We achieve state-of-the-art performance in precipitation nowcasting with a compact latent-space rectified flow transformer on the SEVIR benchmark and show further performance gains by model and test-time scaling.

Usage

Please refer to our GitHub repository for model implementations and usage details.

Setup

1. Clone the repository:

git clone https://github.com/CompVis/weather-rf
cd weather-rf

2. Download model weights:

Download the model weights from 🤗 huggingface:

hf download CompVis/weather-rf --include "*.pt" --local-dir ckpts

3. Create a Python environment and install dependencies:

Conda (recommended):

conda create -n weather-rf python=3.12 -y
conda activate weather-rf
python -m pip install --upgrade pip
pip install -r requirements.txt

Virtual environment:

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
pip install -r requirements.txt

Inference

The notebook notebooks/inference.ipynb contains code for obtaining both

• FREUD reconstructions and
• RaMViD latent-space forecasting (LSM)

Open it and update local paths (dataset + checkpoints) in the config cells.

For script-based evaluation, run:

python eval/eval_forecasting.py \
  --model_path checkpoints/lsm.ckpt \
  --sevir_npy_path <SEVIR_NPY_ROOT_PLACEHOLDER> \
  --txt_path data/test_data.txt

or to evaluate reconstruction quality run:

python eval/eval_freud_recon.py \
  --model_path checkpoints/freud.ckpt \
  --sevir_npy_path <SEVIR_NPY_ROOT_PLACEHOLDER> \
  --txt_path data/test_data.txt

⚠️ Original vs. Clean Implementation

Results in the paper were obtained using models trained with torch==2.5.1. Due to changes in the behavior of flex_attention, we found checkpoints obtained with this version are incompatible with newer PyTorch versions and highly sensitive to implementation details.

Therefore, we provide two implementations of our model:

• Clean: In model/ we provide a clean, easy-to-use, and easy-to-understand implementation of our models compatible with newer PyTorch versions. However, results may differ to results reported in the paper.
• Original: In original_model/ we provide code to run the models we trained for the paper. These models have to be run with torch==2.5.1 (see original_requirements). This implementation can be used to reproduce our results, yet might be fragile.

We provide an example of how to use the original implementation in notebooks/original_inference.ipynb. Some slight modification to the eval scripts is necessary to use them with the original models, yet core logic for evaluation is shared across both model versions.

We recommend using the original_model when exact reproduction/comparison is of essence and model when integrating components of our model into different pipelines.

Citation

If you use our work or parts thereof, please cite us accordingly:

@inproceedings{schusterbauer2026probabilisticprecipitation,
    title = {Probabilistic Precipitation Nowcasting with Rectified Flow Transformers},
    author = {Schusterbauer, Johannes and Wiese, Jannik and Stracke, Nick and Phan, Timy and Ommer, Bj{\"o}rn},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
    year = {2026}
}