/TeCM-CLIP

Primary LanguagePython

TeCM-CLIP: Text-based Controllable Multi-attribute Face Image Manipulation

Our framework supports hairstyle, hair color, emotion, gender and age editing. Select the corresponding attribute mapper and enter the text prompt.

Getting Started

Dependencies

$ Python >= 3.7 (Recommend to use Anaconda or Miniconda)
$ PyTorch >= 1.7
$ CLIP
$ Option: NVIDIA GPU + CUDA

Installation

$ conda install --yes -c pytorch pytorch=1.7.1 torchvision cudatoolkit=11.0
$ pip install ftfy regex tqdm
$ pip install git+https://github.com/openai/CLIP.git

Models

Please download the pre-trained model from the following link. We only provide the pretrained model of Hairstyle here.

Path Description
Hairstyle Our pre-trained Hairstyle model.
Hair Color Our pre-trained Hair Color model.
Emotion Our pre-trained Emotion model.
Gender Our pre-trained Gender model.
Age Our pre-trained Age model.

If you wish to use the pretrained model for training or inference, you may do so using the flag --checkpoint_path.

Auxiliary Models and Latent Codes

In addition, we provide various auxiliary models and latent codes inverted by e4e needed for training your own model from scratch.

Path Description
FFHQ StyleGAN StyleGAN model pretrained on FFHQ taken from rosinality with 1024x1024 output resolution.
IR-SE50 Model Pretrained IR-SE50 model taken from TreB1eN for use in our ID loss during HairCLIP training.
Train Set CelebA-HQ train set latent codes inverted by e4e.
Test Set CelebA-HQ test set latent codes inverted by e4e.

By default, we assume that all auxiliary models are downloaded and saved to the directory pretrained_models.

Training

Training TeCM-CLIP

The main training script can be found in scripts/train.py.
Intermediate training results are saved to opts.exp_dir. This includes checkpoints, train outputs, and test outputs.

Training the TeCM-CLIP Mapper

Example of Using Text to Edit Hairstyle

cd mapper
python scripts/train.py \
--exp_dir=/path/to/experiment \
--change_type="HairStyle" \
--description="description/hairstyle_description.txt" \
--latents_train_path=/path/to/train_faces.pt \
--latents_test_path=/path/to/test_faces.pt \
--batch_size=1  \
--id_lambda=0.2 \
--clip_lambda=0.6 \
--latent_l2_lambda=1.0 \
--face_l2_lambda=2.0 \
--bg_l2_lambda=0.0 \
--learning_rate=0.0001 \
--mapper_type='LevelsMapper'

Testing

Inference

The main inference script can be found in inference.py. Inference results are saved to test_opts.exp_dir.

Example of Using Text to Edit Hairstyle

cd mapper
python inference.py \
--exp_dir=/path/to/experiment \
--checkpoint_path=./pretrained_models/hairclip.pt \
--latents_test_path=/path/to/test_faces.pt \
--description="short curly hair"

Acknowledgements

This code is based on StyleCLIP and HairCLIP.