"""
将One-shot的说话人大模型(os_secc2plane or os_secc2plane_torso)在单一说话人(一张照片或一段视频)上overfit, 实现和GeneFace++类似的效果
"""
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import librosa
import random
import time
import numpy as np
import importlib
import tqdm
import copy
import cv2
import glob
import imageio
# common utils
from utils.commons.hparams import hparams, set_hparams
from utils.commons.tensor_utils import move_to_cuda, convert_to_tensor
from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint
# 3DMM-related utils
from deep_3drecon.deep_3drecon_models.bfm import ParametricFaceModel
from data_util.face3d_helper import Face3DHelper
from data_gen.utils.process_image.fit_3dmm_landmark import fit_3dmm_for_a_image
from data_gen.utils.process_video.fit_3dmm_landmark import fit_3dmm_for_a_video
from data_gen.utils.process_video.extract_segment_imgs import decode_segmap_mask_from_image
from deep_3drecon.secc_renderer import SECC_Renderer
from data_gen.eg3d.convert_to_eg3d_convention import get_eg3d_convention_camera_pose_intrinsic
from data_gen.runs.binarizer_nerf import get_lip_rect
# Face Parsing
from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter
from data_gen.utils.process_video.extract_segment_imgs import inpaint_torso_job, extract_background
# other inference utils
from inference.infer_utils import mirror_index, load_img_to_512_hwc_array, load_img_to_normalized_512_bchw_tensor
from inference.infer_utils import smooth_camera_sequence, smooth_features_xd
from inference.edit_secc import blink_eye_for_secc, hold_eye_opened_for_secc
from modules.commons.loralib.utils import mark_only_lora_as_trainable
from utils.nn.model_utils import num_params
import lpips
from utils.commons.meters import AvgrageMeter
meter = AvgrageMeter()
from torch.utils.tensorboard import SummaryWriter
class LoRATrainer(nn.Module):
def __init__(self, inp):
super().__init__()
self.inp = inp
self.lora_args = {'lora_mode': inp['lora_mode'], 'lora_r': inp['lora_r']}
device = 'cuda' if torch.cuda.is_available() else 'cpu'
head_model_dir = inp['head_ckpt']
torso_model_dir = inp['torso_ckpt']
model_dir = torso_model_dir if torso_model_dir != '' else head_model_dir
cmd = f"cp {os.path.join(model_dir, 'config.yaml')} {self.inp['work_dir']}"
print(cmd)
os.system(cmd)
with open(os.path.join(self.inp['work_dir'], 'config.yaml'), "a") as f:
f.write(f"\nlora_r: {inp['lora_r']}")
f.write(f"\nlora_mode: {inp['lora_mode']}")
f.write(f"\n")
self.secc2video_model = self.load_secc2video(model_dir)
self.secc2video_model.to(device).eval()
self.seg_model = MediapipeSegmenter()
self.secc_renderer = SECC_Renderer(512)
self.face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='lm68')
self.mp_face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='mediapipe')
# self.camera_selector = KNearestCameraSelector()
self.load_training_data(inp)
def load_secc2video(self, model_dir):
inp = self.inp
from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane, OSAvatarSECC_Img2plane_Torso
hp = set_hparams(f"{model_dir}/config.yaml", print_hparams=False, global_hparams=True)
hp['htbsr_head_threshold'] = 1.0
self.neural_rendering_resolution = hp['neural_rendering_resolution']
if 'torso' in hp['task_cls'].lower():
self.torso_mode = True
model = OSAvatarSECC_Img2plane_Torso(hp=hp, lora_args=self.lora_args)
else:
self.torso_mode = False
model = OSAvatarSECC_Img2plane(hp=hp, lora_args=self.lora_args)
mark_only_lora_as_trainable(model, bias='none')
lora_ckpt_path = os.path.join(inp['work_dir'], 'checkpoint.ckpt')
if os.path.exists(lora_ckpt_path):
self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, model.triplane_hid_dim*model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True)
model._last_cano_planes = self.learnable_triplane
load_ckpt(model, lora_ckpt_path, model_name='model', strict=False)
else:
load_ckpt(model, f"{model_dir}", model_name='model', strict=False)
num_params(model)
self.model = model
return model
def load_training_data(self, inp):
video_id = inp['video_id']
if video_id.endswith((".mp4", ".png", ".jpg", ".jpeg")):
# If input video is not GeneFace training videos, convert it into GeneFace convention
video_id_ = video_id
video_id = os.path.basename(video_id)[:-4]
inp['video_id'] = video_id
target_video_path = f'data/raw/videos/{video_id}.mp4'
if not os.path.exists(target_video_path):
print(f"| Copying video to {target_video_path}")
os.makedirs(os.path.dirname(target_video_path), exist_ok=True)
cmd = f"ffmpeg -i {video_id_} -vf fps=25,scale=w=512:h=512 -qmin 1 -q:v 1 -y {target_video_path}"
print(f"| {cmd}")
os.system(cmd)
target_video_path = f'data/raw/videos/{video_id}.mp4'
print(f"| Copy source video into work dir: {self.inp['work_dir']}")
os.system(f"cp {target_video_path} {self.inp['work_dir']}")
# check head_img path
head_img_pattern = f'data/processed/videos/{video_id}/head_imgs/*.png'
head_img_names = sorted(glob.glob(head_img_pattern))
if len(head_img_names) == 0:
# extract head_imgs
head_img_dir = os.path.dirname(head_img_pattern)
print(f"| Pre-extracted head_imgs not found, try to extract and save to {head_img_dir}, this may take a while...")
gt_img_dir = f"data/processed/videos/{video_id}/gt_imgs"
os.makedirs(gt_img_dir, exist_ok=True)
target_video_path = f'data/raw/videos/{video_id}.mp4'
cmd = f"ffmpeg -i {target_video_path} -vf fps=25,scale=w=512:h=512 -qmin 1 -q:v 1 -start_number 0 -y {gt_img_dir}/%08d.jpg"
print(f"| {cmd}")
os.system(cmd)
# extract image, segmap, and background
cmd = f"python data_gen/utils/process_video/extract_segment_imgs.py --ds_name=nerf --vid_dir={target_video_path}"
print(f"| {cmd}")
os.system(cmd)
print("| Head images Extracted!")
num_samples = len(head_img_names)
npy_name = f"data/processed/videos/{video_id}/coeff_fit_mp_for_lora.npy"
if os.path.exists(npy_name):
coeff_dict = np.load(npy_name, allow_pickle=True).tolist()
else:
print(f"| Pre-extracted 3DMM coefficient not found, try to extract and save to {npy_name}, this may take a while...")
coeff_dict = fit_3dmm_for_a_video(f'data/raw/videos/{video_id}.mp4', save=False)
os.makedirs(os.path.dirname(npy_name), exist_ok=True)
np.save(npy_name, coeff_dict)
ids = convert_to_tensor(coeff_dict['id']).reshape([-1,80]).cuda()
exps = convert_to_tensor(coeff_dict['exp']).reshape([-1,64]).cuda()
eulers = convert_to_tensor(coeff_dict['euler']).reshape([-1,3]).cuda()
trans = convert_to_tensor(coeff_dict['trans']).reshape([-1,3]).cuda()
WH = 512 # now we only support 512x512
lm2ds = WH * self.face3d_helper.reconstruct_lm2d(ids, exps, eulers, trans).cpu().numpy()
lip_rects = [get_lip_rect(lm2ds[i], WH, WH) for i in range(len(lm2ds))]
kps = self.face3d_helper.reconstruct_lm2d(ids, exps, eulers, trans).cuda()
kps = (kps-0.5) / 0.5 # rescale to -1~1
kps = torch.cat([kps, torch.zeros([*kps.shape[:-1], 1]).cuda()], dim=-1)
camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler': torch.tensor(coeff_dict['euler']).reshape([-1,3]), 'trans': torch.tensor(coeff_dict['trans']).reshape([-1,3])})
c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics']
cameras = torch.tensor(np.concatenate([c2w.reshape([-1,16]), intrinsics.reshape([-1,9])], axis=-1)).cuda()
camera_smo_ksize = 7
cameras = smooth_camera_sequence(cameras.cpu().numpy(), kernel_size=camera_smo_ksize) # [T, 25]
cameras = torch.tensor(cameras).cuda()
zero_eulers = eulers * 0
zero_trans = trans * 0
_, cano_secc_color = self.secc_renderer(ids[0:1], exps[0:1]*0, zero_eulers[0:1], zero_trans[0:1])
src_idx = 0
_, src_secc_color = self.secc_renderer(ids[0:1], exps[src_idx:src_idx+1], zero_eulers[0:1], zero_trans[0:1])
drv_secc_colors = [None for _ in range(len(exps))]
drv_head_imgs = [None for _ in range(len(exps))]
drv_torso_imgs = [None for _ in range(len(exps))]
drv_com_imgs = [None for _ in range(len(exps))]
segmaps = [None for _ in range(len(exps))]
img_name = f'data/processed/videos/{video_id}/bg.jpg'
bg_img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
ds = {
'id': ids.cuda().float(),
'exps': exps.cuda().float(),
'eulers': eulers.cuda().float(),
'trans': trans.cuda().float(),
'cano_secc_color': cano_secc_color.cuda().float(),
'src_secc_color': src_secc_color.cuda().float(),
'cameras': cameras.float(),
'video_id': video_id,
'lip_rects': lip_rects,
'head_imgs': drv_head_imgs,
'torso_imgs': drv_torso_imgs,
'com_imgs': drv_com_imgs,
'bg_img': bg_img,
'segmaps': segmaps,
'kps': kps,
}
self.ds = ds
return ds
def training_loop(self, inp):
trainer = self
video_id = self.ds['video_id']
lora_params = [p for k, p in self.secc2video_model.named_parameters() if 'lora_' in k]
self.criterion_lpips = lpips.LPIPS(net='alex',lpips=True).cuda()
self.logger = SummaryWriter(log_dir=inp['work_dir'])
if not hasattr(self, 'learnable_triplane'):
src_idx = 0 # init triplane from the first frame's prediction
self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, self.secc2video_model.triplane_hid_dim*self.secc2video_model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True)
img_name = f'data/processed/videos/{video_id}/head_imgs/{format(src_idx, "08d")}.png'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float().cuda().float() # [3, H, W]
cano_plane = self.secc2video_model.cal_cano_plane(img.unsqueeze(0)) # [1, 3, CD, h, w]
self.learnable_triplane.data = cano_plane.data
self.secc2video_model._last_cano_planes = self.learnable_triplane
if len(lora_params) == 0:
self.optimizer = torch.optim.AdamW([self.learnable_triplane], lr=inp['lr_triplane'], weight_decay=0.01, betas=(0.9,0.98))
else:
self.optimizer = torch.optim.Adam(lora_params, lr=inp['lr'], betas=(0.9,0.98))
self.optimizer.add_param_group({
'params': [self.learnable_triplane],
'lr': inp['lr_triplane'],
'betas': (0.9, 0.98)
})
ids = self.ds['id']
exps = self.ds['exps']
zero_eulers = self.ds['eulers']*0
zero_trans = self.ds['trans']*0
num_updates = inp['max_updates']
batch_size = inp['batch_size'] # 1 for lower gpu mem usage
num_samples = len(self.ds['cameras'])
init_plane = self.learnable_triplane.detach().clone()
if num_samples <= 5:
lambda_reg_triplane = 1.0
elif num_samples <= 250:
lambda_reg_triplane = 0.1
else:
lambda_reg_triplane = 0.
for i_step in tqdm.trange(num_updates+1,desc="training lora..."):
milestone_steps = []
# milestone_steps = [100, 200, 500]
if i_step % 2000 == 0 or i_step in milestone_steps:
trainer.test_loop(inp, step=i_step)
if i_step != 0:
filepath = os.path.join(inp['work_dir'], f"model_ckpt_steps_{i_step}.ckpt")
checkpoint = self.dump_checkpoint(inp)
tmp_path = str(filepath) + ".part"
torch.save(checkpoint, tmp_path, _use_new_zipfile_serialization=False)
os.replace(tmp_path, filepath)
drv_idx = [random.randint(0, num_samples-1) for _ in range(batch_size)]
drv_secc_colors = []
gt_imgs = []
head_imgs = []
segmaps_0 = []
segmaps = []
torso_imgs = []
drv_lip_rects = []
kp_src = []
kp_drv = []
for di in drv_idx:
# 读取target image
if self.torso_mode:
if self.ds['com_imgs'][di] is None:
# img_name = f'data/processed/videos/{video_id}/gt_imgs/{format(di, "08d")}.jpg'
img_name = f'data/processed/videos/{video_id}/com_imgs/{format(di, "08d")}.jpg'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['com_imgs'][di] = img
gt_imgs.append(self.ds['com_imgs'][di])
else:
if self.ds['head_imgs'][di] is None:
img_name = f'data/processed/videos/{video_id}/head_imgs/{format(di, "08d")}.png'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['head_imgs'][di] = img
gt_imgs.append(self.ds['head_imgs'][di])
if self.ds['head_imgs'][di] is None:
img_name = f'data/processed/videos/{video_id}/head_imgs/{format(di, "08d")}.png'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['head_imgs'][di] = img
head_imgs.append(self.ds['head_imgs'][di])
# 使用第一帧的torso作为face v2v的输入
if self.ds['torso_imgs'][0] is None:
img_name = f'data/processed/videos/{video_id}/inpaint_torso_imgs/{format(0, "08d")}.png'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['torso_imgs'][0] = img
torso_imgs.append(self.ds['torso_imgs'][0])
# 所以segmap也用第一帧的了
if self.ds['segmaps'][0] is None:
img_name = f'data/processed/videos/{video_id}/segmaps/{format(0, "08d")}.png'
seg_img = cv2.imread(img_name)[:,:, ::-1]
segmap = torch.from_numpy(decode_segmap_mask_from_image(seg_img)) # [6, H, W]
self.ds['segmaps'][0] = segmap
segmaps_0.append(self.ds['segmaps'][0])
if self.ds['segmaps'][di] is None:
img_name = f'data/processed/videos/{video_id}/segmaps/{format(di, "08d")}.png'
seg_img = cv2.imread(img_name)[:,:, ::-1]
segmap = torch.from_numpy(decode_segmap_mask_from_image(seg_img)) # [6, H, W]
self.ds['segmaps'][di] = segmap
segmaps.append(self.ds['segmaps'][di])
_, secc_color = self.secc_renderer(ids[0:1], exps[di:di+1], zero_eulers[0:1], zero_trans[0:1])
drv_secc_colors.append(secc_color)
drv_lip_rects.append(self.ds['lip_rects'][di])
kp_src.append(self.ds['kps'][0])
kp_drv.append(self.ds['kps'][di])
bg_img = self.ds['bg_img'].unsqueeze(0).repeat([batch_size, 1, 1, 1]).cuda()
ref_torso_imgs = torch.stack(torso_imgs).float().cuda()
kp_src = torch.stack(kp_src).float().cuda()
kp_drv = torch.stack(kp_drv).float().cuda()
segmaps = torch.stack(segmaps).float().cuda()
segmaps_0 = torch.stack(segmaps_0).float().cuda()
tgt_imgs = torch.stack(gt_imgs).float().cuda()
head_imgs = torch.stack(head_imgs).float().cuda()
drv_secc_color = torch.cat(drv_secc_colors)
cano_secc_color = self.ds['cano_secc_color'].repeat([batch_size, 1, 1, 1])
src_secc_color = self.ds['src_secc_color'].repeat([batch_size, 1, 1, 1])
cond = {'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_color,
'ref_torso_img': ref_torso_imgs, 'bg_img': bg_img,
'segmap': segmaps_0, # v2v使用第一帧的torso作为source image来warp
'kp_s': kp_src, 'kp_d': kp_drv}
camera = self.ds['cameras'][drv_idx]
gen_output = self.secc2video_model.forward(img=None, camera=camera, cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
pred_imgs = gen_output['image']
pred_imgs_raw = gen_output['image_raw']
losses = {}
loss_weights = {
'v2v_occlusion_reg_l1_loss': 0.001, # loss for face_vid2vid-based torso
'v2v_occlusion_2_reg_l1_loss': 0.001, # loss for face_vid2vid-based torso
'v2v_occlusion_2_weights_entropy_loss': hparams['lam_occlusion_weights_entropy'], # loss for face_vid2vid-based torso
'density_weight_l2_loss': 0.01, # supervised density
'density_weight_entropy_loss': 0.001, # keep the density change sharp
'mse_loss': 1.,
'head_mse_loss': 0.2, # loss on neural rendering low-reso pred_img
'lip_mse_loss': 1.0,
'lpips_loss': 0.5,
'head_lpips_loss': 0.1,
'lip_lpips_loss': 1.0, # make the teeth more clear
'blink_reg_loss': 0.003, # increase it when you find head shake while blinking; decrease it when you find the eye cannot closed.
'triplane_reg_loss': lambda_reg_triplane,
'secc_reg_loss': 0.01, # used to reduce flicking
}
occlusion_reg_l1 = gen_output.get("losses", {}).get('facev2v/occlusion_reg_l1', 0.)
occlusion_2_reg_l1 = gen_output.get("losses", {}).get('facev2v/occlusion_2_reg_l1', 0.)
occlusion_2_weights_entropy = gen_output.get("losses", {}).get('facev2v/occlusion_2_weights_entropy', 0.)
losses['v2v_occlusion_reg_l1_loss'] = occlusion_reg_l1
losses['v2v_occlusion_2_reg_l1_loss'] = occlusion_2_reg_l1
losses['v2v_occlusion_2_weights_entropy_loss'] = occlusion_2_weights_entropy
# Weights Reg loss in torso
neural_rendering_reso = self.neural_rendering_resolution
alphas = gen_output['weights_img'].clamp(1e-5, 1 - 1e-5)
loss_weights_entropy = torch.mean(- alphas * torch.log2(alphas) - (1 - alphas) * torch.log2(1 - alphas))
mv_head_masks = segmaps[:, [1,3,5]].sum(dim=1)
mv_head_masks_raw = F.interpolate(mv_head_masks.unsqueeze(1), size=(neural_rendering_reso,neural_rendering_reso)).squeeze(1)
face_mask = mv_head_masks_raw.bool().unsqueeze(1)
nonface_mask = ~ face_mask
loss_weights_l2_loss = (alphas[nonface_mask]-0).pow(2).mean() + (alphas[face_mask]-1).pow(2).mean()
losses['density_weight_l2_loss'] = loss_weights_l2_loss
losses['density_weight_entropy_loss'] = loss_weights_entropy
mse_loss = (pred_imgs - tgt_imgs).abs().mean()
head_mse_loss = (pred_imgs_raw - F.interpolate(head_imgs, size=(neural_rendering_reso,neural_rendering_reso), mode='bilinear', antialias=True)).abs().mean()
lpips_loss = self.criterion_lpips(pred_imgs, tgt_imgs).mean()
head_lpips_loss = self.criterion_lpips(pred_imgs_raw, F.interpolate(head_imgs, size=(neural_rendering_reso,neural_rendering_reso), mode='bilinear', antialias=True)).mean()
lip_mse_loss = 0
lip_lpips_loss = 0
for i in range(len(drv_idx)):
xmin, xmax, ymin, ymax = drv_lip_rects[i]
lip_tgt_imgs = tgt_imgs[i:i+1,:, ymin:ymax,xmin:xmax].contiguous()
lip_pred_imgs = pred_imgs[i:i+1,:, ymin:ymax,xmin:xmax].contiguous()
try:
lip_mse_loss = lip_mse_loss + (lip_pred_imgs - lip_tgt_imgs).abs().mean()
lip_lpips_loss = lip_lpips_loss + self.criterion_lpips(lip_pred_imgs, lip_tgt_imgs).mean()
except: pass
losses['mse_loss'] = mse_loss
losses['head_mse_loss'] = head_mse_loss
losses['lpips_loss'] = lpips_loss
losses['head_lpips_loss'] = head_lpips_loss
losses['lip_mse_loss'] = lip_mse_loss
losses['lip_lpips_loss'] = lip_lpips_loss
# eye blink reg loss
if i_step % 4 == 0:
blink_secc_lst1 = []
blink_secc_lst2 = []
blink_secc_lst3 = []
for i in range(len(drv_secc_color)):
secc = drv_secc_color[i]
blink_percent1 = random.random() * 0.5 # 0~0.5
blink_percent3 = 0.5 + random.random() * 0.5 # 0.5~1.0
blink_percent2 = (blink_percent1 + blink_percent3)/2
try:
out_secc1 = blink_eye_for_secc(secc, blink_percent1).to(secc.device)
out_secc2 = blink_eye_for_secc(secc, blink_percent2).to(secc.device)
out_secc3 = blink_eye_for_secc(secc, blink_percent3).to(secc.device)
except:
print("blink eye for secc failed, use original secc")
out_secc1 = copy.deepcopy(secc)
out_secc2 = copy.deepcopy(secc)
out_secc3 = copy.deepcopy(secc)
blink_secc_lst1.append(out_secc1)
blink_secc_lst2.append(out_secc2)
blink_secc_lst3.append(out_secc3)
src_secc_color1 = torch.stack(blink_secc_lst1)
src_secc_color2 = torch.stack(blink_secc_lst2)
src_secc_color3 = torch.stack(blink_secc_lst3)
blink_cond1 = {'cond_cano': cano_secc_color, 'cond_src': src_secc_color, 'cond_tgt': src_secc_color1}
blink_cond2 = {'cond_cano': cano_secc_color, 'cond_src': src_secc_color, 'cond_tgt': src_secc_color2}
blink_cond3 = {'cond_cano': cano_secc_color, 'cond_src': src_secc_color, 'cond_tgt': src_secc_color3}
blink_secc_plane1 = self.model.cal_secc_plane(blink_cond1)
blink_secc_plane2 = self.model.cal_secc_plane(blink_cond2)
blink_secc_plane3 = self.model.cal_secc_plane(blink_cond3)
interpolate_blink_secc_plane = (blink_secc_plane1 + blink_secc_plane3) / 2
blink_reg_loss = torch.nn.functional.l1_loss(blink_secc_plane2, interpolate_blink_secc_plane)
losses['blink_reg_loss'] = blink_reg_loss
# Triplane Reg loss
triplane_reg_loss = (self.learnable_triplane - init_plane).abs().mean()
losses['triplane_reg_loss'] = triplane_reg_loss
ref_id = self.ds['id'][0:1]
secc_pertube_randn_scale = hparams['secc_pertube_randn_scale']
perturbed_id = ref_id + torch.randn_like(ref_id) * secc_pertube_randn_scale
drv_exp = self.ds['exps'][drv_idx]
perturbed_exp = drv_exp + torch.randn_like(drv_exp) * secc_pertube_randn_scale
zero_euler = torch.zeros([len(drv_idx), 3], device=ref_id.device, dtype=ref_id.dtype)
zero_trans = torch.zeros([len(drv_idx), 3], device=ref_id.device, dtype=ref_id.dtype)
perturbed_secc = self.secc_renderer(perturbed_id, perturbed_exp, zero_euler, zero_trans)[1]
secc_reg_loss = torch.nn.functional.l1_loss(drv_secc_color, perturbed_secc)
losses['secc_reg_loss'] = secc_reg_loss
total_loss = sum([loss_weights[k] * v for k, v in losses.items() if isinstance(v, torch.Tensor) and v.requires_grad])
# Update weights
self.optimizer.zero_grad()
total_loss.backward()
self.learnable_triplane.grad.data = self.learnable_triplane.grad.data * self.learnable_triplane.numel()
self.optimizer.step()
meter.update(total_loss.item())
if i_step % 10 == 0:
log_line = f"Iter {i_step+1}: total_loss={meter.avg} "
for k, v in losses.items():
log_line = log_line + f" {k}={v.item()}, "
self.logger.add_scalar(f"train/{k}", v.item(), i_step)
print(log_line)
meter.reset()
@torch.no_grad()
def test_loop(self, inp, step=''):
self.model.eval()
# coeff_dict = np.load('data/processed/videos/Lieu/coeff_fit_mp_for_lora.npy', allow_pickle=True).tolist()
# drv_exps = torch.tensor(coeff_dict['exp']).cuda().float()
drv_exps = self.ds['exps']
zero_eulers = self.ds['eulers']*0
zero_trans = self.ds['trans']*0
batch_size = 1
num_samples = len(self.ds['cameras'])
video_writer = imageio.get_writer(os.path.join(inp['work_dir'], f'val_step{step}.mp4'), fps=25)
total_iters = min(num_samples, 250)
video_id = inp['video_id']
for i in tqdm.trange(total_iters,desc="testing lora..."):
drv_idx = [i]
drv_secc_colors = []
gt_imgs = []
segmaps = []
torso_imgs = []
drv_lip_rects = []
kp_src = []
kp_drv = []
for di in drv_idx:
# 读取target image
if self.torso_mode:
if self.ds['com_imgs'][di] is None:
img_name = f'data/processed/videos/{video_id}/com_imgs/{format(di, "08d")}.jpg'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['com_imgs'][di] = img
gt_imgs.append(self.ds['com_imgs'][di])
else:
if self.ds['head_imgs'][di] is None:
img_name = f'data/processed/videos/{video_id}/head_imgs/{format(di, "08d")}.png'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['head_imgs'][di] = img
gt_imgs.append(self.ds['head_imgs'][di])
# 使用第一帧的torso作为face v2v的输入
if self.ds['torso_imgs'][0] is None:
img_name = f'data/processed/videos/{video_id}/inpaint_torso_imgs/{format(0, "08d")}.png'
img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
self.ds['torso_imgs'][0] = img
torso_imgs.append(self.ds['torso_imgs'][0])
# 所以segmap也用第一帧的了
if self.ds['segmaps'][0] is None:
img_name = f'data/processed/videos/{video_id}/segmaps/{format(0, "08d")}.png'
seg_img = cv2.imread(img_name)[:,:, ::-1]
segmap = torch.from_numpy(decode_segmap_mask_from_image(seg_img)) # [6, H, W]
self.ds['segmaps'][0] = segmap
segmaps.append(self.ds['segmaps'][0])
drv_lip_rects.append(self.ds['lip_rects'][di])
kp_src.append(self.ds['kps'][0])
kp_drv.append(self.ds['kps'][di])
bg_img = self.ds['bg_img'].unsqueeze(0).repeat([batch_size, 1, 1, 1]).cuda()
ref_torso_imgs = torch.stack(torso_imgs).float().cuda()
kp_src = torch.stack(kp_src).float().cuda()
kp_drv = torch.stack(kp_drv).float().cuda()
segmaps = torch.stack(segmaps).float().cuda()
tgt_imgs = torch.stack(gt_imgs).float().cuda()
for di in drv_idx:
_, secc_color = self.secc_renderer(self.ds['id'][0:1], drv_exps[di:di+1], zero_eulers[0:1], zero_trans[0:1])
drv_secc_colors.append(secc_color)
drv_secc_color = torch.cat(drv_secc_colors)
cano_secc_color = self.ds['cano_secc_color'].repeat([batch_size, 1, 1, 1])
src_secc_color = self.ds['src_secc_color'].repeat([batch_size, 1, 1, 1])
cond = {'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_color,
'ref_torso_img': ref_torso_imgs, 'bg_img': bg_img, 'segmap': segmaps,
'kp_s': kp_src, 'kp_d': kp_drv}
camera = self.ds['cameras'][drv_idx]
gen_output = self.secc2video_model.forward(img=None, camera=camera, cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
pred_img = gen_output['image']
pred_img = ((pred_img.permute(0, 2, 3, 1) + 1)/2 * 255).int().cpu().numpy().astype(np.uint8)
video_writer.append_data(pred_img[0])
video_writer.close()
self.model.train()
def masked_error_loss(self, img_pred, img_gt, mask, unmasked_weight=0.1, mode='l1'):
# 对raw图像，因为deform的原因背景没法全黑，导致这部分mse过高，我们将其mask掉，只计算人脸部分
masked_weight = 1.0
weight_mask = mask.float() * masked_weight + (~mask).float() * unmasked_weight
if mode == 'l1':
error = (img_pred - img_gt).abs().sum(dim=1) * weight_mask
else:
error = (img_pred - img_gt).pow(2).sum(dim=1) * weight_mask
error.clamp_(0, max(0.5, error.quantile(0.8).item())) # clamp掉较高loss的pixel，避免姿态没对齐的pixel导致的异常值占主导影响训练
loss = error.mean()
return loss
def dilate(self, bin_img, ksize=5, mode='max_pool'):
"""
mode: max_pool or avg_pool
"""
# bin_img, [1, h, w]
pad = (ksize-1)//2
bin_img = F.pad(bin_img, pad=[pad,pad,pad,pad], mode='reflect')
if mode == 'max_pool':
out = F.max_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0)
else:
out = F.avg_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0)
return out
def dilate_mask(self, mask, ksize=21):
mask = self.dilate(mask, ksize=ksize, mode='max_pool')
return mask
def set_unmasked_to_black(self, img, mask):
out_img = img * mask.float() - (~mask).float() # -1 denotes black
return out_img
def dump_checkpoint(self, inp):
checkpoint = {}
# save optimizers
optimizer_states = []
self.optimizers = [self.optimizer]
for i, optimizer in enumerate(self.optimizers):
if optimizer is not None:
state_dict = optimizer.state_dict()
state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
optimizer_states.append(state_dict)
checkpoint['optimizer_states'] = optimizer_states
state_dict = {
'model': self.model.state_dict(),
'learnable_triplane': self.model.state_dict()['_last_cano_planes'],
}
del state_dict['model']['_last_cano_planes']
checkpoint['state_dict'] = state_dict
checkpoint['lora_args'] = self.lora_args
person_ds = {}
video_id = inp['video_id']
img_name = f'data/processed/videos/{video_id}/gt_imgs/{format(0, "08d")}.jpg'
gt_img = torch.tensor(cv2.resize(cv2.imread(img_name), (512, 512))[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
person_ds['gt_img'] = gt_img.reshape([1, 3, 512, 512])
person_ds['id'] = self.ds['id'].cpu().reshape([1, 80])
person_ds['src_kp'] = self.ds['kps'][0].cpu()
person_ds['video_id'] = inp['video_id']
checkpoint['person_ds'] = person_ds
return checkpoint
if __name__ == '__main__':
import argparse, glob, tqdm
parser = argparse.ArgumentParser()
parser.add_argument("--head_ckpt", default='') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
# parser.add_argument("--torso_ckpt", default='checkpoints/240210_real3dportrait_orig/secc2plane_torso_orig') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
parser.add_argument("--torso_ckpt", default='checkpoints/mimictalk_orig/os_secc2plane_torso') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
parser.add_argument("--video_id", default='data/raw/examples/GER.mp4', help="identity source, we support (1) already processed <video_id> of GeneFace, (2) video path, (3) image path")
parser.add_argument("--work_dir", default=None)
parser.add_argument("--max_updates", default=10000, type=int, help="for video, 2000 is good; for an image, 3~10 is good")
parser.add_argument("--test", action='store_true')
parser.add_argument("--batch_size", default=1, type=int, help="batch size during training, 1 needs 8GB, 2 needs 15GB")
parser.add_argument("--lr", default=0.001)
parser.add_argument("--lr_triplane", default=0.005, help="for video, 0.1; for an image, 0.001; for ablation with_triplane, 0.")
parser.add_argument("--lora_r", default=2, type=int, help="width of lora unit")
parser.add_argument("--lora_mode", default='secc2plane_sr', help='for video, full; for an image, none')
args = parser.parse_args()
inp = {
'head_ckpt': args.head_ckpt,
'torso_ckpt': args.torso_ckpt,
'video_id': args.video_id,
'work_dir': args.work_dir,
'max_updates': args.max_updates,
'batch_size': args.batch_size,
'test': args.test,
'lr': float(args.lr),
'lr_triplane': float(args.lr_triplane),
'lora_mode': args.lora_mode,
'lora_r': args.lora_r,
}
if inp['work_dir'] == None:
video_id = os.path.basename(inp['video_id'])[:-4] if inp['video_id'].endswith((".mp4", ".png", ".jpg", ".jpeg")) else inp['video_id']
inp['work_dir'] = f'checkpoints_mimictalk/{video_id}'
os.makedirs(inp['work_dir'], exist_ok=True)
trainer = LoRATrainer(inp)
if inp['test']:
trainer.test_loop(inp)
else:
trainer.training_loop(inp)
trainer.test_loop(inp)
print(" ")