initial commit
This commit is contained in:
kijai
+17
@@ -0,0 +1,17 @@
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
**/__pycache__/
|
||||
*.py[cod]
|
||||
**/*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# Model weights
|
||||
**/*.pth
|
||||
**/*.onnx
|
||||
|
||||
# Ipython notebook
|
||||
*.ipynb
|
||||
|
||||
# Temporary files or benchmark resources
|
||||
animations/*
|
||||
tmp/*
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 Kuaishou Visual Generation and Interaction Center
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,23 @@
|
||||
docs/inference.gif filter=lfs diff=lfs merge=lfs -text
|
||||
docs/showcase2.gif filter=lfs diff=lfs merge=lfs -text
|
||||
docs/showcase.gif filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d5.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d7.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d9.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d6.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d8.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d0.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d1.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d2.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/driving/d3.mp4 filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s5.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s7.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s9.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s4.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s10.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s1.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s2.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s3.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s6.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s8.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
examples/source/s0.jpg filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:63f6f9962e1fdf6e6722172e7a18155204858d5d5ce3b1e0646c150360c33bed
|
||||
size 2958395
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:d8bdb2e3d18db28a7c81bc06070aefe2beebae2b489b976fa009eff5881bc7fe
|
||||
size 48753
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:5ce4a92e80e4d0b44c43af7afba471b3ebdf7dc83894ae98af6d4973f17af484
|
||||
size 47762
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:ef5c86e49b1b43dcb1449b499eb5a7f0cbae2f78aec08b5598193be1e4257099
|
||||
size 1430968
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:021af59d9e1d89c967699efd374e7acc77fd14c30b052abd98665de401d9e511
|
||||
size 135015
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:00e3ea79bbf28cbdc4fbb67ec655d9a0fe876e880ec45af55ae481348d0c0fff
|
||||
size 1967790
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9e3ecc59a7285d52ca3666be6562e59ed45ad8afcef28d0f69a9be0e2850ed37
|
||||
size 185258
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:912a02dcffef9b1a6d6641c9a48e170490080c3715dc9f55d91e168483d27e6e
|
||||
size 312295
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9a414aa1d547be35306d692065a2157434bf40a6025ba8e30ce12e5bb322cc33
|
||||
size 2257929
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:ccd094951787cec391f3d444e056ee9a58d715f84152f49100c03bb3ce2962fc
|
||||
size 116149
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:98931abeba483c27c699197e7c7bab58b7deb4324745b0f84e3c92e9ccef5901
|
||||
size 98411
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:49bdee735bed389d289c20d0bdff9bf0a485068e128cec8eb9e5bb11ae1f0422
|
||||
size 537372
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:4d950dc1c6560e0c7d0bcdca258079ba605d26f0979b7a602af75511a15e4c03
|
||||
size 760146
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:d93728ec60e1c0fceaf0f62b05d0b70283e9383a6eef0f31dc2078cb560f707a
|
||||
size 63626
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:d62af6c554e211942f954237f7b898abb41d9d6166b5079564e76c25ff804c55
|
||||
size 143918
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9c368d9aa2eea0c5319b7d76e3b6c339819e49565ef064933a4511c11b581549
|
||||
size 133517
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:fa226be24fc24c11c4bb276c2d7d789218eba88b8e49155855f5685d1b4d7809
|
||||
size 107520
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7a94b38bb66b58c3d44b98b86be3138cd3fb3ca40f89017294476b879353ebc4
|
||||
size 139855
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:01b6af295ebefcac1291cbe3b84fc6af37e40b168c528113662519aace4bbce5
|
||||
size 227378
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:84a6bfcb3346951cef85c0b0bb07856794e0b04d5b65322c18d77198fbf43f2a
|
||||
size 442850
|
||||
@@ -0,0 +1,44 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
config for user
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
from dataclasses import dataclass
|
||||
#import tyro
|
||||
from typing_extensions import Annotated
|
||||
from .base_config import PrintableConfig, make_abs_path
|
||||
|
||||
|
||||
@dataclass(repr=False) # use repr from PrintableConfig
|
||||
class ArgumentConfig(PrintableConfig):
|
||||
########## input arguments ##########
|
||||
#source_image: Annotated[str, tyro.conf.arg(aliases=["-s"])] = make_abs_path('../../assets/examples/source/s6.jpg') # path to the reference portrait
|
||||
#driving_info: Annotated[str, tyro.conf.arg(aliases=["-d"])] = make_abs_path('../../assets/examples/driving/d0.mp4') # path to driving video or template (.pkl format)
|
||||
#output_dir: Annotated[str, tyro.conf.arg(aliases=["-o"])] = 'animations/' # directory to save output video
|
||||
#####################################
|
||||
|
||||
########## inference arguments ##########
|
||||
device_id: int = 0
|
||||
flag_lip_zero : bool = True # whether let the lip to close state before animation, only take effect when flag_eye_retargeting and flag_lip_retargeting is False
|
||||
flag_eye_retargeting: bool = False
|
||||
flag_lip_retargeting: bool = False
|
||||
flag_stitching: bool = True # we recommend setting it to True!
|
||||
flag_relative: bool = True # whether to use relative pose
|
||||
flag_pasteback: bool = True # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
|
||||
flag_do_crop: bool = True # whether to crop the reference portrait to the face-cropping space
|
||||
flag_do_rot: bool = True # whether to conduct the rotation when flag_do_crop is True
|
||||
#########################################
|
||||
|
||||
########## crop arguments ##########
|
||||
dsize: int = 512
|
||||
scale: float = 2.3
|
||||
vx_ratio: float = 0 # vx ratio
|
||||
vy_ratio: float = -0.125 # vy ratio +up, -down
|
||||
####################################
|
||||
|
||||
########## gradio arguments ##########
|
||||
#server_port: Annotated[int, tyro.conf.arg(aliases=["-p"])] = 8890
|
||||
#share: bool = False
|
||||
#server_name: str = "0.0.0.0"
|
||||
@@ -0,0 +1,29 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
pretty printing class
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import os.path as osp
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
|
||||
class PrintableConfig: # pylint: disable=too-few-public-methods
|
||||
"""Printable Config defining str function"""
|
||||
|
||||
def __repr__(self):
|
||||
lines = [self.__class__.__name__ + ":"]
|
||||
for key, val in vars(self).items():
|
||||
if isinstance(val, Tuple):
|
||||
flattened_val = "["
|
||||
for item in val:
|
||||
flattened_val += str(item) + "\n"
|
||||
flattened_val = flattened_val.rstrip("\n")
|
||||
val = flattened_val + "]"
|
||||
lines += f"{key}: {str(val)}".split("\n")
|
||||
return "\n ".join(lines)
|
||||
@@ -0,0 +1,18 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
parameters used for crop faces
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
from dataclasses import dataclass
|
||||
from typing import Union, List
|
||||
from .base_config import PrintableConfig
|
||||
|
||||
|
||||
@dataclass(repr=False) # use repr from PrintableConfig
|
||||
class CropConfig(PrintableConfig):
|
||||
dsize: int = 512 # crop size
|
||||
scale: float = 2.3 # scale factor
|
||||
vx_ratio: float = 0 # vx ratio
|
||||
vy_ratio: float = -0.125 # vy ratio +up, -down
|
||||
@@ -0,0 +1,49 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
config dataclass used for inference
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, Tuple
|
||||
from .base_config import PrintableConfig, make_abs_path
|
||||
|
||||
|
||||
@dataclass(repr=False) # use repr from PrintableConfig
|
||||
class InferenceConfig(PrintableConfig):
|
||||
models_config: str = make_abs_path('./models.yaml') # portrait animation config
|
||||
checkpoint_F: str = make_abs_path('../../pretrained_weights/liveportrait/base_models/appearance_feature_extractor.pth') # path to checkpoint
|
||||
checkpoint_M: str = make_abs_path('../../pretrained_weights/liveportrait/base_models/motion_extractor.pth') # path to checkpoint
|
||||
checkpoint_G: str = make_abs_path('../../pretrained_weights/liveportrait/base_models/spade_generator.pth') # path to checkpoint
|
||||
checkpoint_W: str = make_abs_path('../../pretrained_weights/liveportrait/base_models/warping_module.pth') # path to checkpoint
|
||||
|
||||
checkpoint_S: str = make_abs_path('../../pretrained_weights/liveportrait/retargeting_models/stitching_retargeting_module.pth') # path to checkpoint
|
||||
flag_use_half_precision: bool = True # whether to use half precision
|
||||
|
||||
flag_lip_zero: bool = True # whether let the lip to close state before animation, only take effect when flag_eye_retargeting and flag_lip_retargeting is False
|
||||
lip_zero_threshold: float = 0.03
|
||||
|
||||
flag_eye_retargeting: bool = False
|
||||
flag_lip_retargeting: bool = False
|
||||
flag_stitching: bool = True # we recommend setting it to True!
|
||||
|
||||
flag_relative: bool = True # whether to use relative pose
|
||||
anchor_frame: int = 0 # set this value if find_best_frame is True
|
||||
|
||||
input_shape: Tuple[int, int] = (256, 256) # input shape
|
||||
output_format: Literal['mp4', 'gif'] = 'mp4' # output video format
|
||||
output_fps: int = 30 # fps for output video
|
||||
crf: int = 15 # crf for output video
|
||||
|
||||
flag_write_result: bool = True # whether to write output video
|
||||
flag_pasteback: bool = True # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
|
||||
mask_crop = None
|
||||
flag_write_gif: bool = False
|
||||
size_gif: int = 256
|
||||
ref_max_shape: int = 1280
|
||||
ref_shape_n: int = 2
|
||||
|
||||
device_id: int = 0
|
||||
flag_do_crop: bool = False # whether to crop the reference portrait to the face-cropping space
|
||||
flag_do_rot: bool = True # whether to conduct the rotation when flag_do_crop is True
|
||||
@@ -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)
|
||||
@@ -0,0 +1,207 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Pipeline of LivePortrait
|
||||
"""
|
||||
|
||||
# TODO:
|
||||
# 1. 当前假定所有的模板都是已经裁好的，需要修改下
|
||||
# 2. pick样例图 source + driving
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pickle
|
||||
import os.path as osp
|
||||
from rich.progress import track
|
||||
|
||||
from .config.argument_config import ArgumentConfig
|
||||
from .config.inference_config import InferenceConfig
|
||||
from .config.crop_config import CropConfig
|
||||
from .utils.cropper import Cropper
|
||||
from .utils.camera import get_rotation_matrix
|
||||
from .utils.video import images2video, concat_frames
|
||||
from .utils.crop import _transform_img
|
||||
from .utils.retargeting_utils import calc_lip_close_ratio
|
||||
from .utils.io import load_image_rgb, load_driving_info
|
||||
from .utils.helper import mkdir, basename, dct2cuda, is_video, is_template, resize_to_limit
|
||||
from .utils.rprint import rlog as log
|
||||
from .live_portrait_wrapper import LivePortraitWrapper
|
||||
|
||||
import comfy.utils
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
|
||||
class LivePortraitPipeline(object):
|
||||
|
||||
def __init__(self, appearance_feature_extractor, motion_extractor, warping_module,
|
||||
spade_generator, stitching_retargeting_module, inference_cfg: InferenceConfig, crop_cfg: CropConfig):
|
||||
|
||||
self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(
|
||||
appearance_feature_extractor, motion_extractor, warping_module,
|
||||
spade_generator, stitching_retargeting_module, cfg=inference_cfg)
|
||||
|
||||
self.cropper = Cropper(crop_cfg=crop_cfg)
|
||||
|
||||
def execute(self, img_rgb, driving_images_np, args: ArgumentConfig):
|
||||
inference_cfg = self.live_portrait_wrapper.cfg # for convenience
|
||||
######## process reference portrait ########
|
||||
#img_rgb = load_image_rgb(args.source_image)
|
||||
img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n)
|
||||
#log(f"Load source image from {args.source_image}")
|
||||
crop_info = self.cropper.crop_single_image(img_rgb)
|
||||
source_lmk = crop_info['lmk_crop']
|
||||
img_crop, img_crop_256x256 = crop_info['img_crop'], crop_info['img_crop_256x256']
|
||||
if inference_cfg.flag_do_crop:
|
||||
I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256)
|
||||
else:
|
||||
I_s = self.live_portrait_wrapper.prepare_source(img_rgb)
|
||||
x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)
|
||||
x_c_s = x_s_info['kp']
|
||||
R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
|
||||
f_s = self.live_portrait_wrapper.extract_feature_3d(I_s)
|
||||
x_s = self.live_portrait_wrapper.transform_keypoint(x_s_info)
|
||||
|
||||
if inference_cfg.flag_lip_zero:
|
||||
# let lip-open scalar to be 0 at first
|
||||
c_d_lip_before_animation = [0.]
|
||||
combined_lip_ratio_tensor_before_animation = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk)
|
||||
if combined_lip_ratio_tensor_before_animation[0][0] < inference_cfg.lip_zero_threshold:
|
||||
inference_cfg.flag_lip_zero = False
|
||||
else:
|
||||
lip_delta_before_animation = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)
|
||||
############################################
|
||||
|
||||
######## process driving info ########
|
||||
#if is_video(args.driving_info):
|
||||
#log(f"Load from video file (mp4 mov avi etc...): {args.driving_info}")
|
||||
# TODO: 这里track一下驱动视频 -> 构建模板
|
||||
#driving_rgb_lst = load_driving_info(args.driving_info)
|
||||
|
||||
driving_rgb_lst = driving_images_np
|
||||
|
||||
driving_rgb_lst_256 = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
|
||||
I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst_256)
|
||||
n_frames = I_d_lst.shape[0]
|
||||
if inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting:
|
||||
driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
|
||||
input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)
|
||||
# elif is_template(args.driving_info):
|
||||
# log(f"Load from video templates {args.driving_info}")
|
||||
# with open(args.driving_info, 'rb') as f:
|
||||
# template_lst, driving_lmk_lst = pickle.load(f)
|
||||
# n_frames = template_lst[0]['n_frames']
|
||||
# input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)
|
||||
# else:
|
||||
# raise Exception("Unsupported driving types!")
|
||||
#########################################
|
||||
|
||||
######## prepare for pasteback ########
|
||||
if inference_cfg.flag_pasteback:
|
||||
if inference_cfg.mask_crop is None:
|
||||
inference_cfg.mask_crop = cv2.imread(make_abs_path('./utils/resources/mask_template.png'), cv2.IMREAD_COLOR)
|
||||
mask_ori = _transform_img(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
|
||||
mask_ori = mask_ori.astype(np.float32) / 255.
|
||||
I_p_paste_lst = []
|
||||
#########################################
|
||||
|
||||
I_p_lst = []
|
||||
R_d_0, x_d_0_info = None, None
|
||||
pbar = comfy.utils.ProgressBar(n_frames)
|
||||
for i in track(range(n_frames), description='Animating...', total=n_frames):
|
||||
#if is_video(args.driving_info):
|
||||
# extract kp info by M
|
||||
I_d_i = I_d_lst[i]
|
||||
x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
|
||||
R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
|
||||
# else:
|
||||
# # from template
|
||||
# x_d_i_info = template_lst[i]
|
||||
# x_d_i_info = dct2cuda(x_d_i_info, inference_cfg.device_id)
|
||||
# R_d_i = x_d_i_info['R_d']
|
||||
|
||||
if i == 0:
|
||||
R_d_0 = R_d_i
|
||||
x_d_0_info = x_d_i_info
|
||||
|
||||
if inference_cfg.flag_relative:
|
||||
R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
|
||||
delta_new = x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp'])
|
||||
scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
|
||||
t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
|
||||
else:
|
||||
R_new = R_d_i
|
||||
delta_new = x_d_i_info['exp']
|
||||
scale_new = x_s_info['scale']
|
||||
t_new = x_d_i_info['t']
|
||||
|
||||
t_new[..., 2].fill_(0) # zero tz
|
||||
x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
|
||||
|
||||
# Algorithm 1:
|
||||
if not inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:
|
||||
# without stitching or retargeting
|
||||
if inference_cfg.flag_lip_zero:
|
||||
x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
|
||||
else:
|
||||
pass
|
||||
elif inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:
|
||||
# with stitching and without retargeting
|
||||
if inference_cfg.flag_lip_zero:
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
|
||||
else:
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
|
||||
else:
|
||||
eyes_delta, lip_delta = None, None
|
||||
if inference_cfg.flag_eye_retargeting:
|
||||
c_d_eyes_i = input_eye_ratio_lst[i]
|
||||
combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(c_d_eyes_i, source_lmk)
|
||||
# ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
|
||||
eyes_delta = self.live_portrait_wrapper.retarget_eye(x_s, combined_eye_ratio_tensor)
|
||||
if inference_cfg.flag_lip_retargeting:
|
||||
c_d_lip_i = input_lip_ratio_lst[i]
|
||||
combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_i, source_lmk)
|
||||
# ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
|
||||
lip_delta = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor)
|
||||
|
||||
if inference_cfg.flag_relative: # use x_s
|
||||
x_d_i_new = x_s + \
|
||||
(eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
|
||||
else: # use x_d,i
|
||||
x_d_i_new = x_d_i_new + \
|
||||
(eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
|
||||
|
||||
if inference_cfg.flag_stitching:
|
||||
x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
|
||||
|
||||
out = self.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)
|
||||
I_p_i = self.live_portrait_wrapper.parse_output(out['out'])[0]
|
||||
I_p_lst.append(I_p_i)
|
||||
pbar.update(1)
|
||||
|
||||
#if inference_cfg.flag_pasteback:
|
||||
I_p_i_to_ori = _transform_img(I_p_i, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
|
||||
I_p_i_to_ori_blend = np.clip(mask_ori * I_p_i_to_ori + (1 - mask_ori) * img_rgb, 0, 255).astype(np.uint8)
|
||||
out = np.hstack([I_p_i_to_ori, I_p_i_to_ori_blend])
|
||||
I_p_paste_lst.append(I_p_i_to_ori_blend)
|
||||
|
||||
#mkdir(args.output_dir)
|
||||
wfp_concat = None
|
||||
#if is_video(args.driving_info):
|
||||
|
||||
#frames_concatenated = concat_frames(I_p_lst, driving_rgb_lst, img_crop_256x256)
|
||||
return I_p_lst, I_p_paste_lst
|
||||
# # save (driving frames, source image, drived frames) result
|
||||
# wfp_concat = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat.mp4')
|
||||
# images2video(frames_concatenated, wfp=wfp_concat)
|
||||
|
||||
# # save drived result
|
||||
# wfp = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}.mp4')
|
||||
# if inference_cfg.flag_pasteback:
|
||||
# images2video(I_p_paste_lst, wfp=wfp)
|
||||
# else:
|
||||
# images2video(I_p_lst, wfp=wfp)
|
||||
|
||||
# return wfp, wfp_concat
|
||||
@@ -0,0 +1,341 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Wrapper for LivePortrait core functions
|
||||
"""
|
||||
|
||||
import os.path as osp
|
||||
import numpy as np
|
||||
import cv2
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from .utils.timer import Timer
|
||||
from .utils.helper import load_model, concat_feat
|
||||
from .utils.retargeting_utils import compute_eye_delta, compute_lip_delta
|
||||
from .utils.camera import headpose_pred_to_degree, get_rotation_matrix
|
||||
from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
|
||||
from .config.inference_config import InferenceConfig
|
||||
from .utils.rprint import rlog as log
|
||||
|
||||
|
||||
class LivePortraitWrapper(object):
|
||||
|
||||
def __init__(self, appearance_feature_extractor, motion_extractor, warping_module,
|
||||
spade_generator, stitching_retargeting_module, 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.appearance_feature_extractor = appearance_feature_extractor
|
||||
self.motion_extractor = motion_extractor
|
||||
self.warping_module = warping_module
|
||||
self.spade_generator = spade_generator
|
||||
self.stitching_retargeting_module = stitching_retargeting_module
|
||||
|
||||
self.cfg = cfg
|
||||
self.device_id = cfg.device_id
|
||||
self.timer = Timer()
|
||||
|
||||
def update_config(self, user_args):
|
||||
for k, v in user_args.items():
|
||||
if hasattr(self.cfg, k):
|
||||
setattr(self.cfg, k, v)
|
||||
|
||||
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 get_pose_dct(self, kp_info: dict) -> dict:
|
||||
pose_dct = dict(
|
||||
pitch=headpose_pred_to_degree(kp_info['pitch']).item(),
|
||||
yaw=headpose_pred_to_degree(kp_info['yaw']).item(),
|
||||
roll=headpose_pred_to_degree(kp_info['roll']).item(),
|
||||
)
|
||||
return pose_dct
|
||||
|
||||
def get_fs_and_kp_info(self, source_prepared, driving_first_frame):
|
||||
|
||||
# get the canonical keypoints of source image by M
|
||||
source_kp_info = self.get_kp_info(source_prepared, flag_refine_info=True)
|
||||
source_rotation = get_rotation_matrix(source_kp_info['pitch'], source_kp_info['yaw'], source_kp_info['roll'])
|
||||
|
||||
# get the canonical keypoints of first driving frame by M
|
||||
driving_first_frame_kp_info = self.get_kp_info(driving_first_frame, flag_refine_info=True)
|
||||
driving_first_frame_rotation = get_rotation_matrix(
|
||||
driving_first_frame_kp_info['pitch'],
|
||||
driving_first_frame_kp_info['yaw'],
|
||||
driving_first_frame_kp_info['roll']
|
||||
)
|
||||
|
||||
# get feature volume by F
|
||||
source_feature_3d = self.extract_feature_3d(source_prepared)
|
||||
|
||||
return source_kp_info, source_rotation, source_feature_3d, driving_first_frame_kp_info, driving_first_frame_rotation
|
||||
|
||||
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 retarget_eye(self, kp_source: torch.Tensor, eye_close_ratio: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
kp_source: BxNx3
|
||||
eye_close_ratio: Bx3
|
||||
Return: Bx(3*num_kp+2)
|
||||
"""
|
||||
feat_eye = concat_feat(kp_source, eye_close_ratio)
|
||||
|
||||
with torch.no_grad():
|
||||
delta = self.stitching_retargeting_module['eye'](feat_eye)
|
||||
|
||||
return delta
|
||||
|
||||
def retarget_lip(self, kp_source: torch.Tensor, lip_close_ratio: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
kp_source: BxNx3
|
||||
lip_close_ratio: Bx2
|
||||
"""
|
||||
feat_lip = concat_feat(kp_source, lip_close_ratio)
|
||||
|
||||
with torch.no_grad():
|
||||
delta = self.stitching_retargeting_module['lip'](feat_lip)
|
||||
|
||||
return delta
|
||||
|
||||
def retarget_keypoints(self, frame_idx, num_keypoints, input_eye_ratios, input_lip_ratios, source_landmarks, portrait_wrapper, kp_source, driving_transformed_kp):
|
||||
# TODO: GPT style, refactor it...
|
||||
if self.cfg.flag_eye_retargeting:
|
||||
# ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
|
||||
eye_delta = compute_eye_delta(frame_idx, input_eye_ratios, source_landmarks, portrait_wrapper, kp_source)
|
||||
else:
|
||||
# α_eyes = 0
|
||||
eye_delta = None
|
||||
|
||||
if self.cfg.flag_lip_retargeting:
|
||||
# ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
|
||||
lip_delta = compute_lip_delta(frame_idx, input_lip_ratios, source_landmarks, portrait_wrapper, kp_source)
|
||||
else:
|
||||
# α_lip = 0
|
||||
lip_delta = None
|
||||
|
||||
if self.cfg.flag_relative: # use x_s
|
||||
new_driving_kp = kp_source + \
|
||||
(eye_delta.reshape(-1, num_keypoints, 3) if eye_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, num_keypoints, 3) if lip_delta is not None else 0)
|
||||
else: # use x_d,i
|
||||
new_driving_kp = driving_transformed_kp + \
|
||||
(eye_delta.reshape(-1, num_keypoints, 3) if eye_delta is not None else 0) + \
|
||||
(lip_delta.reshape(-1, num_keypoints, 3) if lip_delta is not None else 0)
|
||||
|
||||
return new_driving_kp
|
||||
|
||||
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
|
||||
|
||||
def calc_retargeting_ratio(self, source_lmk, driving_lmk_lst):
|
||||
input_eye_ratio_lst = []
|
||||
input_lip_ratio_lst = []
|
||||
for lmk in driving_lmk_lst:
|
||||
# for eyes retargeting
|
||||
input_eye_ratio_lst.append(calc_eye_close_ratio(lmk[None]))
|
||||
# for lip retargeting
|
||||
input_lip_ratio_lst.append(calc_lip_close_ratio(lmk[None]))
|
||||
return input_eye_ratio_lst, input_lip_ratio_lst
|
||||
|
||||
def calc_combined_eye_ratio(self, input_eye_ratio, source_lmk):
|
||||
eye_close_ratio = calc_eye_close_ratio(source_lmk[None])
|
||||
eye_close_ratio_tensor = torch.from_numpy(eye_close_ratio).float().cuda(self.device_id)
|
||||
input_eye_ratio_tensor = torch.Tensor([input_eye_ratio[0][0]]).reshape(1, 1).cuda(self.device_id)
|
||||
# [c_s,eyes, c_d,eyes,i]
|
||||
combined_eye_ratio_tensor = torch.cat([eye_close_ratio_tensor, input_eye_ratio_tensor], dim=1)
|
||||
return combined_eye_ratio_tensor
|
||||
|
||||
def calc_combined_lip_ratio(self, input_lip_ratio, source_lmk):
|
||||
lip_close_ratio = calc_lip_close_ratio(source_lmk[None])
|
||||
lip_close_ratio_tensor = torch.from_numpy(lip_close_ratio).float().cuda(self.device_id)
|
||||
# [c_s,lip, c_d,lip,i]
|
||||
input_lip_ratio_tensor = torch.Tensor([input_lip_ratio[0]]).cuda(self.device_id)
|
||||
if input_lip_ratio_tensor.shape != [1, 1]:
|
||||
input_lip_ratio_tensor = input_lip_ratio_tensor.reshape(1, 1)
|
||||
combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1)
|
||||
return combined_lip_ratio_tensor
|
||||
@@ -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,65 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
Make video template
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pickle
|
||||
from rich.progress import track
|
||||
from .utils.cropper import Cropper
|
||||
|
||||
from .utils.io import load_driving_info
|
||||
from .utils.camera import get_rotation_matrix
|
||||
from .utils.helper import mkdir, basename
|
||||
from .utils.rprint import rlog as log
|
||||
from .config.crop_config import CropConfig
|
||||
from .config.inference_config import InferenceConfig
|
||||
from .live_portrait_wrapper import LivePortraitWrapper
|
||||
|
||||
class TemplateMaker:
|
||||
|
||||
def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):
|
||||
self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)
|
||||
self.cropper = Cropper(crop_cfg=crop_cfg)
|
||||
|
||||
def make_motion_template(self, video_fp: str, output_path: str, **kwargs):
|
||||
""" make video template (.pkl format)
|
||||
video_fp: driving video file path
|
||||
output_path: where to save the pickle file
|
||||
"""
|
||||
|
||||
driving_rgb_lst = load_driving_info(video_fp)
|
||||
driving_rgb_lst = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
|
||||
driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
|
||||
I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst)
|
||||
|
||||
n_frames = I_d_lst.shape[0]
|
||||
|
||||
templates = []
|
||||
|
||||
|
||||
for i in track(range(n_frames), description='Making templates...', total=n_frames):
|
||||
I_d_i = I_d_lst[i]
|
||||
x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
|
||||
R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
|
||||
# collect s_d, R_d, δ_d and t_d for inference
|
||||
template_dct = {
|
||||
'n_frames': n_frames,
|
||||
'frames_index': i,
|
||||
}
|
||||
template_dct['scale'] = x_d_i_info['scale'].cpu().numpy().astype(np.float32)
|
||||
template_dct['R_d'] = R_d_i.cpu().numpy().astype(np.float32)
|
||||
template_dct['exp'] = x_d_i_info['exp'].cpu().numpy().astype(np.float32)
|
||||
template_dct['t'] = x_d_i_info['t'].cpu().numpy().astype(np.float32)
|
||||
|
||||
templates.append(template_dct)
|
||||
|
||||
mkdir(output_path)
|
||||
# Save the dictionary as a pickle file
|
||||
pickle_fp = os.path.join(output_path, f'{basename(video_fp)}.pkl')
|
||||
with open(pickle_fp, 'wb') as f:
|
||||
pickle.dump([templates, driving_lmk_lst], f)
|
||||
log(f"Template saved at {pickle_fp}")
|
||||
@@ -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,393 @@
|
||||
# coding: utf-8
|
||||
|
||||
"""
|
||||
cropping function and the related preprocess functions for cropping
|
||||
"""
|
||||
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) # NOTE: enforce single thread
|
||||
import numpy as np
|
||||
from .rprint import rprint as print
|
||||
from math import sin, cos, acos, degrees
|
||||
|
||||
DTYPE = np.float32
|
||||
CV2_INTERP = cv2.INTER_LINEAR
|
||||
|
||||
|
||||
def _transform_img(img, M, dsize, flags=CV2_INTERP, borderMode=None):
|
||||
""" conduct similarity or affine transformation to the image, do not do border operation!
|
||||
img:
|
||||
M: 2x3 matrix or 3x3 matrix
|
||||
dsize: target shape (width, height)
|
||||
"""
|
||||
if isinstance(dsize, tuple) or isinstance(dsize, list):
|
||||
_dsize = tuple(dsize)
|
||||
else:
|
||||
_dsize = (dsize, dsize)
|
||||
|
||||
if borderMode is not None:
|
||||
return cv2.warpAffine(img, M[:2, :], dsize=_dsize, flags=flags, borderMode=borderMode, borderValue=(0, 0, 0))
|
||||
else:
|
||||
return cv2.warpAffine(img, M[:2, :], dsize=_dsize, flags=flags)
|
||||
|
||||
|
||||
def _transform_pts(pts, M):
|
||||
""" conduct similarity or affine transformation to the pts
|
||||
pts: Nx2 ndarray
|
||||
M: 2x3 matrix or 3x3 matrix
|
||||
return: Nx2
|
||||
"""
|
||||
return pts @ M[:2, :2].T + M[:2, 2]
|
||||
|
||||
|
||||
def parse_pt2_from_pt101(pt101, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 101 points, which cancels the roll
|
||||
"""
|
||||
# the former version use the eye center, but it is not robust, now use interpolation
|
||||
pt_left_eye = np.mean(pt101[[39, 42, 45, 48]], axis=0) # left eye center
|
||||
pt_right_eye = np.mean(pt101[[51, 54, 57, 60]], axis=0) # right eye center
|
||||
|
||||
if use_lip:
|
||||
# use lip
|
||||
pt_center_eye = (pt_left_eye + pt_right_eye) / 2
|
||||
pt_center_lip = (pt101[75] + pt101[81]) / 2
|
||||
pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0)
|
||||
return pt2
|
||||
|
||||
|
||||
def parse_pt2_from_pt106(pt106, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 106 points, which cancels the roll
|
||||
"""
|
||||
pt_left_eye = np.mean(pt106[[33, 35, 40, 39]], axis=0) # left eye center
|
||||
pt_right_eye = np.mean(pt106[[87, 89, 94, 93]], axis=0) # right eye center
|
||||
|
||||
if use_lip:
|
||||
# use lip
|
||||
pt_center_eye = (pt_left_eye + pt_right_eye) / 2
|
||||
pt_center_lip = (pt106[52] + pt106[61]) / 2
|
||||
pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0)
|
||||
return pt2
|
||||
|
||||
|
||||
def parse_pt2_from_pt203(pt203, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 203 points, which cancels the roll
|
||||
"""
|
||||
pt_left_eye = np.mean(pt203[[0, 6, 12, 18]], axis=0) # left eye center
|
||||
pt_right_eye = np.mean(pt203[[24, 30, 36, 42]], axis=0) # right eye center
|
||||
if use_lip:
|
||||
# use lip
|
||||
pt_center_eye = (pt_left_eye + pt_right_eye) / 2
|
||||
pt_center_lip = (pt203[48] + pt203[66]) / 2
|
||||
pt2 = np.stack([pt_center_eye, pt_center_lip], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([pt_left_eye, pt_right_eye], axis=0)
|
||||
return pt2
|
||||
|
||||
|
||||
def parse_pt2_from_pt68(pt68, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 68 points, which cancels the roll
|
||||
"""
|
||||
lm_idx = np.array([31, 37, 40, 43, 46, 49, 55], dtype=np.int32) - 1
|
||||
if use_lip:
|
||||
pt5 = np.stack([
|
||||
np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye
|
||||
np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye
|
||||
pt68[lm_idx[0], :], # nose
|
||||
pt68[lm_idx[5], :], # lip
|
||||
pt68[lm_idx[6], :] # lip
|
||||
], axis=0)
|
||||
|
||||
pt2 = np.stack([
|
||||
(pt5[0] + pt5[1]) / 2,
|
||||
(pt5[3] + pt5[4]) / 2
|
||||
], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([
|
||||
np.mean(pt68[lm_idx[[1, 2]], :], 0), # left eye
|
||||
np.mean(pt68[lm_idx[[3, 4]], :], 0), # right eye
|
||||
], axis=0)
|
||||
|
||||
return pt2
|
||||
|
||||
|
||||
def parse_pt2_from_pt5(pt5, use_lip=True):
|
||||
"""
|
||||
parsing the 2 points according to the 5 points, which cancels the roll
|
||||
"""
|
||||
if use_lip:
|
||||
pt2 = np.stack([
|
||||
(pt5[0] + pt5[1]) / 2,
|
||||
(pt5[3] + pt5[4]) / 2
|
||||
], axis=0)
|
||||
else:
|
||||
pt2 = np.stack([
|
||||
pt5[0],
|
||||
pt5[1]
|
||||
], axis=0)
|
||||
return pt2
|
||||
|
||||
|
||||
def parse_pt2_from_pt_x(pts, use_lip=True):
|
||||
if pts.shape[0] == 101:
|
||||
pt2 = parse_pt2_from_pt101(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] == 106:
|
||||
pt2 = parse_pt2_from_pt106(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] == 68:
|
||||
pt2 = parse_pt2_from_pt68(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] == 5:
|
||||
pt2 = parse_pt2_from_pt5(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] == 203:
|
||||
pt2 = parse_pt2_from_pt203(pts, use_lip=use_lip)
|
||||
elif pts.shape[0] > 101:
|
||||
# take the first 101 points
|
||||
pt2 = parse_pt2_from_pt101(pts[:101], use_lip=use_lip)
|
||||
else:
|
||||
raise Exception(f'Unknow shape: {pts.shape}')
|
||||
|
||||
if not use_lip:
|
||||
# NOTE: to compile with the latter code, need to rotate the pt2 90 degrees clockwise manually
|
||||
v = pt2[1] - pt2[0]
|
||||
pt2[1, 0] = pt2[0, 0] - v[1]
|
||||
pt2[1, 1] = pt2[0, 1] + v[0]
|
||||
|
||||
return pt2
|
||||
|
||||
|
||||
def parse_rect_from_landmark(
|
||||
pts,
|
||||
scale=1.5,
|
||||
need_square=True,
|
||||
vx_ratio=0,
|
||||
vy_ratio=0,
|
||||
use_deg_flag=False,
|
||||
**kwargs
|
||||
):
|
||||
"""parsing center, size, angle from 101/68/5/x landmarks
|
||||
vx_ratio: the offset ratio along the pupil axis x-axis, multiplied by size
|
||||
vy_ratio: the offset ratio along the pupil axis y-axis, multiplied by size, which is used to contain more forehead area
|
||||
|
||||
judge with pts.shape
|
||||
"""
|
||||
pt2 = parse_pt2_from_pt_x(pts, use_lip=kwargs.get('use_lip', True))
|
||||
|
||||
uy = pt2[1] - pt2[0]
|
||||
l = np.linalg.norm(uy)
|
||||
if l <= 1e-3:
|
||||
uy = np.array([0, 1], dtype=DTYPE)
|
||||
else:
|
||||
uy /= l
|
||||
ux = np.array((uy[1], -uy[0]), dtype=DTYPE)
|
||||
|
||||
# the rotation degree of the x-axis, the clockwise is positive, the counterclockwise is negative (image coordinate system)
|
||||
# print(uy)
|
||||
# print(ux)
|
||||
angle = acos(ux[0])
|
||||
if ux[1] < 0:
|
||||
angle = -angle
|
||||
|
||||
# rotation matrix
|
||||
M = np.array([ux, uy])
|
||||
|
||||
# calculate the size which contains the angle degree of the bbox, and the center
|
||||
center0 = np.mean(pts, axis=0)
|
||||
rpts = (pts - center0) @ M.T # (M @ P.T).T = P @ M.T
|
||||
lt_pt = np.min(rpts, axis=0)
|
||||
rb_pt = np.max(rpts, axis=0)
|
||||
center1 = (lt_pt + rb_pt) / 2
|
||||
|
||||
size = rb_pt - lt_pt
|
||||
if need_square:
|
||||
m = max(size[0], size[1])
|
||||
size[0] = m
|
||||
size[1] = m
|
||||
|
||||
size *= scale # scale size
|
||||
center = center0 + ux * center1[0] + uy * center1[1] # counterclockwise rotation, equivalent to M.T @ center1.T
|
||||
center = center + ux * (vx_ratio * size) + uy * \
|
||||
(vy_ratio * size) # considering the offset in vx and vy direction
|
||||
|
||||
if use_deg_flag:
|
||||
angle = degrees(angle)
|
||||
|
||||
return center, size, angle
|
||||
|
||||
|
||||
def parse_bbox_from_landmark(pts, **kwargs):
|
||||
center, size, angle = parse_rect_from_landmark(pts, **kwargs)
|
||||
cx, cy = center
|
||||
w, h = size
|
||||
|
||||
# calculate the vertex positions before rotation
|
||||
bbox = np.array([
|
||||
[cx-w/2, cy-h/2], # left, top
|
||||
[cx+w/2, cy-h/2],
|
||||
[cx+w/2, cy+h/2], # right, bottom
|
||||
[cx-w/2, cy+h/2]
|
||||
], dtype=DTYPE)
|
||||
|
||||
# construct rotation matrix
|
||||
bbox_rot = bbox.copy()
|
||||
R = np.array([
|
||||
[np.cos(angle), -np.sin(angle)],
|
||||
[np.sin(angle), np.cos(angle)]
|
||||
], dtype=DTYPE)
|
||||
|
||||
# calculate the relative position of each vertex from the rotation center, then rotate these positions, and finally add the coordinates of the rotation center
|
||||
bbox_rot = (bbox_rot - center) @ R.T + center
|
||||
|
||||
return {
|
||||
'center': center, # 2x1
|
||||
'size': size, # scalar
|
||||
'angle': angle, # rad, counterclockwise
|
||||
'bbox': bbox, # 4x2
|
||||
'bbox_rot': bbox_rot, # 4x2
|
||||
}
|
||||
|
||||
|
||||
def crop_image_by_bbox(img, bbox, lmk=None, dsize=512, angle=None, flag_rot=False, **kwargs):
|
||||
left, top, right, bot = bbox
|
||||
if int(right - left) != int(bot - top):
|
||||
print(f'right-left {right-left} != bot-top {bot-top}')
|
||||
size = right - left
|
||||
|
||||
src_center = np.array([(left + right) / 2, (top + bot) / 2], dtype=DTYPE)
|
||||
tgt_center = np.array([dsize / 2, dsize / 2], dtype=DTYPE)
|
||||
|
||||
s = dsize / size # scale
|
||||
if flag_rot and angle is not None:
|
||||
costheta, sintheta = cos(angle), sin(angle)
|
||||
cx, cy = src_center[0], src_center[1] # ori center
|
||||
tcx, tcy = tgt_center[0], tgt_center[1] # target center
|
||||
# need to infer
|
||||
M_o2c = np.array(
|
||||
[[s * costheta, s * sintheta, tcx - s * (costheta * cx + sintheta * cy)],
|
||||
[-s * sintheta, s * costheta, tcy - s * (-sintheta * cx + costheta * cy)]],
|
||||
dtype=DTYPE
|
||||
)
|
||||
else:
|
||||
M_o2c = np.array(
|
||||
[[s, 0, tgt_center[0] - s * src_center[0]],
|
||||
[0, s, tgt_center[1] - s * src_center[1]]],
|
||||
dtype=DTYPE
|
||||
)
|
||||
|
||||
if flag_rot and angle is None:
|
||||
print('angle is None, but flag_rotate is True', style="bold yellow")
|
||||
|
||||
img_crop = _transform_img(img, M_o2c, dsize=dsize, borderMode=kwargs.get('borderMode', None))
|
||||
|
||||
lmk_crop = _transform_pts(lmk, M_o2c) if lmk is not None else None
|
||||
|
||||
M_o2c = np.vstack([M_o2c, np.array([0, 0, 1], dtype=DTYPE)])
|
||||
M_c2o = np.linalg.inv(M_o2c)
|
||||
|
||||
# cv2.imwrite('crop.jpg', img_crop)
|
||||
|
||||
return {
|
||||
'img_crop': img_crop,
|
||||
'lmk_crop': lmk_crop,
|
||||
'M_o2c': M_o2c,
|
||||
'M_c2o': M_c2o,
|
||||
}
|
||||
|
||||
|
||||
def _estimate_similar_transform_from_pts(
|
||||
pts,
|
||||
dsize,
|
||||
scale=1.5,
|
||||
vx_ratio=0,
|
||||
vy_ratio=-0.1,
|
||||
flag_do_rot=True,
|
||||
**kwargs
|
||||
):
|
||||
""" calculate the affine matrix of the cropped image from sparse points, the original image to the cropped image, the inverse is the cropped image to the original image
|
||||
pts: landmark, 101 or 68 points or other points, Nx2
|
||||
scale: the larger scale factor, the smaller face ratio
|
||||
vx_ratio: x shift
|
||||
vy_ratio: y shift, the smaller the y shift, the lower the face region
|
||||
rot_flag: if it is true, conduct correction
|
||||
"""
|
||||
center, size, angle = parse_rect_from_landmark(
|
||||
pts, scale=scale, vx_ratio=vx_ratio, vy_ratio=vy_ratio,
|
||||
use_lip=kwargs.get('use_lip', True)
|
||||
)
|
||||
|
||||
s = dsize / size[0] # scale
|
||||
tgt_center = np.array([dsize / 2, dsize / 2], dtype=DTYPE) # center of dsize
|
||||
|
||||
if flag_do_rot:
|
||||
costheta, sintheta = cos(angle), sin(angle)
|
||||
cx, cy = center[0], center[1] # ori center
|
||||
tcx, tcy = tgt_center[0], tgt_center[1] # target center
|
||||
# need to infer
|
||||
M_INV = np.array(
|
||||
[[s * costheta, s * sintheta, tcx - s * (costheta * cx + sintheta * cy)],
|
||||
[-s * sintheta, s * costheta, tcy - s * (-sintheta * cx + costheta * cy)]],
|
||||
dtype=DTYPE
|
||||
)
|
||||
else:
|
||||
M_INV = np.array(
|
||||
[[s, 0, tgt_center[0] - s * center[0]],
|
||||
[0, s, tgt_center[1] - s * center[1]]],
|
||||
dtype=DTYPE
|
||||
)
|
||||
|
||||
M_INV_H = np.vstack([M_INV, np.array([0, 0, 1])])
|
||||
M = np.linalg.inv(M_INV_H)
|
||||
|
||||
# M_INV is from the original image to the cropped image, M is from the cropped image to the original image
|
||||
return M_INV, M[:2, ...]
|
||||
|
||||
|
||||
def crop_image(img, pts: np.ndarray, **kwargs):
|
||||
dsize = kwargs.get('dsize', 224)
|
||||
scale = kwargs.get('scale', 1.5) # 1.5 | 1.6
|
||||
vy_ratio = kwargs.get('vy_ratio', -0.1) # -0.0625 | -0.1
|
||||
|
||||
M_INV, _ = _estimate_similar_transform_from_pts(
|
||||
pts,
|
||||
dsize=dsize,
|
||||
scale=scale,
|
||||
vy_ratio=vy_ratio,
|
||||
flag_do_rot=kwargs.get('flag_do_rot', True),
|
||||
)
|
||||
|
||||
if img is None:
|
||||
M_INV_H = np.vstack([M_INV, np.array([0, 0, 1], dtype=DTYPE)])
|
||||
M = np.linalg.inv(M_INV_H)
|
||||
ret_dct = {
|
||||
'M': M[:2, ...], # from the original image to the cropped image
|
||||
'M_o2c': M[:2, ...], # from the cropped image to the original image
|
||||
'img_crop': None,
|
||||
'pt_crop': None,
|
||||
}
|
||||
return ret_dct
|
||||
|
||||
img_crop = _transform_img(img, M_INV, dsize) # origin to crop
|
||||
pt_crop = _transform_pts(pts, M_INV)
|
||||
|
||||
M_o2c = np.vstack([M_INV, np.array([0, 0, 1], dtype=DTYPE)])
|
||||
M_c2o = np.linalg.inv(M_o2c)
|
||||
|
||||
ret_dct = {
|
||||
'M_o2c': M_o2c, # from the original image to the cropped image 3x3
|
||||
'M_c2o': M_c2o, # from the cropped image to the original image 3x3
|
||||
'img_crop': img_crop, # the cropped image
|
||||
'pt_crop': pt_crop, # the landmarks of the cropped image
|
||||
}
|
||||
|
||||
return ret_dct
|
||||
|
||||
def average_bbox_lst(bbox_lst):
|
||||
if len(bbox_lst) == 0:
|
||||
return None
|
||||
bbox_arr = np.array(bbox_lst)
|
||||
return np.mean(bbox_arr, axis=0).tolist()
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
# coding: utf-8
|
||||
|
||||
import numpy as np
|
||||
import os.path as osp
|
||||
from typing import List, Union, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
|
||||
from .landmark_runner import LandmarkRunner
|
||||
from .face_analysis_diy import FaceAnalysisDIY
|
||||
from .helper import prefix
|
||||
from .crop import crop_image, crop_image_by_bbox, parse_bbox_from_landmark, average_bbox_lst
|
||||
from .timer import Timer
|
||||
from .rprint import rlog as log
|
||||
from .io import load_image_rgb
|
||||
from .video import VideoWriter, get_fps, change_video_fps
|
||||
|
||||
import folder_paths
|
||||
import os
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Trajectory:
|
||||
start: int = -1 # 起始帧 闭区间
|
||||
end: int = -1 # 结束帧 闭区间
|
||||
lmk_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # lmk list
|
||||
bbox_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # bbox list
|
||||
frame_rgb_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # frame list
|
||||
frame_rgb_crop_lst: Union[Tuple, List, np.ndarray] = field(default_factory=list) # frame crop list
|
||||
|
||||
|
||||
class Cropper(object):
|
||||
def __init__(self, **kwargs) -> None:
|
||||
device_id = kwargs.get('device_id', 0)
|
||||
self.landmark_runner = LandmarkRunner(
|
||||
#ckpt_path=make_abs_path('../../pretrained_weights/liveportrait/landmark.onnx'),
|
||||
ckpt_path=os.path.join(folder_paths.models_dir, 'liveportrait', 'landmark.onnx'),
|
||||
onnx_provider='cuda',
|
||||
device_id=device_id
|
||||
)
|
||||
self.landmark_runner.warmup()
|
||||
|
||||
self.face_analysis_wrapper = FaceAnalysisDIY(
|
||||
name='buffalo_l',
|
||||
root=os.path.join(folder_paths.models_dir, 'insightface'),
|
||||
providers=["CUDAExecutionProvider"]
|
||||
)
|
||||
self.face_analysis_wrapper.prepare(ctx_id=device_id, det_size=(512, 512))
|
||||
self.face_analysis_wrapper.warmup()
|
||||
|
||||
self.crop_cfg = kwargs.get('crop_cfg', None)
|
||||
|
||||
def update_config(self, user_args):
|
||||
for k, v in user_args.items():
|
||||
if hasattr(self.crop_cfg, k):
|
||||
setattr(self.crop_cfg, k, v)
|
||||
|
||||
def crop_single_image(self, obj, **kwargs):
|
||||
direction = kwargs.get('direction', 'large-small')
|
||||
|
||||
# crop and align a single image
|
||||
if isinstance(obj, str):
|
||||
img_rgb = load_image_rgb(obj)
|
||||
elif isinstance(obj, np.ndarray):
|
||||
img_rgb = obj
|
||||
|
||||
src_face = self.face_analysis_wrapper.get(
|
||||
img_rgb,
|
||||
flag_do_landmark_2d_106=True,
|
||||
direction=direction
|
||||
)
|
||||
|
||||
if len(src_face) == 0:
|
||||
log('No face detected in the source image.')
|
||||
raise Exception("No face detected in the source image!")
|
||||
elif len(src_face) > 1:
|
||||
log(f'More than one face detected in the image, only pick one face by rule {direction}.')
|
||||
|
||||
src_face = src_face[0]
|
||||
pts = src_face.landmark_2d_106
|
||||
|
||||
# crop the face
|
||||
ret_dct = crop_image(
|
||||
img_rgb, # ndarray
|
||||
pts, # 106x2 or Nx2
|
||||
dsize=kwargs.get('dsize', 512),
|
||||
scale=kwargs.get('scale', 2.3),
|
||||
vy_ratio=kwargs.get('vy_ratio', -0.15),
|
||||
)
|
||||
# update a 256x256 version for network input or else
|
||||
ret_dct['img_crop_256x256'] = cv2.resize(ret_dct['img_crop'], (256, 256), interpolation=cv2.INTER_AREA)
|
||||
ret_dct['pt_crop_256x256'] = ret_dct['pt_crop'] * 256 / kwargs.get('dsize', 512)
|
||||
|
||||
recon_ret = self.landmark_runner.run(img_rgb, pts)
|
||||
lmk = recon_ret['pts']
|
||||
ret_dct['lmk_crop'] = lmk
|
||||
|
||||
return ret_dct
|
||||
|
||||
def get_retargeting_lmk_info(self, driving_rgb_lst):
|
||||
# TODO: implement a tracking-based version
|
||||
driving_lmk_lst = []
|
||||
for driving_image in driving_rgb_lst:
|
||||
ret_dct = self.crop_single_image(driving_image)
|
||||
driving_lmk_lst.append(ret_dct['lmk_crop'])
|
||||
return driving_lmk_lst
|
||||
|
||||
def make_video_clip(self, driving_rgb_lst, output_path, output_fps=30, **kwargs):
|
||||
trajectory = Trajectory()
|
||||
direction = kwargs.get('direction', 'large-small')
|
||||
for idx, driving_image in enumerate(driving_rgb_lst):
|
||||
if idx == 0 or trajectory.start == -1:
|
||||
src_face = self.face_analysis_wrapper.get(
|
||||
driving_image,
|
||||
flag_do_landmark_2d_106=True,
|
||||
direction=direction
|
||||
)
|
||||
if len(src_face) == 0:
|
||||
# No face detected in the driving_image
|
||||
continue
|
||||
elif len(src_face) > 1:
|
||||
log(f'More than one face detected in the driving frame_{idx}, only pick one face by rule {direction}.')
|
||||
src_face = src_face[0]
|
||||
pts = src_face.landmark_2d_106
|
||||
lmk_203 = self.landmark_runner(driving_image, pts)['pts']
|
||||
trajectory.start, trajectory.end = idx, idx
|
||||
else:
|
||||
lmk_203 = self.face_recon_wrapper(driving_image, trajectory.lmk_lst[-1])['pts']
|
||||
trajectory.end = idx
|
||||
|
||||
trajectory.lmk_lst.append(lmk_203)
|
||||
ret_bbox = parse_bbox_from_landmark(lmk_203, scale=self.crop_cfg.globalscale, vy_ratio=elf.crop_cfg.vy_ratio)['bbox']
|
||||
bbox = [ret_bbox[0, 0], ret_bbox[0, 1], ret_bbox[2, 0], ret_bbox[2, 1]] # 4,
|
||||
trajectory.bbox_lst.append(bbox) # bbox
|
||||
trajectory.frame_rgb_lst.append(driving_image)
|
||||
|
||||
global_bbox = average_bbox_lst(trajectory.bbox_lst)
|
||||
for idx, (frame_rgb, lmk) in enumerate(zip(trajectory.frame_rgb_lst, trajectory.lmk_lst)):
|
||||
ret_dct = crop_image_by_bbox(
|
||||
frame_rgb, global_bbox, lmk=lmk,
|
||||
dsize=self.video_crop_cfg.dsize, flag_rot=self.video_crop_cfg.flag_rot, borderValue=self.video_crop_cfg.borderValue
|
||||
)
|
||||
frame_rgb_crop = ret_dct['img_crop']
|
||||
@@ -0,0 +1,79 @@
|
||||
# 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 ..modules.spade_generator import SPADEDecoder
|
||||
from ..modules.warping_network import WarpingNetwork
|
||||
from ..modules.motion_extractor import MotionExtractor
|
||||
from ..modules.appearance_feature_extractor import AppearanceFeatureExtractor
|
||||
from ..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}")
|
||||
@@ -0,0 +1,89 @@
|
||||
# coding: utf-8
|
||||
|
||||
import os.path as osp
|
||||
import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
|
||||
import torch
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
from .timer import Timer
|
||||
from .rprint import rlog
|
||||
from .crop import crop_image, _transform_pts
|
||||
|
||||
|
||||
def make_abs_path(fn):
|
||||
return osp.join(osp.dirname(osp.realpath(__file__)), fn)
|
||||
|
||||
|
||||
def to_ndarray(obj):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.cpu().numpy()
|
||||
elif isinstance(obj, np.ndarray):
|
||||
return obj
|
||||
else:
|
||||
return np.array(obj)
|
||||
|
||||
|
||||
class LandmarkRunner(object):
|
||||
"""landmark runner"""
|
||||
def __init__(self, **kwargs):
|
||||
ckpt_path = kwargs.get('ckpt_path')
|
||||
onnx_provider = kwargs.get('onnx_provider', 'cuda') # 默认用cuda
|
||||
device_id = kwargs.get('device_id', 0)
|
||||
self.dsize = kwargs.get('dsize', 224)
|
||||
self.timer = Timer()
|
||||
|
||||
if onnx_provider.lower() == 'cuda':
|
||||
self.session = onnxruntime.InferenceSession(
|
||||
ckpt_path, providers=[
|
||||
('CUDAExecutionProvider', {'device_id': device_id})
|
||||
]
|
||||
)
|
||||
else:
|
||||
opts = onnxruntime.SessionOptions()
|
||||
opts.intra_op_num_threads = 4 # 默认线程数为 4
|
||||
self.session = onnxruntime.InferenceSession(
|
||||
ckpt_path, providers=['CPUExecutionProvider'],
|
||||
sess_options=opts
|
||||
)
|
||||
|
||||
def _run(self, inp):
|
||||
out = self.session.run(None, {'input': inp})
|
||||
return out
|
||||
|
||||
def run(self, img_rgb: np.ndarray, lmk=None):
|
||||
if lmk is not None:
|
||||
crop_dct = crop_image(img_rgb, lmk, dsize=self.dsize, scale=1.5, vy_ratio=-0.1)
|
||||
img_crop_rgb = crop_dct['img_crop']
|
||||
else:
|
||||
img_crop_rgb = cv2.resize(img_rgb, (self.dsize, self.dsize))
|
||||
scale = max(img_rgb.shape[:2]) / self.dsize
|
||||
crop_dct = {
|
||||
'M_c2o': np.array([
|
||||
[scale, 0., 0.],
|
||||
[0., scale, 0.],
|
||||
[0., 0., 1.],
|
||||
], dtype=np.float32),
|
||||
}
|
||||
|
||||
inp = (img_crop_rgb.astype(np.float32) / 255.).transpose(2, 0, 1)[None, ...] # HxWx3 (BGR) -> 1x3xHxW (RGB!)
|
||||
|
||||
out_lst = self._run(inp)
|
||||
out_pts = out_lst[2]
|
||||
|
||||
pts = to_ndarray(out_pts[0]).reshape(-1, 2) * self.dsize # scale to 0-224
|
||||
pts = _transform_pts(pts, M=crop_dct['M_c2o'])
|
||||
|
||||
return {
|
||||
'pts': pts, # 2d landmarks 203 points
|
||||
}
|
||||
|
||||
def warmup(self):
|
||||
# 构造dummy image进行warmup
|
||||
self.timer.tic()
|
||||
|
||||
dummy_image = np.zeros((1, 3, self.dsize, self.dsize), dtype=np.float32)
|
||||
|
||||
_ = self._run(dummy_image)
|
||||
|
||||
elapse = self.timer.toc()
|
||||
rlog(f'LandmarkRunner warmup time: {elapse:.3f}s')
|
||||
@@ -0,0 +1 @@
|
||||
mask_template.png filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:4c53e64a07ae3056af2b38548ae5b6cceb04d5c3e6514c7a8d2c3aff9da1ee76
|
||||
size 3470
|
||||
@@ -0,0 +1,76 @@
|
||||
|
||||
"""
|
||||
Functions to compute distance ratios between specific pairs of facial landmarks
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def calculate_distance_ratio(lmk: np.ndarray, idx1: int, idx2: int, idx3: int, idx4: int, eps: float = 1e-6) -> np.ndarray:
|
||||
"""
|
||||
Calculate the ratio of the distance between two pairs of landmarks.
|
||||
|
||||
Parameters:
|
||||
lmk (np.ndarray): Landmarks array of shape (B, N, 2).
|
||||
idx1, idx2, idx3, idx4 (int): Indices of the landmarks.
|
||||
eps (float): Small value to avoid division by zero.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Calculated distance ratio.
|
||||
"""
|
||||
return (np.linalg.norm(lmk[:, idx1] - lmk[:, idx2], axis=1, keepdims=True) /
|
||||
(np.linalg.norm(lmk[:, idx3] - lmk[:, idx4], axis=1, keepdims=True) + eps))
|
||||
|
||||
|
||||
def calc_eye_close_ratio(lmk: np.ndarray, target_eye_ratio: np.ndarray = None) -> np.ndarray:
|
||||
"""
|
||||
Calculate the eye-close ratio for left and right eyes.
|
||||
|
||||
Parameters:
|
||||
lmk (np.ndarray): Landmarks array of shape (B, N, 2).
|
||||
target_eye_ratio (np.ndarray, optional): Additional target eye ratio array to include.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Concatenated eye-close ratios.
|
||||
"""
|
||||
lefteye_close_ratio = calculate_distance_ratio(lmk, 6, 18, 0, 12)
|
||||
righteye_close_ratio = calculate_distance_ratio(lmk, 30, 42, 24, 36)
|
||||
if target_eye_ratio is not None:
|
||||
return np.concatenate([lefteye_close_ratio, righteye_close_ratio, target_eye_ratio], axis=1)
|
||||
else:
|
||||
return np.concatenate([lefteye_close_ratio, righteye_close_ratio], axis=1)
|
||||
|
||||
|
||||
def calc_lip_close_ratio(lmk: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Calculate the lip-close ratio.
|
||||
|
||||
Parameters:
|
||||
lmk (np.ndarray): Landmarks array of shape (B, N, 2).
|
||||
|
||||
Returns:
|
||||
np.ndarray: Calculated lip-close ratio.
|
||||
"""
|
||||
return calculate_distance_ratio(lmk, 90, 102, 48, 66)
|
||||
|
||||
|
||||
def compute_eye_delta(frame_idx, input_eye_ratios, source_landmarks, portrait_wrapper, kp_source):
|
||||
input_eye_ratio = input_eye_ratios[frame_idx][0][0]
|
||||
eye_close_ratio = calc_eye_close_ratio(source_landmarks[None])
|
||||
eye_close_ratio_tensor = torch.from_numpy(eye_close_ratio).float().cuda(portrait_wrapper.device_id)
|
||||
input_eye_ratio_tensor = torch.Tensor([input_eye_ratio]).reshape(1, 1).cuda(portrait_wrapper.device_id)
|
||||
combined_eye_ratio_tensor = torch.cat([eye_close_ratio_tensor, input_eye_ratio_tensor], dim=1)
|
||||
# print(combined_eye_ratio_tensor.mean())
|
||||
eye_delta = portrait_wrapper.retarget_eye(kp_source, combined_eye_ratio_tensor)
|
||||
return eye_delta
|
||||
|
||||
|
||||
def compute_lip_delta(frame_idx, input_lip_ratios, source_landmarks, portrait_wrapper, kp_source):
|
||||
input_lip_ratio = input_lip_ratios[frame_idx][0]
|
||||
lip_close_ratio = calc_lip_close_ratio(source_landmarks[None])
|
||||
lip_close_ratio_tensor = torch.from_numpy(lip_close_ratio).float().cuda(portrait_wrapper.device_id)
|
||||
input_lip_ratio_tensor = torch.Tensor([input_lip_ratio]).cuda(portrait_wrapper.device_id)
|
||||
combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1)
|
||||
lip_delta = portrait_wrapper.retarget_lip(kp_source, combined_lip_ratio_tensor)
|
||||
return lip_delta
|
||||
@@ -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,139 @@
|
||||
# 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)
|
||||
for i in track(range(n), description='writing', transient=True):
|
||||
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 = []
|
||||
for idx, _ in 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
|
||||
@@ -0,0 +1,268 @@
|
||||
import os
|
||||
import torch
|
||||
import yaml
|
||||
import folder_paths
|
||||
import comfy.model_management as mm
|
||||
import comfy.utils
|
||||
|
||||
script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
from .liveportrait.config.argument_config import ArgumentConfig
|
||||
from .liveportrait.live_portrait_pipeline import LivePortraitPipeline
|
||||
|
||||
from .liveportrait.modules.spade_generator import SPADEDecoder
|
||||
from .liveportrait.modules.warping_network import WarpingNetwork
|
||||
from .liveportrait.modules.motion_extractor import MotionExtractor
|
||||
from .liveportrait.modules.appearance_feature_extractor import AppearanceFeatureExtractor
|
||||
from .liveportrait.modules.stitching_retargeting_network import StitchingRetargetingNetwork
|
||||
|
||||
class InferenceConfig:
|
||||
def __init__(self,
|
||||
mask_crop = None,
|
||||
flag_use_half_precision=True,
|
||||
flag_lip_zero=True,
|
||||
lip_zero_threshold=0.03,
|
||||
flag_eye_retargeting=False,
|
||||
flag_lip_retargeting=False,
|
||||
flag_stitching=True,
|
||||
flag_relative=True,
|
||||
anchor_frame=0,
|
||||
input_shape=(256, 256),
|
||||
output_format='mp4',
|
||||
output_fps=30,
|
||||
crf=15,
|
||||
flag_write_result=True,
|
||||
flag_pasteback=True,
|
||||
flag_write_gif=False,
|
||||
size_gif=256,
|
||||
ref_max_shape=1280,
|
||||
ref_shape_n=2,
|
||||
device_id=0,
|
||||
flag_do_crop=True,
|
||||
flag_do_rot=True):
|
||||
self.flag_use_half_precision = flag_use_half_precision
|
||||
self.flag_lip_zero = flag_lip_zero
|
||||
self.lip_zero_threshold = lip_zero_threshold
|
||||
self.flag_eye_retargeting = flag_eye_retargeting
|
||||
self.flag_lip_retargeting = flag_lip_retargeting
|
||||
self.flag_stitching = flag_stitching
|
||||
self.flag_relative = flag_relative
|
||||
self.anchor_frame = anchor_frame
|
||||
self.input_shape = input_shape
|
||||
self.output_format = output_format
|
||||
self.output_fps = output_fps
|
||||
self.crf = crf
|
||||
self.flag_write_result = flag_write_result
|
||||
self.flag_pasteback = flag_pasteback
|
||||
self.flag_write_gif = flag_write_gif
|
||||
self.size_gif = size_gif
|
||||
self.ref_max_shape = ref_max_shape
|
||||
self.ref_shape_n = ref_shape_n
|
||||
self.device_id = device_id
|
||||
self.flag_do_crop = flag_do_crop
|
||||
self.flag_do_rot = flag_do_rot
|
||||
self.mask_crop=mask_crop
|
||||
|
||||
class CropConfig:
|
||||
def __init__(self, dsize=512, scale=2.3, vx_ratio=0, vy_ratio=-0.125):
|
||||
self.dsize = dsize
|
||||
self.scale = scale
|
||||
self.vx_ratio = vx_ratio
|
||||
self.vy_ratio = vy_ratio
|
||||
|
||||
class ArgumentConfig:
|
||||
def __init__(self,
|
||||
device_id=0,
|
||||
flag_lip_zero=True,
|
||||
flag_eye_retargeting=False,
|
||||
flag_lip_retargeting=False,
|
||||
flag_stitching=True,
|
||||
flag_relative=True,
|
||||
flag_pasteback=True,
|
||||
flag_do_crop=True,
|
||||
flag_do_rot=True,
|
||||
dsize=512,
|
||||
scale=2.3,
|
||||
vx_ratio=0,
|
||||
vy_ratio=-0.125,
|
||||
):
|
||||
self.device_id = device_id
|
||||
self.flag_lip_zero = flag_lip_zero
|
||||
self.flag_eye_retargeting = flag_eye_retargeting
|
||||
self.flag_lip_retargeting = flag_lip_retargeting
|
||||
self.flag_stitching = flag_stitching
|
||||
self.flag_relative = flag_relative
|
||||
self.flag_pasteback = flag_pasteback
|
||||
self.flag_do_crop = flag_do_crop
|
||||
self.flag_do_rot = flag_do_rot
|
||||
self.dsize = dsize
|
||||
self.scale = scale
|
||||
self.vx_ratio = vx_ratio
|
||||
self.vy_ratio = vy_ratio
|
||||
|
||||
class DownloadAndLoadLivePortraitModels:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIVEPORTRAITPIPE",)
|
||||
RETURN_NAMES = ("live_portrait_pipe",)
|
||||
FUNCTION = "loadmodel"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def loadmodel(self):
|
||||
device = mm.get_torch_device()
|
||||
mm.soft_empty_cache()
|
||||
|
||||
pbar = comfy.utils.ProgressBar(3)
|
||||
|
||||
download_path = os.path.join(folder_paths.models_dir, "liveportrait")
|
||||
model_path = os.path.join(download_path)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
print(f"Downloading model to: {model_path}")
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id="Kijai/LivePortrait_safetensors",
|
||||
local_dir=download_path,
|
||||
local_dir_use_symlinks=False)
|
||||
|
||||
model_config_path = os.path.join(script_directory, 'liveportrait', 'config', 'models.yaml')
|
||||
with open(model_config_path, 'r') as file:
|
||||
model_config = yaml.safe_load(file)
|
||||
print(model_config)
|
||||
feature_extractor_path = os.path.join(model_path, 'appearance_feature_extractor.safetensors')
|
||||
motion_extractor_path = os.path.join(model_path, 'motion_extractor.safetensors')
|
||||
warping_module_path = os.path.join(model_path, 'warping_module.safetensors')
|
||||
spade_generator_path = os.path.join(model_path, 'spade_generator.safetensors')
|
||||
stitching_retargeting_path = os.path.join(model_path, 'stitching_retargeting_module.safetensors')
|
||||
|
||||
# init F
|
||||
model_params = model_config['model_params']['appearance_feature_extractor_params']
|
||||
self.appearance_feature_extractor = AppearanceFeatureExtractor(**model_params).to(device)
|
||||
self.appearance_feature_extractor.load_state_dict(comfy.utils.load_torch_file(feature_extractor_path))
|
||||
self.appearance_feature_extractor.eval()
|
||||
print('Load appearance_feature_extractor done.')
|
||||
pbar.update(1)
|
||||
# init M
|
||||
model_params = model_config['model_params']['motion_extractor_params']
|
||||
self.motion_extractor = MotionExtractor(**model_params).to(device)
|
||||
self.motion_extractor.load_state_dict(comfy.utils.load_torch_file(motion_extractor_path))
|
||||
self.motion_extractor.eval()
|
||||
print('Load motion_extractor done.')
|
||||
pbar.update(1)
|
||||
# init W
|
||||
model_params = model_config['model_params']['warping_module_params']
|
||||
self.warping_module = WarpingNetwork(**model_params).to(device)
|
||||
self.warping_module.load_state_dict(comfy.utils.load_torch_file(warping_module_path))
|
||||
self.warping_module.eval()
|
||||
print('Load warping_module done.')
|
||||
pbar.update(1)
|
||||
# init G
|
||||
model_params = model_config['model_params']['spade_generator_params']
|
||||
self.spade_generator = SPADEDecoder(**model_params).to(device)
|
||||
self.spade_generator.load_state_dict(comfy.utils.load_torch_file(spade_generator_path))
|
||||
self.spade_generator.eval()
|
||||
print('Load spade_generator done.')
|
||||
pbar.update(1)
|
||||
|
||||
def filter_checkpoint_for_model(checkpoint, prefix):
|
||||
"""Filter and adjust the checkpoint dictionary for a specific model based on the prefix."""
|
||||
# Create a new dictionary where keys are adjusted by removing the prefix and the model name
|
||||
filtered_checkpoint = {key.replace(prefix + "_module.", ""): value for key, value in checkpoint.items() if key.startswith(prefix)}
|
||||
return filtered_checkpoint
|
||||
|
||||
config = model_config['model_params']['stitching_retargeting_module_params']
|
||||
checkpoint = comfy.utils.load_torch_file(stitching_retargeting_path)
|
||||
|
||||
# Example usage for the stitcher model
|
||||
stitcher_prefix = 'retarget_shoulder'
|
||||
stitcher_checkpoint = filter_checkpoint_for_model(checkpoint, stitcher_prefix)
|
||||
stitcher = StitchingRetargetingNetwork(**config.get('stitching'))
|
||||
stitcher.load_state_dict(stitcher_checkpoint)
|
||||
stitcher = stitcher.to(device)
|
||||
stitcher.eval()
|
||||
|
||||
# Repeat for other models with their respective prefixes
|
||||
lip_prefix = 'retarget_mouth'
|
||||
lip_checkpoint = filter_checkpoint_for_model(checkpoint, lip_prefix)
|
||||
retargetor_lip = StitchingRetargetingNetwork(**config.get('lip'))
|
||||
retargetor_lip.load_state_dict(lip_checkpoint)
|
||||
retargetor_lip = retargetor_lip.to(device)
|
||||
retargetor_lip.eval()
|
||||
|
||||
eye_prefix = 'retarget_eye'
|
||||
eye_checkpoint = filter_checkpoint_for_model(checkpoint, eye_prefix)
|
||||
retargetor_eye = StitchingRetargetingNetwork(**config.get('eye'))
|
||||
retargetor_eye.load_state_dict(eye_checkpoint)
|
||||
retargetor_eye = retargetor_eye.to(device)
|
||||
retargetor_eye.eval()
|
||||
print('Load stitching_retargeting_module done.')
|
||||
|
||||
self.stich_retargeting_module = {
|
||||
'stitching': stitcher,
|
||||
'lip': retargetor_lip,
|
||||
'eye': retargetor_eye
|
||||
}
|
||||
|
||||
pipeline = LivePortraitPipeline(
|
||||
self.appearance_feature_extractor,
|
||||
self.motion_extractor,
|
||||
self.warping_module,
|
||||
self.spade_generator,
|
||||
self.stich_retargeting_module,
|
||||
InferenceConfig(),
|
||||
CropConfig()
|
||||
)
|
||||
|
||||
return (pipeline,)
|
||||
|
||||
class LivePortraitProcess:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
|
||||
"pipeline": ("LIVEPORTRAITPIPE",),
|
||||
"source_image": ("IMAGE",),
|
||||
"driving_images": ("IMAGE",),
|
||||
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
||||
RETURN_NAMES = ("cropped_images", "full_images",)
|
||||
FUNCTION = "process"
|
||||
CATEGORY = "LivePortrait"
|
||||
|
||||
def process(self, source_image, driving_images, pipeline):
|
||||
device = mm.get_torch_device()
|
||||
|
||||
source_image_np = (source_image.squeeze(0) * 255).byte().numpy()
|
||||
driving_images_np = (driving_images * 255).byte().numpy()
|
||||
|
||||
args = ArgumentConfig()
|
||||
|
||||
cropped_frames, full_frame = pipeline.execute(source_image_np, driving_images_np, args)
|
||||
|
||||
cropped_tensors = [torch.from_numpy(np_array) for np_array in cropped_frames]
|
||||
cropped_tensors_out = torch.stack(cropped_tensors) / 255
|
||||
cropped_tensors_out = cropped_tensors_out.cpu().float()
|
||||
|
||||
full_tensors = [torch.from_numpy(np_array) for np_array in full_frame]
|
||||
full_tensors_out = torch.stack(full_tensors) / 255
|
||||
full_tensors_out = full_tensors_out.cpu().float()
|
||||
|
||||
print(cropped_tensors_out.shape)
|
||||
print(cropped_tensors_out.min(), cropped_tensors_out.max())
|
||||
|
||||
return (cropped_tensors_out, full_tensors_out)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DownloadAndLoadLivePortraitModels": DownloadAndLoadLivePortraitModels,
|
||||
"LivePortraitProcess": LivePortraitProcess,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DownloadAndLoadLivePortraitModels": "(Down)Load LivePortraitModels",
|
||||
"LivePortraitProcess": "LivePortraitProcess",
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
# ComfyUI nodes to use LivePortrait
|
||||
|
||||
I have converted all the pickle files to safetensors, and they
|
||||
|
||||
are automatically downloaded from here to `ComfyUI/models/liveportrait`:
|
||||
|
||||
https://huggingface.co/Kijai/LivePortrait_safetensors/tree/main
|
||||
|
||||
Insightface is required, and it's models are loaded from the usual `ComfyUI/models/insightface`
|
||||
@@ -0,0 +1,3 @@
|
||||
yaml
|
||||
numpy
|
||||
opencv-python
|
||||
Reference in New Issue
Block a user