""" 将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 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(" ")