Update README.md

README.md

@@ -9,7 +9,10 @@ tags:
 
 # Probabilistic Precipitation Nowcasting with Rectified Flow Transformers
 
-**TODO: badges**
+[![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://compvis.github.io/weather-rf/)
+[![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**.
@@ -17,11 +20,118 @@ Our FREUD design enables variable-length inputs, robustness to frame drops, and
 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.
 
-![Reconstruction distributions for different precipitation levels](./docs/teaser_figure_weather.svg)
+<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/teaser_figure_weather.png" alt="Reconstruction distributions for different precipitation levels">
+  <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 --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}
+}
+```