README.md

---
license: cc-by-nc-sa-4.0
pipeline_tag: video-to-video
tags:
- weather
- precipitation
- nowcasting
---

# Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

[![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://compvis.github.io/weather-rf/)
[![arXiv](https://img.shields.io/badge/arXiv-PDF-b31b1b)](https://arxiv.org/abs/2605.31204)
[![Paper PDF](https://img.shields.io/badge/Paper-PDF-orange)](https://openaccess.thecvf.com//content/CVPR2026/papers/Schusterbauer_Probabilistic_Precipitation_Nowcasting_with_Rectified_Flow_Transformers_CVPR_2026_paper.pdf)
[![GitHub](https://img.shields.io/badge/GitHub-Code-black)](https://github.com/CompVis/weather-rf)


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.

<p align="center">
  <img src="./docs/teaser_figure_weather.svg" alt="Reconstruction distributions for different precipitation levels" width=700px>
</p>

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

<p align="center">
  <img src="./docs/qual-forecast.svg" alt="Forecasts with zoom-ins" width=700px>
</p>

*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](https://openaccess.thecvf.com//content/CVPR2026/papers/Schusterbauer_Probabilistic_Precipitation_Nowcasting_with_Rectified_Flow_Transformers_CVPR_2026_paper.pdf), accepted at CVPR 2026.

### Abstract:

<p><em>
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 <strong>Fr</strong>ame-wise <strong>E</strong>ncoder and <strong>U</strong>nited <strong>D</strong>ecoder 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.
</em></p>

## Usage

Please refer to [our GitHub repository](https://github.com/CompVis/weather-rf) for model implementations and usage details.

### Setup

1. Clone the repository:

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

2. Download model weights:

Download the model weights from 🤗 huggingface:

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

3. Create a Python environment and install dependencies:

Conda (recommended):

```bash
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:

```bash
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](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:

```bash
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:

```bash
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: 

```bibtex
@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}
}
```