rebase fixes

This commit is contained in:
Rafael Silva 2024-07-12 23:33:33 -04:00
parent 427a2637cf
commit c27351eb6d

View File

@ -32,20 +32,20 @@ class LivePortraitWrapper(object):
model_config = yaml.load(open(inference_cfg.models_config, 'r'), Loader=yaml.SafeLoader)
# init F
# self.appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device_id, 'appearance_feature_extractor')
self.appearance_feature_extractor = torch.jit.load("build/appearance_feature_extractor.pt")
self.appearance_feature_extractor = load_model(inference_cfg.checkpoint_F, model_config, self.device, 'appearance_feature_extractor')
# self.appearance_feature_extractor = torch.jit.load("build/appearance_feature_extractor.pt")
log(f'Load appearance_feature_extractor done.')
# init M
# self.motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device_id, 'motion_extractor')
self.motion_extractor = torch.jit.load("build/motion_extractor.pt")
self.motion_extractor = load_model(inference_cfg.checkpoint_M, model_config, self.device, 'motion_extractor')
# self.motion_extractor = torch.jit.load("build/motion_extractor.pt")
log(f'Load motion_extractor done.')
# init W
# self.warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module')
self.warping_module = torch.jit.load("build/warping_module.pt")
self.warping_module = load_model(inference_cfg.checkpoint_W, model_config, self.device, 'warping_module')
# self.warping_module = torch.jit.load("build/warping_module.pt")
log(f'Load warping_module done.')
# init G
# self.spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device_id, 'spade_generator')
self.spade_generator = torch.jit.load("build/spade_generator.pt")
self.spade_generator = load_model(inference_cfg.checkpoint_G, model_config, self.device, 'spade_generator')
# self.spade_generator = torch.jit.load("build/spade_generator.pt")
log(f'Load spade_generator done.')
# init S and R
if inference_cfg.checkpoint_S is not None and osp.exists(inference_cfg.checkpoint_S):