milleniums / DifFace

DifFace: Blind Face Restoration with Diffused Error Contraction (PyTorch)

DifFace: Blind Face Restoration with Diffused Error Contraction

Zongsheng Yue, Chen Change Loy

Paper

:star: If DifFace is helpful to your images or projects, please help star this repo. Thanks! :hugs:

Update

• 2022.12.19: Add Colab demo .
• 2022.12.17: Add the .
• 2022.12.13: Create this repo.

Applications

👉 Old Photo Enhancement

👉 Face Restoration

Requirements

A suitable conda environment named DifFace can be created and activated with:

conda env create -f environment.yaml
conda activate DifFace

Inference

👦 Face image restoration (cropped and aligned)

python inference_difface.py --aligned --in_path [image folder/image path] --out_path [result folder] --gpu_id [gpu index]

👫 Whole image enhancement

python inference_difface.py --in_path [image folder/image path] --out_path [result folder] --gpu_id [gpu index]

Training

🐢 Prepare data

1. Download the FFHQ dataset, and resize them into size 512x512.

python datapipe/prepare/face/big2small_face.py --face_dir [Face folder(1024x1024)] --save_dir [Saving folder] --pch_size 512 

2. Extract the image path into 'datapipe/files_txt/ffhq512.txt'

python datapipe/prepare/face/split_train_val.py --face_dir [Face folder(512x512)] --save_dir [Saving folder] 

3. Making the testing dataset

python datapipe/prepare/face/make_testing_data.py --files_txt datapipe/files_txt/ffhq512.txt --save_dir [Saving folder]  

🐬 Train diffusion model

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 --nnodes=1 main_diffusion.py --cfg_path configs/training/diffsuion_ffhq512.yaml --save_dir [Logging Folder]  

🐳 Train diffused estimator (SwinIR)

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 --nnodes=1 main_sr.py --cfg_path configs/training/swinir_ffhq512.yaml --save_dir [Logging Folder]  

License

This project is licensed under NTU S-Lab License 1.0. Redistribution and use should follow this license.

Acknowledgement

This project is based on Improved Diffusion Model. Some codes are brought from BasicSR, YOLOv5-face, and FaceXLib. We also adopt Real-ESRGAN to support background image enhancement. Thanks for their awesome works.

Contact

If you have any questions, please feel free to contact me via [e-mail].