nodes/PowerHouseMan-ComfyUI-AdvancedLivePortrait
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #12 from claussteinmassl/main

Browse Source 
Add support for Mac Users
This commit is contained in:
PowerHouseMan
2024-08-07 17:11:11 +09:00
committed by GitHub
parent 9ecc7d6a63 1891671340
commit 20484fd791
3 changed files with 56 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LivePortrait/live_portrait_wrapper.py
+33 -6
View File
@@ -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
LivePortrait/modules/dense_motion.py
+1 -1
View File
@@ -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
nodes.py
+22 -13
View File
@@ -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