README.md

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

# Probabilistic Precipitation Nowcasting with Rectified Flow Transformers

This repository contains the weights for **FREUD**, as introduced in the paper [Probabilistic Precipitation Nowcasting with Rectified Flow Transformers](https://huggingface.co/papers/2605.31204).

**Authors**: Johannes Schusterbauer, Jannik Wiese, Nick Stracke, Timy Phan, Björn Ommer.

[![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)
[![GitHub](https://img.shields.io/badge/GitHub-Code-black)](https://github.com/CompVis/weather-rf)

We propose FREUD, a **Fr**ame-wise **E**ncoder, **U**nited **D**ecoder rectified flow-based first stage for precipitation nowcasting. Weather forecasting requires probabilistic prediction; our generative decoder allows **uncertainty-aware compression**. Our design enables variable-length inputs, robustness to frame drops, and preserves temporal consistency.

<p align="center">
  <img src="https://huggingface.co/CompVis/weather-rf/resolve/main/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="https://huggingface.co/CompVis/weather-rf/resolve/main/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://huggingface.co/papers/2605.31204), accepted at CVPR 2026.

### Abstract Summary:
In this work, we introduce FREUD, a model based on rectified flow transformers for efficient compression of spatio-temporal weather data. Our uncertainty-preserving first stage allows us to capture aleatoric uncertainty via 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.

## 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:
```bash
hf download CompVis/weather-rf --include "*.pt" --local-dir ckpts
```

3. Create a Python environment and install dependencies:
```bash
pip install -r requirements.txt
```

### Inference

The notebook [notebooks/inference.ipynb](https://github.com/CompVis/weather-rf/blob/main/notebooks/inference.ipynb) contains code for obtaining both FREUD reconstructions and RaMViD latent-space forecasting (LSM).

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

### ⚠️ Original vs. Clean Implementation

Results in the paper were obtained using models trained with `torch==2.5.1`. 
- **Clean**: In `model/` we provide a clean, easy-to-use implementation compatible with newer PyTorch versions.
- **Original**: In `original_model/` we provide code to run the models we trained for the paper (requires `torch==2.5.1`).

## Citation

```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{\"}orn},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
    year = {2026}
}
```