|
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 |
|
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): |
|
|
|
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) |
|
|
|
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) |
|
|
|
|
|
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() |
|
|
|
|
|
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) |
|
|
|
|
|
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 |
|
|
|
|
|
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 = 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(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer, |
|
checkpoint_dir=checkpoint_dir, |
|
checkpoint_interval=hparams.checkpoint_interval, |
|
nepochs=hparams.nepochs) |
|
|
