Update pipeline tag and improve model card metadata

README.md

@@ -1,6 +1,6 @@
 ---
 license: cc-by-nc-sa-4.0
-pipeline_tag: video-to-video
+pipeline_tag: image-to-video
 tags:
 - weather
 - precipitation
@@ -9,39 +9,34 @@ tags:
 
 # 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)
-[![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.
+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="./docs/teaser_figure_weather.svg" alt="Reconstruction distributions for different precipitation levels" width=700px>
+  <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="./docs/qual-forecast.svg" alt="Forecasts with zoom-ins" width=700px>
+  <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://openaccess.thecvf.com//content/CVPR2026/papers/Schusterbauer_Probabilistic_Precipitation_Nowcasting_with_Rectified_Flow_Transformers_CVPR_2026_paper.pdf), accepted at CVPR 2026.
-
-### 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.
 
-<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>
+### 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
 
@@ -50,51 +45,26 @@ Please refer to [our GitHub repository](https://github.com/CompVis/weather-rf) f
 ### 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.
+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 \
@@ -102,37 +72,19 @@ python eval/eval_forecasting.py \
   --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.
+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
 
-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},
+    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}
 }
-```
+```