/pytorch-face-reaging-network

PyTorch Lightning implementation of Disney's face re-aging network (FRAN) paper.

Primary LanguagePythonMIT LicenseMIT

PyTorch Face Re-aging Network (FRAN)

Image Alt Text Image Alt Text

An unofficial PyTorch implementation of Disney's face re-aging network (FRAN) paper.

Paper Link: link

Dataset Link (collected according to the paper): link

Setup

Python 3.12.2

> pip install -r requirements.txt

For full training, download all 2000 subjects from the dataset link.

Training

See Colab Notebook or train a new model by

> python train.py -h
usage: train.py [-h] [--data_dir DATA_DIR]

Train FRAN model.

options:
  -h, --help            show this help message and exit
  --data_dir DATA_DIR, -C DATA_DIR
                        directory for data

> python train.py

Using 16bit Automatic Mixed Precision (AMP)
GPU available: True (mps), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs

  | Name          | Type                  | Params
--------------------------------------------------------
0 | generator     | Generator             | 31.4 M
1 | discriminator | PatchGANDiscriminator | 1.6 M 
--------------------------------------------------------
32.9 M    Trainable params
0         Non-trainable params
32.9 M    Total params
131.800   Total estimated model params size (MB)
Epoch 0:   0%|                                                                       | 0/228 [00:00<?, ?it/s]

Showcase

Applying the model + naive crop and masking. Model was trained on a single V100 for 6 hours in Google Colab.

Hugging Face Spaces

Notes

Why does the image turn grainy/pixelated sometimes?

Model encountered an unfamilar face. More diverse data may help resolve this problem.

Model is noticably worse with turning old people to young.