From 189167134093cdde0cfdbde69faaf1c5fd7f7e73 Mon Sep 17 00:00:00 2001 From: claussteinmassl Date: Tue, 6 Aug 2024 12:09:14 +0200 Subject: [PATCH] Refactor calls to cuda to be more generic, so it works on cuda, mps and cpu, depending on the available device. --- LivePortrait/live_portrait_wrapper.py | 39 ++++++++++++++++++++++----- LivePortrait/modules/dense_motion.py | 2 +- nodes.py | 35 +++++++++++++++--------- 3 files changed, 56 insertions(+), 20 deletions(-) diff --git a/LivePortrait/live_portrait_wrapper.py b/LivePortrait/live_portrait_wrapper.py index 91c2145..173e24a 100644 --- a/LivePortrait/live_portrait_wrapper.py +++ b/LivePortrait/live_portrait_wrapper.py @@ -10,6 +10,18 @@ from .utils.camera import headpose_pred_to_degree, get_rotation_matrix from .config.inference_config import InferenceConfig from .utils.rprint import rlog as log +device_type = 'cpu' +def get_device(): + global device_type + if torch.cuda.is_available(): + device_type = 'cuda' + return torch.device('cuda') + elif torch.backends.mps.is_available(): + device_type = 'mps' + return torch.device('mps') + else: + device_type = 'cpu' + return torch.device('cpu') class LivePortraitWrapper(object): @@ -40,7 +52,7 @@ class LivePortraitWrapper(object): y = _imgs.astype(np.float32) / 255. y = np.clip(y, 0, 1) # clip to 0~1 y = torch.from_numpy(y).permute(0, 4, 3, 1, 2) # TxHxWx3x1 -> Tx1x3xHxW - y = y.cuda(self.device_id) + y = y.to(get_device()) return y @@ -49,8 +61,12 @@ class LivePortraitWrapper(object): x: Bx3xHxW, normalized to 0~1 """ with torch.no_grad(): - with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): - feature_3d = self.appearance_feature_extractor(x) + # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + if device_type == 'cuda' or device_type == 'cpu': + with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + feature_3d = self.appearance_feature_extractor(x) + else: + feature_3d = self.appearance_feature_extractor(x) return feature_3d.float() @@ -61,8 +77,12 @@ class LivePortraitWrapper(object): return: A dict contains keys: 'pitch', 'yaw', 'roll', 't', 'exp', 'scale', 'kp' """ with torch.no_grad(): - with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): - kp_info = self.motion_extractor(x) + # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + if device_type == 'cuda' or device_type == 'cpu': + with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + kp_info = self.motion_extractor(x) + else: + kp_info = self.motion_extractor(x) if self.cfg.flag_use_half_precision: # float the dict @@ -154,7 +174,14 @@ class LivePortraitWrapper(object): """ # The line 18 in Algorithm 1: D(W(f_s; x_s, x′_d,i)） with torch.no_grad(): - with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + if device_type == 'cuda' or device_type == 'cpu': + with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=self.cfg.flag_use_half_precision): + # get decoder input + ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving) + # decode + ret_dct['out'] = self.spade_generator(feature=ret_dct['out']) + else: # get decoder input ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving) # decode diff --git a/LivePortrait/modules/dense_motion.py b/LivePortrait/modules/dense_motion.py index 0eec0c4..c1a7f9a 100644 --- a/LivePortrait/modules/dense_motion.py +++ b/LivePortrait/modules/dense_motion.py @@ -59,7 +59,7 @@ class DenseMotionNetwork(nn.Module): heatmap = gaussian_driving - gaussian_source # (bs, num_kp, d, h, w) # adding background feature - zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type()).to(heatmap.device) + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.dtype).to(heatmap.device) heatmap = torch.cat([zeros, heatmap], dim=1) heatmap = heatmap.unsqueeze(2) # (bs, 1+num_kp, 1, d, h, w) return heatmap diff --git a/nodes.py b/nodes.py index c9f2a19..952d2ee 100644 --- a/nodes.py +++ b/nodes.py @@ -33,6 +33,14 @@ from .LivePortrait.modules.stitching_retargeting_network import StitchingRetarge from collections import OrderedDict +def get_device(): + if torch.cuda.is_available(): + return torch.device('cuda') + elif torch.backends.mps.is_available(): + return torch.device('mps') + else: + return torch.device('cpu') + def tensor2pil(image): return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) def pil2tensor(image): @@ -128,7 +136,9 @@ class LP_Engine: key.startswith(prefix)} return filtered_checkpoint - def load_model(self, model_config, device, model_type): + def load_model(self, model_config, model_type): + + device = get_device() if model_type == 'stitching_retargeting_module': ckpt_path = os.path.join(get_model_dir("liveportrait"), "retargeting_models", model_type + ".pth") @@ -144,13 +154,13 @@ class LP_Engine: "https://huggingface.co/Kijai/LivePortrait_safetensors/resolve/main/" + model_type + ".safetensors") model_params = model_config['model_params'][f'{model_type}_params'] if model_type == 'appearance_feature_extractor': - model = AppearanceFeatureExtractor(**model_params).cuda(device) + model = AppearanceFeatureExtractor(**model_params).to(device) elif model_type == 'motion_extractor': - model = MotionExtractor(**model_params).cuda(device) + model = MotionExtractor(**model_params).to(device) elif model_type == 'warping_module': - model = WarpingNetwork(**model_params).cuda(device) + model = WarpingNetwork(**model_params).to(device) elif model_type == 'spade_generator': - model = SPADEDecoder(**model_params).cuda(device) + model = SPADEDecoder(**model_params).to(device) elif model_type == 'stitching_retargeting_module': # Special handling for stitching and retargeting module config = model_config['model_params']['stitching_retargeting_module_params'] @@ -197,12 +207,11 @@ class LP_Engine: model_config_path = os.path.join(current_directory, 'LivePortrait', 'config', 'models.yaml') model_config = yaml.safe_load(open(model_config_path, 'r')) - device_id = 0 - appearance_feature_extractor = self.load_model(model_config, device_id, 'appearance_feature_extractor') - motion_extractor = self.load_model(model_config, device_id, 'motion_extractor') - warping_module = self.load_model(model_config, device_id, 'warping_module') - spade_generator = self.load_model(model_config, device_id, 'spade_generator') - stitching_retargeting_module = self.load_model(model_config, device_id, 'stitching_retargeting_module') + appearance_feature_extractor = self.load_model(model_config, 'appearance_feature_extractor') + motion_extractor = self.load_model(model_config, 'motion_extractor') + warping_module = self.load_model(model_config, 'warping_module') + spade_generator = self.load_model(model_config, 'spade_generator') + stitching_retargeting_module = self.load_model(model_config, 'stitching_retargeting_module') self.pipeline = LivePortraitWrapper(InferenceConfig(), appearance_feature_extractor, motion_extractor, warping_module, spade_generator, stitching_retargeting_module) @@ -325,7 +334,7 @@ class LP_Engine: raise ValueError(f'img ndim should be 3 or 4: {x.ndim}') x = np.clip(x, 0, 1) # clip to 0~1 x = torch.from_numpy(x).permute(0, 3, 1, 2) # 1xHxWx3 -> 1x3xHxW - x = x.cuda() + x = x.to(get_device()) return x def GetMaskImg(self): @@ -467,7 +476,7 @@ class ExpressionSet: self.s = erst[2] self.t = erst[3] else: - self.e = torch.from_numpy(np.zeros((1, 21, 3))).float().to(device='cuda') + self.e = torch.from_numpy(np.zeros((1, 21, 3))).float().to(get_device()) self.r = torch.Tensor([0, 0, 0]) self.s = 0 self.t = 0