🔥 Update

• [2025.04] Training code is released.
• [2025.03] Inference code and checkpoint is released.

📈 Our model can do both ISR and VSR. Hope you can enjoy it.

Realistic Image SR

Realistic Video SR

🔧 Dependencies and Installation

1. Clone Repo

   git clone https://github.com/ssj9596/SCST.git
   cd SCST
   

2. Create Conda Environment and Install Dependencies

   # create new conda env
   conda create -n SCST python=3.10
   conda activate SCST
   
   # install python dependencies
   pip install -r requirements.txt
   

3. Download Models

   • Download pretrained models from huggingface and put them under the checkpoints folder.
   • Download SD2.1 and put them into checkpoints/stable-diffusion-2-1-base.

   The checkpoints directory structure should be arranged as:

   ├── checkpoints
   │   ├── controlnet
   │   ├── stable-diffusion-2-1-base
   │   ├── localatten_unet.pth
   │   ├── mococtrl_unet.pth
   │   ├── stcm_unet.pth
   

☕️ Quick Inference

We provide several examples in the inputs folder. Run the following commands to try it out:

## Single Image 
## no temporal module
python inference/infer_mococtrl.py 

## Video
## use LocalAttention as temporal module
python inference/infer_localatten.py
## use Mamba as temporal module
python inference/infer_stcm.py

You can enter the script to modify the input path.

🎬 Train A SCST

We divide the training process into several steps to help you reproduce our results from scratch.

1. Download Pretrained Models

   Download Stable Diffusion 2.1 from HuggingFace and place it under the checkpoints/stable-diffusion-2-1-base/ directory

2. Download Training Datasets

   Pre-processed Data (skip Step 3 if using this)

   • REDS Sub

   Raw Data (Download original datasets)

   • REDS (Google Drive)
   • YouHQ (Google Drive)

3. Prepare Dataset Format

   We use .tar packages of video frames as inputs. Follow the steps below:

   • Use dataloader/extract_sub_images.py to crop the high-resolution frames.
   • Pack Each Video Sequence into a .tar File and each file contains frame images named {i%08d}.png.
   • Prepare the Meta Info File Example (datasets_example/reds_meta_info_example.txt).

   Datasets structure:

   datasets_example/
   └── REDS/
       ├── 000_s001.tar 
       ├── xxx.tar
       ...
   

   Example of reds_meta_info_example.txt:

   000_s001 40
   

   The frames inside the .tar file:

   000_s001/00000000.png
   ...
   000_s001/00000039.png
   

4. Modify Training Scripts and Start Training Edit the following parameters in train_stage13.py and train_stage2.py to match your data path:

   meta_path = './datasets_example/reds_meta_info_example.txt'
   hr_root = './datasets_example/REDS/'
   

   We provide training scripts for all three stages:

   • Stage 1: scripts/stage1.sh
   • Stage 2: scripts/stage2.sh
   • Stage 3 with LocalAttention: scripts/stage3_localatten.sh
   • Stage 3 with STCM: scripts/stage3_stcm.sh

   To ensure that Stage 2 and Stage 3 use the pre-trained model from the previous stage, make sure to set the resume_path parameter accordingly.

🎬 Overview

Acknowledgments

Our project is based on diffusers. Some codes are brought from MGLD and PASD. Thanks for their awesome works.

📧 Contact

If you have any questions, please feel free to reach us at ssj180123@gmail.com