Zhiyuan Li, Yanhui Zhou, Hao Wei, Chenyang Ge, Ajmal Mian
🥳 This work is accepted by IEEE Transactions on Circuits and Systems for Video Technology.

⭐ The quantitative metrics for each method presented in our paper can be found in result.xlsx.

# clone this repo
git clone https://github.com/huai-chang/RDEIC.git
cd RDEIC
# create an environment
conda create -n rdeic python=3.8
conda activate rdeic
pip install torch==2.0.1
pip install tb-nightly --index-url https://pypi.org/simple
pip install -r requirements.txt

1. Download LSDIR dataset.

2. Generate file list of training set and validation set.

   python3 make_fire_list.py\
   --train_folder [path_to_train_folder]\
   --test_folder [path_to_test_folder]\
   --save_folder [path_to_save_floder]
   

   After running this script, you will get two file lists in save_folder, each line in a file list contains an absolute path of an image file:

   save_folder
   ├── train.list # training file list
   └── valid.list # validation file list
   

3. Download pretrained Stable Diffusion v2.1 into ./weight.

   wget https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt --no-check-certificate
   

1. Modify the configuration file.

   # ./configs/train_rdeic.yaml
   resume: ~ or path to the initial checkpoint
   default_root_dir: path to save logs and checkpoints
   # ./configs/model/rdeic.yaml
   is_refine: False
   learning_rate: 1e-4 or 2e-5
   l_guide_weight: 2 #{2,1,0.5,0.25,0.1}
   

2. Start training.

   python3 train.py
   

1. Modify the configuration file.

   # ./configs/train_rdeic.yaml
   resume: path to the initial checkpoint
   default_root_dir: path to save logs and checkpoints.
   # ./configs/model/rdeic.yaml
   is_refine: True
   fixed_step: 2 or 5
   learning_rate: 2e-5
   l_guide_weight: 2 #{2,1,0.5,0.25,0.1}
   # ./configs/dataset/lic_train.yaml
   out_size: 256
   

2. Start training.

   python3 train.py
   

1. Download pretrained Stable Diffusion v2.1 into ./weight.

   wget https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt --no-check-certificate
   

2. Download the pre-trained weights for the Compression and Control Module into ./weight.

3. Download test datasets.

4. Run the following command.

   python3 inference_partition.py \
   --ckpt_sd ./weight/v2-1_512-ema-pruned.ckpt \
   --ckpt_lc ./weight/rdeic_2_step2.ckpt \
   --config configs/model/rdeic.yaml \
   --input path to input images \
   --output path to output files \
   --steps 2 \
   --guidance_scale 1.0 \
   --device cuda 
   

• Release code
• Release pretrained models

This work is based on DiffEIC and CVQ-VAE, thanks to their invaluable contributions.

Please cite us if our work is useful for your research.

@article{li2024diffusion,
  title={Diffusion-based Extreme Image Compression with Compressed Feature Initialization},
  author={Li, Zhiyuan and Zhou, Yanhui and Wei, Hao and Ge, Chenyang and Mian, Ajmal},
  journal={arXiv preprint arXiv:2410.02640},
  year={2024}
}