Towards Extreme Image Compression with Latent Feature Guidance and Diffusion Prior

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

📖 Table Of Contents

• 👀 Visual Results
• ⚔️ Quantitative Performance
• 💻 Train
• ⚡ Inference
• 📝 TODO
• ❤️ Acknowledgement
• 📋 Citation

👀 Visual Results

⚔️ Quantitative Performance

🔧 Requirements

- conda create -n diffeic python=3.8
- conda activate diffeic
- pip install torch==2.0.1
- pip install tb-nightly --index-url https://pypi.org/simple
- pip install -r requirements.txt

💻 Train

1. 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
   

2. 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
   

3. Modify the configuration file ./configs/train_diffeic.yaml and ./configs/model/diffeic.yaml accordingly.

4. Start training.

   python3 train.py
   

⚡ Inference

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 LFGCM and Control Module into ./weight.

   Bitrate	Link
   0.12 bpp	1_2_1
   0.09 bpp	1_2_2
   0.06 bpp	1_2_4
   0.04 bpp	1_2_8
   0.02 bpp	1_2_16

3. Run the following command.

   python3 inference_partition.py \
   --ckpt_sd ./weight/v2-1_512-ema-pruned.ckpt \
   --ckpt_lc ./weight/1_2_1/lc.ckpt \
   --config configs/model/diffeic.yaml \
   --input path to input images \
   --output path to output files \
   --steps 50 \
   --devide cuda 
   

📝 TODO

• Release code
• Release pretrained models

❤️ Acknowledgement

This work is based on ControlNet, DiffBIR, and ELIC, thanks to their invaluable contributions.

📋 Citation

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

@article{li2024towards,
  author={Li, Zhiyuan and Zhou, Yanhui and Wei, Hao and Ge, Chenyang and Jiang, Jingwen},
  journal={IEEE Transactions on Circuits and Systems for Video Technology}, 
  title={Towards Extreme Image Compression with Latent Feature Guidance and Diffusion Prior}, 
  year={2024},
  doi={10.1109/TCSVT.2024.3455576}}