README.md

July 12, 2026 · View on GitHub

SATB-VR: Training Few-Step Video Restoration Diffusion Model using SNR-Aware Trajectory Blending

1Alibaba Group   2Harbin Institute of Technology   3Xi'an Jiaotong University
* Equal contribution    † Corresponding author

Paper | Project Page

For more video visualizations, visit our [project page].

🔥 Update

  • [2026.07.06] Inference code is released.
  • [2026.06.27] This repo is created.
  • [2026.06.27] Paper is released at [Arxiv].

🎬 Overview

overall_structure

🔧 Dependencies and Installation

  1. Clone Repo

    git clone https://github.com/chenxx89/SATB-VR.git
cd SATB-VR

  2. Create Conda Environment and Install Dependencies

    # create new conda env
conda create -n SATB-VR python=3.10
conda activate SATB-VR

# install pytorch
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu118

# install python dependencies
pip install -r requirements.txt

  3. Download Models

    # download ckpts from huggingface
python SATB_VR/download_weights.py

    If you prefer manual download, the models are available at:

    The ./ckpts directory structure should be arranged as:

    ├── ckpts
│   ├── CogVideoX1.5-5B
│   │   ├── ...
│   ├── cogvlm2-llama3-caption
│   │   ├── ...
│   ├── SATB-VR
│   │   ├── controlnet
│   │       ├── config.json
│   │       ├── diffusion_pytorch_model.safetensors
│   │   ├── lora_predictor
│   │       ├── pytorch_lora_weights.safetensors
│   │   ├── lora_transformer
│   │       ├── pytorch_lora_weights.safetensors
│   │   ├── connectors.pt
│   │   ├── control_patch_embed.pt
│   │   ├── negative_prompt_embeds.pt
│   │   ├── prompt_embeds.pt

☕️ Quick Inference

Run the following commands to try it out:

python inference.py  \
    --ckpt_path=./ckpts \
    --input_dir=/dir/to/input/videos \
    --output_dir=/dir/to/output/videos \
    --enable_text_encoder \
    --enable_captioner \
    --enable_spatial_tiling \
    --enable_temporal_tiling    \
    --upscale=0 \
    --save_images
  • --enable_text_encoder: Optional, if given, the text encoder will be used
  • --enable_captioner: Optional, if given, the captioner will be used
  • --enable_spatial_tiling: Optional, if given, the spatial tiling will be used
  • --enable_temporal_tiling: Optional, if given, the temporal tiling will be used
  • --upscale: Optional, if set to 0, the short-size of output videos will be 1024
  • --save_images: Optional, if given, the video frames will be saved

GPU memory usage:

  • For a 121-frame video, it requires approximately 40GB GPU memory.
  • If you want to reduce GPU memory usage, replace pipe.enable_model_cpu_offload with pipe.enable_sequential_cpu_offload in inference.py. GPU memory usage is reduced to 25GB, but the inference time is longer.
  • Use --enable_spatial_tiling and --enable_temporal_tiling to process long videos with less GPU memory at the cost of longer inference time.

📧 Citation

If you find our repo useful for your research, please consider citing it:

 @article{bai2026satb,
 title={SATB-VR: Training Few-Step Video Restoration Diffusion Model using SNR-Aware Trajectory Blending},
 author={Bai, Haoran and Chen, Xiaoxu and Liu, Xiaoyu and Yue, Zongsheng and Deng, Sibin and Zuo, Wangmeng and Chen, Ying},
 journal={arXiv preprint arXiv:2606.28677},
 year={2026}
 }

📄 License