Python Version 3.9.23
Required packages contained within requirements.txt. Recommend using miniforge to support packages such as flash-attn.
Ask if in need to help troubleshooting setup.
Visit Drive here to get the MagiCodec model and files and place within root directory and unzip. Will need to set up taming-transformers within the root directory; if you run into trouble I can try to find a way to host the setup for it I have on local, but it's ~21GB so I'll need to look for a storage solution.
For the datasets, we should collaborate to find a solution to transfer as the datasets are too large to store on Drive
The project has so far been written to support:
- HDTF
- TKED
- MEAD
All datasets contain audio files and FLAME representations of the subject, broken up into small sequences
HDTF and TKED are made with EMICA and use 300 shape and 100 expression FLAME features, MEAD is made with EMOCA and uses 100 shape and 50 expression. The FLAME model code will need to be modified to switch between these two cases. The chosen dataset and associating data loading code is configured within train.py.
Right now the code runs simply by executing:
python train.py
Currently the program does not support multi-GPU operation.
The primary loop runs within talkinghead/models/mask_transformer/transformer_trainer_hdtf_stage2_2ndcode.py, and the 2nd stage transformer is located within talkinghead/models/mask_transformer/transformer_trainer_hdtf_stage2_2ndcode.py
In transformer_trainer_hdtf_stage2_2ndcode.py you can change the checkpoint name by modifying lines 538 and 540:
self.save(pjoin(self.opt.model_dir, 'latest_mead_80vq.tar'), epoch, it) #Set checkpoint save name here
self.save(pjoin(self.opt.model_dir, 'latest_mead_80vq'), epoch, it) #Set checkpoint save name here
To load a train checkpoint, modify lines 486 and 488. To disable checkpointing and train a fresh new model, comment out line 486.
self.opt.is_continue = True #Flag whether to load from checkpoint
model_dir = pjoin(self.opt.model_dir, 'latest_mead_80vq.tar') # Set the checkpoint path here
You can switch between training and inference by commenting/uncommenting lines 11 and 12 of train.py:
from models.mask_transformer.transformer_trainer_hdtf_stage2_2ndcode import MaskTransformerTrainer
#from models.mask_transformer.transformer_trainer_hdtf_stage2_2ndcode_inference import MaskTransformerTrainer
Checkpoints are saved within the checkpoints directory.
You can change the stage 1 VQVAE checkpoint from line 253 in train.py from among those currently available within the models directory; the current one is the latest for MEAD.