README.md

metadata

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.

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

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 design enables variable-length inputs, robustness to frame drops, and preserves temporal consistency.

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

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

3. Create a Python environment and install dependencies:

pip install -r requirements.txt

Inference

The notebook notebooks/inference.ipynb contains code for obtaining both FREUD reconstructions and RaMViD latent-space forecasting (LSM).

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

⚠️ 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

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