Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/teaser.png filter=lfs diff=lfs merge=lfs -text
+demo/input/English_WTT5UTZQ9K8_8.wav filter=lfs diff=lfs merge=lfs -text
+demo/input/French_JATq1mUhfiA_8.wav filter=lfs diff=lfs merge=lfs -text
+demo/input/Italian_72pdx3tZwto_4.wav filter=lfs diff=lfs merge=lfs -text
+sample_dataset/wav/Greek_0_38_FbWPEz8NFS8.wav filter=lfs diff=lfs merge=lfs -text
+sample_dataset/wav/Spanish_xyVZDmzt6HY_6.wav filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+/MultiTalk_dataset/raw_video
+/MultiTalk_dataset/multitalk_dataset

MultiTalk_dataset/README.md

@@ -0,0 +1,96 @@
+## Overview
+MultiTalk dataset is a new multilingual 2D video dataset featuring over 420 hours of talking videos across 20 languages. 
+It contains 293,812 clips with a resolution of 512x512, a frame rate of 25 fps, and an average duration of 5.19 seconds per clip.
+The dataset shows a balanced distribution across languages, with each language representing between 2.0% and 9.7% of the total.  
+
+<img alt="statistic" src="../assets/statistic.png" width=560>
+
+
+<details><summary><b>Detailed statistics</b></summary><p>
+
+| Language | Total Duration(h) | #Clips | Avg. Duration(s) |                                                     Annotation                                                      |
+|:---:|:---:|:---:|:---:|:-------------------------------------------------------------------------------------------------------------------:|
+| Arabic | 10.32 | 9048 | 4.11 |     [arabic.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/arabic.json)     |
+| Catalan | 41.0 |  29232 | 5.05 |    [catalan.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/catalan.json)    |
+| Croatian | 41.0 |  25465 | 5.80 |   [croatian.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/croatian.json)   |
+| Czech | 18.9 | 11228 | 6.06 |      [czech.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/czech.json)      |
+| Dutch | 17.05 | 14187 | 4.33 |      [dutch.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/dutch.json)      |
+| English | 15.49 |  11082 | 5.03 |    [english.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/english.json)    |
+| French | 13.17 |  11576 | 4.10 |     [french.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/french.json)     |
+| German | 16.25 | 10856 | 5.39 |     [german.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/german.json)     |
+| Greek | 17.53 | 12698 | 4.97 |      [greek.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/greek.json)      |
+| Hindi | 24.41 | 16120 | 5.45 |      [hindi.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/hindi.json)      |
+| Italian | 13.59 | 9753 | 5.02 |    [italian.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/italian.json)    |
+| Japanese | 8.36 | 5990 | 5.03 |   [japanese.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/japanese.json)   |
+| Mandarin | 8.73 | 6096 | 5.15 |   [mandarin.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/mandarin.json)   |
+| Polish | 21.58 | 15181 | 5.12 |     [polish.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/polish.json)     |
+| Portuguese | 41.0 | 25321 | 5.83 | [portuguese.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/portuguese.json) |
+| Russian | 26.32 | 17811 | 5.32 |    [russian.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/russian.json)    |
+| Spanish | 23.65 | 18758 | 4.54 |    [spanish.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/spanish.json)    |
+| Thai | 10.95 | 7595 | 5.19 |       [thai.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/thai.json)       |
+| Turkish | 12.9 | 11165 | 4.16 |    [turkish.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/turkish.json)    |
+| Ukrainian | 41.0 | 24650 | 5.99 |  [ukrainian.json](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/annotations/ukrainian.json)  |
+</p></details>
+
+## Download
+
+### Usage
+**Prepare the environment:**
+```bash
+pip install pytube
+pip install opencv-python
+```
+
+**Run script:** 
+```bash
+cd MultiTalk_Dataset
+``` 
+You can pass the languages you want to download as arguments to the script. If you want to download all 20 languages, run the following script.  
+```bash
+sh dataset.sh arabic catalan croatian czech dutch english french german greek hindi italian japanese mandarin polish portuguese russian spanish thai turkish ukrainian
+```
+
+After downloading, the folder structure will be as below. Each language folder contains the .mp4 videos.  
+You can change the ${ROOT} folder in the [code](https://github.com/postech-ami/MultiTalk/tree/main/MultiTalk_dataset/download_and_process.py).
+```
+    ${ROOT}
+    ├── multitalk_dataset        # MultiTalk Dataset
+    │   ├── arabic
+    │   │   ├── O-VJXuHb390_0.mp4
+    │   │   ├── O-VJXuHb390_1.mp4
+    │   │   ├── ...
+    │   │   └── ...             
+    │   ├── catalan                                        
+    │   ├── ...          
+    │   └── ...                 
+    └── raw_video              # Original videos (you can remove this directory after downloading)
+        ├── arabic              
+        ├── catalan                          
+        ├── ...          
+        └── ...
+```
+
+### JSON File Structure
+```javascript
+{
+    "QrDZjUeiUwc_0":  // clip 1 
+    {
+        "youtube_id": "QrDZjUeiUwc",                                // youtube id
+        "duration": {"start_sec": 302.0, "end_sec": 305.56},        // start and end times in the original video
+        "bbox": {"top": 0.0, "bottom": 0.8167, "left": 0.4484, "right": 0.9453},  // bounding box
+        "language": "czech",                                        // language
+        "transcript": "já jsem v podstatě obnovil svůj list z minulého roku"      // transcript
+    },
+    "QrDZjUeiUwc_1":  // clip 2 
+    {
+        "youtube_id": "QrDZjUeiUwc",                                
+        "duration": {"start_sec": 0.12, "end_sec": 4.12},        
+        "bbox": {"top": 0.0097, "bottom": 0.55, "left": 0.3406, "right": 0.6398},  
+        "language": "czech",                                       
+        "transcript": "ahoj tady anička a vítejte u dalšího easycheck videa"      
+    }
+    "..."
+    "..."
+
+}
+```

MultiTalk_dataset/dataset.sh

@@ -0,0 +1,4 @@
+for language in "$@"
+do
+python download_and_process.py --language "$language"
+done

MultiTalk_dataset/download_and_process.py

@@ -0,0 +1,147 @@
+import os
+import json
+import cv2
+import argparse
+from pytube import Playlist, YouTube
+from pytube.exceptions import VideoUnavailable
+import os
+import shutil
+import subprocess
+
+def downloadYouTube(yt, videourl, path):
+    video_stream = yt.streams.filter(progressive=False, file_extension='mp4').order_by('resolution').desc().first()
+    audio_stream = yt.streams.filter(only_audio=True).order_by('abr').desc().first()
+    if video_stream.fps >= 25:
+        video_id = videourl.split('=')[-1]
+        video_path = os.path.join(path, f"{video_id}_video.mp4")
+        audio_path = os.path.join(path, f"{video_id}_audio.mp4")
+        final_path = os.path.join(path, f"{video_id}.mp4")
+
+        print("Downloading video...")
+        video_stream.download(filename=video_path)
+        print("Downloading audio...")
+        audio_stream.download(filename=audio_path)
+
+        print("Merging video and audio...")
+        subprocess.run([
+        'ffmpeg', '-i', video_path, '-i', audio_path, '-r', '25',
+        '-c:v', 'copy', '-c:a', 'aac', '-strict', 'experimental',
+        final_path, '-y'
+        ])
+
+        os.remove(video_path)
+        os.remove(audio_path)
+        return True
+
+    else:
+        return False
+
+def process_ffmpeg(raw_vid_path, save_folder, save_vid_name,
+                   bbox, time):
+    """
+    raw_vid_path:
+    save_folder:
+    save_vid_name:
+    bbox: format: top, bottom, left, right. the values are normalized to 0~1
+    time: begin_sec, end_sec
+    """
+
+    def secs_to_timestr(secs):
+        hrs = secs // (60 * 60)
+        min = (secs - hrs * 3600) // 60
+        sec = secs % 60
+        end = (secs - int(secs)) * 100
+        return "{:02d}:{:02d}:{:02d}.{:02d}".format(int(hrs), int(min),
+                                                    int(sec), int(end))
+
+    def expand(bbox, ratio):
+        top, bottom = max(bbox[0] - ratio, 0), min(bbox[1] + ratio, 1)
+        left, right = max(bbox[2] - ratio, 0), min(bbox[3] + ratio, 1)
+
+        return top, bottom, left, right
+
+    def to_square(bbox):
+        top, bottom, left, right = bbox
+        h = bottom - top
+        w = right - left
+        c = min(h, w) // 2
+        c_h = (top + bottom) / 2
+        c_w = (left + right) / 2
+
+        top, bottom = c_h - c, c_h + c
+        left, right = c_w - c, c_w + c
+        return top, bottom, left, right
+
+    def denorm(bbox, height, width):
+        top, bottom, left, right = \
+            round(bbox[0] * height), \
+            round(bbox[1] * height), \
+            round(bbox[2] * width), \
+            round(bbox[3] * width)
+
+        return top, bottom, left, right
+
+    out_path = os.path.join(save_folder, save_vid_name)
+
+    cap = cv2.VideoCapture(raw_vid_path)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    top, bottom, left, right = to_square(
+        denorm(expand(bbox, 0.02), height, width))
+    start_sec, end_sec = time
+    cmd = f"ffmpeg -i {raw_vid_path} -r 25 -vf crop=w={right-left}:h={bottom-top}:x={left}:y={top},scale=512:512 -ss {start_sec} -to {end_sec} -loglevel error {out_path}"
+    os.system(cmd)
+
+
+def load_data(file_path):
+    with open(file_path) as f:
+        data_dict = json.load(f)
+
+    for key, val in data_dict.items():
+        save_name = key+".mp4"
+        ytb_id = val['youtube_id']
+        time = val['duration']['start_sec'], val['duration']['end_sec']
+
+        bbox = [val['bbox']['top'], val['bbox']['bottom'],
+                val['bbox']['left'], val['bbox']['right']]
+        language = val['language']
+        yield ytb_id, save_name, time, bbox, language
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--language', type=str, default="dutch", help='Language')
+    args = parser.parse_args()
+
+    # you can change the root folder
+    root = './'
+    processed_vid_root = os.path.join(root, 'multitalk_dataset') # processed video path
+    raw_vid_root = os.path.join(root, 'raw_video')  # downloaded raw video path
+    os.makedirs(processed_vid_root, exist_ok=True)
+    os.makedirs(raw_vid_root, exist_ok=True)
+
+    json_path = os.path.join('./annotations', f'{args.language}.json')  # json file path
+
+    for vid_id, save_vid_name, time, bbox, language in load_data(json_path):
+        processed_vid_dir = os.path.join(processed_vid_root, language)
+        raw_vid_dir = os.path.join(raw_vid_root, language)
+        raw_vid_path = os.path.join(raw_vid_dir, vid_id + ".mp4")
+
+        os.makedirs(processed_vid_dir, exist_ok=True)
+        os.makedirs(raw_vid_dir, exist_ok=True)
+
+        url = 'https://www.youtube.com/watch?v='+vid_id
+        success = True
+        if not os.path.isfile(raw_vid_path) :
+            while True:
+                try:
+                    yt = YouTube(url, use_oauth=True)
+                    success = downloadYouTube(yt, url, raw_vid_dir)
+                    break
+                except:
+                    continue
+        if success:
+            process_ffmpeg(raw_vid_path, processed_vid_dir, save_vid_name, bbox, time)
+
+    # you can remove this directory after downloading
+    # shutil.rmtree(raw_vid_root)

README.md

@@ -0,0 +1,140 @@
+# MultiTalk (INTERSPEECH 2024)
+
+### [Project Page](https://multi-talk.github.io/) | [Paper](https://arxiv.org/abs/2406.14272) | [Dataset](https://github.com/postech-ami/MultiTalk/blob/main/MultiTalk_dataset/README.md)
+This repository contains a pytorch implementation for the Interspeech 2024 paper, [MultiTalk: Enhancing 3D Talking Head Generation Across Languages with Multilingual Video Dataset](https://multi-talk.github.io/). MultiTalk generates 3D talking head with enhanced multilingual performance.<br><br>
+
+<img width="700" alt="teaser" src="./assets/teaser.png"> 
+
+## Getting started
+This code was developed on Ubuntu 18.04 with Python 3.8, CUDA 11.3 and PyTorch 1.12.0. Later versions should work, but have not been tested.
+
+### Installation
+Create and activate a virtual environment to work in:
+```
+conda create --name multitalk python=3.8
+conda activate multitalk
+```
+
+Install [PyTorch](https://pytorch.org/). For CUDA 11.3 and ffmpeg, this would look like:
+```
+pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu113
+conda install -c conda-forge ffmpeg
+```
+
+Install the remaining requirements with pip:
+```
+pip install -r requirements.txt 
+```
+
+Compile and install `psbody-mesh` package:
+[MPI-IS/mesh](https://github.com/MPI-IS/mesh)
+```
+BOOST_INCLUDE_DIRS=/usr/lib/x86_64-linux-gnu make all
+```
+
+
+### Download models
+To run MultiTalk, you need to download stage1 and stage2 model, and the template file of mean face in FLAME topology,
+Download [stage1 model](https://drive.google.com/file/d/1jI9feFcUuhXst1pM1_xOMvqE8cgUzP_t/view?usp=sharing | [stage2 model](https://drive.google.com/file/d/1zqhzfF-vO_h_0EpkmBS7nO36TRNV4BCr/view?usp=sharing) | [template](https://drive.google.com/file/d/1WuZ87kljz6EK1bAzEKSyBsZ9IlUmiI-i/view?usp=sharing) and download FLAME_sample.ply from [voca](https://github.com/TimoBolkart/voca/tree/master/template).
+
+After downloading the models, place them in `./checkpoints`.
+```
+./checkpoints/stage1.pth.tar
+./checkpoints/stage2.pth.tar
+./checkpoints/FLAME_sample.ply
+```
+
+## Demo
+Run below command to train the model.
+We provide sample audios in **./demo/input**.
+```
+sh scripts/demo.sh multi
+```
+
+To use wav2vec of `facebook/wav2vec2-large-xlsr-53`, please move to `/path/to/conda_environment/lib/python3.8/site-packages/transformers/models/wav2vec2/processing_wav2vec2.py` and change the code as below.
+```
+L105: tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)
+to
+L105: tokenizer=Wav2Vec2CTCTokenizer.from_pretrained("facebook/wav2vec2-base-960h",**kwargs)
+```
+
+## MultiTalk Dataset
+Please follow the instructions in [MultiTalk_dataset/README.md](https://github.com/postech-ami/MultiTalk/blob/main/MultiTalk_dataset/README.md).
+
+## Training and testing
+
+### Training for Discrete Motion Prior
+```
+sh scripts/train_multi.sh MultiTalk_s1 config/multi/stage1.yaml multi s1
+```
+
+### Training for Speech-Driven Motion Synthesis
+Make sure the paths of pre-trained models are correct, i.e.,`vqvae_pretrained_path` and `wav2vec2model_path` in `config/multi/stage2.yaml`.
+```
+sh scripts/train_multi.sh MultiTalk_s2 config/multi/stage2.yaml multi s2
+```
+
+### Testing
+#### Lip Vertex Error (LVE)
+For evaluating the lip vertex error, please run below command.
+
+```
+sh scripts/test.sh MultiTalk_s2 config/multi/stage2.yaml vocaset s2
+```
+
+#### Audio-Visual Lip Reading (AVLR)
+For evaluating lip readability with a pre-trained Audio-Visual Speech Recognition (AVSR), download language specific checkpoint, dictionary, and tokenizer from [muavic](https://github.com/facebookresearch/muavic).   
+Place them in `./avlr/${language}/checkpoints/${language}_avlr`.
+```
+# e.g "Arabic" 
+./avlr/ar/checkpoints/ar_avsr/checkpoint_best.pt
+./avlr/ar/checkpoints/ar_avsr/dict.ar.txt
+./avlr/ar/checkpoints/ar_avsr/tokenizer.model
+```
+And place the rendered videos in `./avlr/${language}/inputs/MultiTalk`, corresponding wav files in `./avlr/${language}/inputs/wav`.
+```
+# e.g "Arabic" 
+./avlr/ar/inputs/MultiTalk
+./avlr/ar/inputs/wav
+```
+
+Run below command to evaluate lip readability.
+```
+python eval_avlr/eval_avlr.py --avhubert-path ./av_hubert/avhubert --work-dir ./avlr --language ${language} --model-name MultiTalk --exp-name ${exp_name}
+```
+
+
+
+[//]: # (## **Citation**)
+
+[//]: # ()
+[//]: # (If you find the code useful for your work, please star this repo and consider citing:)
+
+[//]: # ()
+[//]: # (```)
+
+[//]: # (@inproceedings{xing2023codetalker,)
+
+[//]: # (  title={Codetalker: Speech-driven 3d facial animation with discrete motion prior},)
+
+[//]: # (  author={Xing, Jinbo and Xia, Menghan and Zhang, Yuechen and Cun, Xiaodong and Wang, Jue and Wong, Tien-Tsin},)
+
+[//]: # (  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},)
+
+[//]: # (  pages={12780--12790},)
+
+[//]: # (  year={2023})
+
+[//]: # (})
+
+[//]: # (```)
+
+## **Notes**
+1. Although our codebase allows for training with multi-GPUs, we did not test it and just hardcode the training batch size as one. You may need to change the `data_loader` if needed.
+
+
+## **Acknowledgement**
+We heavily borrow the code from
+[Codetalk](https://doubiiu.github.io/projects/codetalker/).
+We sincerely appreciate those authors.
+

RUN/multi/MultiTalk_s2/test-20240707_000247.log

@@ -0,0 +1,15 @@
+Traceback (most recent call last):
+  File "main/test_multi_pred.py", line 12, in <module>
+    import librosa
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/librosa/__init__.py", line 211, in <module>
+    from . import core
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/librosa/core/__init__.py", line 9, in <module>
+    from .constantq import *  # pylint: disable=wildcard-import
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/librosa/core/constantq.py", line 1059, in <module>
+    dtype=np.complex,
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/numpy/__init__.py", line 305, in __getattr__
+    raise AttributeError(__former_attrs__[attr])
+AttributeError: module 'numpy' has no attribute 'complex'.
+`np.complex` was a deprecated alias for the builtin `complex`. To avoid this error in existing code, use `complex` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.complex128` here.
+The aliases was originally deprecated in NumPy 1.20; for more details and guidance see the original release note at:
+    https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations

RUN/multi/MultiTalk_s2/test-20240707_000302.log

@@ -0,0 +1,110 @@
+[2024-07-07 00:03:03,357 INFO test_multi_pred.py line 63 3703560]=>INaffine: False
+StepLR: False
+adaptive_lr: False
+arch: stage2
+autoencoder: stage1_vocaset
+base_lr: 0.0001
+batch_size: 1
+batch_size_val: 1
+data_root: sample_dataset
+dataset: multi
+device: cuda
+dist_backend: nccl
+dist_url: tcp://127.0.0.1:6701
+epochs: 100
+eval_freq: 5
+evaluate: True
+face_quan_num: 16
+factor: 0.3
+feature_dim: 1024
+gamma: 0.5
+gt_save_folder: demo/gt
+hidden_size: 1024
+in_dim: 15069
+intermediate_size: 1536
+log_dir: None
+loss: MSE
+manual_seed: 131
+measure_lve: False
+model_path: checkpoints/stage2.pth.tar
+momentum: 0.9
+motion_weight: 1.0
+multiprocessing_distributed: True
+n_embed: 256
+n_head: 4
+neg: 0.2
+num_attention_heads: 8
+num_hidden_layers: 6
+num_layers: 6
+patience: 3
+period: 25
+poly_lr: False
+power: 0.9
+print_freq: 10
+quant_factor: 0
+rank: 0
+read_audio: True
+reg_weight: 1.0
+resume: None
+save: True
+save_folder: demo/output
+save_freq: 1
+save_path: None
+start_epoch: 0
+step_size: 100
+style_emb_method: nnemb
+sync_bn: False
+teacher_forcing: True
+template_file: templates.pkl
+test_batch_size: 1
+test_gpu: [0]
+test_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+test_workers: 0
+threshold: 0.0001
+train_gpu: [0]
+train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+use_sgd: False
+val_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+vertice_dim: 15069
+vertices_path: npy
+visualize_mesh: True
+vqvae_pretrained_path: checkpoints/stage1.pth.tar
+warmup_steps: 1
+wav2vec2model_path: facebook/wav2vec2-large-xlsr-53
+wav_path: wav
+weight: None
+weight_decay: 0.002
+window_size: 1
+workers: 10
+world_size: 1
+zquant_dim: 64
+[2024-07-07 00:03:03,357 INFO test_multi_pred.py line 64 3703560]=>=> creating model ...
+
+
+Some weights of the model checkpoint at facebook/wav2vec2-large-xlsr-53 were not used when initializing Wav2Vec2Model: ['quantizer.weight_proj.bias', 'project_hid.bias', 'project_q.bias', 'quantizer.weight_proj.weight', 'project_q.weight', 'quantizer.codevectors', 'project_hid.weight']
+- This IS expected if you are initializing Wav2Vec2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
+- This IS NOT expected if you are initializing Wav2Vec2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
+[2024-07-07 00:04:21,701 INFO test_multi_pred.py line 69 3703560]=>=> loading checkpoint 'checkpoints/stage2.pth.tar'
+[2024-07-07 00:04:22,268 INFO test_multi_pred.py line 72 3703560]=>=> loaded checkpoint 'checkpoints/stage2.pth.tar'
+Loading data...
+
+Traceback (most recent call last):
+  File "main/test_multi_pred.py", line 144, in <module>
+    main()
+  File "main/test_multi_pred.py", line 78, in main
+    dataset = get_dataloaders(cfg, test_config=True)
+  File "/home/rnd/Documents/Ameer/Dream/MultiTalk/dataset/data_loader_multi.py", line 142, in get_dataloaders
+    train_data, valid_data, test_data, subjects_dict = read_data(args, test_config)
+  File "/home/rnd/Documents/Ameer/Dream/MultiTalk/dataset/data_loader_multi.py", line 66, in read_data
+    processor = Wav2Vec2Processor.from_pretrained(args.wav2vec2model_path)
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/transformers/models/wav2vec2/processing_wav2vec2.py", line 105, in from_pretrained
+    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/transformers/tokenization_utils_base.py", line 1708, in from_pretrained
+    raise EnvironmentError(msg)
+OSError: Can't load tokenizer for 'facebook/wav2vec2-large-xlsr-53'. Make sure that:
+
+- 'facebook/wav2vec2-large-xlsr-53' is a correct model identifier listed on 'https://huggingface.co/models'
+
+- or 'facebook/wav2vec2-large-xlsr-53' is the correct path to a directory containing relevant tokenizer files
+
+

RUN/multi/MultiTalk_s2/test-20240707_000539.log

@@ -0,0 +1,107 @@
+[2024-07-07 00:05:39,843 INFO test_multi_pred.py line 63 3704273]=>INaffine: False
+StepLR: False
+adaptive_lr: False
+arch: stage2
+autoencoder: stage1_vocaset
+base_lr: 0.0001
+batch_size: 1
+batch_size_val: 1
+data_root: sample_dataset
+dataset: multi
+device: cuda
+dist_backend: nccl
+dist_url: tcp://127.0.0.1:6701
+epochs: 100
+eval_freq: 5
+evaluate: True
+face_quan_num: 16
+factor: 0.3
+feature_dim: 1024
+gamma: 0.5
+gt_save_folder: demo/gt
+hidden_size: 1024
+in_dim: 15069
+intermediate_size: 1536
+log_dir: None
+loss: MSE
+manual_seed: 131
+measure_lve: False
+model_path: checkpoints/stage2.pth.tar
+momentum: 0.9
+motion_weight: 1.0
+multiprocessing_distributed: True
+n_embed: 256
+n_head: 4
+neg: 0.2
+num_attention_heads: 8
+num_hidden_layers: 6
+num_layers: 6
+patience: 3
+period: 25
+poly_lr: False
+power: 0.9
+print_freq: 10
+quant_factor: 0
+rank: 0
+read_audio: True
+reg_weight: 1.0
+resume: None
+save: True
+save_folder: demo/output
+save_freq: 1
+save_path: None
+start_epoch: 0
+step_size: 100
+style_emb_method: nnemb
+sync_bn: False
+teacher_forcing: True
+template_file: templates.pkl
+test_batch_size: 1
+test_gpu: [0]
+test_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+test_workers: 0
+threshold: 0.0001
+train_gpu: [0]
+train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+use_sgd: False
+val_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+vertice_dim: 15069
+vertices_path: npy
+visualize_mesh: True
+vqvae_pretrained_path: checkpoints/stage1.pth.tar
+warmup_steps: 1
+wav2vec2model_path: facebook/wav2vec2-large-xlsr-53
+wav_path: wav
+weight: None
+weight_decay: 0.002
+window_size: 1
+workers: 10
+world_size: 1
+zquant_dim: 64
+[2024-07-07 00:05:39,843 INFO test_multi_pred.py line 64 3704273]=>=> creating model ...
+Some weights of the model checkpoint at facebook/wav2vec2-large-xlsr-53 were not used when initializing Wav2Vec2Model: ['quantizer.codevectors', 'project_q.weight', 'project_hid.bias', 'project_hid.weight', 'project_q.bias', 'quantizer.weight_proj.weight', 'quantizer.weight_proj.bias']
+- This IS expected if you are initializing Wav2Vec2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
+- This IS NOT expected if you are initializing Wav2Vec2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
+[2024-07-07 00:05:43,925 INFO test_multi_pred.py line 69 3704273]=>=> loading checkpoint 'checkpoints/stage2.pth.tar'
+[2024-07-07 00:05:44,516 INFO test_multi_pred.py line 72 3704273]=>=> loaded checkpoint 'checkpoints/stage2.pth.tar'
+Loading data...
+Traceback (most recent call last):
+  File "main/test_multi_pred.py", line 144, in <module>
+    main()
+  File "main/test_multi_pred.py", line 78, in main
+    dataset = get_dataloaders(cfg, test_config=True)
+  File "/home/rnd/Documents/Ameer/Dream/MultiTalk/dataset/data_loader_multi.py", line 142, in get_dataloaders
+    train_data, valid_data, test_data, subjects_dict = read_data(args, test_config)
+  File "/home/rnd/Documents/Ameer/Dream/MultiTalk/dataset/data_loader_multi.py", line 66, in read_data
+    processor = Wav2Vec2Processor.from_pretrained(args.wav2vec2model_path)
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/transformers/models/wav2vec2/processing_wav2vec2.py", line 105, in from_pretrained
+    tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)
+  File "/home/rnd/miniconda3/envs/multitalk/lib/python3.8/site-packages/transformers/tokenization_utils_base.py", line 1708, in from_pretrained
+    raise EnvironmentError(msg)
+OSError: Can't load tokenizer for 'facebook/wav2vec2-large-xlsr-53'. Make sure that:
+
+- 'facebook/wav2vec2-large-xlsr-53' is a correct model identifier listed on 'https://huggingface.co/models'
+
+- or 'facebook/wav2vec2-large-xlsr-53' is the correct path to a directory containing relevant tokenizer files
+
+

RUN/multi/MultiTalk_s2/test-20240707_000731.log

@@ -0,0 +1,93 @@
+[2024-07-07 00:07:32,456 INFO test_multi_pred.py line 63 3704520]=>INaffine: False
+StepLR: False
+adaptive_lr: False
+arch: stage2
+autoencoder: stage1_vocaset
+base_lr: 0.0001
+batch_size: 1
+batch_size_val: 1
+data_root: sample_dataset
+dataset: multi
+device: cuda
+dist_backend: nccl
+dist_url: tcp://127.0.0.1:6701
+epochs: 100
+eval_freq: 5
+evaluate: True
+face_quan_num: 16
+factor: 0.3
+feature_dim: 1024
+gamma: 0.5
+gt_save_folder: demo/gt
+hidden_size: 1024
+in_dim: 15069
+intermediate_size: 1536
+log_dir: None
+loss: MSE
+manual_seed: 131
+measure_lve: False
+model_path: checkpoints/stage2.pth.tar
+momentum: 0.9
+motion_weight: 1.0
+multiprocessing_distributed: True
+n_embed: 256
+n_head: 4
+neg: 0.2
+num_attention_heads: 8
+num_hidden_layers: 6
+num_layers: 6
+patience: 3
+period: 25
+poly_lr: False
+power: 0.9
+print_freq: 10
+quant_factor: 0
+rank: 0
+read_audio: True
+reg_weight: 1.0
+resume: None
+save: True
+save_folder: demo/output
+save_freq: 1
+save_path: None
+start_epoch: 0
+step_size: 100
+style_emb_method: nnemb
+sync_bn: False
+teacher_forcing: True
+template_file: templates.pkl
+test_batch_size: 1
+test_gpu: [0]
+test_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+test_workers: 0
+threshold: 0.0001
+train_gpu: [0]
+train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+use_sgd: False
+val_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+vertice_dim: 15069
+vertices_path: npy
+visualize_mesh: True
+vqvae_pretrained_path: checkpoints/stage1.pth.tar
+warmup_steps: 1
+wav2vec2model_path: facebook/wav2vec2-large-xlsr-53
+wav_path: wav
+weight: None
+weight_decay: 0.002
+window_size: 1
+workers: 10
+world_size: 1
+zquant_dim: 64
+[2024-07-07 00:07:32,456 INFO test_multi_pred.py line 64 3704520]=>=> creating model ...
+Some weights of the model checkpoint at facebook/wav2vec2-large-xlsr-53 were not used when initializing Wav2Vec2Model: ['project_hid.bias', 'project_q.weight', 'quantizer.weight_proj.weight', 'project_hid.weight', 'quantizer.codevectors', 'quantizer.weight_proj.bias', 'project_q.bias']
+- This IS expected if you are initializing Wav2Vec2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
+- This IS NOT expected if you are initializing Wav2Vec2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
+[2024-07-07 00:07:36,887 INFO test_multi_pred.py line 69 3704520]=>=> loading checkpoint 'checkpoints/stage2.pth.tar'
+[2024-07-07 00:07:37,452 INFO test_multi_pred.py line 72 3704520]=>=> loaded checkpoint 'checkpoints/stage2.pth.tar'
+Loading data...
+
+
+
+
  0%|          | 0/4 [00:00<?, ?it/s]
 50%|█████     | 2/4 [00:00<00:00,  7.68it/s]
+Loaded data: Train-0, Val-0, Test-2
+Lip Vertex Error on test set:  9.532534e-06

RUN/vocaset/MultiTalk_s2/test-20240707_000820.log

@@ -0,0 +1,90 @@
+[2024-07-07 00:08:21,269 INFO test_multi_pred.py line 63 3705029]=>INaffine: False
+StepLR: False
+adaptive_lr: False
+arch: stage2
+autoencoder: stage1_vocaset
+base_lr: 0.0001
+batch_size: 1
+batch_size_val: 1
+data_root: sample_dataset
+dataset: multi
+device: cuda
+dist_backend: nccl
+dist_url: tcp://127.0.0.1:6701
+epochs: 100
+eval_freq: 5
+evaluate: True
+face_quan_num: 16
+factor: 0.3
+feature_dim: 1024
+gamma: 0.5
+gt_save_folder: demo/gt
+hidden_size: 1024
+in_dim: 15069
+intermediate_size: 1536
+log_dir: None
+loss: MSE
+manual_seed: 131
+measure_lve: False
+model_path: checkpoints/stage2.pth.tar
+momentum: 0.9
+motion_weight: 1.0
+multiprocessing_distributed: True
+n_embed: 256
+n_head: 4
+neg: 0.2
+num_attention_heads: 8
+num_hidden_layers: 6
+num_layers: 6
+patience: 3
+period: 25
+poly_lr: False
+power: 0.9
+print_freq: 10
+quant_factor: 0
+rank: 0
+read_audio: True
+reg_weight: 1.0
+resume: None
+save: True
+save_folder: demo/output
+save_freq: 1
+save_path: None
+start_epoch: 0
+step_size: 100
+style_emb_method: nnemb
+sync_bn: False
+teacher_forcing: True
+template_file: templates.pkl
+test_batch_size: 1
+test_gpu: [0]
+test_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+test_workers: 0
+threshold: 0.0001
+train_gpu: [0]
+train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+use_sgd: False
+val_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+vertice_dim: 15069
+vertices_path: npy
+visualize_mesh: True
+vqvae_pretrained_path: checkpoints/stage1.pth.tar
+warmup_steps: 1
+wav2vec2model_path: facebook/wav2vec2-large-xlsr-53
+wav_path: wav
+weight: None
+weight_decay: 0.002
+window_size: 1
+workers: 10
+world_size: 1
+zquant_dim: 64
+[2024-07-07 00:08:21,269 INFO test_multi_pred.py line 64 3705029]=>=> creating model ...
+Some weights of the model checkpoint at facebook/wav2vec2-large-xlsr-53 were not used when initializing Wav2Vec2Model: ['project_q.weight', 'project_hid.bias', 'quantizer.weight_proj.weight', 'quantizer.weight_proj.bias', 'project_q.bias', 'quantizer.codevectors', 'project_hid.weight']
+- This IS expected if you are initializing Wav2Vec2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
+- This IS NOT expected if you are initializing Wav2Vec2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
+[2024-07-07 00:08:25,356 INFO test_multi_pred.py line 69 3705029]=>=> loading checkpoint 'checkpoints/stage2.pth.tar'
+[2024-07-07 00:08:25,922 INFO test_multi_pred.py line 72 3705029]=>=> loaded checkpoint 'checkpoints/stage2.pth.tar'
+Loading data...
+
  0%|          | 0/4 [00:00<?, ?it/s]
 50%|█████     | 2/4 [00:00<00:00,  7.70it/s]
+Loaded data: Train-0, Val-0, Test-2
+Lip Vertex Error on test set:  9.532534e-06

base/__init__.py

@@ -0,0 +1 @@
+from .base_model import *

base/baseTrainer.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+import torch
+from os.path import join
+import torch.distributed as dist
+from .utilities import check_makedirs
+from collections import OrderedDict
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+
+
+def step_learning_rate(base_lr, epoch, step_epoch, multiplier=0.1):
+    lr = base_lr * (multiplier ** (epoch // step_epoch))
+    return lr
+
+
+def poly_learning_rate(base_lr, curr_iter, max_iter, power=0.9):
+    """poly learning rate policy"""
+    lr = base_lr * (1 - float(curr_iter) / max_iter) ** power
+    return lr
+
+
+def adjust_learning_rate(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def save_checkpoint(model, other_state={}, sav_path='', filename='model.pth.tar', stage=1):
+    if isinstance(model, (DistributedDataParallel, DataParallel)):
+        weight = model.module.state_dict()
+    elif isinstance(model, torch.nn.Module):
+        weight = model.state_dict()
+    else:
+        raise ValueError('model must be nn.Module or nn.DataParallel!')
+    check_makedirs(sav_path)
+
+    if stage == 2: # remove vqvae part
+        for k in list(weight.keys()):
+            if 'autoencoder' in k:
+                weight.pop(k)
+
+    other_state['state_dict'] = weight
+    filename = join(sav_path, filename)
+    torch.save(other_state, filename)
+
+
+
+def load_state_dict(model, state_dict, strict=True):
+    if isinstance(model, (DistributedDataParallel, DataParallel)):
+        model.module.load_state_dict(state_dict, strict=strict)
+    else:
+        model.load_state_dict(state_dict, strict=strict)
+
+
+def state_dict_remove_module(state_dict):
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        # name = k[7:]  # remove 'module.' of dataparallel
+        name = k.replace('module.', '')
+        new_state_dict[name] = v
+    return new_state_dict
+
+
+def reduce_tensor(tensor, args):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= args.world_size
+    return rt

base/base_model.py

@@ -0,0 +1,30 @@
+import torch.nn as nn
+import numpy as np
+
+class BaseModel(nn.Module):
+    """
+    Base class for all models
+    """
+
+    def __init__(self):
+        super(BaseModel, self).__init__()
+        # self.logger = logging.getLogger(self.__class__.__name__)
+
+    def forward(self, *x):
+        """
+        Forward pass logic
+
+        :return: Model output
+        """
+        raise NotImplementedError
+
+    def summary(self, logger, writer):
+        """
+        Model summary
+        """
+        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
+        params = sum([np.prod(p.size()) for p in model_parameters]) / 1e6  # Unit is Mega
+        logger.info(self)
+        logger.info('===>Trainable parameters: %.3f M' % params)
+        if writer is not None:
+            writer.add_text('Model Summary', 'Trainable parameters: %.3f M' % params)

base/config.py

@@ -0,0 +1,165 @@
+# -----------------------------------------------------------------------------
+# Functions for parsing args
+# -----------------------------------------------------------------------------
+import yaml
+import os
+from ast import literal_eval
+import copy
+
+
+class CfgNode(dict):
+    """
+    CfgNode represents an internal node in the configuration tree. It's a simple
+    dict-like container that allows for attribute-based access to keys.
+    """
+
+    def __init__(self, init_dict=None, key_list=None, new_allowed=False):
+        # Recursively convert nested dictionaries in init_dict into CfgNodes
+        init_dict = {} if init_dict is None else init_dict
+        key_list = [] if key_list is None else key_list
+        for k, v in init_dict.items():
+            if type(v) is dict:
+                # Convert dict to CfgNode
+                init_dict[k] = CfgNode(v, key_list=key_list + [k])
+        super(CfgNode, self).__init__(init_dict)
+
+    def __getattr__(self, name):
+        if name in self:
+            return self[name]
+        else:
+            raise AttributeError(name)
+
+    def __setattr__(self, name, value):
+        self[name] = value
+
+    def __str__(self):
+        def _indent(s_, num_spaces):
+            s = s_.split("\n")
+            if len(s) == 1:
+                return s_
+            first = s.pop(0)
+            s = [(num_spaces * " ") + line for line in s]
+            s = "\n".join(s)
+            s = first + "\n" + s
+            return s
+
+        r = ""
+        s = []
+        for k, v in sorted(self.items()):
+            seperator = "\n" if isinstance(v, CfgNode) else " "
+            attr_str = "{}:{}{}".format(str(k), seperator, str(v))
+            attr_str = _indent(attr_str, 2)
+            s.append(attr_str)
+        r += "\n".join(s)
+        return r
+
+    def __repr__(self):
+        return "{}({})".format(self.__class__.__name__, super(CfgNode, self).__repr__())
+
+
+def load_cfg_from_cfg_file(file):
+    cfg = {}
+    assert os.path.isfile(file) and file.endswith('.yaml'), \
+        '{} is not a yaml file'.format(file)
+
+    with open(file, 'r') as f:
+        cfg_from_file = yaml.safe_load(f)
+
+    for key in cfg_from_file:
+        for k, v in cfg_from_file[key].items():
+            cfg[k] = v
+
+    cfg = CfgNode(cfg)
+    return cfg
+
+
+def merge_cfg_from_list(cfg, cfg_list):
+    new_cfg = copy.deepcopy(cfg)
+    assert len(cfg_list) % 2 == 0
+    for full_key, v in zip(cfg_list[0::2], cfg_list[1::2]):
+        subkey = full_key.split('.')[-1]
+        assert subkey in cfg, 'Non-existent key: {}'.format(full_key)
+        value = _decode_cfg_value(v)
+        value = _check_and_coerce_cfg_value_type(
+            value, cfg[subkey], subkey, full_key
+        )
+        setattr(new_cfg, subkey, value)
+
+    return new_cfg
+
+
+def _decode_cfg_value(v):
+    """Decodes a raw config value (e.g., from a yaml config files or command
+    line argument) into a Python object.
+    """
+    # All remaining processing is only applied to strings
+    if not isinstance(v, str):
+        return v
+    # Try to interpret `v` as a:
+    #   string, number, tuple, list, dict, boolean, or None
+    try:
+        v = literal_eval(v)
+    # The following two excepts allow v to pass through when it represents a
+    # string.
+    #
+    # Longer explanation:
+    # The type of v is always a string (before calling literal_eval), but
+    # sometimes it *represents* a string and other times a data structure, like
+    # a list. In the case that v represents a string, what we got back from the
+    # yaml parser is 'foo' *without quotes* (so, not '"foo"'). literal_eval is
+    # ok with '"foo"', but will raise a ValueError if given 'foo'. In other
+    # cases, like paths (v = 'foo/bar' and not v = '"foo/bar"'), literal_eval
+    # will raise a SyntaxError.
+    except ValueError:
+        pass
+    except SyntaxError:
+        pass
+    return v
+
+
+def _check_and_coerce_cfg_value_type(replacement, original, key, full_key):
+    """Checks that `replacement`, which is intended to replace `original` is of
+    the right type. The type is correct if it matches exactly or is one of a few
+    cases in which the type can be easily coerced.
+    """
+    original_type = type(original)
+    replacement_type = type(replacement)
+
+    # The types must match (with some exceptions)
+    if replacement_type == original_type or original is None:
+        return replacement
+
+    # Cast replacement from from_type to to_type if the replacement and original
+    # types match from_type and to_type
+    def conditional_cast(from_type, to_type):
+        if replacement_type == from_type and original_type == to_type:
+            return True, to_type(replacement)
+        else:
+            return False, None
+
+    # Conditionally casts
+    # list <-> tuple
+    casts = [(tuple, list), (list, tuple)]
+    # For py2: allow converting from str (bytes) to a unicode string
+    try:
+        casts.append((str, unicode))  # noqa: F821
+    except Exception:
+        pass
+
+    for (from_type, to_type) in casts:
+        converted, converted_value = conditional_cast(from_type, to_type)
+        if converted:
+            return converted_value
+
+    raise ValueError(
+        "Type mismatch ({} vs. {}) with values ({} vs. {}) for config "
+        "key: {}".format(
+            original_type, replacement_type, original, replacement, full_key
+        )
+    )
+
+
+def _assert_with_logging(cond, msg):
+    if not cond:
+        logger.debug(msg)
+    assert cond, msg

base/utilities.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+import argparse
+import os
+import random
+import time
+import logging
+import numpy as np
+from base import config
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description=' ')
+    parser.add_argument('--config', type=str, default='**.yaml', help='config file')
+    parser.add_argument('opts', help=' ', default=None,
+                        nargs=argparse.REMAINDER)
+    args = parser.parse_args()
+    assert args.config is not None
+    cfg = config.load_cfg_from_cfg_file(args.config)
+    if args.opts is not None:
+        cfg = config.merge_cfg_from_list(cfg, args.opts)
+    return cfg
+
+
+def get_logger():
+    logger_name = "main-logger"
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler()
+    fmt = "[%(asctime)s %(levelname)s %(filename)s line %(lineno)d %(process)d]=>%(message)s"
+    handler.setFormatter(logging.Formatter(fmt))
+    logger.addHandler(handler)
+    return logger
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def check_mkdir(dir_name):
+    if not os.path.exists(dir_name):
+        os.mkdir(dir_name)
+
+
+def check_makedirs(dir_name):
+    if not os.path.exists(dir_name):
+        os.makedirs(dir_name)
+
+
+def main_process(args):
+    return not args.multiprocessing_distributed or (
+            args.multiprocessing_distributed and args.rank % args.ngpus_per_node == 0)

checkpoints/stage1.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:204f79b355080035e05ef25b2a69d46a9404aa66ed70c9fc0064eeac4ad95fa2
+size 567634587

checkpoints/stage2.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0777b00137c88384d72c48f58b5f5e929aca7fdbee3c32e08b45aaa099379693
+size 1703000029

checkpoints/templates.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c9fff77287b68808ce18a76e31dfc2d793a6746abbb7620eb4424d1d73919fa
+size 60437

config/multi/demo.yaml

@@ -0,0 +1,47 @@
+DATA:
+  dataset: multi
+  data_root: sample_dataset
+  wav_path: wav
+  vertices_path: npy
+  template_file: templates.pkl
+  train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+
+NETWORK:
+  arch: stage2
+  in_dim: 15069
+  hidden_size: 1024
+  num_hidden_layers: 6
+  num_attention_heads: 8
+  intermediate_size: 1536
+  window_size: 1
+  quant_factor: 0
+  face_quan_num: 16
+  neg: 0.2
+  autoencoder: stage1_vocaset
+  INaffine: False
+  style_emb_method: nnemb # onehot or nnemb
+
+VQuantizer:
+  n_embed: 256
+  zquant_dim: 64
+
+PREDICTOR:
+  feature_dim: 1024
+  vertice_dim: 15069
+  device: cuda
+  period: 25
+  vqvae_pretrained_path: checkpoints/stage1.pth.tar
+  wav2vec2model_path: facebook/wav2vec2-large-xlsr-53
+  teacher_forcing: True
+  num_layers: 6
+  n_head: 4 # not used
+
+DEMO:
+  model_path: checkpoints/stage2.pth.tar
+  #condition: False #if false, the waveform file has the cue for the type of language
+  condition: English
+  subject: id
+  demo_wav_dir_path: demo/input/
+  demo_output_path: demo/output/
+  fps: 25
+  background_black: True # chose the background color of your rendered video

config/multi/stage1.yaml

@@ -0,0 +1,79 @@
+DATA:
+  dataset: multi
+  data_root: sample_dataset
+  wav_path: wav
+  vertices_path: npy
+  template_file: templates.pkl
+  read_audio: False
+  train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+  val_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+  test_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+
+
+LOSS:
+  quant_loss_weight: 1.0
+
+NETWORK:
+  arch: stage1_vocaset
+  in_dim: 15069
+  hidden_size: 1024
+  num_hidden_layers: 6
+  num_attention_heads: 8
+  intermediate_size: 1536
+  window_size: 1
+  quant_factor: 0
+  face_quan_num: 16
+  neg: 0.2
+  INaffine: False
+
+VQuantizer:
+  n_embed: 256
+  zquant_dim: 64
+
+TRAIN:
+  use_sgd: False
+  sync_bn: False  # adopt sync_bn or not
+  train_gpu: [0]
+  workers: 10  # data loader workers
+  batch_size: 1  # batch size for training
+  batch_size_val: 1  # batch size for validation during training, memory and speed tradeoff
+  base_lr: 0.0001
+  StepLR: True
+  warmup_steps: 1
+  adaptive_lr: False
+  factor: 0.3
+  patience: 3
+  threshold: 0.0001
+  poly_lr: False
+  epochs: 200
+  step_size: 20
+  gamma: 0.5
+  start_epoch: 0
+  power: 0.9
+  momentum: 0.9
+  weight_decay: 0.002
+  manual_seed: 131
+  print_freq: 10
+  save_freq: 1
+  save_path:
+#  weight:
+  weight:
+  resume:
+  evaluate: True  # evaluate on validation set, extra gpu memory needed and small batch_size_val is recommend
+  eval_freq: 10
+
+Distributed:
+  dist_url: tcp://127.0.0.1:6701
+  dist_backend: 'nccl'
+  multiprocessing_distributed: True
+  world_size: 1
+  rank: 0
+
+
+TEST:
+  test_workers: 0
+  test_gpu: [0]
+  test_batch_size: 1
+  save: True
+  model_path:
+  save_folder:

config/multi/stage2.yaml

@@ -0,0 +1,97 @@
+DATA:
+  dataset: multi
+  data_root: sample_dataset
+  wav_path: wav
+  vertices_path: npy
+  template_file: templates.pkl
+  read_audio: True
+  train_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+  val_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+  test_subjects: Arabic English French German Greek Italian Portuguese Russian Spanish Korean Mandarin Japanese
+  log_dir:
+
+LOSS:
+  loss: MSE
+  motion_weight: 1.0
+  reg_weight: 1.0
+
+
+NETWORK:
+  arch: stage2
+  in_dim: 15069
+  hidden_size: 1024
+  num_hidden_layers: 6
+  num_attention_heads: 8
+  intermediate_size: 1536
+  window_size: 1
+  quant_factor: 0
+  face_quan_num: 16
+  neg: 0.2
+  autoencoder: stage1_vocaset
+  INaffine: False
+  style_emb_method: nnemb # onehot or nnemb
+
+VQuantizer:
+  n_embed: 256
+  zquant_dim: 64
+
+PREDICTOR:
+  feature_dim: 1024
+  vertice_dim: 15069
+  device: cuda
+  period: 25
+  vqvae_pretrained_path: checkpoints/stage1.pth.tar
+  wav2vec2model_path: facebook/wav2vec2-large-xlsr-53
+  teacher_forcing: True
+  num_layers: 6
+  n_head: 4 # not used
+
+TRAIN:
+  use_sgd: False
+  sync_bn: False  # adopt sync_bn or not
+  train_gpu: [0]
+  workers: 10  # data loader workers
+  batch_size: 1  # batch size for training
+  batch_size_val: 1  # batch size for validation during training, memory and speed tradeoff
+  base_lr: 0.0001
+  StepLR: False
+  warmup_steps: 1
+  adaptive_lr: False
+  factor: 0.3
+  patience: 3
+  threshold: 0.0001
+  poly_lr: False
+  epochs: 100
+  step_size: 100
+  gamma: 0.5
+  start_epoch: 0
+  power: 0.9
+  momentum: 0.9
+  weight_decay: 0.002
+  manual_seed: 131
+  print_freq: 10
+  save_freq: 1
+  save_path:
+#  weight:
+  weight:
+  resume:
+  evaluate: True  # evaluate on validation set, extra gpu memory needed and small batch_size_val is recommend
+  eval_freq: 5
+
+Distributed:
+  dist_url: tcp://127.0.0.1:6701
+  dist_backend: 'nccl'
+  multiprocessing_distributed: True
+  world_size: 1
+  rank: 0
+
+TEST:
+  test_workers: 0
+  test_gpu: [0]
+  test_batch_size: 1
+  save: True
+  model_path: checkpoints/stage2.pth.tar
+  save_folder: demo/output
+  gt_save_folder: demo/gt
+  measure_lve : False
+  visualize_mesh : True