697 lines
33 KiB
Python
697 lines
33 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import io
|
|
import logging
|
|
import os
|
|
import random
|
|
import numpy as np
|
|
import torch
|
|
import torchaudio
|
|
import gc
|
|
from diffusers import AutoencoderKL, DDIMScheduler
|
|
from omegaconf import OmegaConf
|
|
from .src.models.unet_2d_condition import UNet2DConditionModel
|
|
from .src.models.unet_3d_echo import EchoUNet3DConditionModel
|
|
from .src.models.whisper.audio2feature import load_audio_model
|
|
from .src.pipelines.pipeline_echo_mimic import Audio2VideoPipeline
|
|
from .src.pipelines.pipeline_echo_mimic_acc import Audio2VideoPipeline as Audio2VideoACCPipeline
|
|
from .src.pipelines.pipeline_echo_mimic_pose import AudioPose2VideoPipeline
|
|
from .src.pipelines.pipeline_echo_mimic_pose_acc import AudioPose2VideoPipeline as AudioPose2VideoaccPipeline
|
|
from .src.models.face_locator import FaceLocator
|
|
from .src.utils.draw_utils import FaceMeshVisualizer
|
|
from .src.utils.motion_utils import motion_sync
|
|
from .utils import load_images, tensor2cv, find_directories, nomarl_upscale, download_weights, get_instance_path, \
|
|
process_video, narry_list, weight_dtype, cf_tensor2cv,process_video_v2
|
|
from .echomimic_v2.src.models.pose_encoder import PoseEncoder
|
|
from .echomimic_v2.src.pipelines.pipeline_echomimicv2 import EchoMimicV2Pipeline
|
|
from .echomimic_v2.src.models.unet_2d_condition import UNet2DConditionModel as UNet2DConditionModelV2
|
|
from .echomimic_v2.src.models.unet_3d_emo import EMOUNet3DConditionModel as EMOUNet3DConditionModelV2
|
|
from .hallo.video_sr import run_realesrgan, pre_u_loader
|
|
from pathlib import Path
|
|
from huggingface_hub import hf_hub_download
|
|
import folder_paths
|
|
import platform
|
|
import subprocess
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
|
current_path = os.path.dirname(os.path.abspath(__file__))
|
|
inference_config_path = os.path.join(current_path, "configs", "inference", "inference_v2.yaml")
|
|
infer_config = OmegaConf.load(inference_config_path)
|
|
inference_config_path_v2 = os.path.join(current_path, "echomimic_v2/configs/inference/inference_v2.yaml")
|
|
infer_config_v2 = OmegaConf.load(inference_config_path_v2)
|
|
|
|
weigths_hallo_current_path = os.path.join(folder_paths.models_dir, "Hallo")
|
|
if not os.path.exists(weigths_hallo_current_path):
|
|
os.makedirs(weigths_hallo_current_path)
|
|
|
|
try:
|
|
folder_paths.add_model_folder_path("Hallo", weigths_hallo_current_path, False)
|
|
except:
|
|
folder_paths.add_model_folder_path("Hallo", weigths_hallo_current_path)
|
|
|
|
# pre dir
|
|
weigths_current_path = os.path.join(folder_paths.models_dir, "echo_mimic")
|
|
if not os.path.exists(weigths_current_path):
|
|
os.makedirs(weigths_current_path)
|
|
|
|
weigths_uet_current_path = os.path.join(weigths_current_path, "unet")
|
|
if not os.path.exists(weigths_uet_current_path):
|
|
os.makedirs(weigths_uet_current_path)
|
|
|
|
weigths_vae_current_path = os.path.join(weigths_current_path, "vae")
|
|
if not os.path.exists(weigths_vae_current_path):
|
|
os.makedirs(weigths_vae_current_path)
|
|
|
|
weigths_au_current_path = os.path.join(weigths_current_path, "audio_processor")
|
|
if not os.path.exists(weigths_au_current_path):
|
|
os.makedirs(weigths_au_current_path)
|
|
|
|
tensorrt_lite = os.path.join(folder_paths.get_input_directory(), "tensorrt_lite")
|
|
if not os.path.exists(tensorrt_lite):
|
|
os.makedirs(tensorrt_lite)
|
|
|
|
# upscale dir
|
|
weigths_facelib_path = os.path.join(weigths_hallo_current_path, "facelib")
|
|
if not os.path.exists(weigths_hallo_current_path):
|
|
os.makedirs(weigths_facelib_path)
|
|
|
|
weigths_face_analysis_path = os.path.join(weigths_hallo_current_path, "face_analysis/models")
|
|
weigths_face_analysis_dir = os.path.join(weigths_hallo_current_path, "face_analysis")
|
|
if not os.path.exists(weigths_face_analysis_path):
|
|
os.makedirs(weigths_face_analysis_path)
|
|
|
|
# ffmpeg
|
|
ffmpeg_path = os.getenv('FFMPEG_PATH')
|
|
if ffmpeg_path is None and platform.system() in ['Linux', 'Darwin']:
|
|
try:
|
|
result = subprocess.run(['which', 'ffmpeg'], capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
ffmpeg_path = result.stdout.strip()
|
|
print(f"FFmpeg is installed at: {ffmpeg_path}")
|
|
else:
|
|
print("FFmpeg is not installed. Please download ffmpeg-static and export to FFMPEG_PATH.")
|
|
print("For example: export FFMPEG_PATH=/musetalk/ffmpeg-4.4-amd64-static")
|
|
except Exception as e:
|
|
pass
|
|
|
|
if ffmpeg_path is not None and ffmpeg_path not in os.getenv('PATH'):
|
|
print("Adding FFMPEG_PATH to PATH")
|
|
os.environ["PATH"] = f"{ffmpeg_path}:{os.environ['PATH']}"
|
|
|
|
|
|
# *****************mian***************
|
|
|
|
class Echo_LoadModel:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae": ("STRING", {"default": "stabilityai/sd-vae-ft-mse"}),
|
|
"denoising": ("BOOLEAN", {"default": True},),
|
|
"infer_mode": (["audio_drived", "audio_drived_acc", "pose_normal", "pose_acc"],),
|
|
"draw_mouse": ("BOOLEAN", {"default": False},),
|
|
"motion_sync": ("BOOLEAN", {"default": False},),
|
|
"lowvram": ("BOOLEAN", {"default": False},),
|
|
"version": (["V2", "V1", ],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL_PIPE_E", "MODEL_FACE_E", "MODEL_VISUAL_E",)
|
|
RETURN_NAMES = ("model", "face_detector", "visualizer",)
|
|
FUNCTION = "main_loader"
|
|
CATEGORY = "EchoMimic"
|
|
|
|
def main_loader(self, vae, denoising, infer_mode, draw_mouse, motion_sync, lowvram, version):
|
|
|
|
############# model_init started #############
|
|
|
|
## vae init #using local vae first
|
|
try:
|
|
vae = AutoencoderKL.from_pretrained(weigths_vae_current_path).to(device,
|
|
dtype=weight_dtype) # using local vae first
|
|
except:
|
|
try: # try downlaod model ,and load local vae
|
|
download_weights(weigths_vae_current_path, "stabilityai/sd-vae-ft-mse", subfolder="",
|
|
pt_name="diffusion_pytorch_model.safetensors")
|
|
download_weights(weigths_vae_current_path, "stabilityai/sd-vae-ft-mse", subfolder="",
|
|
pt_name="config.json")
|
|
vae = AutoencoderKL.from_pretrained(weigths_vae_current_path).to(device, dtype=weight_dtype)
|
|
except:
|
|
vae = AutoencoderKL.from_pretrained(vae).to(device, dtype=weight_dtype)
|
|
|
|
|
|
## reference net init
|
|
pretrained_base_model_path = get_instance_path(weigths_current_path)
|
|
|
|
# pre models
|
|
download_weights(weigths_current_path, "lambdalabs/sd-image-variations-diffusers", subfolder="unet",
|
|
pt_name="diffusion_pytorch_model.bin")
|
|
download_weights(weigths_current_path, "lambdalabs/sd-image-variations-diffusers", subfolder="unet",
|
|
pt_name="config.json")
|
|
audio_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic", subfolder="audio_processor",
|
|
pt_name="whisper_tiny.pt")
|
|
|
|
# pre pth
|
|
if version == "V1":
|
|
logging.info("****** refer in EchoMimic V1 mode!******")
|
|
if infer_mode == "pose_normal":
|
|
re_ckpt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="reference_unet_pose.pth")
|
|
face_locator_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="face_locator_pose.pth")
|
|
motion_path = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="motion_module_pose.pth")
|
|
denois_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="denoising_unet_pose.pth")
|
|
|
|
elif infer_mode == "pose_acc":
|
|
re_ckpt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="reference_unet_pose.pth")
|
|
motion_path = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="motion_module_pose_acc.pth")
|
|
denois_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="denoising_unet_pose_acc.pth")
|
|
face_locator_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="face_locator_pose.pth")
|
|
elif infer_mode == "audio_drived":
|
|
re_ckpt = download_weights(weigths_current_path, "BadToBest/EchoMimic", pt_name="reference_unet.pth")
|
|
face_locator_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="face_locator.pth")
|
|
motion_path = download_weights(weigths_current_path, "BadToBest/EchoMimic", pt_name="motion_module.pth")
|
|
denois_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic", pt_name="denoising_unet.pth")
|
|
else:
|
|
re_ckpt = download_weights(weigths_current_path, "BadToBest/EchoMimic", pt_name="reference_unet.pth")
|
|
face_locator_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="face_locator.pth")
|
|
motion_path = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="motion_module_acc.pth")
|
|
denois_pt = download_weights(weigths_current_path, "BadToBest/EchoMimic",
|
|
pt_name="denoising_unet_acc.pth")
|
|
else:
|
|
logging.info("****** refer in EchoMimic V2 mode!******")
|
|
weigths_current_path_v2 = os.path.join(weigths_current_path, "v2")
|
|
if not os.path.exists(weigths_current_path_v2):
|
|
os.makedirs(weigths_current_path_v2)
|
|
re_ckpt = download_weights(weigths_current_path_v2, "BadToBest/EchoMimicV2",
|
|
pt_name="reference_unet.pth")
|
|
motion_path = download_weights(weigths_current_path_v2, "BadToBest/EchoMimicV2",
|
|
pt_name="motion_module.pth")
|
|
denois_pt = download_weights(weigths_current_path_v2, "BadToBest/EchoMimicV2",
|
|
pt_name="denoising_unet.pth")
|
|
pose_encoder_pt = download_weights(weigths_current_path_v2, "BadToBest/EchoMimicV2",
|
|
pt_name="pose_encoder.pth")
|
|
|
|
# unet init
|
|
if version == "V1":
|
|
try:
|
|
reference_unet = UNet2DConditionModel.from_config(
|
|
pretrained_base_model_path,
|
|
subfolder="unet",
|
|
).to(dtype=weight_dtype)
|
|
except:
|
|
reference_unet = UNet2DConditionModel.from_pretrained(
|
|
pretrained_base_model_path,
|
|
subfolder="unet",
|
|
).to(dtype=weight_dtype)
|
|
else:
|
|
try:
|
|
reference_unet = UNet2DConditionModelV2.from_config(
|
|
pretrained_base_model_path,
|
|
subfolder="unet",
|
|
).to(dtype=weight_dtype)
|
|
except:
|
|
reference_unet = UNet2DConditionModelV2.from_pretrained(
|
|
pretrained_base_model_path,
|
|
subfolder="unet",
|
|
).to(dtype=weight_dtype)
|
|
|
|
re_state = torch.load(re_ckpt, map_location="cpu")
|
|
reference_unet.load_state_dict(re_state, strict=False)
|
|
del re_state
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
## denoising net init
|
|
if version == "V1":
|
|
if denoising:
|
|
if os.path.exists(motion_path): ### stage1 + stage2
|
|
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
|
pretrained_base_model_path,
|
|
motion_path,
|
|
subfolder="unet",
|
|
unet_additional_kwargs=infer_config.unet_additional_kwargs,
|
|
).to(dtype=weight_dtype)
|
|
else:
|
|
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
|
pretrained_base_model_path,
|
|
"",
|
|
subfolder="unet",
|
|
unet_additional_kwargs={
|
|
"use_motion_module": False,
|
|
"unet_use_temporal_attention": False,
|
|
"cross_attention_dim": infer_config.unet_additional_kwargs.cross_attention_dim
|
|
}
|
|
).to(dtype=weight_dtype)
|
|
else:
|
|
### only stage1
|
|
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
|
pretrained_base_model_path,
|
|
"",
|
|
subfolder="unet",
|
|
unet_additional_kwargs={
|
|
"use_motion_module": False,
|
|
"unet_use_temporal_attention": False,
|
|
"cross_attention_dim": infer_config.unet_additional_kwargs.cross_attention_dim
|
|
}
|
|
).to(dtype=weight_dtype, )
|
|
else: # v2
|
|
denoising_unet = EMOUNet3DConditionModelV2.from_pretrained_2d(
|
|
pretrained_base_model_path,
|
|
motion_path,
|
|
subfolder="unet",
|
|
unet_additional_kwargs=infer_config_v2.unet_additional_kwargs,
|
|
).to(dtype=weight_dtype)
|
|
denoising_state = torch.load(denois_pt, map_location="cpu")
|
|
denoising_unet.load_state_dict(denoising_state, strict=False)
|
|
del denoising_state
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
if version == "V1":
|
|
if infer_mode == "pose_normal" or infer_mode == "pose_acc":
|
|
# face locator init
|
|
face_locator = FaceLocator(320, conditioning_channels=3, block_out_channels=(16, 32, 96, 256)).to(
|
|
dtype=weight_dtype, device=device)
|
|
face_locator.load_state_dict(torch.load(face_locator_pt), strict=False)
|
|
if motion_sync:
|
|
visualizer = FaceMeshVisualizer(draw_iris=False, draw_mouse=True, draw_eye=True, draw_nose=True,
|
|
draw_eyebrow=True, draw_pupil=True)
|
|
else:
|
|
visualizer = FaceMeshVisualizer(draw_iris=False, draw_mouse=draw_mouse)
|
|
else:
|
|
# face locator init
|
|
face_locator = FaceLocator(320, conditioning_channels=1, block_out_channels=(16, 32, 96, 256)).to(
|
|
dtype=weight_dtype, device=device)
|
|
face_locator.load_state_dict(torch.load(face_locator_pt), strict=False)
|
|
visualizer = None
|
|
else: # v2
|
|
# pose net init
|
|
pose_net = PoseEncoder(320, conditioning_channels=3, block_out_channels=(16, 32, 96, 256)).to(device=device,
|
|
dtype=weight_dtype)
|
|
pose_state = torch.load(pose_encoder_pt)
|
|
pose_net.load_state_dict(pose_state)
|
|
del pose_state
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
if infer_mode == "pose_normal":
|
|
from .src.pose import SapiensPoseEstimation
|
|
pose_dir_32=os.path.join(weigths_current_path,"sapiens_1b_goliath_best_goliath_AP_639_torchscript.pt2")
|
|
pose_dir_bf16 = os.path.join(weigths_current_path, "sapiens_1b_goliath_best_goliath_AP_639_torchscript_bf16.pt2")
|
|
dtype=torch.float32
|
|
if os.path.exists(pose_dir_bf16):
|
|
dtype=torch.float16
|
|
pose_dir=pose_dir_bf16
|
|
else:
|
|
if os.path.exists(pose_dir_32):
|
|
pose_dir = pose_dir_32
|
|
else:
|
|
pose_dir = ""
|
|
visualizer = SapiensPoseEstimation(local_pose=pose_dir,model_dir=weigths_current_path,dtype=dtype)
|
|
else:
|
|
visualizer = None
|
|
|
|
|
|
## load audio processor params
|
|
audio_processor = load_audio_model(model_path=audio_pt, device=device)
|
|
|
|
## load face detector params
|
|
if version == "V1":
|
|
from facenet_pytorch import MTCNN
|
|
face_detector = MTCNN(image_size=320, margin=0, min_face_size=20, thresholds=[0.6, 0.7, 0.7], factor=0.709,
|
|
post_process=True, device=device)
|
|
else:
|
|
face_detector = None
|
|
|
|
############# model_init finished #############
|
|
if version == "V1":
|
|
sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
|
|
else:
|
|
sched_kwargs = OmegaConf.to_container(infer_config_v2.noise_scheduler_kwargs)
|
|
scheduler = DDIMScheduler(**sched_kwargs)
|
|
|
|
if version == "V1":
|
|
if infer_mode == "pose_normal":
|
|
pipe = AudioPose2VideoPipeline(
|
|
vae=vae,
|
|
reference_unet=reference_unet,
|
|
denoising_unet=denoising_unet,
|
|
audio_guider=audio_processor,
|
|
face_locator=face_locator,
|
|
scheduler=scheduler,
|
|
).to(dtype=weight_dtype)
|
|
elif infer_mode == "pose_acc":
|
|
pipe = AudioPose2VideoaccPipeline(
|
|
vae=vae,
|
|
reference_unet=reference_unet,
|
|
denoising_unet=denoising_unet,
|
|
audio_guider=audio_processor,
|
|
face_locator=face_locator,
|
|
scheduler=scheduler,
|
|
).to(dtype=weight_dtype)
|
|
elif infer_mode == "audio_drived":
|
|
pipe = Audio2VideoPipeline(
|
|
vae=vae,
|
|
reference_unet=reference_unet,
|
|
denoising_unet=denoising_unet,
|
|
audio_guider=audio_processor,
|
|
face_locator=face_locator,
|
|
scheduler=scheduler,
|
|
).to(dtype=weight_dtype)
|
|
else:
|
|
pipe = Audio2VideoACCPipeline(
|
|
vae=vae,
|
|
reference_unet=reference_unet,
|
|
denoising_unet=denoising_unet,
|
|
audio_guider=audio_processor,
|
|
face_locator=face_locator,
|
|
scheduler=scheduler,
|
|
).to(dtype=weight_dtype)
|
|
else:
|
|
pipe = EchoMimicV2Pipeline(
|
|
vae=vae,
|
|
reference_unet=reference_unet,
|
|
denoising_unet=denoising_unet,
|
|
audio_guider=audio_processor,
|
|
pose_encoder=pose_net,
|
|
scheduler=scheduler, )
|
|
|
|
pipe.enable_vae_slicing()
|
|
if lowvram:
|
|
pipe.enable_sequential_cpu_offload()
|
|
model = {"pipe": pipe, "lowvram": lowvram,"version":version}
|
|
return (model, face_detector, visualizer,)
|
|
|
|
|
|
class Echo_Sampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
pose_path_list = ["none"] + find_directories(tensorrt_lite) if find_directories(tensorrt_lite) else ["none", ]
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",), # [B,H,W,C], C=3
|
|
"audio": ("AUDIO",),
|
|
"model": ("MODEL_PIPE_E",),
|
|
"face_detector": ("MODEL_FACE_E",),
|
|
"pose_dir": (pose_path_list,),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
|
|
"cfg": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"steps": ("INT", {"default": 30, "min": 1, "max": 100}),
|
|
"fps": ("FLOAT", {"default": 25.0, "min": 5.0, "max": 120.0}),
|
|
"sample_rate": ("INT", {"default": 16000, "min": 8000, "max": 48000, "step": 1000, }),
|
|
"facemask_ratio": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
|
|
"facecrop_ratio": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
|
|
"context_frames": ("INT", {"default": 12, "min": 0, "max": 50}),
|
|
"context_overlap": ("INT", {"default": 3, "min": 0, "max": 10}),
|
|
"length": ("INT", {"default": 120, "min": 50, "max": 5000, "step": 1, "display": "number"}),
|
|
"width": ("INT", {"default": 512, "min": 128, "max": 1024, "step": 64, "display": "number"}),
|
|
"height": ("INT", {"default": 512, "min": 128, "max": 1024, "step": 64, "display": "number"}),
|
|
"save_video": ("BOOLEAN", {"default": False},), },
|
|
"optional": {
|
|
"visualizer": ("MODEL_VISUAL_E",),
|
|
"video_images": ("IMAGE",), # [B,H,W,C], C=3,B>1
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT")
|
|
RETURN_NAMES = ("image", "audio", "frame_rate")
|
|
FUNCTION = "em_main"
|
|
CATEGORY = "EchoMimic"
|
|
|
|
def em_main(self, image, audio, model, face_detector, pose_dir, seed, cfg, steps, fps, sample_rate, facemask_ratio,
|
|
facecrop_ratio, context_frames, context_overlap, length,
|
|
width, height, save_video, **kwargs):
|
|
|
|
version= model.get("version")
|
|
pipe = model.get("pipe")
|
|
lowvram = model.get("lowvram")
|
|
if not lowvram:
|
|
pipe.to(device, torch.float16)
|
|
|
|
image = cf_tensor2cv(image, width, height) if version=="V1" else image # v1 cv ,v2 pil
|
|
visualizer = kwargs.get("visualizer")
|
|
video_images = kwargs.get("video_images")
|
|
|
|
audio_file_prefix = ''.join(random.choice("0123456789") for _ in range(6))
|
|
audio_file = os.path.join(folder_paths.get_input_directory(), f"audio_{audio_file_prefix}_temp.wav")
|
|
|
|
# 减少音频数据传递导致的不必要文件存储
|
|
buff = io.BytesIO()
|
|
torchaudio.save(buff, audio["waveform"].squeeze(0), audio["sample_rate"], format="FLAC")
|
|
with open(audio_file, 'wb') as f:
|
|
f.write(buff.getbuffer())
|
|
if version=="V1":
|
|
output_video = process_video(image, audio_file, width, height, length, seed, facemask_ratio,
|
|
facecrop_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps,
|
|
pipe, face_detector, save_video, pose_dir, video_images, audio_file_prefix,
|
|
visualizer)
|
|
|
|
else:
|
|
output_video=process_video_v2(image, audio_file, width, height, length, seed,
|
|
context_frames, context_overlap, cfg, steps, sample_rate, fps, pipe,
|
|
save_video, pose_dir, audio_file_prefix,visualizer,video_images )
|
|
|
|
gen = narry_list(output_video) # pil列表排序
|
|
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
|
|
frame_rate = float(fps)
|
|
if not lowvram: # for upsacle ,need VR
|
|
pipe.to("cpu")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return (images, audio, frame_rate)
|
|
|
|
|
|
class Echo_Upscaleloader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
ckpt_list_filter_U = [i for i in folder_paths.get_filename_list("Hallo") if i.endswith(".pth") and "lib" in i]
|
|
upscale_list = [i for i in folder_paths.get_filename_list("upscale_models") if "x2plus" in i.lower()]
|
|
return {
|
|
"required": {
|
|
"realesrgan": (["none"] + upscale_list,),
|
|
"face_detection_model": (["none"] + ckpt_list_filter_U,),
|
|
"bg_upsampler": (['realesrgan', 'none', ],),
|
|
"face_upsample": ("BOOLEAN", {"default": False},),
|
|
"has_aligned": ("BOOLEAN", {"default": False},),
|
|
"bg_tile": ("INT", {
|
|
"default": 400,
|
|
"min": 200, # Minimum value
|
|
"max": 1000, # Maximum value
|
|
"step": 10, # Slider's step
|
|
"display": "number", # Cosmetic only: display as "number" or "slider"
|
|
}),
|
|
"upscale": ("INT", {
|
|
"default": 2,
|
|
"min": 2, # Minimum value
|
|
"max": 4, # Maximum value
|
|
"step": 2, # Slider's step
|
|
"display": "number", # Cosmetic only: display as "number" or "slider"
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("ECHO_U_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "Upscale_main"
|
|
CATEGORY = "EchoMimic"
|
|
|
|
def Upscale_main(self, realesrgan, face_detection_model, bg_upsampler, face_upsample, has_aligned, bg_tile,
|
|
upscale):
|
|
|
|
if realesrgan == "none":
|
|
if not "RealESRGAN_x2plus.pth" in folder_paths.get_filename_list("upscale_models"):
|
|
logging.info("NO RealESRGAN_x2plus.pth find in upscale_models dir,start auto download")
|
|
model_path = hf_hub_download(
|
|
repo_id="fudan-generative-ai/hallo2",
|
|
subfolder="realesrgan",
|
|
filename="RealESRGAN_x2plus.pth",
|
|
local_dir=os.path.join(folder_paths.models_dir, "upscale_models"),
|
|
)
|
|
else:
|
|
raise "Need choice 'RealESRGAN_x2plus.pth' at 'realesrgan' menu"
|
|
else:
|
|
model_path = folder_paths.get_full_path("upscale_models", realesrgan)
|
|
|
|
parse_model = os.path.join(weigths_facelib_path, "parsing_parsenet.pth")
|
|
if not os.path.exists(parse_model):
|
|
print(f" No 'parsing_parsenet.pth' in {parse_model} ,try download from huggingface!")
|
|
hf_hub_download(
|
|
repo_id="fudan-generative-ai/hallo2",
|
|
subfolder="facelib",
|
|
filename="parsing_parsenet.pth",
|
|
local_dir=weigths_hallo_current_path,
|
|
)
|
|
|
|
if face_detection_model == "none":
|
|
if not "detection_Resnet50_Final.pth" in folder_paths.get_filename_list("Hallo"):
|
|
logging.info("NO detection_Resnet50_Final.pth find in Hallo dir,start auto download")
|
|
face_detection_model = hf_hub_download(
|
|
repo_id="fudan-generative-ai/hallo2",
|
|
subfolder="facelib",
|
|
filename="detection_Resnet50_Final.pth",
|
|
local_dir=weigths_hallo_current_path,
|
|
)
|
|
else:
|
|
raise "need chocie a face_detection_model,resent or yolov5"
|
|
else:
|
|
face_detection_model = folder_paths.get_full_path("Hallo", face_detection_model)
|
|
|
|
hallo_model_path = os.path.join(weigths_hallo_current_path, "hallo2", "net_g.pth")
|
|
if not os.path.exists(hallo_model_path):
|
|
print(f"no net_g.pth in {weigths_hallo_current_path}/hallo2 ,try download from huggingface!")
|
|
hf_hub_download(
|
|
repo_id="fudan-generative-ai/hallo2",
|
|
subfolder="hallo2",
|
|
filename="net_g.pth",
|
|
local_dir=weigths_hallo_current_path,
|
|
)
|
|
|
|
net, face_upsampler, bg_upsampler, face_helper = pre_u_loader(bg_upsampler, model_path, bg_tile, upscale,
|
|
face_upsample, device, hallo_model_path,
|
|
face_detection_model, parse_model, has_aligned)
|
|
model = {"net": net, "face_upsampler": face_upsampler, "bg_upsampler": bg_upsampler, "upscale": upscale,
|
|
"face_helper": face_helper, "has_aligned": has_aligned, "face_upsample": face_upsample}
|
|
return (model,)
|
|
|
|
|
|
class Echo_VideoUpscale:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_path = folder_paths.get_input_directory()
|
|
video_files = [f for f in os.listdir(input_path) if
|
|
os.path.isfile(os.path.join(input_path, f)) and f.split('.')[-1] in ['webm', 'mp4', 'mkv',
|
|
'gif']]
|
|
return {
|
|
"required": {
|
|
"model": ("ECHO_U_MODEL",),
|
|
"video_path": (["none"] + video_files,),
|
|
"fidelity_weight": ("FLOAT", {
|
|
"default": 0.5,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.1,
|
|
"round": 0.01,
|
|
"display": "number",
|
|
}),
|
|
"only_center_face": ("BOOLEAN", {"default": False},),
|
|
"draw_box": ("BOOLEAN", {"default": False},),
|
|
"save_video": ("BOOLEAN", {"default": False},),
|
|
|
|
},
|
|
"optional": {"image": ("IMAGE",),
|
|
"audio": ("AUDIO",),
|
|
"frame_rate": ("FLOAT", {"forceInput": True, "default": 0.5, }),
|
|
"path": ("STRING", {"forceInput": True, "default": "", }),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT",)
|
|
RETURN_NAMES = ("image", "audio", "frame_rate")
|
|
FUNCTION = "Upscale_main"
|
|
CATEGORY = "EchoMimic"
|
|
|
|
def Upscale_main(self, model, video_path, fidelity_weight, only_center_face, draw_box, save_video, **kwargs):
|
|
# pre data
|
|
video_img = kwargs.get("image")
|
|
audio = kwargs.get("audio")
|
|
frame_rate = kwargs.get("frame_rate")
|
|
sampler_path = kwargs.get("path")
|
|
|
|
# pre model
|
|
net_g = model.get("net")
|
|
face_upsampler = model.get("face_upsampler")
|
|
bg_upsampler = model.get("bg_upsampler")
|
|
upscale = model.get("upscale")
|
|
face_helper = model.get("face_helper")
|
|
has_aligned = model.get("has_aligned")
|
|
face_upsample = model.get("face_upsample")
|
|
|
|
front_path = Path(sampler_path) if sampler_path and os.path.exists(Path(sampler_path)) else None
|
|
video_list = []
|
|
if isinstance(video_img, list):
|
|
if isinstance(video_img[0], torch.Tensor):
|
|
video_list = video_img
|
|
elif isinstance(video_img, torch.Tensor):
|
|
b, _, _, _ = video_img.size()
|
|
if b == 1:
|
|
img = [b]
|
|
while img is not []:
|
|
video_list += img
|
|
else:
|
|
video_list = torch.chunk(video_img, chunks=b)
|
|
|
|
print(len(video_list))
|
|
video_list = [tensor2cv(i) for i in video_list] if video_list else [] # tensor to np
|
|
|
|
if video_path != "none":
|
|
if front_path is not None:
|
|
path = front_path
|
|
else:
|
|
path = os.path.join(folder_paths.get_input_directory(), video_path)
|
|
else:
|
|
if front_path is not None:
|
|
path = front_path
|
|
else:
|
|
path = None
|
|
|
|
if video_list: # prior choice
|
|
path = None
|
|
|
|
if not video_list and video_path == "none" and not front_path:
|
|
raise "Need choice a video or link 'path or image' in the front!!!"
|
|
|
|
output_path = folder_paths.get_output_directory()
|
|
|
|
# infer
|
|
print("Start to video upscale processing...")
|
|
video_image, audio_form_v, fps = run_realesrgan(video_list, audio, frame_rate, fidelity_weight, path,
|
|
output_path,
|
|
has_aligned, only_center_face, draw_box, bg_upsampler,
|
|
save_video, net_g, face_upsampler, upscale, face_helper,
|
|
face_upsample, suffix="", )
|
|
if path is not None:
|
|
audio = audio_form_v
|
|
frame_rate = float(fps)
|
|
|
|
img_list = []
|
|
if isinstance(video_image, list):
|
|
for i in video_image:
|
|
for j in i:
|
|
img_list.append(j)
|
|
|
|
image = load_images(img_list)
|
|
return (image, audio, frame_rate,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Echo_LoadModel": Echo_LoadModel,
|
|
"Echo_Sampler": Echo_Sampler,
|
|
"Echo_Upscaleloader": Echo_Upscaleloader,
|
|
"Echo_VideoUpscale": Echo_VideoUpscale
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Echo_LoadModel": "Echo_LoadModel",
|
|
"Echo_Sampler": "Echo_Sampler",
|
|
"Echo_Upscaleloader": "Echo_Upscaleloader",
|
|
"Echo_VideoUpscale": "Echo_VideoUpscale"
|
|
}
|