From 301164ef66a03d8d02c3ce4c917ceb60c6fa6d12 Mon Sep 17 00:00:00 2001 From: PowerHouseMan Date: Fri, 2 Aug 2024 00:19:43 +0900 Subject: [PATCH] main --- LivePortrait/src/live_portrait_wrapper.py | 214 +++++ LivePortrait/src/modules/__init__.py | 0 .../modules/appearance_feature_extractor.py | 48 ++ LivePortrait/src/modules/convnextv2.py | 149 ++++ LivePortrait/src/modules/dense_motion.py | 104 +++ LivePortrait/src/modules/motion_extractor.py | 35 + LivePortrait/src/modules/spade_generator.py | 59 ++ .../modules/stitching_retargeting_network.py | 38 + LivePortrait/src/modules/util.py | 441 ++++++++++ LivePortrait/src/modules/warping_network.py | 77 ++ LivePortrait/src/utils/__init__.py | 0 LivePortrait/src/utils/camera.py | 75 ++ LivePortrait/src/utils/face_analysis_diy.py | 78 ++ LivePortrait/src/utils/helper.py | 175 ++++ LivePortrait/src/utils/io.py | 97 +++ .../src/utils/resources/mask_template.png | Bin 0 -> 3470 bytes LivePortrait/src/utils/rprint.py | 16 + LivePortrait/src/utils/timer.py | 29 + LivePortrait/src/utils/video.py | 142 ++++ __init__.py | 49 ++ config/inference_config.py | 21 + config/models.yaml | 43 + exp_data/입술내밀기.exp | Bin 0 -> 1037 bytes install.bat | 20 + nodes.py | 770 ++++++++++++++++++ requirements.txt | 10 + 26 files changed, 2690 insertions(+) create mode 100644 LivePortrait/src/live_portrait_wrapper.py create mode 100644 LivePortrait/src/modules/__init__.py create mode 100644 LivePortrait/src/modules/appearance_feature_extractor.py create mode 100644 LivePortrait/src/modules/convnextv2.py create mode 100644 LivePortrait/src/modules/dense_motion.py create mode 100644 LivePortrait/src/modules/motion_extractor.py create mode 100644 LivePortrait/src/modules/spade_generator.py create mode 100644 LivePortrait/src/modules/stitching_retargeting_network.py create mode 100644 LivePortrait/src/modules/util.py create mode 100644 LivePortrait/src/modules/warping_network.py create mode 100644 LivePortrait/src/utils/__init__.py create mode 100644 LivePortrait/src/utils/camera.py create mode 100644 LivePortrait/src/utils/face_analysis_diy.py create mode 100644 LivePortrait/src/utils/helper.py create mode 100644 LivePortrait/src/utils/io.py create mode 100644 LivePortrait/src/utils/resources/mask_template.png create mode 100644 LivePortrait/src/utils/rprint.py create mode 100644 LivePortrait/src/utils/timer.py create mode 100644 LivePortrait/src/utils/video.py create mode 100644 __init__.py create mode 100644 config/inference_config.py create mode 100644 config/models.yaml create mode 100644 exp_data/입술내밀기.exp create mode 100644 install.bat create mode 100644 nodes.py create mode 100644 requirements.txt diff --git a/LivePortrait/src/live_portrait_wrapper.py b/LivePortrait/src/live_portrait_wrapper.py new file mode 100644 index 0000000..935b03e --- /dev/null +++ b/LivePortrait/src/live_portrait_wrapper.py @@ -0,0 +1,214 @@ +import os.path as osp +import numpy as np +import cv2 +import torch +import yaml + +from LivePortrait.src.utils.timer import Timer +from LivePortrait.src.utils.helper import load_model, concat_feat +from LivePortrait.src.utils.camera import headpose_pred_to_degree, get_rotation_matrix +from config.inference_config import InferenceConfig +from LivePortrait.src.utils.rprint import rlog as log + + +class LivePortraitWrapper(object): + + def __init__(self, cfg: InferenceConfig): + + model_config = yaml.load(open(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') + log(f'Load appearance_feature_extractor done.') + # init M + self.motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device_id, 'motion_extractor') + log(f'Load motion_extractor done.') + # init W + self.warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module') + log(f'Load warping_module done.') + # init G + self.spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device_id, 'spade_generator') + log(f'Load spade_generator done.') + # init S and R + if cfg.checkpoint_S is not None and osp.exists(cfg.checkpoint_S): + self.stitching_retargeting_module = load_model(cfg.checkpoint_S, model_config, cfg.device_id, 'stitching_retargeting_module') + log(f'Load stitching_retargeting_module done.') + else: + self.stitching_retargeting_module = None + + self.cfg = cfg + self.device_id = cfg.device_id + self.timer = Timer() + + def prepare_source(self, img: np.ndarray) -> torch.Tensor: + """ construct the input as standard + img: HxWx3, uint8, 256x256 + """ + h, w = img.shape[:2] + if h != self.cfg.input_shape[0] or w != self.cfg.input_shape[1]: + x = cv2.resize(img, (self.cfg.input_shape[0], self.cfg.input_shape[1])) + else: + x = img.copy() + + if x.ndim == 3: + x = x[np.newaxis].astype(np.float32) / 255. # HxWx3 -> 1xHxWx3, normalized to 0~1 + elif x.ndim == 4: + x = x.astype(np.float32) / 255. # BxHxWx3, normalized to 0~1 + else: + 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(self.device_id) + return x + + def prepare_driving_videos(self, imgs) -> torch.Tensor: + """ construct the input as standard + imgs: NxBxHxWx3, uint8 + """ + if isinstance(imgs, list): + _imgs = np.array(imgs)[..., np.newaxis] # TxHxWx3x1 + elif isinstance(imgs, np.ndarray): + _imgs = imgs + else: + raise ValueError(f'imgs type error: {type(imgs)}') + + 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) + + return y + + def extract_feature_3d(self, x: torch.Tensor) -> torch.Tensor: + """ get the appearance feature of the image by F + 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) + + return feature_3d.float() + + def get_kp_info(self, x: torch.Tensor, **kwargs) -> dict: + """ get the implicit keypoint information + x: Bx3xHxW, normalized to 0~1 + flag_refine_info: whether to trandform the pose to degrees and the dimention of the reshape + 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) + + if self.cfg.flag_use_half_precision: + # float the dict + for k, v in kp_info.items(): + if isinstance(v, torch.Tensor): + kp_info[k] = v.float() + + flag_refine_info: bool = kwargs.get('flag_refine_info', True) + if flag_refine_info: + bs = kp_info['kp'].shape[0] + kp_info['pitch'] = headpose_pred_to_degree(kp_info['pitch'])[:, None] # Bx1 + kp_info['yaw'] = headpose_pred_to_degree(kp_info['yaw'])[:, None] # Bx1 + kp_info['roll'] = headpose_pred_to_degree(kp_info['roll'])[:, None] # Bx1 + kp_info['kp'] = kp_info['kp'].reshape(bs, -1, 3) # BxNx3 + kp_info['exp'] = kp_info['exp'].reshape(bs, -1, 3) # BxNx3 + + return kp_info + def transform_keypoint(self, kp_info: dict): + """ + transform the implicit keypoints with the pose, shift, and expression deformation + kp: BxNx3 + """ + kp = kp_info['kp'] # (bs, k, 3) + pitch, yaw, roll = kp_info['pitch'], kp_info['yaw'], kp_info['roll'] + + t, exp = kp_info['t'], kp_info['exp'] + scale = kp_info['scale'] + + pitch = headpose_pred_to_degree(pitch) + yaw = headpose_pred_to_degree(yaw) + roll = headpose_pred_to_degree(roll) + + bs = kp.shape[0] + if kp.ndim == 2: + num_kp = kp.shape[1] // 3 # Bx(num_kpx3) + else: + num_kp = kp.shape[1] # Bxnum_kpx3 + + rot_mat = get_rotation_matrix(pitch, yaw, roll) # (bs, 3, 3) + + # Eqn.2: s * (R * x_c,s + exp) + t + kp_transformed = kp.view(bs, num_kp, 3) @ rot_mat + exp.view(bs, num_kp, 3) + kp_transformed *= scale[..., None] # (bs, k, 3) * (bs, 1, 1) = (bs, k, 3) + kp_transformed[:, :, 0:2] += t[:, None, 0:2] # remove z, only apply tx ty + + return kp_transformed + + def stitch(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor: + """ + kp_source: BxNx3 + kp_driving: BxNx3 + Return: Bx(3*num_kp+2) + """ + feat_stiching = concat_feat(kp_source, kp_driving) + + with torch.no_grad(): + delta = self.stitching_retargeting_module['stitching'](feat_stiching) + + return delta + + def stitching(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor: + """ conduct the stitching + kp_source: Bxnum_kpx3 + kp_driving: Bxnum_kpx3 + """ + + if self.stitching_retargeting_module is not None: + + bs, num_kp = kp_source.shape[:2] + + kp_driving_new = kp_driving.clone() + delta = self.stitch(kp_source, kp_driving_new) + + delta_exp = delta[..., :3*num_kp].reshape(bs, num_kp, 3) # 1x20x3 + delta_tx_ty = delta[..., 3*num_kp:3*num_kp+2].reshape(bs, 1, 2) # 1x1x2 + + kp_driving_new += delta_exp + kp_driving_new[..., :2] += delta_tx_ty + + return kp_driving_new + + return kp_driving + + def warp_decode(self, feature_3d: torch.Tensor, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor: + """ get the image after the warping of the implicit keypoints + feature_3d: Bx32x16x64x64, feature volume + kp_source: BxNx3 + kp_driving: BxNx3 + """ + # 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): + # 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']) + + # float the dict + if self.cfg.flag_use_half_precision: + for k, v in ret_dct.items(): + if isinstance(v, torch.Tensor): + ret_dct[k] = v.float() + + return ret_dct + + def parse_output(self, out: torch.Tensor) -> np.ndarray: + """ construct the output as standard + return: 1xHxWx3, uint8 + """ + out = np.transpose(out.data.cpu().numpy(), [0, 2, 3, 1]) # 1x3xHxW -> 1xHxWx3 + out = np.clip(out, 0, 1) # clip to 0~1 + out = np.clip(out * 255, 0, 255).astype(np.uint8) # 0~1 -> 0~255 + + return out diff --git a/LivePortrait/src/modules/__init__.py b/LivePortrait/src/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/LivePortrait/src/modules/appearance_feature_extractor.py b/LivePortrait/src/modules/appearance_feature_extractor.py new file mode 100644 index 0000000..8d89e4f --- /dev/null +++ b/LivePortrait/src/modules/appearance_feature_extractor.py @@ -0,0 +1,48 @@ +# coding: utf-8 + +""" +Appearance extractor(F) defined in paper, which maps the source image s to a 3D appearance feature volume. +""" + +import torch +from torch import nn +from .util import SameBlock2d, DownBlock2d, ResBlock3d + + +class AppearanceFeatureExtractor(nn.Module): + + def __init__(self, image_channel, block_expansion, num_down_blocks, max_features, reshape_channel, reshape_depth, num_resblocks): + super(AppearanceFeatureExtractor, self).__init__() + self.image_channel = image_channel + self.block_expansion = block_expansion + self.num_down_blocks = num_down_blocks + self.max_features = max_features + self.reshape_channel = reshape_channel + self.reshape_depth = reshape_depth + + self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(3, 3), padding=(1, 1)) + + down_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.down_blocks = nn.ModuleList(down_blocks) + + self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1) + + self.resblocks_3d = torch.nn.Sequential() + for i in range(num_resblocks): + self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1)) + + def forward(self, source_image): + out = self.first(source_image) # Bx3x256x256 -> Bx64x256x256 + + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + out = self.second(out) + bs, c, h, w = out.shape # ->Bx512x64x64 + + f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w) # ->Bx32x16x64x64 + f_s = self.resblocks_3d(f_s) # ->Bx32x16x64x64 + return f_s diff --git a/LivePortrait/src/modules/convnextv2.py b/LivePortrait/src/modules/convnextv2.py new file mode 100644 index 0000000..83ea126 --- /dev/null +++ b/LivePortrait/src/modules/convnextv2.py @@ -0,0 +1,149 @@ +# coding: utf-8 + +""" +This moudle is adapted to the ConvNeXtV2 version for the extraction of implicit keypoints, poses, and expression deformation. +""" + +import torch +import torch.nn as nn +# from timm.models.layers import trunc_normal_, DropPath +from .util import LayerNorm, DropPath, trunc_normal_, GRN + +__all__ = ['convnextv2_tiny'] + + +class Block(nn.Module): + """ ConvNeXtV2 Block. + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + """ + + def __init__(self, dim, drop_path=0.): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.norm = LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + + +class ConvNeXtV2(nn.Module): + """ ConvNeXt V2 + + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + + def __init__( + self, + in_chans=3, + depths=[3, 3, 9, 3], + dims=[96, 192, 384, 768], + drop_path_rate=0., + **kwargs + ): + super().__init__() + self.depths = depths + self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format="channels_first") + ) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(4): + stage = nn.Sequential( + *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + + # NOTE: the output semantic items + num_bins = kwargs.get('num_bins', 66) + num_kp = kwargs.get('num_kp', 24) # the number of implicit keypoints + self.fc_kp = nn.Linear(dims[-1], 3 * num_kp) # implicit keypoints + + # print('dims[-1]: ', dims[-1]) + self.fc_scale = nn.Linear(dims[-1], 1) # scale + self.fc_pitch = nn.Linear(dims[-1], num_bins) # pitch bins + self.fc_yaw = nn.Linear(dims[-1], num_bins) # yaw bins + self.fc_roll = nn.Linear(dims[-1], num_bins) # roll bins + self.fc_t = nn.Linear(dims[-1], 3) # translation + self.fc_exp = nn.Linear(dims[-1], 3 * num_kp) # expression / delta + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x) + x = self.stages[i](x) + return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + + # implicit keypoints + kp = self.fc_kp(x) + + # pose and expression deformation + pitch = self.fc_pitch(x) + yaw = self.fc_yaw(x) + roll = self.fc_roll(x) + t = self.fc_t(x) + exp = self.fc_exp(x) + scale = self.fc_scale(x) + + ret_dct = { + 'pitch': pitch, + 'yaw': yaw, + 'roll': roll, + 't': t, + 'exp': exp, + 'scale': scale, + + 'kp': kp, # canonical keypoint + } + + return ret_dct + + +def convnextv2_tiny(**kwargs): + model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs) + return model diff --git a/LivePortrait/src/modules/dense_motion.py b/LivePortrait/src/modules/dense_motion.py new file mode 100644 index 0000000..0eec0c4 --- /dev/null +++ b/LivePortrait/src/modules/dense_motion.py @@ -0,0 +1,104 @@ +# coding: utf-8 + +""" +The module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving +""" + +from torch import nn +import torch.nn.functional as F +import torch +from .util import Hourglass, make_coordinate_grid, kp2gaussian + + +class DenseMotionNetwork(nn.Module): + def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, estimate_occlusion_map=True): + super(DenseMotionNetwork, self).__init__() + self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks) # ~60+G + + self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3) # 65G! NOTE: computation cost is large + self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1) # 0.8G + self.norm = nn.BatchNorm3d(compress, affine=True) + self.num_kp = num_kp + self.flag_estimate_occlusion_map = estimate_occlusion_map + + if self.flag_estimate_occlusion_map: + self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3) + else: + self.occlusion = None + + def create_sparse_motions(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape # (bs, 4, 16, 64, 64) + identity_grid = make_coordinate_grid((d, h, w), ref=kp_source) # (16, 64, 64, 3) + identity_grid = identity_grid.view(1, 1, d, h, w, 3) # (1, 1, d=16, h=64, w=64, 3) + coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3) + + k = coordinate_grid.shape[1] + + # NOTE: there lacks an one-order flow + driving_to_source = coordinate_grid + kp_source.view(bs, self.num_kp, 1, 1, 1, 3) # (bs, num_kp, d, h, w, 3) + + # adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1) + sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) # (bs, 1+num_kp, d, h, w, 3) + return sparse_motions + + def create_deformed_feature(self, feature, sparse_motions): + bs, _, d, h, w = feature.shape + feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1) # (bs, num_kp+1, 1, c, d, h, w) + feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w) # (bs*(num_kp+1), c, d, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1)) # (bs*(num_kp+1), d, h, w, 3) + sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False) + sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w) + + return sparse_deformed + + def create_heatmap_representations(self, feature, kp_driving, kp_source): + spatial_size = feature.shape[3:] # (d=16, h=64, w=64) + gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w) + gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w) + 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) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) # (bs, 1+num_kp, 1, d, h, w) + return heatmap + + def forward(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape # (bs, 32, 16, 64, 64) + + feature = self.compress(feature) # (bs, 4, 16, 64, 64) + feature = self.norm(feature) # (bs, 4, 16, 64, 64) + feature = F.relu(feature) # (bs, 4, 16, 64, 64) + + out_dict = dict() + + # 1. deform 3d feature + sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source) # (bs, 1+num_kp, d, h, w, 3) + deformed_feature = self.create_deformed_feature(feature, sparse_motion) # (bs, 1+num_kp, c=4, d=16, h=64, w=64) + + # 2. (bs, 1+num_kp, d, h, w) + heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source) # (bs, 1+num_kp, 1, d, h, w) + + input = torch.cat([heatmap, deformed_feature], dim=2) # (bs, 1+num_kp, c=5, d=16, h=64, w=64) + input = input.view(bs, -1, d, h, w) # (bs, (1+num_kp)*c=105, d=16, h=64, w=64) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) # (bs, 1+num_kp, d=16, h=64, w=64) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w) + sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4) # (bs, num_kp+1, 3, d, h, w) + deformation = (sparse_motion * mask).sum(dim=1) # (bs, 3, d, h, w) mask take effect in this place + deformation = deformation.permute(0, 2, 3, 4, 1) # (bs, d, h, w, 3) + + out_dict['deformation'] = deformation + + if self.flag_estimate_occlusion_map: + bs, _, d, h, w = prediction.shape + prediction_reshape = prediction.view(bs, -1, h, w) + occlusion_map = torch.sigmoid(self.occlusion(prediction_reshape)) # Bx1x64x64 + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/LivePortrait/src/modules/motion_extractor.py b/LivePortrait/src/modules/motion_extractor.py new file mode 100644 index 0000000..b2982e5 --- /dev/null +++ b/LivePortrait/src/modules/motion_extractor.py @@ -0,0 +1,35 @@ +# coding: utf-8 + +""" +Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image +""" + +from torch import nn +import torch + +from .convnextv2 import convnextv2_tiny +from .util import filter_state_dict + +model_dict = { + 'convnextv2_tiny': convnextv2_tiny, +} + + +class MotionExtractor(nn.Module): + def __init__(self, **kwargs): + super(MotionExtractor, self).__init__() + + # default is convnextv2_base + backbone = kwargs.get('backbone', 'convnextv2_tiny') + self.detector = model_dict.get(backbone)(**kwargs) + + def load_pretrained(self, init_path: str): + if init_path not in (None, ''): + state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model'] + state_dict = filter_state_dict(state_dict, remove_name='head') + ret = self.detector.load_state_dict(state_dict, strict=False) + print(f'Load pretrained model from {init_path}, ret: {ret}') + + def forward(self, x): + out = self.detector(x) + return out diff --git a/LivePortrait/src/modules/spade_generator.py b/LivePortrait/src/modules/spade_generator.py new file mode 100644 index 0000000..147a9ae --- /dev/null +++ b/LivePortrait/src/modules/spade_generator.py @@ -0,0 +1,59 @@ +# coding: utf-8 + +""" +Spade decoder(G) defined in the paper, which input the warped feature to generate the animated image. +""" + +import torch +from torch import nn +import torch.nn.functional as F +from .util import SPADEResnetBlock + + +class SPADEDecoder(nn.Module): + def __init__(self, upscale=1, max_features=256, block_expansion=64, out_channels=64, num_down_blocks=2): + for i in range(num_down_blocks): + input_channels = min(max_features, block_expansion * (2 ** (i + 1))) + self.upscale = upscale + super().__init__() + norm_G = 'spadespectralinstance' + label_num_channels = input_channels # 256 + + self.fc = nn.Conv2d(input_channels, 2 * input_channels, 3, padding=1) + self.G_middle_0 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_1 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_2 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_3 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_4 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_5 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.up_0 = SPADEResnetBlock(2 * input_channels, input_channels, norm_G, label_num_channels) + self.up_1 = SPADEResnetBlock(input_channels, out_channels, norm_G, label_num_channels) + self.up = nn.Upsample(scale_factor=2) + + if self.upscale is None or self.upscale <= 1: + self.conv_img = nn.Conv2d(out_channels, 3, 3, padding=1) + else: + self.conv_img = nn.Sequential( + nn.Conv2d(out_channels, 3 * (2 * 2), kernel_size=3, padding=1), + nn.PixelShuffle(upscale_factor=2) + ) + + def forward(self, feature): + seg = feature # Bx256x64x64 + x = self.fc(feature) # Bx512x64x64 + x = self.G_middle_0(x, seg) + x = self.G_middle_1(x, seg) + x = self.G_middle_2(x, seg) + x = self.G_middle_3(x, seg) + x = self.G_middle_4(x, seg) + x = self.G_middle_5(x, seg) + + x = self.up(x) # Bx512x64x64 -> Bx512x128x128 + x = self.up_0(x, seg) # Bx512x128x128 -> Bx256x128x128 + x = self.up(x) # Bx256x128x128 -> Bx256x256x256 + x = self.up_1(x, seg) # Bx256x256x256 -> Bx64x256x256 + + x = self.conv_img(F.leaky_relu(x, 2e-1)) # Bx64x256x256 -> Bx3xHxW + x = torch.sigmoid(x) # Bx3xHxW + + return x \ No newline at end of file diff --git a/LivePortrait/src/modules/stitching_retargeting_network.py b/LivePortrait/src/modules/stitching_retargeting_network.py new file mode 100644 index 0000000..5f50b7c --- /dev/null +++ b/LivePortrait/src/modules/stitching_retargeting_network.py @@ -0,0 +1,38 @@ +# coding: utf-8 + +""" +Stitching module(S) and two retargeting modules(R) defined in the paper. + +- The stitching module pastes the animated portrait back into the original image space without pixel misalignment, such as in +the stitching region. + +- The eyes retargeting module is designed to address the issue of incomplete eye closure during cross-id reenactment, especially +when a person with small eyes drives a person with larger eyes. + +- The lip retargeting module is designed similarly to the eye retargeting module, and can also normalize the input by ensuring that +the lips are in a closed state, which facilitates better animation driving. +""" +from torch import nn + + +class StitchingRetargetingNetwork(nn.Module): + def __init__(self, input_size, hidden_sizes, output_size): + super(StitchingRetargetingNetwork, self).__init__() + layers = [] + for i in range(len(hidden_sizes)): + if i == 0: + layers.append(nn.Linear(input_size, hidden_sizes[i])) + else: + layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i])) + layers.append(nn.ReLU(inplace=True)) + layers.append(nn.Linear(hidden_sizes[-1], output_size)) + self.mlp = nn.Sequential(*layers) + + def initialize_weights_to_zero(self): + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.zeros_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, x): + return self.mlp(x) diff --git a/LivePortrait/src/modules/util.py b/LivePortrait/src/modules/util.py new file mode 100644 index 0000000..f83980b --- /dev/null +++ b/LivePortrait/src/modules/util.py @@ -0,0 +1,441 @@ +# coding: utf-8 + +""" +This file defines various neural network modules and utility functions, including convolutional and residual blocks, +normalizations, and functions for spatial transformation and tensor manipulation. +""" + +from torch import nn +import torch.nn.functional as F +import torch +import torch.nn.utils.spectral_norm as spectral_norm +import math +import warnings + + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp + + coordinate_grid = make_coordinate_grid(spatial_size, mean) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance) + + return out + + +def make_coordinate_grid(spatial_size, ref, **kwargs): + d, h, w = spatial_size + x = torch.arange(w).type(ref.dtype).to(ref.device) + y = torch.arange(h).type(ref.dtype).to(ref.device) + z = torch.arange(d).type(ref.dtype).to(ref.device) + + # NOTE: must be right-down-in + x = (2 * (x / (w - 1)) - 1) # the x axis faces to the right + y = (2 * (y / (h - 1)) - 1) # the y axis faces to the bottom + z = (2 * (z / (d - 1)) - 1) # the z axis faces to the inner + + yy = y.view(1, -1, 1).repeat(d, 1, w) + xx = x.view(1, 1, -1).repeat(d, h, 1) + zz = z.view(-1, 1, 1).repeat(1, h, w) + + meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3) + + return meshed + + +class ConvT2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1): + super(ConvT2d, self).__init__() + + self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride, + padding=padding, output_padding=output_padding) + self.norm = nn.InstanceNorm2d(out_features) + + def forward(self, x): + out = self.convT(x) + out = self.norm(out) + out = F.leaky_relu(out) + return out + + +class ResBlock3d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock3d, self).__init__() + self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.norm1 = nn.BatchNorm3d(in_features, affine=True) + self.norm2 = nn.BatchNorm3d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock3d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock3d, self).__init__() + + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=(1, 2, 2)) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class DownBlock3d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock3d, self).__init__() + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups, stride=(1, 2, 2)) + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + if lrelu: + self.ac = nn.LeakyReLU() + else: + self.ac = nn.ReLU() + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = self.ac(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1) + self.norm = nn.BatchNorm3d(self.out_filters, affine=True) + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class SPADE(nn.Module): + def __init__(self, norm_nc, label_nc): + super().__init__() + + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + nhidden = 128 + + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1), + nn.ReLU()) + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + + def forward(self, x, segmap): + normalized = self.param_free_norm(x) + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + actv = self.mlp_shared(segmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + out = normalized * (1 + gamma) + beta + return out + + +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1): + super().__init__() + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + self.use_se = use_se + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + # apply spectral norm if specified + if 'spectral' in norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + # define normalization layers + self.norm_0 = SPADE(fin, label_nc) + self.norm_1 = SPADE(fmiddle, label_nc) + if self.learned_shortcut: + self.norm_s = SPADE(fin, label_nc) + + def forward(self, x, seg1): + x_s = self.shortcut(x, seg1) + dx = self.conv_0(self.actvn(self.norm_0(x, seg1))) + dx = self.conv_1(self.actvn(self.norm_1(dx, seg1))) + out = x_s + dx + return out + + def shortcut(self, x, seg1): + if self.learned_shortcut: + x_s = self.conv_s(self.norm_s(x, seg1)) + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + + +def filter_state_dict(state_dict, remove_name='fc'): + new_state_dict = {} + for key in state_dict: + if remove_name in key: + continue + new_state_dict[key] = state_dict[key] + return new_state_dict + + +class GRN(nn.Module): + """ GRN (Global Response Normalization) layer + """ + + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +class LayerNorm(nn.Module): + r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + self.normalized_shape = (normalized_shape, ) + + def forward(self, x): + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + elif self.data_format == "channels_first": + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1. + math.erf(x / math.sqrt(2.))) / 2. + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2) + + with torch.no_grad(): + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + tensor.uniform_(2 * l - 1, 2 * u - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + return tensor + + +def drop_path(x, drop_prob=0., training=False, scale_by_keep=True): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + + def __init__(self, drop_prob=None, scale_by_keep=True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + +def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): + return _no_grad_trunc_normal_(tensor, mean, std, a, b) diff --git a/LivePortrait/src/modules/warping_network.py b/LivePortrait/src/modules/warping_network.py new file mode 100644 index 0000000..9191a19 --- /dev/null +++ b/LivePortrait/src/modules/warping_network.py @@ -0,0 +1,77 @@ +# coding: utf-8 + +""" +Warping field estimator(W) defined in the paper, which generates a warping field using the implicit +keypoint representations x_s and x_d, and employs this flow field to warp the source feature volume f_s. +""" + +from torch import nn +import torch.nn.functional as F +from .util import SameBlock2d +from .dense_motion import DenseMotionNetwork + + +class WarpingNetwork(nn.Module): + def __init__( + self, + num_kp, + block_expansion, + max_features, + num_down_blocks, + reshape_channel, + estimate_occlusion_map=False, + dense_motion_params=None, + **kwargs + ): + super(WarpingNetwork, self).__init__() + + self.upscale = kwargs.get('upscale', 1) + self.flag_use_occlusion_map = kwargs.get('flag_use_occlusion_map', True) + + if dense_motion_params is not None: + self.dense_motion_network = DenseMotionNetwork( + num_kp=num_kp, + feature_channel=reshape_channel, + estimate_occlusion_map=estimate_occlusion_map, + **dense_motion_params + ) + else: + self.dense_motion_network = None + + self.third = SameBlock2d(max_features, block_expansion * (2 ** num_down_blocks), kernel_size=(3, 3), padding=(1, 1), lrelu=True) + self.fourth = nn.Conv2d(in_channels=block_expansion * (2 ** num_down_blocks), out_channels=block_expansion * (2 ** num_down_blocks), kernel_size=1, stride=1) + + self.estimate_occlusion_map = estimate_occlusion_map + + def deform_input(self, inp, deformation): + return F.grid_sample(inp, deformation, align_corners=False) + + def forward(self, feature_3d, kp_driving, kp_source): + if self.dense_motion_network is not None: + # Feature warper, Transforming feature representation according to deformation and occlusion + dense_motion = self.dense_motion_network( + feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source + ) + if 'occlusion_map' in dense_motion: + occlusion_map = dense_motion['occlusion_map'] # Bx1x64x64 + else: + occlusion_map = None + + deformation = dense_motion['deformation'] # Bx16x64x64x3 + out = self.deform_input(feature_3d, deformation) # Bx32x16x64x64 + + bs, c, d, h, w = out.shape # Bx32x16x64x64 + out = out.view(bs, c * d, h, w) # -> Bx512x64x64 + out = self.third(out) # -> Bx256x64x64 + out = self.fourth(out) # -> Bx256x64x64 + + if self.flag_use_occlusion_map and (occlusion_map is not None): + out = out * occlusion_map + + ret_dct = { + 'occlusion_map': occlusion_map, + 'deformation': deformation, + 'out': out, + } + + return ret_dct diff --git a/LivePortrait/src/utils/__init__.py b/LivePortrait/src/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/LivePortrait/src/utils/camera.py b/LivePortrait/src/utils/camera.py new file mode 100644 index 0000000..8bbfc90 --- /dev/null +++ b/LivePortrait/src/utils/camera.py @@ -0,0 +1,75 @@ +# coding: utf-8 + +""" +functions for processing and transforming 3D facial keypoints +""" + +import numpy as np +import torch +import torch.nn.functional as F + +PI = np.pi + + +def headpose_pred_to_degree(pred): + """ + pred: (bs, 66) or (bs, 1) or others + """ + if pred.ndim > 1 and pred.shape[1] == 66: + # NOTE: note that the average is modified to 97.5 + device = pred.device + idx_tensor = [idx for idx in range(0, 66)] + idx_tensor = torch.FloatTensor(idx_tensor).to(device) + pred = F.softmax(pred, dim=1) + degree = torch.sum(pred*idx_tensor, axis=1) * 3 - 97.5 + + return degree + + return pred + + +def get_rotation_matrix(pitch_, yaw_, roll_): + """ the input is in degree + """ + # calculate the rotation matrix: vps @ rot + + # transform to radian + pitch = pitch_ / 180 * PI + yaw = yaw_ / 180 * PI + roll = roll_ / 180 * PI + + device = pitch.device + + if pitch.ndim == 1: + pitch = pitch.unsqueeze(1) + if yaw.ndim == 1: + yaw = yaw.unsqueeze(1) + if roll.ndim == 1: + roll = roll.unsqueeze(1) + + # calculate the euler matrix + bs = pitch.shape[0] + ones = torch.ones([bs, 1]).to(device) + zeros = torch.zeros([bs, 1]).to(device) + x, y, z = pitch, yaw, roll + + rot_x = torch.cat([ + ones, zeros, zeros, + zeros, torch.cos(x), -torch.sin(x), + zeros, torch.sin(x), torch.cos(x) + ], dim=1).reshape([bs, 3, 3]) + + rot_y = torch.cat([ + torch.cos(y), zeros, torch.sin(y), + zeros, ones, zeros, + -torch.sin(y), zeros, torch.cos(y) + ], dim=1).reshape([bs, 3, 3]) + + rot_z = torch.cat([ + torch.cos(z), -torch.sin(z), zeros, + torch.sin(z), torch.cos(z), zeros, + zeros, zeros, ones + ], dim=1).reshape([bs, 3, 3]) + + rot = rot_z @ rot_y @ rot_x + return rot.permute(0, 2, 1) # transpose diff --git a/LivePortrait/src/utils/face_analysis_diy.py b/LivePortrait/src/utils/face_analysis_diy.py new file mode 100644 index 0000000..8be4c15 --- /dev/null +++ b/LivePortrait/src/utils/face_analysis_diy.py @@ -0,0 +1,78 @@ +# coding: utf-8 + +""" +face detectoin and alignment using InsightFace +""" + +import numpy as np +from .rprint import rlog as log +from insightface.app import FaceAnalysis +from insightface.app.common import Face +from .timer import Timer + + +def sort_by_direction(faces, direction: str = 'large-small', face_center=None): + if len(faces) <= 0: + return faces + if direction == 'left-right': + return sorted(faces, key=lambda face: face['bbox'][0]) + if direction == 'right-left': + return sorted(faces, key=lambda face: face['bbox'][0], reverse=True) + if direction == 'top-bottom': + return sorted(faces, key=lambda face: face['bbox'][1]) + if direction == 'bottom-top': + return sorted(faces, key=lambda face: face['bbox'][1], reverse=True) + if direction == 'small-large': + return sorted(faces, key=lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1])) + if direction == 'large-small': + return sorted(faces, key=lambda face: (face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]), reverse=True) + if direction == 'distance-from-retarget-face': + return sorted(faces, key=lambda face: (((face['bbox'][2]+face['bbox'][0])/2-face_center[0])**2+((face['bbox'][3]+face['bbox'][1])/2-face_center[1])**2)**0.5) + return faces + + +class FaceAnalysisDIY(FaceAnalysis): + def __init__(self, name='buffalo_l', root='~/.insightface', allowed_modules=None, **kwargs): + super().__init__(name=name, root=root, allowed_modules=allowed_modules, **kwargs) + + self.timer = Timer() + + def get(self, img_bgr, **kwargs): + max_num = kwargs.get('max_num', 0) # the number of the detected faces, 0 means no limit + flag_do_landmark_2d_106 = kwargs.get('flag_do_landmark_2d_106', True) # whether to do 106-point detection + direction = kwargs.get('direction', 'large-small') # sorting direction + face_center = None + + bboxes, kpss = self.det_model.detect(img_bgr, max_num=max_num, metric='default') + if bboxes.shape[0] == 0: + return [] + ret = [] + for i in range(bboxes.shape[0]): + bbox = bboxes[i, 0:4] + det_score = bboxes[i, 4] + kps = None + if kpss is not None: + kps = kpss[i] + face = Face(bbox=bbox, kps=kps, det_score=det_score) + for taskname, model in self.models.items(): + if taskname == 'detection': + continue + + if (not flag_do_landmark_2d_106) and taskname == 'landmark_2d_106': + continue + + # print(f'taskname: {taskname}') + model.get(img_bgr, face) + ret.append(face) + + ret = sort_by_direction(ret, direction, face_center) + return ret + + def warmup(self): + self.timer.tic() + + img_bgr = np.zeros((512, 512, 3), dtype=np.uint8) + self.get(img_bgr) + + elapse = self.timer.toc() + log(f'FaceAnalysisDIY warmup time: {elapse:.3f}s') diff --git a/LivePortrait/src/utils/helper.py b/LivePortrait/src/utils/helper.py new file mode 100644 index 0000000..009c3df --- /dev/null +++ b/LivePortrait/src/utils/helper.py @@ -0,0 +1,175 @@ +# coding: utf-8 + +""" +utility functions and classes to handle feature extraction and model loading +""" + +import os +import os.path as osp +import cv2 +import torch +from rich.console import Console +from collections import OrderedDict + +from LivePortrait.src.modules.spade_generator import SPADEDecoder +from LivePortrait.src.modules.warping_network import WarpingNetwork +from LivePortrait.src.modules.motion_extractor import MotionExtractor +from LivePortrait.src.modules.appearance_feature_extractor import AppearanceFeatureExtractor +from LivePortrait.src.modules.stitching_retargeting_network import StitchingRetargetingNetwork +from .rprint import rlog as log + + +def suffix(filename): + """a.jpg -> jpg""" + pos = filename.rfind(".") + if pos == -1: + return "" + return filename[pos + 1:] + + +def prefix(filename): + """a.jpg -> a""" + pos = filename.rfind(".") + if pos == -1: + return filename + return filename[:pos] + + +def basename(filename): + """a/b/c.jpg -> c""" + return prefix(osp.basename(filename)) + + +def is_video(file_path): + if file_path.lower().endswith((".mp4", ".mov", ".avi", ".webm")) or osp.isdir(file_path): + return True + return False + +def is_template(file_path): + if file_path.endswith(".pkl"): + return True + return False + + +def mkdir(d, log=False): + # return self-assined `d`, for one line code + if not osp.exists(d): + os.makedirs(d, exist_ok=True) + if log: + print(f"Make dir: {d}") + return d + + +def squeeze_tensor_to_numpy(tensor): + out = tensor.data.squeeze(0).cpu().numpy() + return out + + +def dct2cuda(dct: dict, device_id: int): + for key in dct: + dct[key] = torch.tensor(dct[key]).cuda(device_id) + return dct + + +def concat_feat(kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor: + """ + kp_source: (bs, k, 3) + kp_driving: (bs, k, 3) + Return: (bs, 2k*3) + """ + bs_src = kp_source.shape[0] + bs_dri = kp_driving.shape[0] + assert bs_src == bs_dri, 'batch size must be equal' + + feat = torch.cat([kp_source.view(bs_src, -1), kp_driving.view(bs_dri, -1)], dim=1) + return feat + + +def remove_ddp_dumplicate_key(state_dict): + state_dict_new = OrderedDict() + for key in state_dict.keys(): + state_dict_new[key.replace('module.', '')] = state_dict[key] + return state_dict_new + + +def load_model(ckpt_path, model_config, device, model_type): + model_params = model_config['model_params'][f'{model_type}_params'] + + if model_type == 'appearance_feature_extractor': + model = AppearanceFeatureExtractor(**model_params).cuda(device) + elif model_type == 'motion_extractor': + model = MotionExtractor(**model_params).cuda(device) + elif model_type == 'warping_module': + model = WarpingNetwork(**model_params).cuda(device) + elif model_type == 'spade_generator': + model = SPADEDecoder(**model_params).cuda(device) + elif model_type == 'stitching_retargeting_module': + # Special handling for stitching and retargeting module + config = model_config['model_params']['stitching_retargeting_module_params'] + checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage) + + stitcher = StitchingRetargetingNetwork(**config.get('stitching')) + stitcher.load_state_dict(remove_ddp_dumplicate_key(checkpoint['retarget_shoulder'])) + stitcher = stitcher.cuda(device) + stitcher.eval() + + retargetor_lip = StitchingRetargetingNetwork(**config.get('lip')) + retargetor_lip.load_state_dict(remove_ddp_dumplicate_key(checkpoint['retarget_mouth'])) + retargetor_lip = retargetor_lip.cuda(device) + retargetor_lip.eval() + + retargetor_eye = StitchingRetargetingNetwork(**config.get('eye')) + retargetor_eye.load_state_dict(remove_ddp_dumplicate_key(checkpoint['retarget_eye'])) + retargetor_eye = retargetor_eye.cuda(device) + retargetor_eye.eval() + + return { + 'stitching': stitcher, + 'lip': retargetor_lip, + 'eye': retargetor_eye + } + else: + raise ValueError(f"Unknown model type: {model_type}") + + model.load_state_dict(torch.load(ckpt_path, map_location=lambda storage, loc: storage)) + model.eval() + return model + + +# get coefficients of Eqn. 7 +def calculate_transformation(config, s_kp_info, t_0_kp_info, t_i_kp_info, R_s, R_t_0, R_t_i): + if config.relative: + new_rotation = (R_t_i @ R_t_0.permute(0, 2, 1)) @ R_s + new_expression = s_kp_info['exp'] + (t_i_kp_info['exp'] - t_0_kp_info['exp']) + else: + new_rotation = R_t_i + new_expression = t_i_kp_info['exp'] + new_translation = s_kp_info['t'] + (t_i_kp_info['t'] - t_0_kp_info['t']) + new_translation[..., 2].fill_(0) # Keep the z-axis unchanged + new_scale = s_kp_info['scale'] * (t_i_kp_info['scale'] / t_0_kp_info['scale']) + return new_rotation, new_expression, new_translation, new_scale + +def load_description(fp): + with open(fp, 'r', encoding='utf-8') as f: + content = f.read() + return content + + +def resize_to_limit(img, max_dim=1280, n=2): + h, w = img.shape[:2] + if max_dim > 0 and max(h, w) > max_dim: + if h > w: + new_h = max_dim + new_w = int(w * (max_dim / h)) + else: + new_w = max_dim + new_h = int(h * (max_dim / w)) + img = cv2.resize(img, (new_w, new_h)) + n = max(n, 1) + new_h = img.shape[0] - (img.shape[0] % n) + new_w = img.shape[1] - (img.shape[1] % n) + if new_h == 0 or new_w == 0: + return img + if new_h != img.shape[0] or new_w != img.shape[1]: + img = img[:new_h, :new_w] + return img diff --git a/LivePortrait/src/utils/io.py b/LivePortrait/src/utils/io.py new file mode 100644 index 0000000..f930c48 --- /dev/null +++ b/LivePortrait/src/utils/io.py @@ -0,0 +1,97 @@ +# coding: utf-8 + +import os +from glob import glob +import os.path as osp +import imageio +import numpy as np +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) + + +def load_image_rgb(image_path: str): + if not osp.exists(image_path): + raise FileNotFoundError(f"Image not found: {image_path}") + img = cv2.imread(image_path, cv2.IMREAD_COLOR) + return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + +def load_driving_info(driving_info): + driving_video_ori = [] + + def load_images_from_directory(directory): + image_paths = sorted(glob(osp.join(directory, '*.png')) + glob(osp.join(directory, '*.jpg'))) + return [load_image_rgb(im_path) for im_path in image_paths] + + def load_images_from_video(file_path): + reader = imageio.get_reader(file_path) + return [image for idx, image in enumerate(reader)] + + if osp.isdir(driving_info): + driving_video_ori = load_images_from_directory(driving_info) + elif osp.isfile(driving_info): + driving_video_ori = load_images_from_video(driving_info) + + return driving_video_ori + + +def contiguous(obj): + if not obj.flags.c_contiguous: + obj = obj.copy(order="C") + return obj + + +def _resize_to_limit(img: np.ndarray, max_dim=1920, n=2): + """ + ajust the size of the image so that the maximum dimension does not exceed max_dim, and the width and the height of the image are multiples of n. + :param img: the image to be processed. + :param max_dim: the maximum dimension constraint. + :param n: the number that needs to be multiples of. + :return: the adjusted image. + """ + h, w = img.shape[:2] + + # ajust the size of the image according to the maximum dimension + if max_dim > 0 and max(h, w) > max_dim: + if h > w: + new_h = max_dim + new_w = int(w * (max_dim / h)) + else: + new_w = max_dim + new_h = int(h * (max_dim / w)) + img = cv2.resize(img, (new_w, new_h)) + + # ensure that the image dimensions are multiples of n + n = max(n, 1) + new_h = img.shape[0] - (img.shape[0] % n) + new_w = img.shape[1] - (img.shape[1] % n) + + if new_h == 0 or new_w == 0: + # when the width or height is less than n, no need to process + return img + + if new_h != img.shape[0] or new_w != img.shape[1]: + img = img[:new_h, :new_w] + + return img + + +def load_img_online(obj, mode="bgr", **kwargs): + max_dim = kwargs.get("max_dim", 1920) + n = kwargs.get("n", 2) + if isinstance(obj, str): + if mode.lower() == "gray": + img = cv2.imread(obj, cv2.IMREAD_GRAYSCALE) + else: + img = cv2.imread(obj, cv2.IMREAD_COLOR) + else: + img = obj + + # Resize image to satisfy constraints + img = _resize_to_limit(img, max_dim=max_dim, n=n) + + if mode.lower() == "bgr": + return contiguous(img) + elif mode.lower() == "rgb": + return contiguous(img[..., ::-1]) + else: + raise Exception(f"Unknown mode {mode}") diff --git a/LivePortrait/src/utils/resources/mask_template.png b/LivePortrait/src/utils/resources/mask_template.png new file mode 100644 index 0000000000000000000000000000000000000000..bca6ca5977ba820d0d2c05b3793c6231cc82e715 GIT binary patch literal 3470 zcmeH~?@JqZ9LF=M*mY>y?tvJpX<^-HH{ydL-40tfQCCMX>u76PqgPs5g@DGH&8W%U zXvVgKB-?gFLA9kP6yHliAbB5Isc9IW=`>NRN0Xx4U?E=Euc-Z+$kle_-5` z2&$E=c>Na>2knpj&47A>K5O<|33JBskZ8~i;lNq zQ+YohxtG6hYD@giu!kPM6||kH-rB5OgunmkL~yT+uV4xbo0^7b1o8X*NzX175tnCX zl10|>0S?D8*_=I=ZA04IU9QeHiNziTrqPDbtI)cop~2EO@fWPNu|rxi6C9snqL%v`K-#f1Pj0@icF}qvkQSH zkB*lgDZntjULG5xJcZYLCH1lt?qL|OcPYAa!twF0&Ea%T1OmS2FAiEqH(+BcIM-}8 zdzbi~j}~Hg^ll;?CPE=`rpP)vfXfv1_o%(R7~5QpR1~-gf>7p{Rfvs&y-QJhW)Rza z6%k0@##7HC>joP8dk8?e4g636Ym!yH#brE%SQY~7w M)pb?FmDg_n1%!;u`2YX_ literal 0 HcmV?d00001 diff --git a/LivePortrait/src/utils/rprint.py b/LivePortrait/src/utils/rprint.py new file mode 100644 index 0000000..c43a42f --- /dev/null +++ b/LivePortrait/src/utils/rprint.py @@ -0,0 +1,16 @@ +# coding: utf-8 + +""" +custom print and log functions +""" + +__all__ = ['rprint', 'rlog'] + +try: + from rich.console import Console + console = Console() + rprint = console.print + rlog = console.log +except: + rprint = print + rlog = print diff --git a/LivePortrait/src/utils/timer.py b/LivePortrait/src/utils/timer.py new file mode 100644 index 0000000..3570fa4 --- /dev/null +++ b/LivePortrait/src/utils/timer.py @@ -0,0 +1,29 @@ +# coding: utf-8 + +""" +tools to measure elapsed time +""" + +import time + +class Timer(object): + """A simple timer.""" + + def __init__(self): + self.total_time = 0. + self.calls = 0 + self.start_time = 0. + self.diff = 0. + + def tic(self): + # using time.time instead of time.clock because time time.clock + # does not normalize for multithreading + self.start_time = time.time() + + def toc(self, average=True): + self.diff = time.time() - self.start_time + return self.diff + + def clear(self): + self.start_time = 0. + self.diff = 0. diff --git a/LivePortrait/src/utils/video.py b/LivePortrait/src/utils/video.py new file mode 100644 index 0000000..3efbbbb --- /dev/null +++ b/LivePortrait/src/utils/video.py @@ -0,0 +1,142 @@ +# coding: utf-8 + +""" +functions for processing video +""" + +import os.path as osp +import numpy as np +import subprocess +import imageio +import cv2 + +# from rich.progress import track +from .helper import prefix +from .rprint import rprint as print + + +def exec_cmd(cmd): + subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + + +def images2video(images, wfp, **kwargs): + fps = kwargs.get('fps', 30) + video_format = kwargs.get('format', 'mp4') # default is mp4 format + codec = kwargs.get('codec', 'libx264') # default is libx264 encoding + quality = kwargs.get('quality') # video quality + pixelformat = kwargs.get('pixelformat', 'yuv420p') # video pixel format + image_mode = kwargs.get('image_mode', 'rgb') + macro_block_size = kwargs.get('macro_block_size', 2) + ffmpeg_params = ['-crf', str(kwargs.get('crf', 18))] + + writer = imageio.get_writer( + wfp, fps=fps, format=video_format, + codec=codec, quality=quality, ffmpeg_params=ffmpeg_params, pixelformat=pixelformat, macro_block_size=macro_block_size + ) + + n = len(images) + print('writing',n) + for i in range(n): + if image_mode.lower() == 'bgr': + writer.append_data(images[i][..., ::-1]) + else: + writer.append_data(images[i]) + + writer.close() + + # print(f':smiley: Dump to {wfp}\n', style="bold green") + print(f'Dump to {wfp}\n') + + +def video2gif(video_fp, fps=30, size=256): + if osp.exists(video_fp): + d = osp.split(video_fp)[0] + fn = prefix(osp.basename(video_fp)) + palette_wfp = osp.join(d, 'palette.png') + gif_wfp = osp.join(d, f'{fn}.gif') + # generate the palette + cmd = f'ffmpeg -i {video_fp} -vf "fps={fps},scale={size}:-1:flags=lanczos,palettegen" {palette_wfp} -y' + exec_cmd(cmd) + # use the palette to generate the gif + cmd = f'ffmpeg -i {video_fp} -i {palette_wfp} -filter_complex "fps={fps},scale={size}:-1:flags=lanczos[x];[x][1:v]paletteuse" {gif_wfp} -y' + exec_cmd(cmd) + else: + print(f'video_fp: {video_fp} not exists!') + + +def merge_audio_video(video_fp, audio_fp, wfp): + if osp.exists(video_fp) and osp.exists(audio_fp): + cmd = f'ffmpeg -i {video_fp} -i {audio_fp} -c:v copy -c:a aac {wfp} -y' + exec_cmd(cmd) + print(f'merge {video_fp} and {audio_fp} to {wfp}') + else: + print(f'video_fp: {video_fp} or audio_fp: {audio_fp} not exists!') + + +def blend(img: np.ndarray, mask: np.ndarray, background_color=(255, 255, 255)): + mask_float = mask.astype(np.float32) / 255. + background_color = np.array(background_color).reshape([1, 1, 3]) + bg = np.ones_like(img) * background_color + img = np.clip(mask_float * img + (1 - mask_float) * bg, 0, 255).astype(np.uint8) + return img + + +def concat_frames(I_p_lst, driving_rgb_lst, img_rgb): + # TODO: add more concat style, e.g., left-down corner driving + out_lst = [] + print('Concatenating result...',len(I_p_lst)) + for idx, _ in enumerate(I_p_lst): + # track(enumerate(I_p_lst), total=len(I_p_lst), description='Concatenating result...'): + source_image_drived = I_p_lst[idx] + image_drive = driving_rgb_lst[idx] + + # resize images to match source_image_drived shape + h, w, _ = source_image_drived.shape + image_drive_resized = cv2.resize(image_drive, (w, h)) + img_rgb_resized = cv2.resize(img_rgb, (w, h)) + + # concatenate images horizontally + frame = np.concatenate((image_drive_resized, img_rgb_resized, source_image_drived), axis=1) + out_lst.append(frame) + return out_lst + + +class VideoWriter: + def __init__(self, **kwargs): + self.fps = kwargs.get('fps', 30) + self.wfp = kwargs.get('wfp', 'video.mp4') + self.video_format = kwargs.get('format', 'mp4') + self.codec = kwargs.get('codec', 'libx264') + self.quality = kwargs.get('quality') + self.pixelformat = kwargs.get('pixelformat', 'yuv420p') + self.image_mode = kwargs.get('image_mode', 'rgb') + self.ffmpeg_params = kwargs.get('ffmpeg_params') + + self.writer = imageio.get_writer( + self.wfp, fps=self.fps, format=self.video_format, + codec=self.codec, quality=self.quality, + ffmpeg_params=self.ffmpeg_params, pixelformat=self.pixelformat + ) + + def write(self, image): + if self.image_mode.lower() == 'bgr': + self.writer.append_data(image[..., ::-1]) + else: + self.writer.append_data(image) + + def close(self): + if self.writer is not None: + self.writer.close() + + +def change_video_fps(input_file, output_file, fps=20, codec='libx264', crf=5): + cmd = f"ffmpeg -i {input_file} -c:v {codec} -crf {crf} -r {fps} {output_file} -y" + exec_cmd(cmd) + + +def get_fps(filepath): + import ffmpeg + probe = ffmpeg.probe(filepath) + video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None) + fps = eval(video_stream['avg_frame_rate']) + return fps diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..476fdb7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,49 @@ +import requests +import os, sys +import subprocess +from tqdm import tqdm +from pip._internal import main as pip_main +from pathlib import Path +from folder_paths import models_dir + +def download_model(file_path, model_url): + print('AdvancedLivePortrait: Downloading model...') + response = requests.get(model_url, stream=True) + try: + if response.status_code == 200: + total_size = int(response.headers.get('content-length', 0)) + block_size = 1024 # 1 Kibibyte + + # tqdm will display a progress bar + with open(file_path, 'wb') as file, tqdm( + desc='Downloading', + total=total_size, + unit='iB', + unit_scale=True, + unit_divisor=1024, + ) as bar: + for data in response.iter_content(block_size): + bar.update(len(data)) + file.write(data) + + except requests.exceptions.RequestException as err: + print('AdvancedLivePortrait: Model download failed: {err}') + print(f'AdvancedLivePortrait: Download it manually from: {model_url}') + print(f'AdvancedLivePortrait: And put it in {file_path}') + except Exception as e: + print(f'AdvancedLivePortrait: An unexpected error occurred: {e}') + +save_path = os.path.join(models_dir, "ultralytics") +if not os.path.exists(save_path): + os.makedirs(save_path, exist_ok=True) + file_path = os.path.join(save_path, "face_yolov8n.pt") + if not Path().is_file(): + download_model(file_path, "https://huggingface.co/Bingsu/adetailer/resolve/main/face_yolov8n.pt") + + + + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + + diff --git a/config/inference_config.py b/config/inference_config.py new file mode 100644 index 0000000..3a735f6 --- /dev/null +++ b/config/inference_config.py @@ -0,0 +1,21 @@ + +import os +from typing import Literal, Tuple +import folder_paths + +current_file_path = os.path.abspath(__file__) +current_directory = os.path.dirname(current_file_path) +class InferenceConfig: + def __init__(self): + self.models_config: str = os.path.join(current_directory, "models.yaml") + self.checkpoint_F: str = os.path.join(folder_paths.models_dir, "liveportrait", "base_models", "appearance_feature_extractor.pth") + self.checkpoint_M: str = os.path.join(folder_paths.models_dir, "liveportrait", "base_models", "motion_extractor.pth") + self.checkpoint_G: str = os.path.join(folder_paths.models_dir, "liveportrait", "base_models", "spade_generator.pth") + self.checkpoint_W: str = os.path.join(folder_paths.models_dir, "liveportrait", "base_models", "warping_module.pth") + self.checkpoint_S: str = os.path.join(folder_paths.models_dir, "liveportrait", "retargeting_models", "stitching_retargeting_module.pth") + self.flag_use_half_precision: bool = True # whether to use half precision + + self.input_shape: Tuple[int, int] = (256, 256) # input shape + + self.mask_crop = None + self.device_id: int = 0 \ No newline at end of file diff --git a/config/models.yaml b/config/models.yaml new file mode 100644 index 0000000..131d1c6 --- /dev/null +++ b/config/models.yaml @@ -0,0 +1,43 @@ +model_params: + appearance_feature_extractor_params: # the F in the paper + image_channel: 3 + block_expansion: 64 + num_down_blocks: 2 + max_features: 512 + reshape_channel: 32 + reshape_depth: 16 + num_resblocks: 6 + motion_extractor_params: # the M in the paper + num_kp: 21 + backbone: convnextv2_tiny + warping_module_params: # the W in the paper + num_kp: 21 + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + reshape_channel: 32 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 32 + max_features: 1024 + num_blocks: 5 + reshape_depth: 16 + compress: 4 + spade_generator_params: # the G in the paper + upscale: 2 # represents upsample factor 256x256 -> 512x512 + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + stitching_retargeting_module_params: # the S in the paper + stitching: + input_size: 126 # (21*3)*2 + hidden_sizes: [128, 128, 64] + output_size: 65 # (21*3)+2(tx,ty) + lip: + input_size: 65 # (21*3)+2 + hidden_sizes: [128, 128, 64] + output_size: 63 # (21*3) + eye: + input_size: 66 # (21*3)+3 + hidden_sizes: [256, 256, 128, 128, 64] + output_size: 63 # (21*3) diff --git a/exp_data/입술내밀기.exp b/exp_data/입술내밀기.exp new file mode 100644 index 0000000000000000000000000000000000000000..66c75aaa3ebcabbab5f49ac3c3663898f5f0e1d7 GIT binary patch literal 1037 zcmdT?-%Aux6uz^+j5STN0-*=-Ay%?w+sR$s5hW`j#%xoPF^pk%=3aMVoN;FEt{Dj1 zgPE}};q_rb`auu{Jw)Q_ItD^u2onWDgwR7TMo=WwLxK<8T@iYyKcMq)?!D)J=X~d$ z`(-^dE)OVMuHlS6jK(gXKXF>lBs2-i-AV?=brYKj1qU@5<@htMpX5L zEFx1M7L#L`vP_fCUkm_Q*B!sQGuyRYf9*-UE|_&)_~NF1+yDWghW0#S>R6X_Rm?!M z)R4hYfKWwmR0U%dicpi4gl4#I3w|2JVs>X#U(ie<*g@eXWBYoRo!NCRpZ``szQ{5=w_!UDo=1gV zoBc7m48YDjIu_(u2`;g7Ul-7OS6hDVgPm`DytrEdfc(ItMX>TY<~(~z>#Uv15Bywm zj((zb;sthnb_r}nV~*3lfR0qRf~{H(Z8k(fb~FZ_Cpolwzz)5*U-US~ZIB&YK&uan zZ+5DS2+b5g_K;1&?Jeg(^HSaWSwyxa0Lb+Z?QNO@} z_|q8jZP@5mcMO2HO9&Ms`Ry$lQ(M_}KgYiEa?aoh4m`NWkqOdEc)&CKVctDS_zvF9 zGm~U0S4lrxg%YNJP4!iI(UhSH<+F-Z7Gf{KG%v@PB(*P)s_gXs_Yw+6T3V=tO3IW* s(o{e0f9j_UrTQsh|FCRnzc)|!1mT$pGL@=r+?$m}=bU3O_332rH*4))Q~&?~ literal 0 HcmV?d00001 diff --git a/install.bat b/install.bat new file mode 100644 index 0000000..9b5cc7b --- /dev/null +++ b/install.bat @@ -0,0 +1,20 @@ +@echo off + +set "requirements_txt=%~dp0\requirements.txt" +set "python_exec=..\..\..\python_embeded\python.exe" + +echo Installing ComfyUI-AdvancedLivePortrait.. + +if exist "%python_exec%" ( + echo Installing with ComfyUI Portable + for /f "delims=" %%i in (%requirements_txt%) do ( + %python_exec% -s -m pip install "%%i" + ) +) else ( + echo Installing with system Python + for /f "delims=" %%i in (%requirements_txt%) do ( + pip install "%%i" + ) +) + +pause \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..eaf3a5f --- /dev/null +++ b/nodes.py @@ -0,0 +1,770 @@ +import os +import sys +import numpy as np +import torch +import cv2 +from PIL import Image +import folder_paths +import comfy.utils +import time +import copy +import math +import dill +import torch.nn.functional as torchfn +from torchvision import transforms +from ultralytics import YOLO +from scipy.spatial.transform import Rotation as R + + +current_file_path = os.path.abspath(__file__) +current_directory = os.path.dirname(current_file_path) +sys.path.append(current_directory) +from .LivePortrait.src.live_portrait_wrapper import LivePortraitWrapper +from .LivePortrait.src.utils.rprint import rlog as log +from .LivePortrait.src.utils.camera import get_rotation_matrix +from .config.inference_config import InferenceConfig + +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) +def rgb_crop(rgb, region): + return rgb[region[1]:region[3], region[0]:region[2]] + +def rgb_crop_batch(rgbs, region): + return rgbs[:, region[1]:region[3], region[0]:region[2]] +def get_rgb_size(rgb): + return rgb.shape[1], rgb.shape[0] +def create_transform_matrix(x, y, scale=1): + return np.float32([[scale, 0, x], [0, scale, y]]) + +def get_model_dir(m): + try: + return folder_paths.get_folder_paths(m)[0] + except: + return os.path.join(folder_paths.models_dir, m) + +def calc_crop_limit(center, img_size, crop_size): + pos = center - crop_size / 2 + if pos < 0: + crop_size += pos * 2 + pos = 0 + + pos2 = pos + crop_size + + if img_size < pos2: + crop_size -= (pos2 - img_size) * 2 + pos2 = img_size + pos = pos2 - crop_size + + return pos, pos2, crop_size + +def retargeting(delta_out, driving_exp, factor, idxes): + for idx in idxes: + #delta_out[0, idx] -= src_exp[0, idx] * factor + delta_out[0, idx] += driving_exp[0, idx] * factor + +class PreparedSrcImg: + def __init__(self, src_rgb, crop_trans_m, x_s_info, f_s_user, x_s_user, mask_ori): + self.src_rgb = src_rgb + self.crop_trans_m = crop_trans_m + self.x_s_info = x_s_info + self.f_s_user = f_s_user + self.x_s_user = x_s_user + self.mask_ori = mask_ori + +class LP_Engine: + pipeline = None + bbox_model = None + mask_img = None + + def detect_face(self, image_rgb): + + crop_factor = 1.7 + bbox_drop_size = 10 + + if self.bbox_model == None: + bbox_model_path = os.path.join(get_model_dir("ultralytics"), "face_yolov8n.pt") + self.bbox_model = YOLO(bbox_model_path) + + pred = self.bbox_model(image_rgb, conf=0.7, device="") + bboxes = pred[0].boxes.xyxy.cpu().numpy() + + w, h = get_rgb_size(image_rgb) + + # for x, label in zip(segmasks, detected_results[0]): + for x1, y1, x2, y2 in bboxes: + bbox_w = x2 - x1 + bbox_h = y2 - y1 + + crop_w = bbox_w * crop_factor + crop_h = bbox_h * crop_factor + + crop_w = max(crop_h, crop_w) + crop_h = crop_w + + kernel_x = x1 + bbox_w / 2 + kernel_y = y1 + bbox_h / 2 + + new_x1, new_x2, crop_w = calc_crop_limit(kernel_x, w, crop_w) + + if crop_w < crop_h: + crop_h = crop_w + + new_y1, new_y2, crop_h = calc_crop_limit(kernel_y, h, crop_h) + + if crop_h < crop_w: + crop_w = crop_h + new_x1, new_x2, crop_w = calc_crop_limit(kernel_x, w, crop_w) + + return [int(new_x1), int(new_y1), int(new_x2), int(new_y2)] + + print("Failed to detect face!!") + return [0, 0, w, h] + + def crop_face(self, rgb_img): + region = self.detect_face(rgb_img) + face_image = rgb_crop(rgb_img, region) + return face_image, region + + def get_pipeline(self): + if self.pipeline == None: + print("Load pipeline...") + self.pipeline = LivePortraitWrapper(cfg=InferenceConfig()) + + return self.pipeline + + def prepare_src_image(self, img): + h, w = img.shape[:2] + input_shape = [256,256] + if h != input_shape[0] or w != input_shape[1]: + x = cv2.resize(img, (input_shape[0], input_shape[1]), interpolation = cv2.INTER_LINEAR) + else: + x = img.copy() + + if x.ndim == 3: + x = x[np.newaxis].astype(np.float32) / 255. # HxWx3 -> 1xHxWx3, normalized to 0~1 + elif x.ndim == 4: + x = x.astype(np.float32) / 255. # BxHxWx3, normalized to 0~1 + else: + 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() + return x + + def GetMask(self): + if self.mask_img is None: + path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "./LivePortrait/src/utils/resources/mask_template.png") + self.mask_img = cv2.imread(path, cv2.IMREAD_COLOR) + return self.mask_img + + def prepare_source(self, source_image, is_video = False): + print("Prepare source...") + engine = self.get_pipeline() + source_image_np = (source_image * 255).byte().numpy() + img_rgb = source_image_np[0] + face_img, crop_region = self.crop_face(img_rgb) + + scale = face_img.shape[0] / 512. + crop_trans_m = create_transform_matrix(crop_region[0], crop_region[1], scale) + mask_ori = cv2.warpAffine(self.GetMask(), crop_trans_m, get_rgb_size(img_rgb), cv2.INTER_LINEAR) + mask_ori = mask_ori.astype(np.float32) / 255. + + psi_list = [] + for img_rgb in source_image_np: + face_img = rgb_crop(img_rgb, crop_region) + i_s = self.prepare_src_image(face_img) + x_s_info = engine.get_kp_info(i_s) + f_s_user = engine.extract_feature_3d(i_s) + x_s_user = engine.transform_keypoint(x_s_info) + psi = PreparedSrcImg(img_rgb, crop_trans_m, x_s_info, f_s_user, x_s_user, mask_ori) + if is_video == False: + return psi + psi_list.append(psi) + + return psi_list + + def prepare_driving_video(self, face_images): + print("Prepare driving video...") + pipeline = self.get_pipeline() + f_img_np = (face_images * 255).byte().numpy() + + out_list = [] + for f_img in f_img_np: + i_d = pipeline.prepare_source(f_img) + d_info = pipeline.get_kp_info(i_d) + #out_list.append((d_info, get_rotation_matrix(d_info['pitch'], d_info['yaw'], d_info['roll']))) + out_list.append(d_info) + + return out_list + + def calc_fe(_, x_d_new, eyes, eyebrow, wink, pupil_x, pupil_y, mouth, eee, woo, smile, + rotate_pitch, rotate_yaw, rotate_roll): + + x_d_new[0, 20, 1] += smile * -0.01 + x_d_new[0, 14, 1] += smile * -0.02 + x_d_new[0, 17, 1] += smile * 0.0065 + x_d_new[0, 17, 2] += smile * 0.003 + x_d_new[0, 13, 1] += smile * -0.00275 + x_d_new[0, 16, 1] += smile * -0.00275 + x_d_new[0, 3, 1] += smile * -0.0035 + x_d_new[0, 7, 1] += smile * -0.0035 + + x_d_new[0, 19, 1] += mouth * 0.001 + x_d_new[0, 19, 2] += mouth * 0.0001 + x_d_new[0, 17, 1] += mouth * -0.0001 + rotate_pitch -= mouth * 0.05 + + x_d_new[0, 20, 2] += eee * -0.001 + x_d_new[0, 20, 1] += eee * -0.001 + #x_d_new[0, 19, 1] += eee * 0.0006 + x_d_new[0, 14, 1] += eee * -0.001 + + x_d_new[0, 14, 1] += woo * 0.001 + x_d_new[0, 3, 1] += woo * -0.0005 + x_d_new[0, 7, 1] += woo * -0.0005 + x_d_new[0, 17, 2] += woo * -0.0005 + + x_d_new[0, 11, 1] += wink * 0.001 + x_d_new[0, 13, 1] += wink * -0.0003 + x_d_new[0, 17, 0] += wink * 0.0003 + x_d_new[0, 17, 1] += wink * 0.0003 + x_d_new[0, 3, 1] += wink * -0.0003 + rotate_roll -= wink * 0.1 + rotate_yaw -= wink * 0.1 + + if 0 < pupil_x: + x_d_new[0, 11, 0] += pupil_x * 0.0007 + x_d_new[0, 15, 0] += pupil_x * 0.001 + else: + x_d_new[0, 11, 0] += pupil_x * 0.001 + x_d_new[0, 15, 0] += pupil_x * 0.0007 + + x_d_new[0, 11, 1] += pupil_y * -0.001 + x_d_new[0, 15, 1] += pupil_y * -0.001 + eyes -= pupil_y / 2. + + x_d_new[0, 11, 1] += eyes * -0.001 + x_d_new[0, 13, 1] += eyes * 0.0003 + x_d_new[0, 15, 1] += eyes * -0.001 + x_d_new[0, 16, 1] += eyes * 0.0003 + + + if 0 < eyebrow: + x_d_new[0, 1, 1] += eyebrow * 0.001 + x_d_new[0, 2, 1] += eyebrow * -0.001 + else: + x_d_new[0, 1, 0] += eyebrow * -0.001 + x_d_new[0, 2, 0] += eyebrow * 0.001 + x_d_new[0, 1, 1] += eyebrow * 0.0003 + x_d_new[0, 2, 1] += eyebrow * -0.0003 + + + return torch.Tensor([rotate_pitch, rotate_yaw, rotate_roll]) +g_engine = LP_Engine() + +class ExpressionSet: + def __init__(self, erst = None, es = None): + if es != None: + self.e = copy.deepcopy(es.e) # [:, :, :] + self.r = copy.deepcopy(es.r) # [:] + self.s = copy.deepcopy(es.s) + self.t = copy.deepcopy(es.t) + elif erst != None: + self.e = erst[0] + self.r = erst[1] + self.s = erst[2] + self.t = erst[3] + else: + self.e = torch.from_numpy(np.zeros((1, 21, 3))).float().to(device='cuda') + self.r = torch.Tensor([0, 0, 0]) + self.s = 0 + self.t = 0 + def div(self, value): + self.e /= value + self.r /= value + self.s /= value + self.t /= value + def add(self, other): + self.e += other.e + self.r += other.r + self.s += other.s + self.t += other.t + def sub(self, other): + self.e -= other.e + self.r -= other.r + self.s -= other.s + self.t -= other.t + def mul(self, value): + self.e *= value + self.r *= value + self.s *= value + self.t *= value + + #def apply_ratio(self, ratio): self.exp *= ratio + +def logging_time(original_fn): + def wrapper_fn(*args, **kwargs): + start_time = time.time() + result = original_fn(*args, **kwargs) + end_time = time.time() + print("WorkingTime[{}]: {} sec".format(original_fn.__name__, end_time - start_time)) + return result + + return wrapper_fn + + +exp_data_dir = os.path.join(current_directory, "exp_data") +class SaveExpData: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "file_name": ("STRING", {"multiline": False, "default": ""}), + }, + "optional": {"save_exp": ("EXP_DATA",), } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("file_name",) + FUNCTION = "run" + CATEGORY = "파워집돌이" + OUTPUT_NODE = True + + def run(self, file_name, save_exp:ExpressionSet=None): + if save_exp == None or file_name == "": + return file_name + + with open(os.path.join(exp_data_dir, file_name + ".exp"), "wb") as f: + dill.dump(save_exp, f) + + return file_name + +class LoadExpData: + @classmethod + def INPUT_TYPES(s): + file_list = [os.path.splitext(file)[0] for file in os.listdir(exp_data_dir) if file.endswith('.exp')] + return {"required": { + "file_name": (sorted(file_list, key=str.lower),), + "ratio": ("FLOAT", {"default": 1, "min": 0, "max": 1, "step": 0.01}), + }, + } + + RETURN_TYPES = ("EXP_DATA",) + RETURN_NAMES = ("exp",) + FUNCTION = "run" + CATEGORY = "파워집돌이" + + def run(self, file_name, ratio): + # es = ExpressionSet() + with open(os.path.join(exp_data_dir, file_name + ".exp"), 'rb') as f: + es = dill.load(f) + es.mul(ratio) + return (es,) + +class ExpData: + @classmethod + def INPUT_TYPES(s): + return {"required":{ + #"code": ("STRING", {"multiline": False, "default": ""}), + "code1": ("INT", {"default": 0}), + "value1": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "code2": ("INT", {"default": 0}), + "value2": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "code3": ("INT", {"default": 0}), + "value3": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "code4": ("INT", {"default": 0}), + "value4": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + "code5": ("INT", {"default": 0}), + "value5": ("FLOAT", {"default": 0, "min": -100, "max": 100, "step": 0.1}), + }, + "optional":{"add_exp": ("EXP_DATA",),} + } + + RETURN_TYPES = ("EXP_DATA",) + RETURN_NAMES = ("exp",) + FUNCTION = "run" + CATEGORY = "파워집돌이" + + def run(self, code1, value1, code2, value2, code3, value3, code4, value4, code5, value5, add_exp=None): + #print(f"type(None):{type(None)}") + #if type(add_exp) == type(None): + if add_exp == None: + es = ExpressionSet() + log(f"exp11:{es.exp[0,1,1]}") + else: + es = ExpressionSet(es = add_exp) + if id(es.exp) == id(add_exp.exp): + log("id(es.exp) == id(add_exp.exp) is True") + + codes = [code1, code2, code3, code4, code5] + values = [value1, value2, value3, value4, value5] + for i in range(5): + idx = int(codes[i] / 10) + r = codes[i] % 10 + es.exp[0, idx, r] += values[i] * 0.001 + + return (es,) + +class PrintExpData: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "cut_noise": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 0.1}), + }, + "optional": {"exp": ("EXP_DATA",), } + } + + RETURN_TYPES = ("EXP_DATA",) + RETURN_NAMES = ("exp",) + FUNCTION = "run" + CATEGORY = "파워집돌이" + OUTPUT_NODE = True + + def run(self, cut_noise, exp = None): + if exp == None: return (exp,) + + cuted_list = [] + e = exp.exp * 1000 + for idx in range(21): + for r in range(3): + a = abs(e[0, idx, r]) + if(cut_noise < a): cuted_list.append((a, e[0, idx, r], idx*10+r)) + + sorted_list = sorted(cuted_list, reverse=True, key=lambda item: item[0]) + print(f"sorted_list: {[[item[2], round(float(item[1]),1)] for item in sorted_list]}") + return (exp,) + +class Command: + def __init__(self, es, change, keep): + self.es:ExpressionSet = es + self.change = change + self.keep = keep +class AdvancedLivePortrait: + def __init__(s): + s.src_images = None + s.driving_images = None + s.pbar = comfy.utils.ProgressBar(1) + + @classmethod + def INPUT_TYPES(s): + + return { + "required": { + "retargeting_eyes": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.01}), + "retargeting_mouth": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.01}), + "turn_on": ("BOOLEAN", {"default": True}), + "command": ("STRING", {"multiline": True, "default": ""}), + }, + "optional": { + "src_images": ("IMAGE",), + "motion_link": ("EDITOR_LINK",), + "driving_images": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "run" + OUTPUT_NODE = True + CATEGORY = "파워집돌이" + + # INPUT_IS_LIST = False + # OUTPUT_IS_LIST = (False,) + + def parsing_command(self, command, motoin_link): + command.replace(' ', '') + # if command == '': return + lines = command.split('\n') + + cmd_list = [] + + total_length = 0 + + i = 0 + #old_es = None + for line in lines: + i += 1 + if line == '': continue + try: + cmds = line.split('=') + idx = int(cmds[0]) + if idx == 0: es = ExpressionSet() + else: es = ExpressionSet(es = motoin_link[idx]) + cmds = cmds[1].split(':') + change = int(cmds[0]) + keep = int(cmds[1]) + except: + log(f"(파워집돌이) 명령어오류 {i}번줄: {line}") + return None, None + + total_length += change + keep + es.div(change) + cmd_list.append(Command(es, change, keep)) + + return cmd_list, total_length + + + def run(self, retargeting_eyes, retargeting_mouth, turn_on, command, + src_images=None, driving_images=None, motion_link=None): + if turn_on == False: return (None,None) + src_length = 1 + + if src_images == None: + if motion_link != None: + self.psi_list = [motion_link[0]] + else: return (None,None) + + if src_images != None: + src_length = len(src_images) + + if id(src_images) != id(self.src_images): + self.src_images = src_images + if 1 < src_length: + self.psi_list = g_engine.prepare_source(src_images, True) + else: + self.psi_list = [g_engine.prepare_source(src_images)] + + cmd_list, cmd_length = self.parsing_command(command, motion_link) + if cmd_list == None: return (None,None) + cmd_idx = 0 + + driving_length = 0 + if driving_images is not None: + if id(driving_images) != id(self.driving_images): + self.driving_images = driving_images + self.driving_values = g_engine.prepare_driving_video(driving_images) + driving_length = len(self.driving_values) + + #total_length = max(driving_length, cmd_length, src_length) + total_length = max(driving_length, src_length) + + c_i_es = ExpressionSet() + c_o_es = ExpressionSet() + d_0_es = None + out_list = [] + + psi = None + pipeline = g_engine.get_pipeline() + for i in range(total_length): + + if i < src_length: + psi = self.psi_list[i] + s_info = psi.x_s_info + s_es = ExpressionSet(erst=(s_info['kp'] + s_info['exp'], torch.Tensor([0, 0, 0]), s_info['scale'], s_info['t'])) + + new_es = ExpressionSet(es = s_es) + + if i < cmd_length: + cmd = cmd_list[cmd_idx] + #cmd = Command()#지울거 + if 0 < cmd.change: + cmd.change -= 1 + c_i_es.add(cmd.es) + c_i_es.sub(c_o_es) + elif 0 < cmd.keep: + cmd.keep -= 1 + + new_es.add(c_i_es) + + if cmd.change == 0 and cmd.keep == 0: + cmd_idx += 1 + if cmd_idx < len(cmd_list): + c_o_es = ExpressionSet(es = c_i_es) + cmd = cmd_list[cmd_idx] + c_o_es.div(cmd.change) + elif 0 < cmd_length: + new_es.add(c_i_es) + + if i < driving_length: + d_i_info = self.driving_values[i] + d_i_r = torch.Tensor([d_i_info['pitch'], d_i_info['yaw'], d_i_info['roll']])#.float().to(device="cuda:0") + + if d_0_es is None: + d_0_es = ExpressionSet(erst = (d_i_info['exp'], d_i_r, d_i_info['scale'], d_i_info['t'])) + #d_i_es = d_0_es + retargeting(s_es.e, d_0_es.e, retargeting_eyes, (11, 13, 15, 16)) + retargeting(s_es.e, d_0_es.e, retargeting_mouth, (14, 17, 19, 20)) + + #r_new = (r_d_i @ d_0_es.r.permute(0, 2, 1)) @ r_new + new_es.e += d_i_info['exp'] - d_0_es.e + new_es.r += d_i_r - d_0_es.r + new_es.t += d_i_info['t'] - d_0_es.t + + r_new = get_rotation_matrix( + s_info['pitch'] + new_es.r[0], s_info['yaw'] + new_es.r[1], s_info['roll'] + new_es.r[2]) + d_new = new_es.s * (new_es.e @ r_new) + new_es.t + d_new = pipeline.stitching(psi.x_s_user, d_new) + crop_out = pipeline.warp_decode(psi.f_s_user, psi.x_s_user, d_new) + crop_out = pipeline.parse_output(crop_out['out'])[0] + + crop_with_fullsize = cv2.warpAffine(crop_out, psi.crop_trans_m, get_rgb_size(psi.src_rgb), + cv2.INTER_LINEAR) + out = np.clip(psi.mask_ori * crop_with_fullsize + (1 - psi.mask_ori) * psi.src_rgb, 0, 255).astype( + np.uint8) + out_list.append(out) + + self.pbar.update_absolute(i+1, total_length, ("PNG", Image.fromarray(crop_out), None)) + + if len(out_list) == 0: return (None,) + + out_imgs = torch.cat([pil2tensor(img_rgb) for img_rgb in out_list]) + return (out_imgs,) + +class ExpressionEditor: + def __init__(self): + self.sample_image = None + self.src_image = None + + @classmethod + def INPUT_TYPES(s): + display = "number" + #display = "slider" + return { + "required": { + + "rotate_pitch": ("FLOAT", {"default": 0, "min": -20, "max": 20, "step": 0.5, "display": display}), + "rotate_yaw": ("FLOAT", {"default": 0, "min": -20, "max": 20, "step": 0.5, "display": display}), + "rotate_roll": ("FLOAT", {"default": 0, "min": -20, "max": 20, "step": 0.5, "display": display}), + + "blink": ("FLOAT", {"default": 0, "min": -20, "max": 5, "step": 0.5, "display": display}), + "eyebrow": ("FLOAT", {"default": 0, "min": -10, "max": 15, "step": 0.5, "display": display}), + "wink": ("FLOAT", {"default": 0, "min": 0, "max": 25, "step": 0.5, "display": display}), + "pupil_x": ("FLOAT", {"default": 0, "min": -15, "max": 15, "step": 0.5, "display": display}), + "pupil_y": ("FLOAT", {"default": 0, "min": -15, "max": 15, "step": 0.5, "display": display}), + "aaa": ("FLOAT", {"default": 0, "min": 0, "max": 120, "step": 1, "display": display}), + "eee": ("FLOAT", {"default": 0, "min": 0, "max": 15, "step": 0.2, "display": display}), + "woo": ("FLOAT", {"default": 0, "min": 0, "max": 15, "step": 0.2, "display": display}), + "smile": ("FLOAT", {"default": 0, "min": 0, "max": 1.3, "step": 0.01, "display": display}), + + "src_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 1, "step": 0.01, "display": display}), + "sample_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 1, "step": 0.01, "display": display}), + }, + + "optional": {"src_image": ("IMAGE",), "motion_link": ("EDITOR_LINK",), + "sample_image": ("IMAGE",), "add_exp": ("EXP_DATA",), + }, + } + + RETURN_TYPES = ("IMAGE", "EDITOR_LINK", "EXP_DATA") + RETURN_NAMES = ("image", "motion_link", "save_exp") + + FUNCTION = "run" + + OUTPUT_NODE = True + + CATEGORY = "파워집돌이" + + # INPUT_IS_LIST = False + # OUTPUT_IS_LIST = (False,) + + def run(self, rotate_pitch, rotate_yaw, rotate_roll, blink, eyebrow, wink, pupil_x, pupil_y, aaa, eee, woo, smile, + src_ratio, sample_ratio, src_image=None, sample_image=None, motion_link=None, add_exp=None): + rotate_yaw = -rotate_yaw + + new_editor_link = None + if motion_link != None: + self.psi = motion_link[0] + new_editor_link = motion_link.copy() + elif src_image != None: + if id(src_image) != id(self.src_image): + self.psi = g_engine.prepare_source(src_image) + self.src_image = src_image + new_editor_link = [] + new_editor_link.append(self.psi) + else: + return (None,None) + + pipeline = g_engine.get_pipeline() + + psi = self.psi + s_info = psi.x_s_info + #delta_new = copy.deepcopy() + s_exp = s_info['exp'] * src_ratio + s_exp[0, 5] = s_info['exp'][0, 5] + s_exp += s_info['kp'] + + es = ExpressionSet() + + if sample_image != None: + if id(self.sample_image) != id(sample_image): + self.sample_image = sample_image + d_image_np = (sample_image * 255).byte().numpy() + d_face, _ = g_engine.crop_face(d_image_np[0]) + i_d = pipeline.prepare_source(d_face) + self.d_info = pipeline.get_kp_info(i_d) + self.d_info['exp'][0, 5, 0] = 0 + self.d_info['exp'][0, 5, 1] = 0 + + # delta_new += s_exp * (1 - sample_ratio) + self.d_info['exp'] * sample_ratio + es.e += self.d_info['exp'] * sample_ratio + + es.r = g_engine.calc_fe(es.e, blink, eyebrow, wink, pupil_x, pupil_y, aaa, eee, woo, smile, + rotate_pitch, rotate_yaw, rotate_roll) + + if add_exp != None: + es.add(add_exp) + + new_rotate = get_rotation_matrix(s_info['pitch'] + es.r[0], s_info['yaw'] + es.r[1], + s_info['roll'] + es.r[2]) + x_d_new = (s_info['scale'] * (1 + es.s)) * ((s_exp + es.e) @ new_rotate) + s_info['t'] + + x_d_new = pipeline.stitching(psi.x_s_user, x_d_new) + + crop_out = pipeline.warp_decode(psi.f_s_user, psi.x_s_user, x_d_new) + crop_out = pipeline.parse_output(crop_out['out'])[0] + + crop_with_fullsize = cv2.warpAffine(crop_out, psi.crop_trans_m, get_rgb_size(psi.src_rgb), cv2.INTER_LINEAR) + out = np.clip(psi.mask_ori * crop_with_fullsize + (1 - psi.mask_ori) * psi.src_rgb, 0, 255).astype(np.uint8) + + out_img = pil2tensor(out) + + filename = "fe_edit_preview.png" + folder_paths.get_save_image_path(filename, folder_paths.get_temp_directory()) + img = Image.fromarray(crop_out) + img.save(os.path.join(folder_paths.get_temp_directory(), filename), compress_level=1) + results = list() + results.append({"filename": filename, "type": "temp"}) + + new_editor_link.append(es) + + return {"ui": {"images": results}, "result": (out_img, new_editor_link, es)} + +class TestNode: + + def __init__(s): + s.pbar = comfy.utils.ProgressBar(1) + + @classmethod + def INPUT_TYPES(s): + return {"required": { + #"images": ("IMAGE",), + #"test_value": ("INT", {"default": 1, "min": 1}), + "command": ("STRING", {"multiline": True, "default": ""}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "run" + OUTPUT_NODE = True + + def run(self, command): + self.parsing_command(command) + + return (None,) + +NODE_CLASS_MAPPINGS = { + "AdvancedLivePortrait": AdvancedLivePortrait, + "ExpressionEditor": ExpressionEditor, + "ExpData": ExpData, + "LoadExpData": LoadExpData, + "SaveExpData": SaveExpData, + "PrintExpData:": PrintExpData, + #"TestNode": TestNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AdvancedLivePortrait": "Advanced Live Portrait (파워집돌이)", + "ExpressionEditor": "Expression Editor (파워집돌이)", +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..03f0383 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +numpy>=1.26.4 +opencv-python-headless +imageio-ffmpeg>=0.5.1 +lmdb>=1.4.1 +timm>=1.0.7 +rich>=13.7.1 +albumentations>=1.4.10 +ultralytics +tyro==0.8.5 +dill \ No newline at end of file