Upload 10 files

README.md

@@ -1,12 +1,166 @@
----
-title: Wav2lip1
-emoji: 💻
-colorFrom: green
-colorTo: purple
-sdk: gradio
-sdk_version: 4.39.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# **Wav2Lip**: *Accurately Lip-syncing Videos In The Wild* 
+
+### Wav2Lip is hosted for free at [Sync Labs](https://synclabs.so/)
+
+Are you looking to integrate this into a product? We have a turn-key hosted API with new and improved lip-syncing models here: https://synclabs.so/
+
+For any other commercial / enterprise requests, please contact us at pavan@synclabs.so and prady@synclabs.so
+
+To reach out to the authors directly you can reach us at [email protected], [email protected].
+
+This code is part of the paper: _A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild_ published at ACM Multimedia 2020. 
+
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/a-lip-sync-expert-is-all-you-need-for-speech/lip-sync-on-lrs2)](https://paperswithcode.com/sota/lip-sync-on-lrs2?p=a-lip-sync-expert-is-all-you-need-for-speech)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/a-lip-sync-expert-is-all-you-need-for-speech/lip-sync-on-lrs3)](https://paperswithcode.com/sota/lip-sync-on-lrs3?p=a-lip-sync-expert-is-all-you-need-for-speech)
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/a-lip-sync-expert-is-all-you-need-for-speech/lip-sync-on-lrw)](https://paperswithcode.com/sota/lip-sync-on-lrw?p=a-lip-sync-expert-is-all-you-need-for-speech)
+
+|📑 Original Paper|📰 Project Page|🌀 Demo|⚡ Live Testing|📔 Colab Notebook
+|:-:|:-:|:-:|:-:|:-:|
+[Paper](http://arxiv.org/abs/2008.10010) | [Project Page](http://cvit.iiit.ac.in/research/projects/cvit-projects/a-lip-sync-expert-is-all-you-need-for-speech-to-lip-generation-in-the-wild/) | [Demo Video](https://youtu.be/0fXaDCZNOJc) | [Interactive Demo](https://synclabs.so/) | [Colab Notebook](https://colab.research.google.com/drive/1tZpDWXz49W6wDcTprANRGLo2D_EbD5J8?usp=sharing) /[Updated Collab Notebook](https://colab.research.google.com/drive/1IjFW1cLevs6Ouyu4Yht4mnR4yeuMqO7Y#scrollTo=MH1m608OymLH)
+ 
+![Logo](https://drive.google.com/uc?export=view&id=1Wn0hPmpo4GRbCIJR8Tf20Akzdi1qjjG9)
+
+----------
+**Highlights**
+----------
+ - Weights of the visual quality disc has been updated in readme!
+ - Lip-sync videos to any target speech with high accuracy :100:. Try our [interactive demo](https://synclabs.so/).
+ - :sparkles: Works for any identity, voice, and language. Also works for CGI faces and synthetic voices.
+ - Complete training code, inference code, and pretrained models are available :boom:
+ - Or, quick-start with the Google Colab Notebook: [Link](https://colab.research.google.com/drive/1tZpDWXz49W6wDcTprANRGLo2D_EbD5J8?usp=sharing). Checkpoints and samples are available in a Google Drive [folder](https://drive.google.com/drive/folders/1I-0dNLfFOSFwrfqjNa-SXuwaURHE5K4k?usp=sharing) as well. There is also a [tutorial video](https://www.youtube.com/watch?v=Ic0TBhfuOrA) on this, courtesy of [What Make Art](https://www.youtube.com/channel/UCmGXH-jy0o2CuhqtpxbaQgA). Also, thanks to [Eyal Gruss](https://eyalgruss.com), there is a more accessible [Google Colab notebook](https://j.mp/wav2lip) with more useful features. A tutorial collab notebook is present at this [link](https://colab.research.google.com/drive/1IjFW1cLevs6Ouyu4Yht4mnR4yeuMqO7Y#scrollTo=MH1m608OymLH).  
+ - :fire: :fire: Several new, reliable evaluation benchmarks and metrics [[`evaluation/` folder of this repo]](https://github.com/Rudrabha/Wav2Lip/tree/master/evaluation) released. Instructions to calculate the metrics reported in the paper are also present.
+
+--------
+**Disclaimer**
+--------
+All results from this open-source code or our [demo website](https://bhaasha.iiit.ac.in/lipsync) should only be used for research/academic/personal purposes only. As the models are trained on the <a href="http://www.robots.ox.ac.uk/~vgg/data/lip_reading/lrs2.html">LRS2 dataset</a>, any form of commercial use is strictly prohibited. For commercial requests please contact us directly!
+
+Prerequisites
+-------------
+- `Python 3.6` 
+- ffmpeg: `sudo apt-get install ffmpeg`
+- Install necessary packages using `pip install -r requirements.txt`. Alternatively, instructions for using a docker image is provided [here](https://gist.github.com/xenogenesi/e62d3d13dadbc164124c830e9c453668). Have a look at [this comment](https://github.com/Rudrabha/Wav2Lip/issues/131#issuecomment-725478562) and comment on [the gist](https://gist.github.com/xenogenesi/e62d3d13dadbc164124c830e9c453668) if you encounter any issues. 
+- Face detection [pre-trained model](https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth) should be downloaded to `face_detection/detection/sfd/s3fd.pth`. Alternative [link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/prajwal_k_research_iiit_ac_in/EZsy6qWuivtDnANIG73iHjIBjMSoojcIV0NULXV-yiuiIg?e=qTasa8) if the above does not work.
+
+Getting the weights
+----------
+| Model  | Description |  Link to the model | 
+| :-------------: | :---------------: | :---------------: |
+| Wav2Lip  | Highly accurate lip-sync | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/Eb3LEzbfuKlJiR600lQWRxgBIY27JZg80f7V9jtMfbNDaQ?e=TBFBVW)  |
+| Wav2Lip + GAN  | Slightly inferior lip-sync, but better visual quality | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EdjI7bZlgApMqsVoEUUXpLsBxqXbn5z8VTmoxp55YNDcIA?e=n9ljGW) |
+| Expert Discriminator  | Weights of the expert discriminator | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EQRvmiZg-HRAjvI6zqN9eTEBP74KefynCwPWVmF57l-AYA?e=ZRPHKP) |
+| Visual Quality Discriminator  | Weights of the visual disc trained in a GAN setup | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EQVqH88dTm1HjlK11eNba5gBbn15WMS0B0EZbDBttqrqkg?e=ic0ljo) |
+
+Lip-syncing videos using the pre-trained models (Inference)
+-------
+You can lip-sync any video to any audio:
+```bash
+python inference.py --checkpoint_path <ckpt> --face <video.mp4> --audio <an-audio-source> 
+```
+The result is saved (by default) in `results/result_voice.mp4`. You can specify it as an argument,  similar to several other available options. The audio source can be any file supported by `FFMPEG` containing audio data: `*.wav`, `*.mp3` or even a video file, from which the code will automatically extract the audio.
+
+##### Tips for better results:
+- Experiment with the `--pads` argument to adjust the detected face bounding box. Often leads to improved results. You might need to increase the bottom padding to include the chin region. E.g. `--pads 0 20 0 0`.
+- If you see the mouth position dislocated or some weird artifacts such as two mouths, then it can be because of over-smoothing the face detections. Use the `--nosmooth` argument and give it another try. 
+- Experiment with the `--resize_factor` argument, to get a lower-resolution video. Why? The models are trained on faces that were at a lower resolution. You might get better, visually pleasing results for 720p videos than for 1080p videos (in many cases, the latter works well too). 
+- The Wav2Lip model without GAN usually needs more experimenting with the above two to get the most ideal results, and sometimes, can give you a better result as well.
+
+Preparing LRS2 for training
+----------
+Our models are trained on LRS2. See [here](#training-on-datasets-other-than-lrs2) for a few suggestions regarding training on other datasets.
+##### LRS2 dataset folder structure
+
+```
+data_root (mvlrs_v1)
+├── main, pretrain (we use only main folder in this work)
+|	├── list of folders
+|	│   ├── five-digit numbered video IDs ending with (.mp4)
+```
+
+Place the LRS2 filelists (train, val, test) `.txt` files in the `filelists/` folder.
+
+##### Preprocess the dataset for fast training
+
+```bash
+python preprocess.py --data_root data_root/main --preprocessed_root lrs2_preprocessed/
+```
+Additional options like `batch_size` and the number of GPUs to use in parallel to use can also be set.
+
+##### Preprocessed LRS2 folder structure
+```
+preprocessed_root (lrs2_preprocessed)
+├── list of folders
+|	├── Folders with five-digit numbered video IDs
+|	│   ├── *.jpg
+|	│   ├── audio.wav
+```
+
+Train!
+----------
+There are two major steps: (i) Train the expert lip-sync discriminator, (ii) Train the Wav2Lip model(s).
+
+##### Training the expert discriminator
+You can download [the pre-trained weights](#getting-the-weights) if you want to skip this step. To train it:
+```bash
+python color_syncnet_train.py --data_root lrs2_preprocessed/ --checkpoint_dir <folder_to_save_checkpoints>
+```
+##### Training the Wav2Lip models
+You can either train the model without the additional visual quality discriminator (< 1 day of training) or use the discriminator (~2 days). For the former, run: 
+```bash
+python wav2lip_train.py --data_root lrs2_preprocessed/ --checkpoint_dir <folder_to_save_checkpoints> --syncnet_checkpoint_path <path_to_expert_disc_checkpoint>
+```
+
+To train with the visual quality discriminator, you should run `hq_wav2lip_train.py` instead. The arguments for both files are similar. In both cases, you can resume training as well. Look at `python wav2lip_train.py --help` for more details. You can also set additional less commonly-used hyper-parameters at the bottom of the `hparams.py` file.
+
+Training on datasets other than LRS2
+------------------------------------
+Training on other datasets might require modifications to the code. Please read the following before you raise an issue:
+
+- You might not get good results by training/fine-tuning on a few minutes of a single speaker. This is a separate research problem, to which we do not have a solution yet. Thus, we would most likely not be able to resolve your issue. 
+- You must train the expert discriminator for your own dataset before training Wav2Lip.
+- If it is your own dataset downloaded from the web, in most cases, needs to be sync-corrected.
+- Be mindful of the FPS of the videos of your dataset. Changes to FPS would need significant code changes. 
+- The expert discriminator's eval loss should go down to ~0.25 and the Wav2Lip eval sync loss should go down to ~0.2 to get good results. 
+
+When raising an issue on this topic, please let us know that you are aware of all these points.
+
+We have an HD model trained on a dataset allowing commercial usage. The size of the generated face will be 192 x 288 in our new model.
+
+Evaluation
+----------
+Please check the `evaluation/` folder for the instructions.
+
+License and Citation
+----------
+This repository can only be used for personal/research/non-commercial purposes. However, for commercial requests, please contact us directly at [email protected] or [email protected]. We have a turn-key hosted API with new and improved lip-syncing models here: https://synclabs.so/
+The size of the generated face will be 192 x 288 in our new models. Please cite the following paper if you use this repository:
+```
+@inproceedings{10.1145/3394171.3413532,
+author = {Prajwal, K R and Mukhopadhyay, Rudrabha and Namboodiri, Vinay P. and Jawahar, C.V.},
+title = {A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild},
+year = {2020},
+isbn = {9781450379885},
+publisher = {Association for Computing Machinery},
+address = {New York, NY, USA},
+url = {https://doi.org/10.1145/3394171.3413532},
+doi = {10.1145/3394171.3413532},
+booktitle = {Proceedings of the 28th ACM International Conference on Multimedia},
+pages = {484–492},
+numpages = {9},
+keywords = {lip sync, talking face generation, video generation},
+location = {Seattle, WA, USA},
+series = {MM '20}
+}
+```
+
+
+Acknowledgments
+----------
+Parts of the code structure are inspired by this [TTS repository](https://github.com/r9y9/deepvoice3_pytorch). We thank the author for this wonderful code. The code for Face Detection has been taken from the [face_alignment](https://github.com/1adrianb/face-alignment) repository. We thank the authors for releasing their code and models. We thank [zabique](https://github.com/zabique) for the tutorial collab notebook.
+
+## Acknowledgements
+
+ - [Awesome Readme Templates](https://awesomeopensource.com/project/elangosundar/awesome-README-templates)
+ - [Awesome README](https://github.com/matiassingers/awesome-readme)
+ - [How to write a Good readme](https://bulldogjob.com/news/449-how-to-write-a-good-readme-for-your-github-project)
+

app.py

@@ -0,0 +1,48 @@
+
+import gradio as gr
+import subprocess
+from subprocess import call
+
+with gr.Blocks() as ui:
+    with gr.Row():
+        video = gr.File(label="Video or Image")
+        audio = gr.File(label="Audio")
+        with gr.Column():
+            checkpoint = gr.Radio(["wav2lip", "wav2lip_gan"], label="Checkpoint")
+            no_smooth = gr.Checkbox(label="No Smooth")
+            resize_factor = gr.Slider(minimum=1, maximum=4, step=1, label="Resize Factor")
+    with gr.Row():
+        with gr.Column():
+            pad_top = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Pad Top")
+            pad_bottom = gr.Slider(minimum=0, maximum=50, step=1, value=10, label="Pad Bottom (Often increasing this to 20 allows chin to be included)")
+            pad_left = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Pad Left")
+            pad_right = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Pad Right")
+            generate_btn = gr.Button("Generate")
+        with gr.Column():
+           result = gr.Video()
+
+    def generate(video, audio, checkpoint, no_smooth, resize_factor, pad_top, pad_bottom, pad_left, pad_right):
+        if video is None or audio is None or checkpoint is None:
+            return
+        
+        smooth = "--nosmooth" if no_smooth else ""
+        
+        
+        cmd = [
+            "python",
+            "inference.py",
+            "--checkpoint_path", f"checkpoints/{checkpoint}.pth",
+            "--face", video.name,
+            "--audio", audio.name,
+            "--outfile", "results/output.mp4",
+        ]
+
+        call(cmd)
+        return "results/output.mp4"
+
+    generate_btn.click(
+        generate, 
+        [video, audio, checkpoint, pad_top, pad_bottom, pad_left, pad_right, resize_factor], 
+        result)
+    
+ui.queue().launch(share=True,debug=True)

audio.py

@@ -0,0 +1,136 @@
+import librosa
+import librosa.filters
+import numpy as np
+# import tensorflow as tf
+from scipy import signal
+from scipy.io import wavfile
+from hparams import hparams as hp
+
+def load_wav(path, sr):
+    return librosa.core.load(path, sr=sr)[0]
+
+def save_wav(wav, path, sr):
+    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
+    #proposed by @dsmiller
+    wavfile.write(path, sr, wav.astype(np.int16))
+
+def save_wavenet_wav(wav, path, sr):
+    librosa.output.write_wav(path, wav, sr=sr)
+
+def preemphasis(wav, k, preemphasize=True):
+    if preemphasize:
+        return signal.lfilter([1, -k], [1], wav)
+    return wav
+
+def inv_preemphasis(wav, k, inv_preemphasize=True):
+    if inv_preemphasize:
+        return signal.lfilter([1], [1, -k], wav)
+    return wav
+
+def get_hop_size():
+    hop_size = hp.hop_size
+    if hop_size is None:
+        assert hp.frame_shift_ms is not None
+        hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate)
+    return hop_size
+
+def linearspectrogram(wav):
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(np.abs(D)) - hp.ref_level_db
+    
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+def melspectrogram(wav):
+    D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize))
+    S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db
+    
+    if hp.signal_normalization:
+        return _normalize(S)
+    return S
+
+def _lws_processor():
+    import lws
+    return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech")
+
+def _stft(y):
+    if hp.use_lws:
+        return _lws_processor(hp).stft(y).T
+    else:
+        return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=get_hop_size(), win_length=hp.win_size)
+
+##########################################################
+#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
+def num_frames(length, fsize, fshift):
+    """Compute number of time frames of spectrogram
+    """
+    pad = (fsize - fshift)
+    if length % fshift == 0:
+        M = (length + pad * 2 - fsize) // fshift + 1
+    else:
+        M = (length + pad * 2 - fsize) // fshift + 2
+    return M
+
+
+def pad_lr(x, fsize, fshift):
+    """Compute left and right padding
+    """
+    M = num_frames(len(x), fsize, fshift)
+    pad = (fsize - fshift)
+    T = len(x) + 2 * pad
+    r = (M - 1) * fshift + fsize - T
+    return pad, pad + r
+##########################################################
+#Librosa correct padding
+def librosa_pad_lr(x, fsize, fshift):
+    return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
+
+# Conversions
+_mel_basis = None
+
+def _linear_to_mel(spectogram):
+    global _mel_basis
+    if _mel_basis is None:
+        _mel_basis = _build_mel_basis()
+    return np.dot(_mel_basis, spectogram)
+    
+def _build_mel_basis():
+    assert hp.fmax <= hp.sample_rate // 2
+    return librosa.filters.mel(sr=hp.sample_rate, n_fft= hp.n_fft, n_mels=hp.num_mels,
+                               fmin=hp.fmin, fmax=hp.fmax)
+
+def _amp_to_db(x):
+    min_level = np.exp(hp.min_level_db / 20 * np.log(10))
+    return 20 * np.log10(np.maximum(min_level, x))
+
+def _db_to_amp(x):
+    return np.power(10.0, (x) * 0.05)
+
+def _normalize(S):
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return np.clip((2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value,
+                           -hp.max_abs_value, hp.max_abs_value)
+        else:
+            return np.clip(hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)), 0, hp.max_abs_value)
+    
+    assert S.max() <= 0 and S.min() - hp.min_level_db >= 0
+    if hp.symmetric_mels:
+        return (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value
+    else:
+        return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db))
+
+def _denormalize(D):
+    if hp.allow_clipping_in_normalization:
+        if hp.symmetric_mels:
+            return (((np.clip(D, -hp.max_abs_value,
+                              hp.max_abs_value) + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value))
+                    + hp.min_level_db)
+        else:
+            return ((np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)
+    
+    if hp.symmetric_mels:
+        return (((D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + hp.min_level_db)
+    else:
+        return ((D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db)

color_syncnet_train.py

@@ -0,0 +1,279 @@
+from os.path import dirname, join, basename, isfile
+from tqdm import tqdm
+
+from models import SyncNet_color as SyncNet
+import audio
+
+import torch
+from torch import nn
+from torch import optim
+import torch.backends.cudnn as cudnn
+from torch.utils import data as data_utils
+import numpy as np
+
+from glob import glob
+
+import os, random, cv2, argparse
+from hparams import hparams, get_image_list
+
+parser = argparse.ArgumentParser(description='Code to train the expert lip-sync discriminator')
+
+parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True)
+
+parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str)
+parser.add_argument('--checkpoint_path', help='Resumed from this checkpoint', default=None, type=str)
+
+args = parser.parse_args()
+
+
+global_step = 0
+global_epoch = 0
+use_cuda = torch.cuda.is_available()
+print('use_cuda: {}'.format(use_cuda))
+
+syncnet_T = 5
+syncnet_mel_step_size = 16
+
+class Dataset(object):
+    def __init__(self, split):
+        self.all_videos = get_image_list(args.data_root, split)
+
+    def get_frame_id(self, frame):
+        return int(basename(frame).split('.')[0])
+
+    def get_window(self, start_frame):
+        start_id = self.get_frame_id(start_frame)
+        vidname = dirname(start_frame)
+
+        window_fnames = []
+        for frame_id in range(start_id, start_id + syncnet_T):
+            frame = join(vidname, '{}.jpg'.format(frame_id))
+            if not isfile(frame):
+                return None
+            window_fnames.append(frame)
+        return window_fnames
+
+    def crop_audio_window(self, spec, start_frame):
+        # num_frames = (T x hop_size * fps) / sample_rate
+        start_frame_num = self.get_frame_id(start_frame)
+        start_idx = int(80. * (start_frame_num / float(hparams.fps)))
+
+        end_idx = start_idx + syncnet_mel_step_size
+
+        return spec[start_idx : end_idx, :]
+
+
+    def __len__(self):
+        return len(self.all_videos)
+
+    def __getitem__(self, idx):
+        while 1:
+            idx = random.randint(0, len(self.all_videos) - 1)
+            vidname = self.all_videos[idx]
+
+            img_names = list(glob(join(vidname, '*.jpg')))
+            if len(img_names) <= 3 * syncnet_T:
+                continue
+            img_name = random.choice(img_names)
+            wrong_img_name = random.choice(img_names)
+            while wrong_img_name == img_name:
+                wrong_img_name = random.choice(img_names)
+
+            if random.choice([True, False]):
+                y = torch.ones(1).float()
+                chosen = img_name
+            else:
+                y = torch.zeros(1).float()
+                chosen = wrong_img_name
+
+            window_fnames = self.get_window(chosen)
+            if window_fnames is None:
+                continue
+
+            window = []
+            all_read = True
+            for fname in window_fnames:
+                img = cv2.imread(fname)
+                if img is None:
+                    all_read = False
+                    break
+                try:
+                    img = cv2.resize(img, (hparams.img_size, hparams.img_size))
+                except Exception as e:
+                    all_read = False
+                    break
+
+                window.append(img)
+
+            if not all_read: continue
+
+            try:
+                wavpath = join(vidname, "audio.wav")
+                wav = audio.load_wav(wavpath, hparams.sample_rate)
+
+                orig_mel = audio.melspectrogram(wav).T
+            except Exception as e:
+                continue
+
+            mel = self.crop_audio_window(orig_mel.copy(), img_name)
+
+            if (mel.shape[0] != syncnet_mel_step_size):
+                continue
+
+            # H x W x 3 * T
+            x = np.concatenate(window, axis=2) / 255.
+            x = x.transpose(2, 0, 1)
+            x = x[:, x.shape[1]//2:]
+
+            x = torch.FloatTensor(x)
+            mel = torch.FloatTensor(mel.T).unsqueeze(0)
+
+            return x, mel, y
+
+logloss = nn.BCELoss()
+def cosine_loss(a, v, y):
+    d = nn.functional.cosine_similarity(a, v)
+    loss = logloss(d.unsqueeze(1), y)
+
+    return loss
+
+def train(device, model, train_data_loader, test_data_loader, optimizer,
+          checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
+
+    global global_step, global_epoch
+    resumed_step = global_step
+    
+    while global_epoch < nepochs:
+        running_loss = 0.
+        prog_bar = tqdm(enumerate(train_data_loader))
+        for step, (x, mel, y) in prog_bar:
+            model.train()
+            optimizer.zero_grad()
+
+            # Transform data to CUDA device
+            x = x.to(device)
+
+            mel = mel.to(device)
+
+            a, v = model(mel, x)
+            y = y.to(device)
+
+            loss = cosine_loss(a, v, y)
+            loss.backward()
+            optimizer.step()
+
+            global_step += 1
+            cur_session_steps = global_step - resumed_step
+            running_loss += loss.item()
+
+            if global_step == 1 or global_step % checkpoint_interval == 0:
+                save_checkpoint(
+                    model, optimizer, global_step, checkpoint_dir, global_epoch)
+
+            if global_step % hparams.syncnet_eval_interval == 0:
+                with torch.no_grad():
+                    eval_model(test_data_loader, global_step, device, model, checkpoint_dir)
+
+            prog_bar.set_description('Loss: {}'.format(running_loss / (step + 1)))
+
+        global_epoch += 1
+
+def eval_model(test_data_loader, global_step, device, model, checkpoint_dir):
+    eval_steps = 1400
+    print('Evaluating for {} steps'.format(eval_steps))
+    losses = []
+    while 1:
+        for step, (x, mel, y) in enumerate(test_data_loader):
+
+            model.eval()
+
+            # Transform data to CUDA device
+            x = x.to(device)
+
+            mel = mel.to(device)
+
+            a, v = model(mel, x)
+            y = y.to(device)
+
+            loss = cosine_loss(a, v, y)
+            losses.append(loss.item())
+
+            if step > eval_steps: break
+
+        averaged_loss = sum(losses) / len(losses)
+        print(averaged_loss)
+
+        return
+
+def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch):
+
+    checkpoint_path = join(
+        checkpoint_dir, "checkpoint_step{:09d}.pth".format(global_step))
+    optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None
+    torch.save({
+        "state_dict": model.state_dict(),
+        "optimizer": optimizer_state,
+        "global_step": step,
+        "global_epoch": epoch,
+    }, checkpoint_path)
+    print("Saved checkpoint:", checkpoint_path)
+
+def _load(checkpoint_path):
+    if use_cuda:
+        checkpoint = torch.load(checkpoint_path)
+    else:
+        checkpoint = torch.load(checkpoint_path,
+                                map_location=lambda storage, loc: storage)
+    return checkpoint
+
+def load_checkpoint(path, model, optimizer, reset_optimizer=False):
+    global global_step
+    global global_epoch
+
+    print("Load checkpoint from: {}".format(path))
+    checkpoint = _load(path)
+    model.load_state_dict(checkpoint["state_dict"])
+    if not reset_optimizer:
+        optimizer_state = checkpoint["optimizer"]
+        if optimizer_state is not None:
+            print("Load optimizer state from {}".format(path))
+            optimizer.load_state_dict(checkpoint["optimizer"])
+    global_step = checkpoint["global_step"]
+    global_epoch = checkpoint["global_epoch"]
+
+    return model
+
+if __name__ == "__main__":
+    checkpoint_dir = args.checkpoint_dir
+    checkpoint_path = args.checkpoint_path
+
+    if not os.path.exists(checkpoint_dir): os.mkdir(checkpoint_dir)
+
+    # Dataset and Dataloader setup
+    train_dataset = Dataset('train')
+    test_dataset = Dataset('val')
+
+    train_data_loader = data_utils.DataLoader(
+        train_dataset, batch_size=hparams.syncnet_batch_size, shuffle=True,
+        num_workers=hparams.num_workers)
+
+    test_data_loader = data_utils.DataLoader(
+        test_dataset, batch_size=hparams.syncnet_batch_size,
+        num_workers=8)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    # Model
+    model = SyncNet().to(device)
+    print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad)))
+
+    optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
+                           lr=hparams.syncnet_lr)
+
+    if checkpoint_path is not None:
+        load_checkpoint(checkpoint_path, model, optimizer, reset_optimizer=False)
+
+    train(device, model, train_data_loader, test_data_loader, optimizer,
+          checkpoint_dir=checkpoint_dir,
+          checkpoint_interval=hparams.syncnet_checkpoint_interval,
+          nepochs=hparams.nepochs)

hparams.py

@@ -0,0 +1,101 @@
+from glob import glob
+import os
+
+def get_image_list(data_root, split):
+	filelist = []
+
+	with open('filelists/{}.txt'.format(split)) as f:
+		for line in f:
+			line = line.strip()
+			if ' ' in line: line = line.split()[0]
+			filelist.append(os.path.join(data_root, line))
+
+	return filelist
+
+class HParams:
+	def __init__(self, **kwargs):
+		self.data = {}
+
+		for key, value in kwargs.items():
+			self.data[key] = value
+
+	def __getattr__(self, key):
+		if key not in self.data:
+			raise AttributeError("'HParams' object has no attribute %s" % key)
+		return self.data[key]
+
+	def set_hparam(self, key, value):
+		self.data[key] = value
+
+
+# Default hyperparameters
+hparams = HParams(
+	num_mels=80,  # Number of mel-spectrogram channels and local conditioning dimensionality
+	#  network
+	rescale=True,  # Whether to rescale audio prior to preprocessing
+	rescaling_max=0.9,  # Rescaling value
+	
+	# Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
+	# It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
+	# Does not work if n_ffit is not multiple of hop_size!!
+	use_lws=False,
+	
+	n_fft=800,  # Extra window size is filled with 0 paddings to match this parameter
+	hop_size=200,  # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
+	win_size=800,  # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
+	sample_rate=16000,  # 16000Hz (corresponding to librispeech) (sox --i <filename>)
+	
+	frame_shift_ms=None,  # Can replace hop_size parameter. (Recommended: 12.5)
+	
+	# Mel and Linear spectrograms normalization/scaling and clipping
+	signal_normalization=True,
+	# Whether to normalize mel spectrograms to some predefined range (following below parameters)
+	allow_clipping_in_normalization=True,  # Only relevant if mel_normalization = True
+	symmetric_mels=True,
+	# Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2, 
+	# faster and cleaner convergence)
+	max_abs_value=4.,
+	# max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not 
+	# be too big to avoid gradient explosion, 
+	# not too small for fast convergence)
+	# Contribution by @begeekmyfriend
+	# Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude 
+	# levels. Also allows for better G&L phase reconstruction)
+	preemphasize=True,  # whether to apply filter
+	preemphasis=0.97,  # filter coefficient.
+	
+	# Limits
+	min_level_db=-100,
+	ref_level_db=20,
+	fmin=55,
+	# Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To 
+	# test depending on dataset. Pitch info: male~[65, 260], female~[100, 525])
+	fmax=7600,  # To be increased/reduced depending on data.
+
+	###################### Our training parameters #################################
+	img_size=96,
+	fps=25,
+	
+	batch_size=16,
+	initial_learning_rate=1e-4,
+	nepochs=200000000000000000,  ### ctrl + c, stop whenever eval loss is consistently greater than train loss for ~10 epochs
+	num_workers=16,
+	checkpoint_interval=3000,
+	eval_interval=3000,
+    save_optimizer_state=True,
+
+    syncnet_wt=0.0, # is initially zero, will be set automatically to 0.03 later. Leads to faster convergence. 
+	syncnet_batch_size=64,
+	syncnet_lr=1e-4,
+	syncnet_eval_interval=10000,
+	syncnet_checkpoint_interval=10000,
+
+	disc_wt=0.07,
+	disc_initial_learning_rate=1e-4,
+)
+
+
+def hparams_debug_string():
+	values = hparams.values()
+	hp = ["  %s: %s" % (name, values[name]) for name in sorted(values) if name != "sentences"]
+	return "Hyperparameters:\n" + "\n".join(hp)

hq_wav2lip_train.py

@@ -0,0 +1,443 @@
+from os.path import dirname, join, basename, isfile
+from tqdm import tqdm
+
+from models import SyncNet_color as SyncNet
+from models import Wav2Lip, Wav2Lip_disc_qual
+import audio
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch import optim
+import torch.backends.cudnn as cudnn
+from torch.utils import data as data_utils
+import numpy as np
+
+from glob import glob
+
+import os, random, cv2, argparse
+from hparams import hparams, get_image_list
+
+parser = argparse.ArgumentParser(description='Code to train the Wav2Lip model WITH the visual quality discriminator')
+
+parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True, type=str)
+
+parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str)
+parser.add_argument('--syncnet_checkpoint_path', help='Load the pre-trained Expert discriminator', required=True, type=str)
+
+parser.add_argument('--checkpoint_path', help='Resume generator from this checkpoint', default=None, type=str)
+parser.add_argument('--disc_checkpoint_path', help='Resume quality disc from this checkpoint', default=None, type=str)
+
+args = parser.parse_args()
+
+
+global_step = 0
+global_epoch = 0
+use_cuda = torch.cuda.is_available()
+print('use_cuda: {}'.format(use_cuda))
+
+syncnet_T = 5
+syncnet_mel_step_size = 16
+
+class Dataset(object):
+    def __init__(self, split):
+        self.all_videos = get_image_list(args.data_root, split)
+
+    def get_frame_id(self, frame):
+        return int(basename(frame).split('.')[0])
+
+    def get_window(self, start_frame):
+        start_id = self.get_frame_id(start_frame)
+        vidname = dirname(start_frame)
+
+        window_fnames = []
+        for frame_id in range(start_id, start_id + syncnet_T):
+            frame = join(vidname, '{}.jpg'.format(frame_id))
+            if not isfile(frame):
+                return None
+            window_fnames.append(frame)
+        return window_fnames
+
+    def read_window(self, window_fnames):
+        if window_fnames is None: return None
+        window = []
+        for fname in window_fnames:
+            img = cv2.imread(fname)
+            if img is None:
+                return None
+            try:
+                img = cv2.resize(img, (hparams.img_size, hparams.img_size))
+            except Exception as e:
+                return None
+
+            window.append(img)
+
+        return window
+
+    def crop_audio_window(self, spec, start_frame):
+        if type(start_frame) == int:
+            start_frame_num = start_frame
+        else:
+            start_frame_num = self.get_frame_id(start_frame)
+        start_idx = int(80. * (start_frame_num / float(hparams.fps)))
+        
+        end_idx = start_idx + syncnet_mel_step_size
+
+        return spec[start_idx : end_idx, :]
+
+    def get_segmented_mels(self, spec, start_frame):
+        mels = []
+        assert syncnet_T == 5
+        start_frame_num = self.get_frame_id(start_frame) + 1 # 0-indexing ---> 1-indexing
+        if start_frame_num - 2 < 0: return None
+        for i in range(start_frame_num, start_frame_num + syncnet_T):
+            m = self.crop_audio_window(spec, i - 2)
+            if m.shape[0] != syncnet_mel_step_size:
+                return None
+            mels.append(m.T)
+
+        mels = np.asarray(mels)
+
+        return mels
+
+    def prepare_window(self, window):
+        # 3 x T x H x W
+        x = np.asarray(window) / 255.
+        x = np.transpose(x, (3, 0, 1, 2))
+
+        return x
+
+    def __len__(self):
+        return len(self.all_videos)
+
+    def __getitem__(self, idx):
+        while 1:
+            idx = random.randint(0, len(self.all_videos) - 1)
+            vidname = self.all_videos[idx]
+            img_names = list(glob(join(vidname, '*.jpg')))
+            if len(img_names) <= 3 * syncnet_T:
+                continue
+            
+            img_name = random.choice(img_names)
+            wrong_img_name = random.choice(img_names)
+            while wrong_img_name == img_name:
+                wrong_img_name = random.choice(img_names)
+
+            window_fnames = self.get_window(img_name)
+            wrong_window_fnames = self.get_window(wrong_img_name)
+            if window_fnames is None or wrong_window_fnames is None:
+                continue
+
+            window = self.read_window(window_fnames)
+            if window is None:
+                continue
+
+            wrong_window = self.read_window(wrong_window_fnames)
+            if wrong_window is None:
+                continue
+
+            try:
+                wavpath = join(vidname, "audio.wav")
+                wav = audio.load_wav(wavpath, hparams.sample_rate)
+
+                orig_mel = audio.melspectrogram(wav).T
+            except Exception as e:
+                continue
+
+            mel = self.crop_audio_window(orig_mel.copy(), img_name)
+            
+            if (mel.shape[0] != syncnet_mel_step_size):
+                continue
+
+            indiv_mels = self.get_segmented_mels(orig_mel.copy(), img_name)
+            if indiv_mels is None: continue
+
+            window = self.prepare_window(window)
+            y = window.copy()
+            window[:, :, window.shape[2]//2:] = 0.
+
+            wrong_window = self.prepare_window(wrong_window)
+            x = np.concatenate([window, wrong_window], axis=0)
+
+            x = torch.FloatTensor(x)
+            mel = torch.FloatTensor(mel.T).unsqueeze(0)
+            indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1)
+            y = torch.FloatTensor(y)
+            return x, indiv_mels, mel, y
+
+def save_sample_images(x, g, gt, global_step, checkpoint_dir):
+    x = (x.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
+    g = (g.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
+    gt = (gt.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
+
+    refs, inps = x[..., 3:], x[..., :3]
+    folder = join(checkpoint_dir, "samples_step{:09d}".format(global_step))
+    if not os.path.exists(folder): os.mkdir(folder)
+    collage = np.concatenate((refs, inps, g, gt), axis=-2)
+    for batch_idx, c in enumerate(collage):
+        for t in range(len(c)):
+            cv2.imwrite('{}/{}_{}.jpg'.format(folder, batch_idx, t), c[t])
+
+logloss = nn.BCELoss()
+def cosine_loss(a, v, y):
+    d = nn.functional.cosine_similarity(a, v)
+    loss = logloss(d.unsqueeze(1), y)
+
+    return loss
+
+device = torch.device("cuda" if use_cuda else "cpu")
+syncnet = SyncNet().to(device)
+for p in syncnet.parameters():
+    p.requires_grad = False
+
+recon_loss = nn.L1Loss()
+def get_sync_loss(mel, g):
+    g = g[:, :, :, g.size(3)//2:]
+    g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1)
+    # B, 3 * T, H//2, W
+    a, v = syncnet(mel, g)
+    y = torch.ones(g.size(0), 1).float().to(device)
+    return cosine_loss(a, v, y)
+
+def train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer,
+          checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
+    global global_step, global_epoch
+    resumed_step = global_step
+
+    while global_epoch < nepochs:
+        print('Starting Epoch: {}'.format(global_epoch))
+        running_sync_loss, running_l1_loss, disc_loss, running_perceptual_loss = 0., 0., 0., 0.
+        running_disc_real_loss, running_disc_fake_loss = 0., 0.
+        prog_bar = tqdm(enumerate(train_data_loader))
+        for step, (x, indiv_mels, mel, gt) in prog_bar:
+            disc.train()
+            model.train()
+
+            x = x.to(device)
+            mel = mel.to(device)
+            indiv_mels = indiv_mels.to(device)
+            gt = gt.to(device)
+
+            ### Train generator now. Remove ALL grads. 
+            optimizer.zero_grad()
+            disc_optimizer.zero_grad()
+
+            g = model(indiv_mels, x)
+
+            if hparams.syncnet_wt > 0.:
+                sync_loss = get_sync_loss(mel, g)
+            else:
+                sync_loss = 0.
+
+            if hparams.disc_wt > 0.:
+                perceptual_loss = disc.perceptual_forward(g)
+            else:
+                perceptual_loss = 0.
+
+            l1loss = recon_loss(g, gt)
+
+            loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \
+                                    (1. - hparams.syncnet_wt - hparams.disc_wt) * l1loss
+
+            loss.backward()
+            optimizer.step()
+
+            ### Remove all gradients before Training disc
+            disc_optimizer.zero_grad()
+
+            pred = disc(gt)
+            disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device))
+            disc_real_loss.backward()
+
+            pred = disc(g.detach())
+            disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device))
+            disc_fake_loss.backward()
+
+            disc_optimizer.step()
+
+            running_disc_real_loss += disc_real_loss.item()
+            running_disc_fake_loss += disc_fake_loss.item()
+
+            if global_step % checkpoint_interval == 0:
+                save_sample_images(x, g, gt, global_step, checkpoint_dir)
+
+            # Logs
+            global_step += 1
+            cur_session_steps = global_step - resumed_step
+
+            running_l1_loss += l1loss.item()
+            if hparams.syncnet_wt > 0.:
+                running_sync_loss += sync_loss.item()
+            else:
+                running_sync_loss += 0.
+
+            if hparams.disc_wt > 0.:
+                running_perceptual_loss += perceptual_loss.item()
+            else:
+                running_perceptual_loss += 0.
+
+            if global_step == 1 or global_step % checkpoint_interval == 0:
+                save_checkpoint(
+                    model, optimizer, global_step, checkpoint_dir, global_epoch)
+                save_checkpoint(disc, disc_optimizer, global_step, checkpoint_dir, global_epoch, prefix='disc_')
+
+
+            if global_step % hparams.eval_interval == 0:
+                with torch.no_grad():
+                    average_sync_loss = eval_model(test_data_loader, global_step, device, model, disc)
+
+                    if average_sync_loss < .75:
+                        hparams.set_hparam('syncnet_wt', 0.03)
+
+            prog_bar.set_description('L1: {}, Sync: {}, Percep: {} | Fake: {}, Real: {}'.format(running_l1_loss / (step + 1),
+                                                                                        running_sync_loss / (step + 1),
+                                                                                        running_perceptual_loss / (step + 1),
+                                                                                        running_disc_fake_loss / (step + 1),
+                                                                                        running_disc_real_loss / (step + 1)))
+
+        global_epoch += 1
+
+def eval_model(test_data_loader, global_step, device, model, disc):
+    eval_steps = 300
+    print('Evaluating for {} steps'.format(eval_steps))
+    running_sync_loss, running_l1_loss, running_disc_real_loss, running_disc_fake_loss, running_perceptual_loss = [], [], [], [], []
+    while 1:
+        for step, (x, indiv_mels, mel, gt) in enumerate((test_data_loader)):
+            model.eval()
+            disc.eval()
+
+            x = x.to(device)
+            mel = mel.to(device)
+            indiv_mels = indiv_mels.to(device)
+            gt = gt.to(device)
+
+            pred = disc(gt)
+            disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device))
+
+            g = model(indiv_mels, x)
+            pred = disc(g)
+            disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device))
+
+            running_disc_real_loss.append(disc_real_loss.item())
+            running_disc_fake_loss.append(disc_fake_loss.item())
+
+            sync_loss = get_sync_loss(mel, g)
+            
+            if hparams.disc_wt > 0.:
+                perceptual_loss = disc.perceptual_forward(g)
+            else:
+                perceptual_loss = 0.
+
+            l1loss = recon_loss(g, gt)
+
+            loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \
+                                    (1. - hparams.syncnet_wt - hparams.disc_wt) * l1loss
+
+            running_l1_loss.append(l1loss.item())
+            running_sync_loss.append(sync_loss.item())
+            
+            if hparams.disc_wt > 0.:
+                running_perceptual_loss.append(perceptual_loss.item())
+            else:
+                running_perceptual_loss.append(0.)
+
+            if step > eval_steps: break
+
+        print('L1: {}, Sync: {}, Percep: {} | Fake: {}, Real: {}'.format(sum(running_l1_loss) / len(running_l1_loss),
+                                                            sum(running_sync_loss) / len(running_sync_loss),
+                                                            sum(running_perceptual_loss) / len(running_perceptual_loss),
+                                                            sum(running_disc_fake_loss) / len(running_disc_fake_loss),
+                                                             sum(running_disc_real_loss) / len(running_disc_real_loss)))
+        return sum(running_sync_loss) / len(running_sync_loss)
+
+
+def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch, prefix=''):
+    checkpoint_path = join(
+        checkpoint_dir, "{}checkpoint_step{:09d}.pth".format(prefix, global_step))
+    optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None
+    torch.save({
+        "state_dict": model.state_dict(),
+        "optimizer": optimizer_state,
+        "global_step": step,
+        "global_epoch": epoch,
+    }, checkpoint_path)
+    print("Saved checkpoint:", checkpoint_path)
+
+def _load(checkpoint_path):
+    if use_cuda:
+        checkpoint = torch.load(checkpoint_path)
+    else:
+        checkpoint = torch.load(checkpoint_path,
+                                map_location=lambda storage, loc: storage)
+    return checkpoint
+
+
+def load_checkpoint(path, model, optimizer, reset_optimizer=False, overwrite_global_states=True):
+    global global_step
+    global global_epoch
+
+    print("Load checkpoint from: {}".format(path))
+    checkpoint = _load(path)
+    s = checkpoint["state_dict"]
+    new_s = {}
+    for k, v in s.items():
+        new_s[k.replace('module.', '')] = v
+    model.load_state_dict(new_s)
+    if not reset_optimizer:
+        optimizer_state = checkpoint["optimizer"]
+        if optimizer_state is not None:
+            print("Load optimizer state from {}".format(path))
+            optimizer.load_state_dict(checkpoint["optimizer"])
+    if overwrite_global_states:
+        global_step = checkpoint["global_step"]
+        global_epoch = checkpoint["global_epoch"]
+
+    return model
+
+if __name__ == "__main__":
+    checkpoint_dir = args.checkpoint_dir
+
+    # Dataset and Dataloader setup
+    train_dataset = Dataset('train')
+    test_dataset = Dataset('val')
+
+    train_data_loader = data_utils.DataLoader(
+        train_dataset, batch_size=hparams.batch_size, shuffle=True,
+        num_workers=hparams.num_workers)
+
+    test_data_loader = data_utils.DataLoader(
+        test_dataset, batch_size=hparams.batch_size,
+        num_workers=4)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+     # Model
+    model = Wav2Lip().to(device)
+    disc = Wav2Lip_disc_qual().to(device)
+
+    print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad)))
+    print('total DISC trainable params {}'.format(sum(p.numel() for p in disc.parameters() if p.requires_grad)))
+
+    optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
+                           lr=hparams.initial_learning_rate, betas=(0.5, 0.999))
+    disc_optimizer = optim.Adam([p for p in disc.parameters() if p.requires_grad],
+                           lr=hparams.disc_initial_learning_rate, betas=(0.5, 0.999))
+
+    if args.checkpoint_path is not None:
+        load_checkpoint(args.checkpoint_path, model, optimizer, reset_optimizer=False)
+
+    if args.disc_checkpoint_path is not None:
+        load_checkpoint(args.disc_checkpoint_path, disc, disc_optimizer, 
+                                reset_optimizer=False, overwrite_global_states=False)
+        
+    load_checkpoint(args.syncnet_checkpoint_path, syncnet, None, reset_optimizer=True, 
+                                overwrite_global_states=False)
+
+    if not os.path.exists(checkpoint_dir):
+        os.mkdir(checkpoint_dir)
+
+    # Train!
+    train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer,
+              checkpoint_dir=checkpoint_dir,
+              checkpoint_interval=hparams.checkpoint_interval,
+              nepochs=hparams.nepochs)

inference.py

@@ -0,0 +1,280 @@
+from os import listdir, path
+import numpy as np
+import scipy, cv2, os, sys, argparse, audio
+import json, subprocess, random, string
+from tqdm import tqdm
+from glob import glob
+import torch, face_detection
+from models import Wav2Lip
+import platform
+
+parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models')
+
+parser.add_argument('--checkpoint_path', type=str, 
+					help='Name of saved checkpoint to load weights from', required=True)
+
+parser.add_argument('--face', type=str, 
+					help='Filepath of video/image that contains faces to use', required=True)
+parser.add_argument('--audio', type=str, 
+					help='Filepath of video/audio file to use as raw audio source', required=True)
+parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', 
+								default='results/result_voice.mp4')
+
+parser.add_argument('--static', type=bool, 
+					help='If True, then use only first video frame for inference', default=False)
+parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)', 
+					default=25., required=False)
+
+parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], 
+					help='Padding (top, bottom, left, right). Please adjust to include chin at least')
+
+parser.add_argument('--face_det_batch_size', type=int, 
+					help='Batch size for face detection', default=16)
+parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128)
+
+parser.add_argument('--resize_factor', default=1, type=int, 
+			help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p')
+
+parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], 
+					help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' 
+					'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width')
+
+parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], 
+					help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.'
+					'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).')
+
+parser.add_argument('--rotate', default=False, action='store_true',
+					help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.'
+					'Use if you get a flipped result, despite feeding a normal looking video')
+
+parser.add_argument('--nosmooth', default=False, action='store_true',
+					help='Prevent smoothing face detections over a short temporal window')
+
+args = parser.parse_args()
+args.img_size = 96
+
+if os.path.isfile(args.face) and args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
+	args.static = True
+
+def get_smoothened_boxes(boxes, T):
+	for i in range(len(boxes)):
+		if i + T > len(boxes):
+			window = boxes[len(boxes) - T:]
+		else:
+			window = boxes[i : i + T]
+		boxes[i] = np.mean(window, axis=0)
+	return boxes
+
+def face_detect(images):
+	detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, 
+											flip_input=False, device=device)
+
+	batch_size = args.face_det_batch_size
+	
+	while 1:
+		predictions = []
+		try:
+			for i in tqdm(range(0, len(images), batch_size)):
+				predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
+		except RuntimeError:
+			if batch_size == 1: 
+				raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
+			batch_size //= 2
+			print('Recovering from OOM error; New batch size: {}'.format(batch_size))
+			continue
+		break
+
+	results = []
+	pady1, pady2, padx1, padx2 = args.pads
+	for rect, image in zip(predictions, images):
+		if rect is None:
+			cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
+			raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
+
+		y1 = max(0, rect[1] - pady1)
+		y2 = min(image.shape[0], rect[3] + pady2)
+		x1 = max(0, rect[0] - padx1)
+		x2 = min(image.shape[1], rect[2] + padx2)
+		
+		results.append([x1, y1, x2, y2])
+
+	boxes = np.array(results)
+	if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
+	results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+
+	del detector
+	return results 
+
+def datagen(frames, mels):
+	img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
+
+	if args.box[0] == -1:
+		if not args.static:
+			face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
+		else:
+			face_det_results = face_detect([frames[0]])
+	else:
+		print('Using the specified bounding box instead of face detection...')
+		y1, y2, x1, x2 = args.box
+		face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
+
+	for i, m in enumerate(mels):
+		idx = 0 if args.static else i%len(frames)
+		frame_to_save = frames[idx].copy()
+		face, coords = face_det_results[idx].copy()
+
+		face = cv2.resize(face, (args.img_size, args.img_size))
+			
+		img_batch.append(face)
+		mel_batch.append(m)
+		frame_batch.append(frame_to_save)
+		coords_batch.append(coords)
+
+		if len(img_batch) >= args.wav2lip_batch_size:
+			img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+
+			img_masked = img_batch.copy()
+			img_masked[:, args.img_size//2:] = 0
+
+			img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
+			mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+
+			yield img_batch, mel_batch, frame_batch, coords_batch
+			img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
+
+	if len(img_batch) > 0:
+		img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
+
+		img_masked = img_batch.copy()
+		img_masked[:, args.img_size//2:] = 0
+
+		img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
+		mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
+
+		yield img_batch, mel_batch, frame_batch, coords_batch
+
+mel_step_size = 16
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('Using {} for inference.'.format(device))
+
+def _load(checkpoint_path):
+	if device == 'cuda':
+		checkpoint = torch.load(checkpoint_path)
+	else:
+		checkpoint = torch.load(checkpoint_path,
+								map_location=lambda storage, loc: storage)
+	return checkpoint
+
+def load_model(path):
+	model = Wav2Lip()
+	print("Load checkpoint from: {}".format(path))
+	checkpoint = _load(path)
+	s = checkpoint["state_dict"]
+	new_s = {}
+	for k, v in s.items():
+		new_s[k.replace('module.', '')] = v
+	model.load_state_dict(new_s)
+
+	model = model.to(device)
+	return model.eval()
+
+def main():
+	if not os.path.isfile(args.face):
+		raise ValueError('--face argument must be a valid path to video/image file')
+
+	elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']:
+		full_frames = [cv2.imread(args.face)]
+		fps = args.fps
+
+	else:
+		video_stream = cv2.VideoCapture(args.face)
+		fps = video_stream.get(cv2.CAP_PROP_FPS)
+
+		print('Reading video frames...')
+
+		full_frames = []
+		while 1:
+			still_reading, frame = video_stream.read()
+			if not still_reading:
+				video_stream.release()
+				break
+			if args.resize_factor > 1:
+				frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor))
+
+			if args.rotate:
+				frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE)
+
+			y1, y2, x1, x2 = args.crop
+			if x2 == -1: x2 = frame.shape[1]
+			if y2 == -1: y2 = frame.shape[0]
+
+			frame = frame[y1:y2, x1:x2]
+
+			full_frames.append(frame)
+
+	print ("Number of frames available for inference: "+str(len(full_frames)))
+
+	if not args.audio.endswith('.wav'):
+		print('Extracting raw audio...')
+		command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, 'temp/temp.wav')
+
+		subprocess.call(command, shell=True)
+		args.audio = 'temp/temp.wav'
+
+	wav = audio.load_wav(args.audio, 16000)
+	mel = audio.melspectrogram(wav)
+	print(mel.shape)
+
+	if np.isnan(mel.reshape(-1)).sum() > 0:
+		raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again')
+
+	mel_chunks = []
+	mel_idx_multiplier = 80./fps 
+	i = 0
+	while 1:
+		start_idx = int(i * mel_idx_multiplier)
+		if start_idx + mel_step_size > len(mel[0]):
+			mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:])
+			break
+		mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])
+		i += 1
+
+	print("Length of mel chunks: {}".format(len(mel_chunks)))
+
+	full_frames = full_frames[:len(mel_chunks)]
+
+	batch_size = args.wav2lip_batch_size
+	gen = datagen(full_frames.copy(), mel_chunks)
+
+	for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, 
+											total=int(np.ceil(float(len(mel_chunks))/batch_size)))):
+		if i == 0:
+			model = load_model(args.checkpoint_path)
+			print ("Model loaded")
+
+			frame_h, frame_w = full_frames[0].shape[:-1]
+			out = cv2.VideoWriter('temp/result.avi', 
+									cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h))
+
+		img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
+		mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
+
+		with torch.no_grad():
+			pred = model(mel_batch, img_batch)
+
+		pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
+		
+		for p, f, c in zip(pred, frames, coords):
+			y1, y2, x1, x2 = c
+			p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))
+
+			f[y1:y2, x1:x2] = p
+			out.write(f)
+
+	out.release()
+
+	command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/result.avi', args.outfile)
+	subprocess.call(command, shell=platform.system() != 'Windows')
+
+if __name__ == '__main__':
+	main()

preprocess.py

@@ -0,0 +1,113 @@
+import sys
+
+if sys.version_info[0] < 3 and sys.version_info[1] < 2:
+	raise Exception("Must be using >= Python 3.2")
+
+from os import listdir, path
+
+if not path.isfile('face_detection/detection/sfd/s3fd.pth'):
+	raise FileNotFoundError('Save the s3fd model to face_detection/detection/sfd/s3fd.pth \
+							before running this script!')
+
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import numpy as np
+import argparse, os, cv2, traceback, subprocess
+from tqdm import tqdm
+from glob import glob
+import audio
+from hparams import hparams as hp
+
+import face_detection
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('--ngpu', help='Number of GPUs across which to run in parallel', default=1, type=int)
+parser.add_argument('--batch_size', help='Single GPU Face detection batch size', default=32, type=int)
+parser.add_argument("--data_root", help="Root folder of the LRS2 dataset", required=True)
+parser.add_argument("--preprocessed_root", help="Root folder of the preprocessed dataset", required=True)
+
+args = parser.parse_args()
+
+fa = [face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False, 
+									device='cuda:{}'.format(id)) for id in range(args.ngpu)]
+
+template = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'
+# template2 = 'ffmpeg -hide_banner -loglevel panic -threads 1 -y -i {} -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 {}'
+
+def process_video_file(vfile, args, gpu_id):
+	video_stream = cv2.VideoCapture(vfile)
+	
+	frames = []
+	while 1:
+		still_reading, frame = video_stream.read()
+		if not still_reading:
+			video_stream.release()
+			break
+		frames.append(frame)
+	
+	vidname = os.path.basename(vfile).split('.')[0]
+	dirname = vfile.split('/')[-2]
+
+	fulldir = path.join(args.preprocessed_root, dirname, vidname)
+	os.makedirs(fulldir, exist_ok=True)
+
+	batches = [frames[i:i + args.batch_size] for i in range(0, len(frames), args.batch_size)]
+
+	i = -1
+	for fb in batches:
+		preds = fa[gpu_id].get_detections_for_batch(np.asarray(fb))
+
+		for j, f in enumerate(preds):
+			i += 1
+			if f is None:
+				continue
+
+			x1, y1, x2, y2 = f
+			cv2.imwrite(path.join(fulldir, '{}.jpg'.format(i)), fb[j][y1:y2, x1:x2])
+
+def process_audio_file(vfile, args):
+	vidname = os.path.basename(vfile).split('.')[0]
+	dirname = vfile.split('/')[-2]
+
+	fulldir = path.join(args.preprocessed_root, dirname, vidname)
+	os.makedirs(fulldir, exist_ok=True)
+
+	wavpath = path.join(fulldir, 'audio.wav')
+
+	command = template.format(vfile, wavpath)
+	subprocess.call(command, shell=True)
+
+	
+def mp_handler(job):
+	vfile, args, gpu_id = job
+	try:
+		process_video_file(vfile, args, gpu_id)
+	except KeyboardInterrupt:
+		exit(0)
+	except:
+		traceback.print_exc()
+		
+def main(args):
+	print('Started processing for {} with {} GPUs'.format(args.data_root, args.ngpu))
+
+	filelist = glob(path.join(args.data_root, '*/*.mp4'))
+
+	jobs = [(vfile, args, i%args.ngpu) for i, vfile in enumerate(filelist)]
+	p = ThreadPoolExecutor(args.ngpu)
+	futures = [p.submit(mp_handler, j) for j in jobs]
+	_ = [r.result() for r in tqdm(as_completed(futures), total=len(futures))]
+
+	print('Dumping audios...')
+
+	for vfile in tqdm(filelist):
+		try:
+			process_audio_file(vfile, args)
+		except KeyboardInterrupt:
+			exit(0)
+		except:
+			traceback.print_exc()
+			continue
+
+if __name__ == '__main__':
+	main(args)

requirements.txt

@@ -0,0 +1,8 @@
+librosa==0.7.0
+numpy==1.17.1
+opencv-contrib-python>=4.2.0.34
+opencv-python==4.3.0.38
+torch==1.11.0
+torchvision==0.12.0
+tqdm==4.45.0
+numba==0.48

wav2lip_train.py

@@ -0,0 +1,374 @@
+from os.path import dirname, join, basename, isfile
+from tqdm import tqdm
+
+from models import SyncNet_color as SyncNet
+from models import Wav2Lip as Wav2Lip
+import audio
+
+import torch
+from torch import nn
+from torch import optim
+import torch.backends.cudnn as cudnn
+from torch.utils import data as data_utils
+import numpy as np
+
+from glob import glob
+
+import os, random, cv2, argparse
+from hparams import hparams, get_image_list
+
+parser = argparse.ArgumentParser(description='Code to train the Wav2Lip model without the visual quality discriminator')
+
+parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True, type=str)
+
+parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str)
+parser.add_argument('--syncnet_checkpoint_path', help='Load the pre-trained Expert discriminator', required=True, type=str)
+
+parser.add_argument('--checkpoint_path', help='Resume from this checkpoint', default=None, type=str)
+
+args = parser.parse_args()
+
+
+global_step = 0
+global_epoch = 0
+use_cuda = torch.cuda.is_available()
+print('use_cuda: {}'.format(use_cuda))
+
+syncnet_T = 5
+syncnet_mel_step_size = 16
+
+class Dataset(object):
+    def __init__(self, split):
+        self.all_videos = get_image_list(args.data_root, split)
+
+    def get_frame_id(self, frame):
+        return int(basename(frame).split('.')[0])
+
+    def get_window(self, start_frame):
+        start_id = self.get_frame_id(start_frame)
+        vidname = dirname(start_frame)
+
+        window_fnames = []
+        for frame_id in range(start_id, start_id + syncnet_T):
+            frame = join(vidname, '{}.jpg'.format(frame_id))
+            if not isfile(frame):
+                return None
+            window_fnames.append(frame)
+        return window_fnames
+
+    def read_window(self, window_fnames):
+        if window_fnames is None: return None
+        window = []
+        for fname in window_fnames:
+            img = cv2.imread(fname)
+            if img is None:
+                return None
+            try:
+                img = cv2.resize(img, (hparams.img_size, hparams.img_size))
+            except Exception as e:
+                return None
+
+            window.append(img)
+
+        return window
+
+    def crop_audio_window(self, spec, start_frame):
+        if type(start_frame) == int:
+            start_frame_num = start_frame
+        else:
+            start_frame_num = self.get_frame_id(start_frame) # 0-indexing ---> 1-indexing
+        start_idx = int(80. * (start_frame_num / float(hparams.fps)))
+        
+        end_idx = start_idx + syncnet_mel_step_size
+
+        return spec[start_idx : end_idx, :]
+
+    def get_segmented_mels(self, spec, start_frame):
+        mels = []
+        assert syncnet_T == 5
+        start_frame_num = self.get_frame_id(start_frame) + 1 # 0-indexing ---> 1-indexing
+        if start_frame_num - 2 < 0: return None
+        for i in range(start_frame_num, start_frame_num + syncnet_T):
+            m = self.crop_audio_window(spec, i - 2)
+            if m.shape[0] != syncnet_mel_step_size:
+                return None
+            mels.append(m.T)
+
+        mels = np.asarray(mels)
+
+        return mels
+
+    def prepare_window(self, window):
+        # 3 x T x H x W
+        x = np.asarray(window) / 255.
+        x = np.transpose(x, (3, 0, 1, 2))
+
+        return x
+
+    def __len__(self):
+        return len(self.all_videos)
+
+    def __getitem__(self, idx):
+        while 1:
+            idx = random.randint(0, len(self.all_videos) - 1)
+            vidname = self.all_videos[idx]
+            img_names = list(glob(join(vidname, '*.jpg')))
+            if len(img_names) <= 3 * syncnet_T:
+                continue
+            
+            img_name = random.choice(img_names)
+            wrong_img_name = random.choice(img_names)
+            while wrong_img_name == img_name:
+                wrong_img_name = random.choice(img_names)
+
+            window_fnames = self.get_window(img_name)
+            wrong_window_fnames = self.get_window(wrong_img_name)
+            if window_fnames is None or wrong_window_fnames is None:
+                continue
+
+            window = self.read_window(window_fnames)
+            if window is None:
+                continue
+
+            wrong_window = self.read_window(wrong_window_fnames)
+            if wrong_window is None:
+                continue
+
+            try:
+                wavpath = join(vidname, "audio.wav")
+                wav = audio.load_wav(wavpath, hparams.sample_rate)
+
+                orig_mel = audio.melspectrogram(wav).T
+            except Exception as e:
+                continue
+
+            mel = self.crop_audio_window(orig_mel.copy(), img_name)
+            
+            if (mel.shape[0] != syncnet_mel_step_size):
+                continue
+
+            indiv_mels = self.get_segmented_mels(orig_mel.copy(), img_name)
+            if indiv_mels is None: continue
+
+            window = self.prepare_window(window)
+            y = window.copy()
+            window[:, :, window.shape[2]//2:] = 0.
+
+            wrong_window = self.prepare_window(wrong_window)
+            x = np.concatenate([window, wrong_window], axis=0)
+
+            x = torch.FloatTensor(x)
+            mel = torch.FloatTensor(mel.T).unsqueeze(0)
+            indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1)
+            y = torch.FloatTensor(y)
+            return x, indiv_mels, mel, y
+
+def save_sample_images(x, g, gt, global_step, checkpoint_dir):
+    x = (x.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
+    g = (g.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
+    gt = (gt.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8)
+
+    refs, inps = x[..., 3:], x[..., :3]
+    folder = join(checkpoint_dir, "samples_step{:09d}".format(global_step))
+    if not os.path.exists(folder): os.mkdir(folder)
+    collage = np.concatenate((refs, inps, g, gt), axis=-2)
+    for batch_idx, c in enumerate(collage):
+        for t in range(len(c)):
+            cv2.imwrite('{}/{}_{}.jpg'.format(folder, batch_idx, t), c[t])
+
+logloss = nn.BCELoss()
+def cosine_loss(a, v, y):
+    d = nn.functional.cosine_similarity(a, v)
+    loss = logloss(d.unsqueeze(1), y)
+
+    return loss
+
+device = torch.device("cuda" if use_cuda else "cpu")
+syncnet = SyncNet().to(device)
+for p in syncnet.parameters():
+    p.requires_grad = False
+
+recon_loss = nn.L1Loss()
+def get_sync_loss(mel, g):
+    g = g[:, :, :, g.size(3)//2:]
+    g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1)
+    # B, 3 * T, H//2, W
+    a, v = syncnet(mel, g)
+    y = torch.ones(g.size(0), 1).float().to(device)
+    return cosine_loss(a, v, y)
+
+def train(device, model, train_data_loader, test_data_loader, optimizer,
+          checkpoint_dir=None, checkpoint_interval=None, nepochs=None):
+
+    global global_step, global_epoch
+    resumed_step = global_step
+ 
+    while global_epoch < nepochs:
+        print('Starting Epoch: {}'.format(global_epoch))
+        running_sync_loss, running_l1_loss = 0., 0.
+        prog_bar = tqdm(enumerate(train_data_loader))
+        for step, (x, indiv_mels, mel, gt) in prog_bar:
+            model.train()
+            optimizer.zero_grad()
+
+            # Move data to CUDA device
+            x = x.to(device)
+            mel = mel.to(device)
+            indiv_mels = indiv_mels.to(device)
+            gt = gt.to(device)
+
+            g = model(indiv_mels, x)
+
+            if hparams.syncnet_wt > 0.:
+                sync_loss = get_sync_loss(mel, g)
+            else:
+                sync_loss = 0.
+
+            l1loss = recon_loss(g, gt)
+
+            loss = hparams.syncnet_wt * sync_loss + (1 - hparams.syncnet_wt) * l1loss
+            loss.backward()
+            optimizer.step()
+
+            if global_step % checkpoint_interval == 0:
+                save_sample_images(x, g, gt, global_step, checkpoint_dir)
+
+            global_step += 1
+            cur_session_steps = global_step - resumed_step
+
+            running_l1_loss += l1loss.item()
+            if hparams.syncnet_wt > 0.:
+                running_sync_loss += sync_loss.item()
+            else:
+                running_sync_loss += 0.
+
+            if global_step == 1 or global_step % checkpoint_interval == 0:
+                save_checkpoint(
+                    model, optimizer, global_step, checkpoint_dir, global_epoch)
+
+            if global_step == 1 or global_step % hparams.eval_interval == 0:
+                with torch.no_grad():
+                    average_sync_loss = eval_model(test_data_loader, global_step, device, model, checkpoint_dir)
+
+                    if average_sync_loss < .75:
+                        hparams.set_hparam('syncnet_wt', 0.01) # without image GAN a lesser weight is sufficient
+
+            prog_bar.set_description('L1: {}, Sync Loss: {}'.format(running_l1_loss / (step + 1),
+                                                                    running_sync_loss / (step + 1)))
+
+        global_epoch += 1
+        
+
+def eval_model(test_data_loader, global_step, device, model, checkpoint_dir):
+    eval_steps = 700
+    print('Evaluating for {} steps'.format(eval_steps))
+    sync_losses, recon_losses = [], []
+    step = 0
+    while 1:
+        for x, indiv_mels, mel, gt in test_data_loader:
+            step += 1
+            model.eval()
+
+            # Move data to CUDA device
+            x = x.to(device)
+            gt = gt.to(device)
+            indiv_mels = indiv_mels.to(device)
+            mel = mel.to(device)
+
+            g = model(indiv_mels, x)
+
+            sync_loss = get_sync_loss(mel, g)
+            l1loss = recon_loss(g, gt)
+
+            sync_losses.append(sync_loss.item())
+            recon_losses.append(l1loss.item())
+
+            if step > eval_steps: 
+                averaged_sync_loss = sum(sync_losses) / len(sync_losses)
+                averaged_recon_loss = sum(recon_losses) / len(recon_losses)
+
+                print('L1: {}, Sync loss: {}'.format(averaged_recon_loss, averaged_sync_loss))
+
+                return averaged_sync_loss
+
+def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch):
+
+    checkpoint_path = join(
+        checkpoint_dir, "checkpoint_step{:09d}.pth".format(global_step))
+    optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None
+    torch.save({
+        "state_dict": model.state_dict(),
+        "optimizer": optimizer_state,
+        "global_step": step,
+        "global_epoch": epoch,
+    }, checkpoint_path)
+    print("Saved checkpoint:", checkpoint_path)
+
+
+def _load(checkpoint_path):
+    if use_cuda:
+        checkpoint = torch.load(checkpoint_path)
+    else:
+        checkpoint = torch.load(checkpoint_path,
+                                map_location=lambda storage, loc: storage)
+    return checkpoint
+
+def load_checkpoint(path, model, optimizer, reset_optimizer=False, overwrite_global_states=True):
+    global global_step
+    global global_epoch
+
+    print("Load checkpoint from: {}".format(path))
+    checkpoint = _load(path)
+    s = checkpoint["state_dict"]
+    new_s = {}
+    for k, v in s.items():
+        new_s[k.replace('module.', '')] = v
+    model.load_state_dict(new_s)
+    if not reset_optimizer:
+        optimizer_state = checkpoint["optimizer"]
+        if optimizer_state is not None:
+            print("Load optimizer state from {}".format(path))
+            optimizer.load_state_dict(checkpoint["optimizer"])
+    if overwrite_global_states:
+        global_step = checkpoint["global_step"]
+        global_epoch = checkpoint["global_epoch"]
+
+    return model
+
+if __name__ == "__main__":
+    checkpoint_dir = args.checkpoint_dir
+
+    # Dataset and Dataloader setup
+    train_dataset = Dataset('train')
+    test_dataset = Dataset('val')
+
+    train_data_loader = data_utils.DataLoader(
+        train_dataset, batch_size=hparams.batch_size, shuffle=True,
+        num_workers=hparams.num_workers)
+
+    test_data_loader = data_utils.DataLoader(
+        test_dataset, batch_size=hparams.batch_size,
+        num_workers=4)
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    # Model
+    model = Wav2Lip().to(device)
+    print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad)))
+
+    optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad],
+                           lr=hparams.initial_learning_rate)
+
+    if args.checkpoint_path is not None:
+        load_checkpoint(args.checkpoint_path, model, optimizer, reset_optimizer=False)
+        
+    load_checkpoint(args.syncnet_checkpoint_path, syncnet, None, reset_optimizer=True, overwrite_global_states=False)
+
+    if not os.path.exists(checkpoint_dir):
+        os.mkdir(checkpoint_dir)
+
+    # Train!
+    train(device, model, train_data_loader, test_data_loader, optimizer,
+              checkpoint_dir=checkpoint_dir,
+              checkpoint_interval=hparams.checkpoint_interval,
+              nepochs=hparams.nepochs)