cnywt/SyncTalk
0
1import argparse2 3from nerf_triplane.provider import NeRFDataset4from nerf_triplane.utils import *5from nerf_triplane.network import NeRFNetwork6 7# torch.autograd.set_detect_anomaly(True)8# Close tf32 features. Fix low numerical accuracy on rtx30xx gpu.9try:10 torch.backends.cuda.matmul.allow_tf32 = False11 torch.backends.cudnn.allow_tf32 = False12except AttributeError as e:13 print('Info. This pytorch version is not support with tf32.')14 15if __name__ == '__main__':16 17 parser = argparse.ArgumentParser()18 parser.add_argument('path', type=str)19 parser.add_argument('-O', action='store_true', help="equals --fp16 --cuda_ray --exp_eye")20 parser.add_argument('--test', action='store_true', help="test mode (load model and test dataset)")21 parser.add_argument('--test_train', action='store_true', help="test mode (load model and train dataset)")22 parser.add_argument('--data_range', type=int, nargs='*', default=[0, -1], help="data range to use")23 parser.add_argument('--workspace', type=str, default='workspace')24 parser.add_argument('--seed', type=int, default=0)25 26 ### training options27 parser.add_argument('--iters', type=int, default=200000, help="training iters")28 parser.add_argument('--lr', type=float, default=1e-2, help="initial learning rate")29 parser.add_argument('--lr_net', type=float, default=1e-3, help="initial learning rate")30 parser.add_argument('--ckpt', type=str, default='latest')31 parser.add_argument('--num_rays', type=int, default=4096 * 16, help="num rays sampled per image for each training step")32 parser.add_argument('--cuda_ray', action='store_true', help="use CUDA raymarching instead of pytorch")33 parser.add_argument('--max_steps', type=int, default=16, help="max num steps sampled per ray (only valid when using --cuda_ray)")34 parser.add_argument('--num_steps', type=int, default=16, help="num steps sampled per ray (only valid when NOT using --cuda_ray)")35 parser.add_argument('--upsample_steps', type=int, default=0, help="num steps up-sampled per ray (only valid when NOT using --cuda_ray)")36 parser.add_argument('--update_extra_interval', type=int, default=16, help="iter interval to update extra status (only valid when using --cuda_ray)")37 parser.add_argument('--max_ray_batch', type=int, default=4096, help="batch size of rays at inference to avoid OOM (only valid when NOT using --cuda_ray)")38 39 ### loss set40 parser.add_argument('--warmup_step', type=int, default=10000, help="warm up steps")41 parser.add_argument('--amb_aud_loss', type=int, default=1, help="use ambient aud loss")42 parser.add_argument('--amb_eye_loss', type=int, default=1, help="use ambient eye loss")43 parser.add_argument('--unc_loss', type=int, default=1, help="use uncertainty loss")44 parser.add_argument('--lambda_amb', type=float, default=1e-4, help="lambda for ambient loss")45 parser.add_argument('--pyramid_loss', type=int, default=0, help="use perceptual loss")46 47 ### network backbone options48 parser.add_argument('--fp16', action='store_true', help="use amp mixed precision training")49 50 parser.add_argument('--bg_img', type=str, default='', help="background image")51 parser.add_argument('--fbg', action='store_true', help="frame-wise bg")52 parser.add_argument('--exp_eye', action='store_true', help="explicitly control the eyes")53 parser.add_argument('--fix_eye', type=float, default=-1, help="fixed eye area, negative to disable, set to 0-0.3 for a reasonable eye")54 parser.add_argument('--smooth_eye', action='store_true', help="smooth the eye area sequence")55 parser.add_argument('--bs_area', type=str, default="upper", help="upper or eye")56 parser.add_argument('--au45', action='store_true', help="use openface au45")57 parser.add_argument('--torso_shrink', type=float, default=0.8, help="shrink bg coords to allow more flexibility in deform")58 59 ### dataset options60 parser.add_argument('--color_space', type=str, default='srgb', help="Color space, supports (linear, srgb)")61 parser.add_argument('--preload', type=int, default=0, help="0 means load data from disk on-the-fly, 1 means preload to CPU, 2 means GPU.")62 # (the default value is for the fox dataset)63 parser.add_argument('--bound', type=float, default=1, help="assume the scene is bounded in box[-bound, bound]^3, if > 1, will invoke adaptive ray marching.")64 parser.add_argument('--scale', type=float, default=4, help="scale camera location into box[-bound, bound]^3")65 parser.add_argument('--offset', type=float, nargs='*', default=[0, 0, 0], help="offset of camera location")66 parser.add_argument('--dt_gamma', type=float, default=1/256, help="dt_gamma (>=0) for adaptive ray marching. set to 0 to disable, >0 to accelerate rendering (but usually with worse quality)")67 parser.add_argument('--min_near', type=float, default=0.05, help="minimum near distance for camera")68 parser.add_argument('--density_thresh', type=float, default=10, help="threshold for density grid to be occupied (sigma)")69 parser.add_argument('--density_thresh_torso', type=float, default=0.01, help="threshold for density grid to be occupied (alpha)")70 parser.add_argument('--patch_size', type=int, default=1, help="[experimental] render patches in training, so as to apply LPIPS loss. 1 means disabled, use [64, 32, 16] to enable")71 72 parser.add_argument('--init_lips', action='store_true', help="init lips region")73 parser.add_argument('--finetune_lips', action='store_true', help="use LPIPS and landmarks to fine tune lips region")74 parser.add_argument('--smooth_lips', action='store_true', help="smooth the enc_a in a exponential decay way...")75 76 parser.add_argument('--torso', action='store_true', help="fix head and train torso")77 parser.add_argument('--head_ckpt', type=str, default='', help="head model")78 79 ### GUI options80 parser.add_argument('--gui', action='store_true', help="start a GUI")81 parser.add_argument('--W', type=int, default=450, help="GUI width")82 parser.add_argument('--H', type=int, default=450, help="GUI height")83 parser.add_argument('--radius', type=float, default=3.35, help="default GUI camera radius from center")84 parser.add_argument('--fovy', type=float, default=21.24, help="default GUI camera fovy")85 parser.add_argument('--max_spp', type=int, default=1, help="GUI rendering max sample per pixel")86 87 ### else88 parser.add_argument('--att', type=int, default=2, help="audio attention mode (0 = turn off, 1 = left-direction, 2 = bi-direction)")89 parser.add_argument('--aud', type=str, default='', help="audio source (empty will load the default, else should be a path to a npy file)")90 parser.add_argument('--emb', action='store_true', help="use audio class + embedding instead of logits")91 parser.add_argument('--portrait', action='store_true', help="only render face")92 parser.add_argument('--ind_dim', type=int, default=4, help="individual code dim, 0 to turn off")93 parser.add_argument('--ind_num', type=int, default=20000, help="number of individual codes, should be larger than training dataset size")94 95 parser.add_argument('--ind_dim_torso', type=int, default=8, help="individual code dim, 0 to turn off")96 97 parser.add_argument('--amb_dim', type=int, default=2, help="ambient dimension")98 parser.add_argument('--part', action='store_true', help="use partial training data (1/10)")99 parser.add_argument('--part2', action='store_true', help="use partial training data (first 15s)")100 101 parser.add_argument('--train_camera', action='store_true', help="optimize camera pose")102 parser.add_argument('--smooth_path', action='store_true', help="brute-force smooth camera pose trajectory with a window size")103 parser.add_argument('--smooth_path_window', type=int, default=7, help="smoothing window size")104 105 # asr106 parser.add_argument('--asr', action='store_true', help="load asr for real-time app")107 parser.add_argument('--asr_wav', type=str, default='', help="load the wav and use as input")108 parser.add_argument('--asr_play', action='store_true', help="play out the audio")109 110 parser.add_argument('--asr_model', type=str, default='deepspeech')111 112 parser.add_argument('--asr_save_feats', action='store_true')113 # audio FPS114 parser.add_argument('--fps', type=int, default=50)115 # sliding window left-middle-right length (unit: 20ms)116 parser.add_argument('-l', type=int, default=10)117 parser.add_argument('-m', type=int, default=50)118 parser.add_argument('-r', type=int, default=10)119 120 opt = parser.parse_args()121 122 if opt.O:123 opt.fp16 = True124 opt.exp_eye = True125 126 if opt.test and False:127 opt.smooth_path = True128 opt.smooth_eye = True129 opt.smooth_lips = True130 131 opt.cuda_ray = True132 # assert opt.cuda_ray, "Only support CUDA ray mode."133 134 if opt.patch_size > 1:135 # assert opt.patch_size > 16, "patch_size should > 16 to run LPIPS loss."136 assert opt.num_rays % (opt.patch_size ** 2) == 0, "patch_size ** 2 should be dividable by num_rays."137 138 # if opt.finetune_lips:139 # # do not update density grid in finetune stage140 # opt.update_extra_interval = 1e9141 142 print(opt)143 144 seed_everything(opt.seed)145 146 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')147 148 model = NeRFNetwork(opt)149 150 # manually load state dict for head151 if opt.torso and opt.head_ckpt != '':152 153 model_dict = torch.load(opt.head_ckpt, map_location='cpu')['model']154 155 missing_keys, unexpected_keys = model.load_state_dict(model_dict, strict=False)156 157 if len(missing_keys) > 0:158 print(f"[WARN] missing keys: {missing_keys}")159 if len(unexpected_keys) > 0:160 print(f"[WARN] unexpected keys: {unexpected_keys}") 161 162 # freeze these keys163 for k, v in model.named_parameters():164 if k in model_dict:165 print(f'[INFO] freeze {k}, {v.shape}')166 v.requires_grad = False167 168 169 # print(model)170 171 # criterion = torch.nn.MSELoss(reduction='none')172 criterion = torch.nn.L1Loss(reduction='none')173 174 175 if opt.test:176 177 if opt.gui:178 metrics = [] # use no metric in GUI for faster initialization...179 else:180 # metrics = [PSNRMeter(), LPIPSMeter(device=device)]181 metrics = [PSNRMeter(), LPIPSMeter(device=device), LMDMeter(backend='fan')]182 183 trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, criterion=criterion, fp16=opt.fp16, metrics=metrics, use_checkpoint=opt.ckpt)184 185 if opt.test_train:186 test_set = NeRFDataset(opt, device=device, type='train')187 # a manual fix to test on the training dataset188 test_set.training = False 189 test_set.num_rays = -1190 test_loader = test_set.dataloader()191 else:192 test_loader = NeRFDataset(opt, device=device, type='test').dataloader()193 194 195 # temp fix: for update_extra_states196 model.aud_features = test_loader._data.auds197 model.eye_areas = test_loader._data.eye_area198 199 if opt.gui:200 from nerf_triplane.gui import NeRFGUI201 # we still need test_loader to provide audio features for testing.202 with NeRFGUI(opt, trainer, test_loader) as gui:203 gui.render()204 205 else:206 ### test and save video (fast) 207 trainer.test(test_loader)208 209 ### evaluate metrics (slow)210 if test_loader.has_gt:211 trainer.evaluate(test_loader)212 213 214 215 else:216 217 optimizer = lambda model: torch.optim.AdamW(model.get_params(opt.lr, opt.lr_net), betas=(0, 0.99), eps=1e-8)218 219 train_loader = NeRFDataset(opt, device=device, type='train').dataloader()220 221 assert len(train_loader) < opt.ind_num, f"[ERROR] dataset too many frames: {len(train_loader)}, please increase --ind_num to this number!"222 223 # temp fix: for update_extra_states224 model.aud_features = train_loader._data.auds225 model.eye_area = train_loader._data.eye_area226 model.poses = train_loader._data.poses227 228 # decay to 0.1 * init_lr at last iter step229 if opt.finetune_lips:230 scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer, lambda iter: 0.05 ** (iter / opt.iters))231 else:232 scheduler = lambda optimizer: optim.lr_scheduler.LambdaLR(optimizer, lambda iter: 0.5 ** (iter / opt.iters))233 234 metrics = [PSNRMeter(), LPIPSMeter(device=device),LMDMeter(backend='fan')]235 236 eval_interval = max(1, int(5000 / len(train_loader)))237 trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, optimizer=optimizer, criterion=criterion, ema_decay=0.95, fp16=opt.fp16, lr_scheduler=scheduler, scheduler_update_every_step=True, metrics=metrics, use_checkpoint=opt.ckpt, eval_interval=eval_interval)238 with open(os.path.join(opt.workspace, 'opt.txt'), 'a') as f:239 f.write(str(opt))240 if opt.gui:241 with NeRFGUI(opt, trainer, train_loader) as gui:242 gui.render()243 244 else:245 valid_loader = NeRFDataset(opt, device=device, type='val', downscale=1).dataloader()246 247 max_epochs = np.ceil(opt.iters / len(train_loader)).astype(np.int32)248 print(f'[INFO] max_epoch = {max_epochs}')249 trainer.train(train_loader, valid_loader, max_epochs)250 251 # free some mem252 del train_loader, valid_loader253 torch.cuda.empty_cache()254 255 # also test256 test_loader = NeRFDataset(opt, device=device, type='test').dataloader()257 258 if test_loader.has_gt:259 trainer.evaluate(test_loader) # blender has gt, so evaluate it.260 261 trainer.test(test_loader)262