main
This commit is contained in:
PowerHouseMan
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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')
|
||||
@@ -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
|
||||
@@ -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}")
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 3.4 KiB |
@@ -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
|
||||
@@ -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.
|
||||
@@ -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
|
||||
+49
@@ -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"]
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Binary file not shown.
+20
@@ -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
|
||||
@@ -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 (파워집돌이)",
|
||||
}
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user