Ditto Talking Head Synthesis

Download the model checkpoints from HuggingFace and place them in a checkpoints directory. You can use Git LFS to clone them directly:

git lfs install
git clone https://huggingface.co/digital-avatar/ditto-talkinghead checkpoints

Checkpoint Structure

The directory should contain subdirectories for configurations (ditto_cfg), ONNX models (ditto_onnx), and TensorRT engines (ditto_trt_Ampere_Plus or ditto_pytorch).

Configuration Types:

• ditto_cfg/v0.4_hubert_cfg_trt_online.pkl: Online configuration.
• ditto_cfg/v0.4_hubert_cfg_trt.pkl: Offline configuration.
• ditto_cfg/v0.4_hubert_cfg_pytorch.pkl: PyTorch configuration.

To use Ditto, first clone the repository:

git clone https://github.com/antgroup/ditto-talkinghead
cd ditto-talkinghead

Option 1: Conda (Recommended)

Create and activate the environment using the provided environment.yaml:

conda env create -f environment.yaml
conda activate ditto

Option 2: Pip

If you have pytorch, cuda, and cudnn already installed, you can install the required dependencies via pip:

pip install \
    tensorrt==8.6.1 \
    librosa \
    tqdm \
    filetype \
    imageio \
    opencv_python_headless \
    scikit-image \
    cython \
    cuda-python \
    imageio-ffmpeg \
    colored \
    polygraphy \
    numpy==2.0.1

Note: If you are not using Conda, you may also need to install ffmpeg manually.

To run inference using the TensorRT models, use the inference.py script.

Note on Hardware Compatibility: The provided TensorRT models are optimized for hardware-compatibility-level=Ampere_Plus (found in checkpoints/ditto_trt_Ampere_Plus/). If your GPU does not support this, you must convert the ONNX models to TensorRT using the cvt_onnx_to_trt.py script first:

python scripts/cvt_onnx_to_trt.py --onnx_dir "./checkpoints/ditto_onnx" --trt_dir "./checkpoints/ditto_trt_custom"

Then, run inference using the custom directory with --data_root="./checkpoints/ditto_trt_custom".

python inference.py \
    --data_root "./checkpoints/ditto_trt_Ampere_Plus" \
    --cfg_pkl "./checkpoints/ditto_cfg/v0.4_hubert_cfg_trt.pkl" \
    --audio_path "./example/audio.wav" \
    --source_path "./example/image.png" \
    --output_path "./tmp/result.mp4"

If you prefer using the PyTorch version of the model, ensure you have downloaded the ditto_pytorch checkpoints. The directory structure should include ditto_pytorch/models/ and ditto_pytorch/aux_models/.

Run the following command:

python inference.py \
    --data_root "./checkpoints/ditto_pytorch" \
    --cfg_pkl "./checkpoints/ditto_cfg/v0.4_hubert_cfg_pytorch.pkl" \
    --audio_path "./example/audio.wav" \
    --source_path "./example/image.png" \
    --output_path "./tmp/result.mp4"

The inference.py script accepts the following arguments: