From 56245461271f8bccf2270668d97566ac49b5bb2c Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Thu, 4 Jul 2024 18:23:26 +0300 Subject: [PATCH] initial commit --- .gitignore | 17 + LICENSE | 21 + __init__.py | 3 + assets/.gitattributes | 23 + assets/examples/driving/d0.mp4 | 3 + assets/examples/driving/d1.mp4 | 3 + assets/examples/driving/d2.mp4 | 3 + assets/examples/driving/d3.mp4 | 3 + assets/examples/driving/d5.mp4 | 3 + assets/examples/driving/d6.mp4 | 3 + assets/examples/driving/d7.mp4 | 3 + assets/examples/driving/d8.mp4 | 3 + assets/examples/driving/d9.mp4 | 3 + assets/examples/source/s0.jpg | 3 + assets/examples/source/s1.jpg | 3 + assets/examples/source/s10.jpg | 3 + assets/examples/source/s2.jpg | 3 + assets/examples/source/s3.jpg | 3 + assets/examples/source/s4.jpg | 3 + assets/examples/source/s5.jpg | 3 + assets/examples/source/s6.jpg | 3 + assets/examples/source/s7.jpg | 3 + assets/examples/source/s8.jpg | 3 + assets/examples/source/s9.jpg | 3 + liveportrait/config/__init__.py | 0 liveportrait/config/argument_config.py | 44 ++ liveportrait/config/base_config.py | 29 ++ liveportrait/config/crop_config.py | 18 + liveportrait/config/inference_config.py | 49 ++ liveportrait/config/models.yaml | 43 ++ liveportrait/live_portrait_pipeline.py | 207 ++++++++ liveportrait/live_portrait_wrapper.py | 341 ++++++++++++++ liveportrait/modules/__init__.py | 0 .../modules/appearance_feature_extractor.py | 48 ++ liveportrait/modules/convnextv2.py | 149 ++++++ liveportrait/modules/dense_motion.py | 104 +++++ liveportrait/modules/motion_extractor.py | 35 ++ liveportrait/modules/spade_generator.py | 59 +++ .../modules/stitching_retargeting_network.py | 38 ++ liveportrait/modules/util.py | 441 ++++++++++++++++++ liveportrait/modules/warping_network.py | 77 +++ liveportrait/template_maker.py | 65 +++ liveportrait/utils/__init__.py | 0 liveportrait/utils/camera.py | 75 +++ liveportrait/utils/crop.py | 393 ++++++++++++++++ liveportrait/utils/cropper.py | 147 ++++++ liveportrait/utils/face_analysis_diy.py | 79 ++++ liveportrait/utils/helper.py | 175 +++++++ liveportrait/utils/io.py | 97 ++++ liveportrait/utils/landmark_runner.py | 89 ++++ liveportrait/utils/resources/.gitattributes | 1 + .../utils/resources/mask_template.png | 3 + liveportrait/utils/retargeting_utils.py | 76 +++ liveportrait/utils/rprint.py | 16 + liveportrait/utils/timer.py | 29 ++ liveportrait/utils/video.py | 139 ++++++ nodes.py | 268 +++++++++++ readme.md | 9 + requirements.txt | 3 + 59 files changed, 3470 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 __init__.py create mode 100644 assets/.gitattributes create mode 100644 assets/examples/driving/d0.mp4 create mode 100644 assets/examples/driving/d1.mp4 create mode 100644 assets/examples/driving/d2.mp4 create mode 100644 assets/examples/driving/d3.mp4 create mode 100644 assets/examples/driving/d5.mp4 create mode 100644 assets/examples/driving/d6.mp4 create mode 100644 assets/examples/driving/d7.mp4 create mode 100644 assets/examples/driving/d8.mp4 create mode 100644 assets/examples/driving/d9.mp4 create mode 100644 assets/examples/source/s0.jpg create mode 100644 assets/examples/source/s1.jpg create mode 100644 assets/examples/source/s10.jpg create mode 100644 assets/examples/source/s2.jpg create mode 100644 assets/examples/source/s3.jpg create mode 100644 assets/examples/source/s4.jpg create mode 100644 assets/examples/source/s5.jpg create mode 100644 assets/examples/source/s6.jpg create mode 100644 assets/examples/source/s7.jpg create mode 100644 assets/examples/source/s8.jpg create mode 100644 assets/examples/source/s9.jpg create mode 100644 liveportrait/config/__init__.py create mode 100644 liveportrait/config/argument_config.py create mode 100644 liveportrait/config/base_config.py create mode 100644 liveportrait/config/crop_config.py create mode 100644 liveportrait/config/inference_config.py create mode 100644 liveportrait/config/models.yaml create mode 100644 liveportrait/live_portrait_pipeline.py create mode 100644 liveportrait/live_portrait_wrapper.py create mode 100644 liveportrait/modules/__init__.py create mode 100644 liveportrait/modules/appearance_feature_extractor.py create mode 100644 liveportrait/modules/convnextv2.py create mode 100644 liveportrait/modules/dense_motion.py create mode 100644 liveportrait/modules/motion_extractor.py create mode 100644 liveportrait/modules/spade_generator.py create mode 100644 liveportrait/modules/stitching_retargeting_network.py create mode 100644 liveportrait/modules/util.py create mode 100644 liveportrait/modules/warping_network.py create mode 100644 liveportrait/template_maker.py create mode 100644 liveportrait/utils/__init__.py create mode 100644 liveportrait/utils/camera.py create mode 100644 liveportrait/utils/crop.py create mode 100644 liveportrait/utils/cropper.py create mode 100644 liveportrait/utils/face_analysis_diy.py create mode 100644 liveportrait/utils/helper.py create mode 100644 liveportrait/utils/io.py create mode 100644 liveportrait/utils/landmark_runner.py create mode 100644 liveportrait/utils/resources/.gitattributes create mode 100644 liveportrait/utils/resources/mask_template.png create mode 100644 liveportrait/utils/retargeting_utils.py create mode 100644 liveportrait/utils/rprint.py create mode 100644 liveportrait/utils/timer.py create mode 100644 liveportrait/utils/video.py create mode 100644 nodes.py create mode 100644 readme.md create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..07050fd --- /dev/null +++ b/.gitignore @@ -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/* diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..9e8f502 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2e96bd6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/assets/.gitattributes b/assets/.gitattributes new file mode 100644 index 0000000..3b11779 --- /dev/null +++ b/assets/.gitattributes @@ -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 diff --git a/assets/examples/driving/d0.mp4 b/assets/examples/driving/d0.mp4 new file mode 100644 index 0000000..92391dd --- /dev/null +++ b/assets/examples/driving/d0.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:63f6f9962e1fdf6e6722172e7a18155204858d5d5ce3b1e0646c150360c33bed +size 2958395 diff --git a/assets/examples/driving/d1.mp4 b/assets/examples/driving/d1.mp4 new file mode 100644 index 0000000..8f58096 --- /dev/null +++ b/assets/examples/driving/d1.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d8bdb2e3d18db28a7c81bc06070aefe2beebae2b489b976fa009eff5881bc7fe +size 48753 diff --git a/assets/examples/driving/d2.mp4 b/assets/examples/driving/d2.mp4 new file mode 100644 index 0000000..5372520 --- /dev/null +++ b/assets/examples/driving/d2.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ce4a92e80e4d0b44c43af7afba471b3ebdf7dc83894ae98af6d4973f17af484 +size 47762 diff --git a/assets/examples/driving/d3.mp4 b/assets/examples/driving/d3.mp4 new file mode 100644 index 0000000..8b70b6a --- /dev/null +++ b/assets/examples/driving/d3.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ef5c86e49b1b43dcb1449b499eb5a7f0cbae2f78aec08b5598193be1e4257099 +size 1430968 diff --git a/assets/examples/driving/d5.mp4 b/assets/examples/driving/d5.mp4 new file mode 100644 index 0000000..0d34edb --- /dev/null +++ b/assets/examples/driving/d5.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:021af59d9e1d89c967699efd374e7acc77fd14c30b052abd98665de401d9e511 +size 135015 diff --git a/assets/examples/driving/d6.mp4 b/assets/examples/driving/d6.mp4 new file mode 100644 index 0000000..44f3513 --- /dev/null +++ b/assets/examples/driving/d6.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:00e3ea79bbf28cbdc4fbb67ec655d9a0fe876e880ec45af55ae481348d0c0fff +size 1967790 diff --git a/assets/examples/driving/d7.mp4 b/assets/examples/driving/d7.mp4 new file mode 100644 index 0000000..d50ec4e --- /dev/null +++ b/assets/examples/driving/d7.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9e3ecc59a7285d52ca3666be6562e59ed45ad8afcef28d0f69a9be0e2850ed37 +size 185258 diff --git a/assets/examples/driving/d8.mp4 b/assets/examples/driving/d8.mp4 new file mode 100644 index 0000000..de9c9d1 --- /dev/null +++ b/assets/examples/driving/d8.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:912a02dcffef9b1a6d6641c9a48e170490080c3715dc9f55d91e168483d27e6e +size 312295 diff --git a/assets/examples/driving/d9.mp4 b/assets/examples/driving/d9.mp4 new file mode 100644 index 0000000..7803b3b --- /dev/null +++ b/assets/examples/driving/d9.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9a414aa1d547be35306d692065a2157434bf40a6025ba8e30ce12e5bb322cc33 +size 2257929 diff --git a/assets/examples/source/s0.jpg b/assets/examples/source/s0.jpg new file mode 100644 index 0000000..7cef081 --- /dev/null +++ b/assets/examples/source/s0.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ccd094951787cec391f3d444e056ee9a58d715f84152f49100c03bb3ce2962fc +size 116149 diff --git a/assets/examples/source/s1.jpg b/assets/examples/source/s1.jpg new file mode 100644 index 0000000..40f5043 --- /dev/null +++ b/assets/examples/source/s1.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:98931abeba483c27c699197e7c7bab58b7deb4324745b0f84e3c92e9ccef5901 +size 98411 diff --git a/assets/examples/source/s10.jpg b/assets/examples/source/s10.jpg new file mode 100644 index 0000000..1806f0d --- /dev/null +++ b/assets/examples/source/s10.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:49bdee735bed389d289c20d0bdff9bf0a485068e128cec8eb9e5bb11ae1f0422 +size 537372 diff --git a/assets/examples/source/s2.jpg b/assets/examples/source/s2.jpg new file mode 100644 index 0000000..d172b2c --- /dev/null +++ b/assets/examples/source/s2.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4d950dc1c6560e0c7d0bcdca258079ba605d26f0979b7a602af75511a15e4c03 +size 760146 diff --git a/assets/examples/source/s3.jpg b/assets/examples/source/s3.jpg new file mode 100644 index 0000000..3b508a9 --- /dev/null +++ b/assets/examples/source/s3.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d93728ec60e1c0fceaf0f62b05d0b70283e9383a6eef0f31dc2078cb560f707a +size 63626 diff --git a/assets/examples/source/s4.jpg b/assets/examples/source/s4.jpg new file mode 100644 index 0000000..058e1f5 --- /dev/null +++ b/assets/examples/source/s4.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d62af6c554e211942f954237f7b898abb41d9d6166b5079564e76c25ff804c55 +size 143918 diff --git a/assets/examples/source/s5.jpg b/assets/examples/source/s5.jpg new file mode 100644 index 0000000..bd8dc95 --- /dev/null +++ b/assets/examples/source/s5.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c368d9aa2eea0c5319b7d76e3b6c339819e49565ef064933a4511c11b581549 +size 133517 diff --git a/assets/examples/source/s6.jpg b/assets/examples/source/s6.jpg new file mode 100644 index 0000000..494b6bd --- /dev/null +++ b/assets/examples/source/s6.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa226be24fc24c11c4bb276c2d7d789218eba88b8e49155855f5685d1b4d7809 +size 107520 diff --git a/assets/examples/source/s7.jpg b/assets/examples/source/s7.jpg new file mode 100644 index 0000000..d87d7b5 --- /dev/null +++ b/assets/examples/source/s7.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7a94b38bb66b58c3d44b98b86be3138cd3fb3ca40f89017294476b879353ebc4 +size 139855 diff --git a/assets/examples/source/s8.jpg b/assets/examples/source/s8.jpg new file mode 100644 index 0000000..398cd85 --- /dev/null +++ b/assets/examples/source/s8.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:01b6af295ebefcac1291cbe3b84fc6af37e40b168c528113662519aace4bbce5 +size 227378 diff --git a/assets/examples/source/s9.jpg b/assets/examples/source/s9.jpg new file mode 100644 index 0000000..83191d9 --- /dev/null +++ b/assets/examples/source/s9.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:84a6bfcb3346951cef85c0b0bb07856794e0b04d5b65322c18d77198fbf43f2a +size 442850 diff --git a/liveportrait/config/__init__.py b/liveportrait/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/liveportrait/config/argument_config.py b/liveportrait/config/argument_config.py new file mode 100644 index 0000000..8a91c3e --- /dev/null +++ b/liveportrait/config/argument_config.py @@ -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" diff --git a/liveportrait/config/base_config.py b/liveportrait/config/base_config.py new file mode 100644 index 0000000..216b8be --- /dev/null +++ b/liveportrait/config/base_config.py @@ -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) diff --git a/liveportrait/config/crop_config.py b/liveportrait/config/crop_config.py new file mode 100644 index 0000000..d3c79be --- /dev/null +++ b/liveportrait/config/crop_config.py @@ -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 diff --git a/liveportrait/config/inference_config.py b/liveportrait/config/inference_config.py new file mode 100644 index 0000000..0da3e3c --- /dev/null +++ b/liveportrait/config/inference_config.py @@ -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 diff --git a/liveportrait/config/models.yaml b/liveportrait/config/models.yaml new file mode 100644 index 0000000..131d1c6 --- /dev/null +++ b/liveportrait/config/models.yaml @@ -0,0 +1,43 @@ +model_params: + appearance_feature_extractor_params: # the F in the paper + image_channel: 3 + block_expansion: 64 + num_down_blocks: 2 + max_features: 512 + reshape_channel: 32 + reshape_depth: 16 + num_resblocks: 6 + motion_extractor_params: # the M in the paper + num_kp: 21 + backbone: convnextv2_tiny + warping_module_params: # the W in the paper + num_kp: 21 + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + reshape_channel: 32 + estimate_occlusion_map: True + dense_motion_params: + block_expansion: 32 + max_features: 1024 + num_blocks: 5 + reshape_depth: 16 + compress: 4 + spade_generator_params: # the G in the paper + upscale: 2 # represents upsample factor 256x256 -> 512x512 + block_expansion: 64 + max_features: 512 + num_down_blocks: 2 + stitching_retargeting_module_params: # the S in the paper + stitching: + input_size: 126 # (21*3)*2 + hidden_sizes: [128, 128, 64] + output_size: 65 # (21*3)+2(tx,ty) + lip: + input_size: 65 # (21*3)+2 + hidden_sizes: [128, 128, 64] + output_size: 63 # (21*3) + eye: + input_size: 66 # (21*3)+3 + hidden_sizes: [256, 256, 128, 128, 64] + output_size: 63 # (21*3) diff --git a/liveportrait/live_portrait_pipeline.py b/liveportrait/live_portrait_pipeline.py new file mode 100644 index 0000000..197e062 --- /dev/null +++ b/liveportrait/live_portrait_pipeline.py @@ -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 diff --git a/liveportrait/live_portrait_wrapper.py b/liveportrait/live_portrait_wrapper.py new file mode 100644 index 0000000..eb1e4f5 --- /dev/null +++ b/liveportrait/live_portrait_wrapper.py @@ -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 diff --git a/liveportrait/modules/__init__.py b/liveportrait/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/liveportrait/modules/appearance_feature_extractor.py b/liveportrait/modules/appearance_feature_extractor.py new file mode 100644 index 0000000..8d89e4f --- /dev/null +++ b/liveportrait/modules/appearance_feature_extractor.py @@ -0,0 +1,48 @@ +# coding: utf-8 + +""" +Appearance extractor(F) defined in paper, which maps the source image s to a 3D appearance feature volume. +""" + +import torch +from torch import nn +from .util import SameBlock2d, DownBlock2d, ResBlock3d + + +class AppearanceFeatureExtractor(nn.Module): + + def __init__(self, image_channel, block_expansion, num_down_blocks, max_features, reshape_channel, reshape_depth, num_resblocks): + super(AppearanceFeatureExtractor, self).__init__() + self.image_channel = image_channel + self.block_expansion = block_expansion + self.num_down_blocks = num_down_blocks + self.max_features = max_features + self.reshape_channel = reshape_channel + self.reshape_depth = reshape_depth + + self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(3, 3), padding=(1, 1)) + + down_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.down_blocks = nn.ModuleList(down_blocks) + + self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1) + + self.resblocks_3d = torch.nn.Sequential() + for i in range(num_resblocks): + self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1)) + + def forward(self, source_image): + out = self.first(source_image) # Bx3x256x256 -> Bx64x256x256 + + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + out = self.second(out) + bs, c, h, w = out.shape # ->Bx512x64x64 + + f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w) # ->Bx32x16x64x64 + f_s = self.resblocks_3d(f_s) # ->Bx32x16x64x64 + return f_s diff --git a/liveportrait/modules/convnextv2.py b/liveportrait/modules/convnextv2.py new file mode 100644 index 0000000..83ea126 --- /dev/null +++ b/liveportrait/modules/convnextv2.py @@ -0,0 +1,149 @@ +# coding: utf-8 + +""" +This moudle is adapted to the ConvNeXtV2 version for the extraction of implicit keypoints, poses, and expression deformation. +""" + +import torch +import torch.nn as nn +# from timm.models.layers import trunc_normal_, DropPath +from .util import LayerNorm, DropPath, trunc_normal_, GRN + +__all__ = ['convnextv2_tiny'] + + +class Block(nn.Module): + """ ConvNeXtV2 Block. + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + """ + + def __init__(self, dim, drop_path=0.): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.norm = LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + + +class ConvNeXtV2(nn.Module): + """ ConvNeXt V2 + + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + + def __init__( + self, + in_chans=3, + depths=[3, 3, 9, 3], + dims=[96, 192, 384, 768], + drop_path_rate=0., + **kwargs + ): + super().__init__() + self.depths = depths + self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format="channels_first") + ) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(4): + stage = nn.Sequential( + *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + + # NOTE: the output semantic items + num_bins = kwargs.get('num_bins', 66) + num_kp = kwargs.get('num_kp', 24) # the number of implicit keypoints + self.fc_kp = nn.Linear(dims[-1], 3 * num_kp) # implicit keypoints + + # print('dims[-1]: ', dims[-1]) + self.fc_scale = nn.Linear(dims[-1], 1) # scale + self.fc_pitch = nn.Linear(dims[-1], num_bins) # pitch bins + self.fc_yaw = nn.Linear(dims[-1], num_bins) # yaw bins + self.fc_roll = nn.Linear(dims[-1], num_bins) # roll bins + self.fc_t = nn.Linear(dims[-1], 3) # translation + self.fc_exp = nn.Linear(dims[-1], 3 * num_kp) # expression / delta + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x) + x = self.stages[i](x) + return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + + # implicit keypoints + kp = self.fc_kp(x) + + # pose and expression deformation + pitch = self.fc_pitch(x) + yaw = self.fc_yaw(x) + roll = self.fc_roll(x) + t = self.fc_t(x) + exp = self.fc_exp(x) + scale = self.fc_scale(x) + + ret_dct = { + 'pitch': pitch, + 'yaw': yaw, + 'roll': roll, + 't': t, + 'exp': exp, + 'scale': scale, + + 'kp': kp, # canonical keypoint + } + + return ret_dct + + +def convnextv2_tiny(**kwargs): + model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs) + return model diff --git a/liveportrait/modules/dense_motion.py b/liveportrait/modules/dense_motion.py new file mode 100644 index 0000000..0eec0c4 --- /dev/null +++ b/liveportrait/modules/dense_motion.py @@ -0,0 +1,104 @@ +# coding: utf-8 + +""" +The module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving +""" + +from torch import nn +import torch.nn.functional as F +import torch +from .util import Hourglass, make_coordinate_grid, kp2gaussian + + +class DenseMotionNetwork(nn.Module): + def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, estimate_occlusion_map=True): + super(DenseMotionNetwork, self).__init__() + self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks) # ~60+G + + self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3) # 65G! NOTE: computation cost is large + self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1) # 0.8G + self.norm = nn.BatchNorm3d(compress, affine=True) + self.num_kp = num_kp + self.flag_estimate_occlusion_map = estimate_occlusion_map + + if self.flag_estimate_occlusion_map: + self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3) + else: + self.occlusion = None + + def create_sparse_motions(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape # (bs, 4, 16, 64, 64) + identity_grid = make_coordinate_grid((d, h, w), ref=kp_source) # (16, 64, 64, 3) + identity_grid = identity_grid.view(1, 1, d, h, w, 3) # (1, 1, d=16, h=64, w=64, 3) + coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3) + + k = coordinate_grid.shape[1] + + # NOTE: there lacks an one-order flow + driving_to_source = coordinate_grid + kp_source.view(bs, self.num_kp, 1, 1, 1, 3) # (bs, num_kp, d, h, w, 3) + + # adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1) + sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) # (bs, 1+num_kp, d, h, w, 3) + return sparse_motions + + def create_deformed_feature(self, feature, sparse_motions): + bs, _, d, h, w = feature.shape + feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1) # (bs, num_kp+1, 1, c, d, h, w) + feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w) # (bs*(num_kp+1), c, d, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1)) # (bs*(num_kp+1), d, h, w, 3) + sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False) + sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w) + + return sparse_deformed + + def create_heatmap_representations(self, feature, kp_driving, kp_source): + spatial_size = feature.shape[3:] # (d=16, h=64, w=64) + gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w) + gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01) # (bs, num_kp, d, h, w) + heatmap = gaussian_driving - gaussian_source # (bs, num_kp, d, h, w) + + # adding background feature + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type()).to(heatmap.device) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) # (bs, 1+num_kp, 1, d, h, w) + return heatmap + + def forward(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape # (bs, 32, 16, 64, 64) + + feature = self.compress(feature) # (bs, 4, 16, 64, 64) + feature = self.norm(feature) # (bs, 4, 16, 64, 64) + feature = F.relu(feature) # (bs, 4, 16, 64, 64) + + out_dict = dict() + + # 1. deform 3d feature + sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source) # (bs, 1+num_kp, d, h, w, 3) + deformed_feature = self.create_deformed_feature(feature, sparse_motion) # (bs, 1+num_kp, c=4, d=16, h=64, w=64) + + # 2. (bs, 1+num_kp, d, h, w) + heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source) # (bs, 1+num_kp, 1, d, h, w) + + input = torch.cat([heatmap, deformed_feature], dim=2) # (bs, 1+num_kp, c=5, d=16, h=64, w=64) + input = input.view(bs, -1, d, h, w) # (bs, (1+num_kp)*c=105, d=16, h=64, w=64) + + prediction = self.hourglass(input) + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) # (bs, 1+num_kp, d=16, h=64, w=64) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w) + sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4) # (bs, num_kp+1, 3, d, h, w) + deformation = (sparse_motion * mask).sum(dim=1) # (bs, 3, d, h, w) mask take effect in this place + deformation = deformation.permute(0, 2, 3, 4, 1) # (bs, d, h, w, 3) + + out_dict['deformation'] = deformation + + if self.flag_estimate_occlusion_map: + bs, _, d, h, w = prediction.shape + prediction_reshape = prediction.view(bs, -1, h, w) + occlusion_map = torch.sigmoid(self.occlusion(prediction_reshape)) # Bx1x64x64 + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/liveportrait/modules/motion_extractor.py b/liveportrait/modules/motion_extractor.py new file mode 100644 index 0000000..b2982e5 --- /dev/null +++ b/liveportrait/modules/motion_extractor.py @@ -0,0 +1,35 @@ +# coding: utf-8 + +""" +Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image +""" + +from torch import nn +import torch + +from .convnextv2 import convnextv2_tiny +from .util import filter_state_dict + +model_dict = { + 'convnextv2_tiny': convnextv2_tiny, +} + + +class MotionExtractor(nn.Module): + def __init__(self, **kwargs): + super(MotionExtractor, self).__init__() + + # default is convnextv2_base + backbone = kwargs.get('backbone', 'convnextv2_tiny') + self.detector = model_dict.get(backbone)(**kwargs) + + def load_pretrained(self, init_path: str): + if init_path not in (None, ''): + state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model'] + state_dict = filter_state_dict(state_dict, remove_name='head') + ret = self.detector.load_state_dict(state_dict, strict=False) + print(f'Load pretrained model from {init_path}, ret: {ret}') + + def forward(self, x): + out = self.detector(x) + return out diff --git a/liveportrait/modules/spade_generator.py b/liveportrait/modules/spade_generator.py new file mode 100644 index 0000000..147a9ae --- /dev/null +++ b/liveportrait/modules/spade_generator.py @@ -0,0 +1,59 @@ +# coding: utf-8 + +""" +Spade decoder(G) defined in the paper, which input the warped feature to generate the animated image. +""" + +import torch +from torch import nn +import torch.nn.functional as F +from .util import SPADEResnetBlock + + +class SPADEDecoder(nn.Module): + def __init__(self, upscale=1, max_features=256, block_expansion=64, out_channels=64, num_down_blocks=2): + for i in range(num_down_blocks): + input_channels = min(max_features, block_expansion * (2 ** (i + 1))) + self.upscale = upscale + super().__init__() + norm_G = 'spadespectralinstance' + label_num_channels = input_channels # 256 + + self.fc = nn.Conv2d(input_channels, 2 * input_channels, 3, padding=1) + self.G_middle_0 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_1 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_2 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_3 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_4 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.G_middle_5 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels) + self.up_0 = SPADEResnetBlock(2 * input_channels, input_channels, norm_G, label_num_channels) + self.up_1 = SPADEResnetBlock(input_channels, out_channels, norm_G, label_num_channels) + self.up = nn.Upsample(scale_factor=2) + + if self.upscale is None or self.upscale <= 1: + self.conv_img = nn.Conv2d(out_channels, 3, 3, padding=1) + else: + self.conv_img = nn.Sequential( + nn.Conv2d(out_channels, 3 * (2 * 2), kernel_size=3, padding=1), + nn.PixelShuffle(upscale_factor=2) + ) + + def forward(self, feature): + seg = feature # Bx256x64x64 + x = self.fc(feature) # Bx512x64x64 + x = self.G_middle_0(x, seg) + x = self.G_middle_1(x, seg) + x = self.G_middle_2(x, seg) + x = self.G_middle_3(x, seg) + x = self.G_middle_4(x, seg) + x = self.G_middle_5(x, seg) + + x = self.up(x) # Bx512x64x64 -> Bx512x128x128 + x = self.up_0(x, seg) # Bx512x128x128 -> Bx256x128x128 + x = self.up(x) # Bx256x128x128 -> Bx256x256x256 + x = self.up_1(x, seg) # Bx256x256x256 -> Bx64x256x256 + + x = self.conv_img(F.leaky_relu(x, 2e-1)) # Bx64x256x256 -> Bx3xHxW + x = torch.sigmoid(x) # Bx3xHxW + + return x \ No newline at end of file diff --git a/liveportrait/modules/stitching_retargeting_network.py b/liveportrait/modules/stitching_retargeting_network.py new file mode 100644 index 0000000..5f50b7c --- /dev/null +++ b/liveportrait/modules/stitching_retargeting_network.py @@ -0,0 +1,38 @@ +# coding: utf-8 + +""" +Stitching module(S) and two retargeting modules(R) defined in the paper. + +- The stitching module pastes the animated portrait back into the original image space without pixel misalignment, such as in +the stitching region. + +- The eyes retargeting module is designed to address the issue of incomplete eye closure during cross-id reenactment, especially +when a person with small eyes drives a person with larger eyes. + +- The lip retargeting module is designed similarly to the eye retargeting module, and can also normalize the input by ensuring that +the lips are in a closed state, which facilitates better animation driving. +""" +from torch import nn + + +class StitchingRetargetingNetwork(nn.Module): + def __init__(self, input_size, hidden_sizes, output_size): + super(StitchingRetargetingNetwork, self).__init__() + layers = [] + for i in range(len(hidden_sizes)): + if i == 0: + layers.append(nn.Linear(input_size, hidden_sizes[i])) + else: + layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i])) + layers.append(nn.ReLU(inplace=True)) + layers.append(nn.Linear(hidden_sizes[-1], output_size)) + self.mlp = nn.Sequential(*layers) + + def initialize_weights_to_zero(self): + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.zeros_(m.weight) + nn.init.zeros_(m.bias) + + def forward(self, x): + return self.mlp(x) diff --git a/liveportrait/modules/util.py b/liveportrait/modules/util.py new file mode 100644 index 0000000..f83980b --- /dev/null +++ b/liveportrait/modules/util.py @@ -0,0 +1,441 @@ +# coding: utf-8 + +""" +This file defines various neural network modules and utility functions, including convolutional and residual blocks, +normalizations, and functions for spatial transformation and tensor manipulation. +""" + +from torch import nn +import torch.nn.functional as F +import torch +import torch.nn.utils.spectral_norm as spectral_norm +import math +import warnings + + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp + + coordinate_grid = make_coordinate_grid(spatial_size, mean) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance) + + return out + + +def make_coordinate_grid(spatial_size, ref, **kwargs): + d, h, w = spatial_size + x = torch.arange(w).type(ref.dtype).to(ref.device) + y = torch.arange(h).type(ref.dtype).to(ref.device) + z = torch.arange(d).type(ref.dtype).to(ref.device) + + # NOTE: must be right-down-in + x = (2 * (x / (w - 1)) - 1) # the x axis faces to the right + y = (2 * (y / (h - 1)) - 1) # the y axis faces to the bottom + z = (2 * (z / (d - 1)) - 1) # the z axis faces to the inner + + yy = y.view(1, -1, 1).repeat(d, 1, w) + xx = x.view(1, 1, -1).repeat(d, h, 1) + zz = z.view(-1, 1, 1).repeat(1, h, w) + + meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3) + + return meshed + + +class ConvT2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1): + super(ConvT2d, self).__init__() + + self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride, + padding=padding, output_padding=output_padding) + self.norm = nn.InstanceNorm2d(out_features) + + def forward(self, x): + out = self.convT(x) + out = self.norm(out) + out = F.leaky_relu(out) + return out + + +class ResBlock3d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock3d, self).__init__() + self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.norm1 = nn.BatchNorm3d(in_features, affine=True) + self.norm2 = nn.BatchNorm3d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock3d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock3d, self).__init__() + + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=(1, 2, 2)) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class DownBlock3d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock3d, self).__init__() + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups, stride=(1, 2, 2)) + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + if lrelu: + self.ac = nn.LeakyReLU() + else: + self.ac = nn.ReLU() + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = self.ac(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1) + self.norm = nn.BatchNorm3d(self.out_filters, affine=True) + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class SPADE(nn.Module): + def __init__(self, norm_nc, label_nc): + super().__init__() + + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + nhidden = 128 + + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1), + nn.ReLU()) + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + + def forward(self, x, segmap): + normalized = self.param_free_norm(x) + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + actv = self.mlp_shared(segmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + out = normalized * (1 + gamma) + beta + return out + + +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1): + super().__init__() + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + self.use_se = use_se + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + # apply spectral norm if specified + if 'spectral' in norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + # define normalization layers + self.norm_0 = SPADE(fin, label_nc) + self.norm_1 = SPADE(fmiddle, label_nc) + if self.learned_shortcut: + self.norm_s = SPADE(fin, label_nc) + + def forward(self, x, seg1): + x_s = self.shortcut(x, seg1) + dx = self.conv_0(self.actvn(self.norm_0(x, seg1))) + dx = self.conv_1(self.actvn(self.norm_1(dx, seg1))) + out = x_s + dx + return out + + def shortcut(self, x, seg1): + if self.learned_shortcut: + x_s = self.conv_s(self.norm_s(x, seg1)) + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + + +def filter_state_dict(state_dict, remove_name='fc'): + new_state_dict = {} + for key in state_dict: + if remove_name in key: + continue + new_state_dict[key] = state_dict[key] + return new_state_dict + + +class GRN(nn.Module): + """ GRN (Global Response Normalization) layer + """ + + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +class LayerNorm(nn.Module): + r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + self.normalized_shape = (normalized_shape, ) + + def forward(self, x): + if self.data_format == "channels_last": + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + elif self.data_format == "channels_first": + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1. + math.erf(x / math.sqrt(2.))) / 2. + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2) + + with torch.no_grad(): + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + tensor.uniform_(2 * l - 1, 2 * u - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + return tensor + + +def drop_path(x, drop_prob=0., training=False, scale_by_keep=True): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + + def __init__(self, drop_prob=None, scale_by_keep=True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + +def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): + return _no_grad_trunc_normal_(tensor, mean, std, a, b) diff --git a/liveportrait/modules/warping_network.py b/liveportrait/modules/warping_network.py new file mode 100644 index 0000000..9191a19 --- /dev/null +++ b/liveportrait/modules/warping_network.py @@ -0,0 +1,77 @@ +# coding: utf-8 + +""" +Warping field estimator(W) defined in the paper, which generates a warping field using the implicit +keypoint representations x_s and x_d, and employs this flow field to warp the source feature volume f_s. +""" + +from torch import nn +import torch.nn.functional as F +from .util import SameBlock2d +from .dense_motion import DenseMotionNetwork + + +class WarpingNetwork(nn.Module): + def __init__( + self, + num_kp, + block_expansion, + max_features, + num_down_blocks, + reshape_channel, + estimate_occlusion_map=False, + dense_motion_params=None, + **kwargs + ): + super(WarpingNetwork, self).__init__() + + self.upscale = kwargs.get('upscale', 1) + self.flag_use_occlusion_map = kwargs.get('flag_use_occlusion_map', True) + + if dense_motion_params is not None: + self.dense_motion_network = DenseMotionNetwork( + num_kp=num_kp, + feature_channel=reshape_channel, + estimate_occlusion_map=estimate_occlusion_map, + **dense_motion_params + ) + else: + self.dense_motion_network = None + + self.third = SameBlock2d(max_features, block_expansion * (2 ** num_down_blocks), kernel_size=(3, 3), padding=(1, 1), lrelu=True) + self.fourth = nn.Conv2d(in_channels=block_expansion * (2 ** num_down_blocks), out_channels=block_expansion * (2 ** num_down_blocks), kernel_size=1, stride=1) + + self.estimate_occlusion_map = estimate_occlusion_map + + def deform_input(self, inp, deformation): + return F.grid_sample(inp, deformation, align_corners=False) + + def forward(self, feature_3d, kp_driving, kp_source): + if self.dense_motion_network is not None: + # Feature warper, Transforming feature representation according to deformation and occlusion + dense_motion = self.dense_motion_network( + feature=feature_3d, kp_driving=kp_driving, kp_source=kp_source + ) + if 'occlusion_map' in dense_motion: + occlusion_map = dense_motion['occlusion_map'] # Bx1x64x64 + else: + occlusion_map = None + + deformation = dense_motion['deformation'] # Bx16x64x64x3 + out = self.deform_input(feature_3d, deformation) # Bx32x16x64x64 + + bs, c, d, h, w = out.shape # Bx32x16x64x64 + out = out.view(bs, c * d, h, w) # -> Bx512x64x64 + out = self.third(out) # -> Bx256x64x64 + out = self.fourth(out) # -> Bx256x64x64 + + if self.flag_use_occlusion_map and (occlusion_map is not None): + out = out * occlusion_map + + ret_dct = { + 'occlusion_map': occlusion_map, + 'deformation': deformation, + 'out': out, + } + + return ret_dct diff --git a/liveportrait/template_maker.py b/liveportrait/template_maker.py new file mode 100644 index 0000000..7f3ce06 --- /dev/null +++ b/liveportrait/template_maker.py @@ -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}") diff --git a/liveportrait/utils/__init__.py b/liveportrait/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/liveportrait/utils/camera.py b/liveportrait/utils/camera.py new file mode 100644 index 0000000..8bbfc90 --- /dev/null +++ b/liveportrait/utils/camera.py @@ -0,0 +1,75 @@ +# coding: utf-8 + +""" +functions for processing and transforming 3D facial keypoints +""" + +import numpy as np +import torch +import torch.nn.functional as F + +PI = np.pi + + +def headpose_pred_to_degree(pred): + """ + pred: (bs, 66) or (bs, 1) or others + """ + if pred.ndim > 1 and pred.shape[1] == 66: + # NOTE: note that the average is modified to 97.5 + device = pred.device + idx_tensor = [idx for idx in range(0, 66)] + idx_tensor = torch.FloatTensor(idx_tensor).to(device) + pred = F.softmax(pred, dim=1) + degree = torch.sum(pred*idx_tensor, axis=1) * 3 - 97.5 + + return degree + + return pred + + +def get_rotation_matrix(pitch_, yaw_, roll_): + """ the input is in degree + """ + # calculate the rotation matrix: vps @ rot + + # transform to radian + pitch = pitch_ / 180 * PI + yaw = yaw_ / 180 * PI + roll = roll_ / 180 * PI + + device = pitch.device + + if pitch.ndim == 1: + pitch = pitch.unsqueeze(1) + if yaw.ndim == 1: + yaw = yaw.unsqueeze(1) + if roll.ndim == 1: + roll = roll.unsqueeze(1) + + # calculate the euler matrix + bs = pitch.shape[0] + ones = torch.ones([bs, 1]).to(device) + zeros = torch.zeros([bs, 1]).to(device) + x, y, z = pitch, yaw, roll + + rot_x = torch.cat([ + ones, zeros, zeros, + zeros, torch.cos(x), -torch.sin(x), + zeros, torch.sin(x), torch.cos(x) + ], dim=1).reshape([bs, 3, 3]) + + rot_y = torch.cat([ + torch.cos(y), zeros, torch.sin(y), + zeros, ones, zeros, + -torch.sin(y), zeros, torch.cos(y) + ], dim=1).reshape([bs, 3, 3]) + + rot_z = torch.cat([ + torch.cos(z), -torch.sin(z), zeros, + torch.sin(z), torch.cos(z), zeros, + zeros, zeros, ones + ], dim=1).reshape([bs, 3, 3]) + + rot = rot_z @ rot_y @ rot_x + return rot.permute(0, 2, 1) # transpose diff --git a/liveportrait/utils/crop.py b/liveportrait/utils/crop.py new file mode 100644 index 0000000..c061ef4 --- /dev/null +++ b/liveportrait/utils/crop.py @@ -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() + diff --git a/liveportrait/utils/cropper.py b/liveportrait/utils/cropper.py new file mode 100644 index 0000000..6245d14 --- /dev/null +++ b/liveportrait/utils/cropper.py @@ -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'] diff --git a/liveportrait/utils/face_analysis_diy.py b/liveportrait/utils/face_analysis_diy.py new file mode 100644 index 0000000..376334f --- /dev/null +++ b/liveportrait/utils/face_analysis_diy.py @@ -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') diff --git a/liveportrait/utils/helper.py b/liveportrait/utils/helper.py new file mode 100644 index 0000000..454050c --- /dev/null +++ b/liveportrait/utils/helper.py @@ -0,0 +1,175 @@ +# coding: utf-8 + +""" +utility functions and classes to handle feature extraction and model loading +""" + +import os +import os.path as osp +import cv2 +import torch +from rich.console import Console +from collections import OrderedDict + +from ..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 diff --git a/liveportrait/utils/io.py b/liveportrait/utils/io.py new file mode 100644 index 0000000..f930c48 --- /dev/null +++ b/liveportrait/utils/io.py @@ -0,0 +1,97 @@ +# coding: utf-8 + +import os +from glob import glob +import os.path as osp +import imageio +import numpy as np +import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False) + + +def load_image_rgb(image_path: str): + if not osp.exists(image_path): + raise FileNotFoundError(f"Image not found: {image_path}") + img = cv2.imread(image_path, cv2.IMREAD_COLOR) + return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + +def load_driving_info(driving_info): + driving_video_ori = [] + + def load_images_from_directory(directory): + image_paths = sorted(glob(osp.join(directory, '*.png')) + glob(osp.join(directory, '*.jpg'))) + return [load_image_rgb(im_path) for im_path in image_paths] + + def load_images_from_video(file_path): + reader = imageio.get_reader(file_path) + return [image for idx, image in enumerate(reader)] + + if osp.isdir(driving_info): + driving_video_ori = load_images_from_directory(driving_info) + elif osp.isfile(driving_info): + driving_video_ori = load_images_from_video(driving_info) + + return driving_video_ori + + +def contiguous(obj): + if not obj.flags.c_contiguous: + obj = obj.copy(order="C") + return obj + + +def _resize_to_limit(img: np.ndarray, max_dim=1920, n=2): + """ + ajust the size of the image so that the maximum dimension does not exceed max_dim, and the width and the height of the image are multiples of n. + :param img: the image to be processed. + :param max_dim: the maximum dimension constraint. + :param n: the number that needs to be multiples of. + :return: the adjusted image. + """ + h, w = img.shape[:2] + + # ajust the size of the image according to the maximum dimension + if max_dim > 0 and max(h, w) > max_dim: + if h > w: + new_h = max_dim + new_w = int(w * (max_dim / h)) + else: + new_w = max_dim + new_h = int(h * (max_dim / w)) + img = cv2.resize(img, (new_w, new_h)) + + # ensure that the image dimensions are multiples of n + n = max(n, 1) + new_h = img.shape[0] - (img.shape[0] % n) + new_w = img.shape[1] - (img.shape[1] % n) + + if new_h == 0 or new_w == 0: + # when the width or height is less than n, no need to process + return img + + if new_h != img.shape[0] or new_w != img.shape[1]: + img = img[:new_h, :new_w] + + return img + + +def load_img_online(obj, mode="bgr", **kwargs): + max_dim = kwargs.get("max_dim", 1920) + n = kwargs.get("n", 2) + if isinstance(obj, str): + if mode.lower() == "gray": + img = cv2.imread(obj, cv2.IMREAD_GRAYSCALE) + else: + img = cv2.imread(obj, cv2.IMREAD_COLOR) + else: + img = obj + + # Resize image to satisfy constraints + img = _resize_to_limit(img, max_dim=max_dim, n=n) + + if mode.lower() == "bgr": + return contiguous(img) + elif mode.lower() == "rgb": + return contiguous(img[..., ::-1]) + else: + raise Exception(f"Unknown mode {mode}") diff --git a/liveportrait/utils/landmark_runner.py b/liveportrait/utils/landmark_runner.py new file mode 100644 index 0000000..7b0dcbe --- /dev/null +++ b/liveportrait/utils/landmark_runner.py @@ -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') diff --git a/liveportrait/utils/resources/.gitattributes b/liveportrait/utils/resources/.gitattributes new file mode 100644 index 0000000..e9c5856 --- /dev/null +++ b/liveportrait/utils/resources/.gitattributes @@ -0,0 +1 @@ +mask_template.png filter=lfs diff=lfs merge=lfs -text diff --git a/liveportrait/utils/resources/mask_template.png b/liveportrait/utils/resources/mask_template.png new file mode 100644 index 0000000..d352221 --- /dev/null +++ b/liveportrait/utils/resources/mask_template.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4c53e64a07ae3056af2b38548ae5b6cceb04d5c3e6514c7a8d2c3aff9da1ee76 +size 3470 diff --git a/liveportrait/utils/retargeting_utils.py b/liveportrait/utils/retargeting_utils.py new file mode 100644 index 0000000..2028590 --- /dev/null +++ b/liveportrait/utils/retargeting_utils.py @@ -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 diff --git a/liveportrait/utils/rprint.py b/liveportrait/utils/rprint.py new file mode 100644 index 0000000..c43a42f --- /dev/null +++ b/liveportrait/utils/rprint.py @@ -0,0 +1,16 @@ +# coding: utf-8 + +""" +custom print and log functions +""" + +__all__ = ['rprint', 'rlog'] + +try: + from rich.console import Console + console = Console() + rprint = console.print + rlog = console.log +except: + rprint = print + rlog = print diff --git a/liveportrait/utils/timer.py b/liveportrait/utils/timer.py new file mode 100644 index 0000000..3570fa4 --- /dev/null +++ b/liveportrait/utils/timer.py @@ -0,0 +1,29 @@ +# coding: utf-8 + +""" +tools to measure elapsed time +""" + +import time + +class Timer(object): + """A simple timer.""" + + def __init__(self): + self.total_time = 0. + self.calls = 0 + self.start_time = 0. + self.diff = 0. + + def tic(self): + # using time.time instead of time.clock because time time.clock + # does not normalize for multithreading + self.start_time = time.time() + + def toc(self, average=True): + self.diff = time.time() - self.start_time + return self.diff + + def clear(self): + self.start_time = 0. + self.diff = 0. diff --git a/liveportrait/utils/video.py b/liveportrait/utils/video.py new file mode 100644 index 0000000..720e082 --- /dev/null +++ b/liveportrait/utils/video.py @@ -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 diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..b10dc4c --- /dev/null +++ b/nodes.py @@ -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", + } diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..2435fa8 --- /dev/null +++ b/readme.md @@ -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` diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ce6d039 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +yaml +numpy +opencv-python