inference/train_mimictalk_on_a_video.py

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