# !/usr/bin/env python # -*- coding: UTF-8 -*- import io import os import random import numpy as np import torch import torchaudio import gc import platform import subprocess from .utils import find_directories,process_video, cf_tensor2cv,process_video_v2,load_images,tensor_to_pil,nomarl_upscale from .origin_infer import Echo_v1_load_model,Echo_v2_load_model,Echo_v1_predata,Echo_v2_predata from .echomimic_v3.infer import load_v3_model,infer_v3,Config,Echo_v3_predata import folder_paths 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__)) # 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_unet_current_path = os.path.join(weigths_current_path, "unet") if not os.path.exists(weigths_unet_current_path): os.makedirs(weigths_unet_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) weigths_trans_current_path = os.path.join(weigths_current_path, "transformer") if not os.path.exists(weigths_trans_current_path): os.makedirs(weigths_trans_current_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']}" # *****************main*************** class Echo_LoadModel: def __init__(self): pass @classmethod def INPUT_TYPES(cls): return { "required": { "vae": (folder_paths.get_filename_list("vae"),), "denoising": ("BOOLEAN", {"default": True},), "infer_mode": (["audio_drived", "audio_drived_acc", "pose_normal_dwpose","pose_normal_sapiens", "pose_acc"],), "lowvram": ("BOOLEAN", {"default": True},), "teacache_offload": ("BOOLEAN", {"default": True},), "use_mmgp": (["LowRAM_LowVRAM","None", "VerylowRAM_LowVRAM","LowRAM_HighVRAM","HighRAM_LowVRAM","HighRAM_HighVRAM" ],), "version": (["V3","V2", "V1", ],), }, } RETURN_TYPES = ("MODEL_PIPE_E", "MODEL_INFO_E") RETURN_NAMES = ("model", "info") FUNCTION = "main_loader" CATEGORY = "EchoMimic" def main_loader(self, vae, denoising, infer_mode, lowvram,teacache_offload,use_mmgp, version): config_v3,audio_pt,face_locator_pt,pose_encoder_pt,tokenizer,temporal_compression_ratio=None,None,None,None,None,None if "V1"==version : model, audio_pt, face_locator_pt= Echo_v1_load_model(vae,weigths_current_path,version, infer_mode, denoising, current_path, device,lowvram) elif "V2"==version : model,pose_encoder_pt, audio_pt= Echo_v2_load_model(vae,weigths_current_path, version, infer_mode,current_path,device,lowvram) else:#v3 config_v3=Config() config_v3.config_path=os.path.join(current_path, "echomimic_v3/config/config.yaml") config_v3.teacache_offload=teacache_offload model, temporal_compression_ratio,tokenizer=load_v3_model(current_path,weigths_current_path,config_v3, device,use_mmgp,vae) print("model loaded") info = {"lowvram": lowvram,"version":version,"tokenizer":tokenizer, "infer_mode":infer_mode,"audio_pt":audio_pt, "face_locator_pt":face_locator_pt,"pose_encoder_pt":pose_encoder_pt,"temporal_compression_ratio":temporal_compression_ratio} info.update({"config":config_v3}) return (model,info) class Echo_Predata: def __init__(self): pass @classmethod def INPUT_TYPES(cls): pose_path_list = ["pose_01","pose_02","pose_03","pose_04","pose_fight","pose_good","pose_salute","pose_ultraman"] + find_directories(tensorrt_lite) if find_directories(tensorrt_lite) else ["pose_01","pose_02","pose_03","pose_04","pose_fight","pose_good","pose_salute","pose_ultraman"] return { "required": { "info": ("MODEL_INFO_E",), "image": ("IMAGE",), # [B,H,W,C], C=3 "audio": ("AUDIO",), "prompt": ("STRING", {"multiline": True,"default":" Gesture, Body, Face Expressions and Movements: Hands and Fingers Movements: The person's hands are not visible in the image, so there are no specific hand or finger movements to describe. Body Positions and Posture: The person is standing upright with a straight posture. Body Coverage in Frame: The person is fully visible from the waist up. Face Expressions Change: The person appears to have a neutral expression with a slight smile. Eyes Movement: The eyes are looking directly at the camera. Head Movement: The head is slightly tilted forward. Overall Description: The character is standing in a front-facing shot, wearing a pink knitted vest over a white collared shirt and a white pleated skirt. The background appears to be a studio setting with soft lighting and some blurred elements that suggest a modern, clean environment. The person is adorned with pearl earrings, adding a touch of elegance to their appearance. The overall impression is one of a professional or casual presentation, possibly for a broadcast or a photoshoot."}), "negative_prompt" :("STRING", {"multiline": True,"default":" Gesture is bad. Gesture is unclear. Strange and twisted hands. Bad hands. Bad fingers. Unclear and blurry hands. 手部快速摆动, 手指频繁抽搐, 夸张手势, 重复机械性动作." }), "pose_dir": (pose_path_list,), "width": ("INT", {"default": 512, "min": 128, "max": 1024, "step": 64, "display": "number"}), "height": ("INT", {"default": 512, "min": 128, "max": 1024, "step": 64, "display": "number"}), "fps": ("FLOAT", {"default": 25.0, "min": 5.0, "max": 120.0}), "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}), "length": ("INT", {"default": 100, "min": 12, "max": 5000, "step": 1, "display": "number"}), "partial_video_length": ([33,65,97,113, 129,193],), "draw_mouse": ("BOOLEAN", {"default": False},), "motion_sync_": ("BOOLEAN", {"default": False},), }, "optional": { "clip": ("CLIP",), "clip_vision": ("CLIP_VISION",), "video_images": ("IMAGE",), # [B,H,W,C], C=3,B>1 } } RETURN_TYPES = ("MODEL_EMB_E", ) RETURN_NAMES = ("emb",) FUNCTION = "main_loader" CATEGORY = "EchoMimic" def main_loader(self, info, image,audio,prompt,negative_prompt,pose_dir,width,height,fps,facemask_ratio,facecrop_ratio,length,partial_video_length,draw_mouse,motion_sync_,**kwargs): version=info.get("version", "V3") lowvram=info.get("lowvram", False) text_encoder=kwargs.get("clip", None) image_encoder=kwargs.get("clip_vision", None) video_images = kwargs.get("video_images") ip_mask_path="" # pre img if version=="V1" or version=="V2": face_img = cf_tensor2cv(image, width, height) if version=="V1" else image # v1 cv ,v2 tensor else: face_img = nomarl_upscale(image, width, height) #pre audio 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()) # pre data if "V1"==version : emb = Echo_v1_predata(face_img,audio_file,fps,info.get("audio_pt"),info.get("face_locator_pt"),device,info.get("infer_mode"),draw_mouse,motion_sync_,lowvram, width,height,facemask_ratio,facecrop_ratio,video_images,audio_file_prefix,current_path,tensorrt_lite,length,pose_dir) elif "V2"==version : emb = Echo_v2_predata(face_img,audio_file,height,width,info.get("pose_encoder_pt"),info.get("audio_pt"), current_path,video_images,tensorrt_lite,device,fps,length,info.get("infer_mode"),weigths_current_path,pose_dir) else:#v3 config=info.get("config") config.fps=fps config.partial_video_length=int(partial_video_length) emb= Echo_v3_predata(image_encoder,text_encoder,info.get("tokenizer"),face_img,audio_file,ip_mask_path,config,device,info.get("temporal_compression_ratio"),prompt,negative_prompt,weigths_current_path,current_path,length) emb["config"]=config emb.update(info) emb.update({"audio_file_prefix":audio_file_prefix,"fps":fps,"height":height,"width":width}) return (emb,) class Echo_Sampler: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL_PIPE_E",), "emb": ("MODEL_EMB_E",), "seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}), "cfg": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}), "steps": ("INT", {"default": 25, "min": 1, "max": 100}), "sample_rate": ("INT", {"default": 16000, "min": 8000, "max": 48000, "step": 1000, }), "context_frames": ("INT", {"default": 12, "min": 0, "max": 50}), "context_overlap": ("INT", {"default": 3, "min": 0, "max": 10}), "save_video": ("BOOLEAN", {"default": False},), }, } RETURN_TYPES = ("IMAGE", "FLOAT") RETURN_NAMES = ("image", "frame_rate") FUNCTION = "em_main" CATEGORY = "EchoMimic" def em_main(self,model, emb, seed, cfg, steps, sample_rate,context_frames, context_overlap,save_video,): version= emb.get("version") lowvram = emb.get("lowvram") if version=="V1": model.to(device) output_video = process_video(emb.get("ref_image_pil"), emb.get("audio_path"), emb.get("width"), emb.get("height"), emb.get("length"), seed, emb.get("face_locator_tensor"), context_frames, context_overlap, cfg, steps, sample_rate, emb.get("fps"), model,save_video,emb.get("mask_len"), emb.get("audio_file_prefix"),emb.get("whisper_chunks")) elif version=="V2": model.to(device) output_video=process_video_v2(emb.get("infer_image_pil"),emb.get("ref_image_pil"), emb.get("audio_path"),emb.get("face_locator_tensor"), emb.get("W_change"), emb.get("H_change"), emb.get("start_idx"), seed, context_frames, context_overlap, cfg, steps, sample_rate, emb.get("fps"), model,emb.get("mask_len"),emb.get("height"),emb.get("width"), save_video, emb.get("audio_file_prefix"),emb.get("LEN"),emb.get("whisper_chunks") ) else: #v3 config = emb.get("config") config.num_inference_steps=steps config.guidance_scale=cfg #config.sample_size=(emb.get("height"),emb.get("width")) output_video=infer_v3(model, config, device,emb.get("video_length"),emb.get("prompt_embeds"),emb.get("negative_prompt_embeds"), emb.get("temporal_compression_ratio"),seed,emb.get("partial_video_length"),emb.get("audio_embeds"),emb.get("ip_mask"), emb.get("sample_height"),emb.get("sample_width"),emb.get("clip_context"), emb.get("latent_frames"),emb.get("ref_image_pil"),emb.get("audio_file_prefix")) frame_rate = float(emb.get("fps")) if not lowvram and version!="V3": # for upsacle ,need VR model.to("cpu") gc.collect() torch.cuda.empty_cache() return (load_images(output_video) , frame_rate) NODE_CLASS_MAPPINGS = { "Echo_LoadModel": Echo_LoadModel, "Echo_Predata":Echo_Predata, "Echo_Sampler": Echo_Sampler, } NODE_DISPLAY_NAME_MAPPINGS = { "Echo_LoadModel": "Echo_LoadModel", "Echo_Predata": "Echo_Predata", "Echo_Sampler": "Echo_Sampler", }