387 lines
13 KiB
Python
387 lines
13 KiB
Python
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.utils.cropper import Cropper
|
|
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),
|
|
flag_write_result=True,
|
|
flag_pasteback=True,
|
|
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.flag_write_result = flag_write_result
|
|
self.flag_pasteback = flag_pasteback
|
|
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": {},
|
|
"optional": {
|
|
"precision": (
|
|
[
|
|
"fp16",
|
|
"fp32",
|
|
],
|
|
{"default": "fp16"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LIVEPORTRAITPIPE",)
|
|
RETURN_NAMES = ("live_portrait_pipe",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "LivePortrait"
|
|
|
|
def loadmodel(self, precision="fp16"):
|
|
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)
|
|
|
|
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)
|
|
|
|
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()
|
|
|
|
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(
|
|
device_id=device,
|
|
flag_use_half_precision=True if precision == "fp16" else False,
|
|
),
|
|
)
|
|
|
|
return (pipeline,)
|
|
|
|
|
|
# OUR CURRENT NODE
|
|
class LivePortraitProcess:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("LIVEPORTRAITPIPE",),
|
|
"source_image": ("IMAGE",),
|
|
"driving_images": ("IMAGE",),
|
|
"dsize": ("INT", {"default": 512, "min": 64, "max": 2048}),
|
|
"scale": (
|
|
"FLOAT",
|
|
{"default": 2.3, "min": 1.0, "max": 4.0, "step": 0.01},
|
|
),
|
|
"vx_ratio": (
|
|
"FLOAT",
|
|
{"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"vy_ratio": (
|
|
"FLOAT",
|
|
{"default": -0.125, "min": -1.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"lip_zero": ("BOOLEAN", {"default": True}),
|
|
"eye_retargeting": ("BOOLEAN", {"default": False}),
|
|
"eyes_retargeting_multiplier": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.01, "max": 10.0, "step": 0.001},
|
|
),
|
|
"lip_retargeting": ("BOOLEAN", {"default": False}),
|
|
"lip_retargeting_multiplier": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.01, "max": 10.0, "step": 0.001},
|
|
),
|
|
"stitching": ("BOOLEAN", {"default": True}),
|
|
"relative": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"mismatch_method": (
|
|
["repeat", "cycle", "mirror", "nearest"],
|
|
{"default": "repeat"},
|
|
),
|
|
"onnx_device": (
|
|
[
|
|
"CPU",
|
|
"CUDA",
|
|
],
|
|
{"default": "CPU"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
)
|
|
RETURN_NAMES = (
|
|
"cropped_images",
|
|
"full_images",
|
|
)
|
|
FUNCTION = "process"
|
|
CATEGORY = "LivePortrait"
|
|
|
|
def process(
|
|
self,
|
|
source_image: torch.Tensor,
|
|
driving_images: torch.Tensor,
|
|
dsize: int,
|
|
scale: float,
|
|
vx_ratio: float,
|
|
vy_ratio: float,
|
|
pipeline: LivePortraitPipeline,
|
|
lip_zero: bool,
|
|
eye_retargeting: bool,
|
|
lip_retargeting: bool,
|
|
stitching: bool,
|
|
relative: bool,
|
|
eyes_retargeting_multiplier: float,
|
|
lip_retargeting_multiplier: float,
|
|
mismatch_method: str = "repeat",
|
|
onnx_device="CUDA",
|
|
):
|
|
source_np = (source_image * 255).byte().numpy()
|
|
driving_images_np = (driving_images * 255).byte().numpy()
|
|
|
|
crop_cfg = CropConfig(
|
|
dsize=dsize,
|
|
scale=scale,
|
|
vx_ratio=vx_ratio,
|
|
vy_ratio=vy_ratio,
|
|
)
|
|
|
|
cropper = Cropper(crop_cfg=crop_cfg, provider=onnx_device)
|
|
pipeline.cropper = cropper
|
|
pipeline.live_portrait_wrapper.cfg.flag_eye_retargeting = eye_retargeting
|
|
pipeline.live_portrait_wrapper.cfg.eyes_retargeting_multiplier = (
|
|
eyes_retargeting_multiplier
|
|
)
|
|
pipeline.live_portrait_wrapper.cfg.flag_lip_retargeting = lip_retargeting
|
|
pipeline.live_portrait_wrapper.cfg.lip_retargeting_multiplier = (
|
|
lip_retargeting_multiplier
|
|
)
|
|
pipeline.live_portrait_wrapper.cfg.flag_stitching = stitching
|
|
pipeline.live_portrait_wrapper.cfg.flag_relative = relative
|
|
pipeline.live_portrait_wrapper.cfg.flag_lip_zero = lip_zero
|
|
|
|
cropped_out_list = []
|
|
full_out_list = []
|
|
|
|
cropped_out_list, full_out_list = pipeline.execute(
|
|
source_np, driving_images_np, mismatch_method
|
|
)
|
|
cropped_tensors_out = (
|
|
torch.stack([torch.from_numpy(np_array) for np_array in cropped_out_list])
|
|
/ 255
|
|
)
|
|
full_tensors_out = (
|
|
torch.stack([torch.from_numpy(np_array) for np_array in full_out_list])
|
|
/ 255
|
|
)
|
|
|
|
return (cropped_tensors_out.cpu().float(), full_tensors_out.cpu().float())
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DownloadAndLoadLivePortraitModels": DownloadAndLoadLivePortraitModels,
|
|
"LivePortraitProcess": LivePortraitProcess,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DownloadAndLoadLivePortraitModels": "(Down)Load LivePortraitModels",
|
|
}
|