-
Prepare training data:
- download CelebAMask-HQ dataset
- change file path in the
preprocess_data.py
script and run
python preprocess_data.py
-
Train the model using CelebAMask-HQ dataset: Just run the train script:
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
If you do not wish to train the model, you can download our pre-trained model and save it in res/cp
.
- Evaluate the trained model using:
# evaluate using GPU
python test.py
Hair | Lip | |
---|---|---|
Original Input | ||
Color |