Improve model card: Add pipeline tag, paper link, GitHub link, and usage details

README.md

@@ -1,16 +1,34 @@
 ---
 license: apache-2.0
+pipeline_tag: image-to-image
 ---
 
-## Checkpoints for [OFTSR](https://github.com/yuanzhi-zhu/OFTSR)
+# OFTSR: One-Step Flow for Image Super-Resolution with Tunable Fidelity-Realism Trade-offs
+
+This repository contains the official implementation and checkpoints for **OFTSR** (One-Step Flow for Image Super-Resolution with Tunable Fidelity-Realism Trade-offs).
+
+OFTSR introduces a novel flow-based framework for one-step image super-resolution, capable of producing outputs with tunable levels of fidelity and realism. It achieves state-of-the-art performance for one-step image super-resolution by distilling a conditional flow-based teacher model, aligning its single-step predictions from initial states to match the teacher's predictions from a closer intermediate state. This approach significantly reduces computational overhead while offering flexibility in the fidelity-realism trade-off.
+
+The model was presented in the paper:
+[**OFTSR: One-Step Flow for Image Super-Resolution with Tunable Fidelity-Realism Trade-offs**](https://huggingface.co/papers/2412.09465)
+
+For the official codebase and more details, visit the [GitHub repository](https://github.com/yuanzhi-zhu/OFTSR).
+
+## Sample Usage
+
+### Checkpoints and Setup
+
+Checkpoints can be found here on Hugging Face: [Yuanzhi/OFTSR](https://huggingface.co/Yuanzhi/OFTSR).
+To download these checkpoints and prepare your environment, follow these steps:
 
 ```bash
 cd OFTSR
 huggingface-cli download Yuanzhi/OFTSR --local-dir ckpts
 ```
-put data at ./val_data
+Please ensure your validation data is placed at `./val_data`.
+
+### FFHQ Generation
 
-########################## FFHQ Generation ##########################
 ```bash
 python sample_fm.py --opt configs/ir_fm_ffhq.yml \
         --overrides \
@@ -21,9 +39,9 @@ python sample_fm.py --opt configs/ir_fm_ffhq.yml \
         fm_model.use_cond=false
 ```
 
-########################## FFHQ Noiseless Restoration ##########################
+### FFHQ Noiseless Restoration
 
-#--------------------- multi-step ---------------------
+**Multi-step:**
 ```bash
 python sample_fm.py --opt configs/ir_fm_ffhq.yml \
         --overrides \
@@ -35,7 +53,7 @@ python sample_fm.py --opt configs/ir_fm_ffhq.yml \
         fm_model.use_cond=true
 ```
 
-#--------------------- one-step ---------------------
+**One-step:**
 ```bash
 python sample_fm.py --opt configs/dis_fm_ffhq.yml \
         --overrides \
@@ -44,13 +62,13 @@ python sample_fm.py --opt configs/dis_fm_ffhq.yml \
         sample.num_sample=100 \
         ir.sigma_y=0. \
         ir.sigma_pertubation=0.1 \
-        fm_model.use_cond=true
+        fm_model.use_cond=true \
         sample.one_step_t=0.99
 ```
 
-########################## FFHQ Noisy Restoration ##########################
+### FFHQ Noisy Restoration
 
-#--------------------- multi-step ---------------------
+**Multi-step:**
 ```bash
 python sample_fm.py --opt configs/ir_fm_ffhq.yml \
         --overrides \
@@ -62,7 +80,7 @@ python sample_fm.py --opt configs/ir_fm_ffhq.yml \
         fm_model.use_cond=true
 ```
 
-#--------------------- one-step ---------------------
+**One-step:**
 ```bash
 python sample_fm.py --opt configs/dis_fm_ffhq.yml \
         --overrides \
@@ -71,13 +89,13 @@ python sample_fm.py --opt configs/dis_fm_ffhq.yml \
         sample.num_sample=100 \
         ir.sigma_y=0.05 \
         ir.sigma_pertubation=0.5 \
-        fm_model.use_cond=true
+        fm_model.use_cond=true \
         sample.one_step_t=0.99
 ```
 
-########################## DIV2K Noiseless Restoration ##########################
+### DIV2K Noiseless Restoration
 
-#--------------------- multi-step ---------------------
+**Multi-step:**
 ```bash
 python sample_fm.py --opt configs/ir_fm_DIV2K.yml \
         --overrides \
@@ -90,7 +108,7 @@ python sample_fm.py --opt configs/ir_fm_DIV2K.yml \
         fm_model.use_cond=true
 ```
 
-#--------------------- one-step ---------------------
+**One-step:**
 ```bash
 python sample_fm.py --opt configs/dis_fm_DIV2K.yml \
         --overrides \
@@ -100,15 +118,13 @@ python sample_fm.py --opt configs/dis_fm_DIV2K.yml \
         sample.psnr_batch_size=1 \
         ir.sigma_y=0. \
         ir.sigma_pertubation=0.2 \
-        fm_model.use_cond=true
+        fm_model.use_cond=true \
         sample.one_step_t=0.99
 ```
 
+### ImageNet Noiseless Restoration
 
-
-########################## ImageNet Noiseless Restoration ##########################
-
-#--------------------- multi-step ---------------------
+**Multi-step:**
 ```bash
 python sample_fm.py --opt configs/ir_fm_imagenet.yml \
         --overrides \
@@ -120,7 +136,7 @@ python sample_fm.py --opt configs/ir_fm_imagenet.yml \
         fm_model.use_cond=true
 ```
 
-#--------------------- one-step ---------------------
+**One-step:**
 ```bash
 python sample_fm.py --opt configs/dis_fm_imagenet.yml \
         --overrides \
@@ -129,6 +145,28 @@ python sample_fm.py --opt configs/dis_fm_imagenet.yml \
         sample.num_sample=100 \
         ir.sigma_y=0. \
         ir.sigma_pertubation=0.2 \
-        fm_model.use_cond=true
+        fm_model.use_cond=true \
         sample.one_step_t=0.99
+```
+
+## Evaluation
+
+To evaluate the model, use the following commands:
+
+```bash
+python evaluate.py --fdir1 hr_folder --fdir2 sr_folder
+python fid.py --fdir1 hr_folder --fdir2 sr_folder
+```
+
+## Citation
+
+If you find this repository helpful, please cite the paper:
+
+```bibtex
+@article{zhu2024oftsr,
+  title={OFTSR: One-Step Flow for Image Super-Resolution with Tunable Fidelity-Realism Trade-offs},
+  author={Zhu, Yuanzhi and Wang, Ruiqing and Lu, Shilin and Li, Junnan and Yan, Hanshu and Zhang, Kai},
+  journal={arXiv preprint arXiv:2412.09465},
+  year={2024}
+}
 ```