-
Prepare training data: -- download CelebAMask-HQ dataset
-- change file path in the
prepropess_data.py
and run
python prepropess_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 |