multimodalart/EchoMimic-zero
8
1#!/usr/bin/env python2# -*- coding: UTF-8 -*-3'''4webui5'''6import spaces7import os8 9os.system('pip install scikit-image')10os.system('pip install IPython')11import random12from datetime import datetime13from pathlib import Path14 15import cv216import numpy as np17import torch18from diffusers import AutoencoderKL, DDIMScheduler19from omegaconf import OmegaConf20from PIL import Image21from src.models.unet_2d_condition import UNet2DConditionModel22from src.models.unet_3d_echo import EchoUNet3DConditionModel23from src.models.whisper.audio2feature import load_audio_model24from src.pipelines.pipeline_echo_mimic import Audio2VideoPipeline25from src.utils.util import save_videos_grid, crop_and_pad26from src.models.face_locator import FaceLocator27from moviepy.editor import VideoFileClip, AudioFileClip28from facenet_pytorch import MTCNN29import argparse30 31import gradio as gr32 33import huggingface_hub34 35import pickle36from src.utils.draw_utils import FaceMeshVisualizer37from src.utils.motion_utils import motion_sync38from src.utils.mp_utils import LMKExtractor39 40 41huggingface_hub.snapshot_download(42 repo_id='BadToBest/EchoMimic',43 local_dir='./pretrained_weights',44 local_dir_use_symlinks=False,45)46 47is_shared_ui = True if "fffiloni/EchoMimic" in os.environ['SPACE_ID'] else False48available_property = False if is_shared_ui else True49advanced_settings_label = "Advanced Configuration (only for duplicated spaces)" if is_shared_ui else "Advanced Configuration"50 51default_values = {52 "width": 512,53 "height": 512,54 "length": 1200,55 "seed": 420,56 "facemask_dilation_ratio": 0.1,57 "facecrop_dilation_ratio": 0.5,58 "context_frames": 12,59 "context_overlap": 3,60 "cfg": 2.5,61 "steps": 30,62 "sample_rate": 16000,63 "fps": 24,64 "device": "cuda"65}66 67ffmpeg_path = os.getenv('FFMPEG_PATH')68if ffmpeg_path is None:69 print("please download ffmpeg-static and export to FFMPEG_PATH. \nFor example: export FFMPEG_PATH=/musetalk/ffmpeg-4.4-amd64-static")70elif ffmpeg_path not in os.getenv('PATH'):71 print("add ffmpeg to path")72 os.environ["PATH"] = f"{ffmpeg_path}:{os.environ['PATH']}"73 74 75config_path = "./configs/prompts/animation.yaml"76config = OmegaConf.load(config_path)77if config.weight_dtype == "fp16":78 weight_dtype = torch.float1679else:80 weight_dtype = torch.float3281 82device = "cuda"83if not torch.cuda.is_available():84 device = "cpu"85 86inference_config_path = config.inference_config87infer_config = OmegaConf.load(inference_config_path)88 89############# model_init started #############90## vae init91vae = AutoencoderKL.from_pretrained(config.pretrained_vae_path).to("cuda", dtype=weight_dtype)92 93## reference net init94reference_unet = UNet2DConditionModel.from_pretrained(95 config.pretrained_base_model_path,96 subfolder="unet",97).to(dtype=weight_dtype, device=device)98reference_unet.load_state_dict(torch.load(config.reference_unet_path, map_location="cpu"))99 100## denoising net init101if os.path.exists(config.motion_module_path):102 ### stage1 + stage2103 denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(104 config.pretrained_base_model_path,105 config.motion_module_path,106 subfolder="unet",107 unet_additional_kwargs=infer_config.unet_additional_kwargs,108 ).to(dtype=weight_dtype, device=device)109else:110 ### only stage1111 denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(112 config.pretrained_base_model_path,113 "",114 subfolder="unet",115 unet_additional_kwargs={116 "use_motion_module": False,117 "unet_use_temporal_attention": False,118 "cross_attention_dim": infer_config.unet_additional_kwargs.cross_attention_dim119 }120 ).to(dtype=weight_dtype, device=device)121 122denoising_unet.load_state_dict(torch.load(config.denoising_unet_path, map_location="cpu"), strict=False)123 124## face locator init125face_locator = FaceLocator(320, conditioning_channels=1, block_out_channels=(16, 32, 96, 256)).to("cuda", dtype=weight_dtype)126face_locator.load_state_dict(torch.load(config.face_locator_path, map_location='cpu'))127 128## load audio processor params129audio_processor = load_audio_model(model_path=config.audio_model_path, device=device)130 131## load face detector params132face_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="cpu")133 134############# model_init finished #############135 136sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)137scheduler = DDIMScheduler(**sched_kwargs)138 139pipe = Audio2VideoPipeline(140 vae=vae,141 reference_unet=reference_unet,142 denoising_unet=denoising_unet,143 audio_guider=audio_processor,144 face_locator=face_locator,145 scheduler=scheduler,146).to("cuda", dtype=weight_dtype)147 148def select_face(det_bboxes, probs):149 ## max face from faces that the prob is above 0.8150 ## box: xyxy151 if det_bboxes is None or probs is None:152 return None153 filtered_bboxes = []154 for bbox_i in range(len(det_bboxes)):155 if probs[bbox_i] > 0.8:156 filtered_bboxes.append(det_bboxes[bbox_i])157 if len(filtered_bboxes) == 0:158 return None159 sorted_bboxes = sorted(filtered_bboxes, key=lambda x:(x[3]-x[1]) * (x[2] - x[0]), reverse=True)160 return sorted_bboxes[0]161 162lmk_extractor = LMKExtractor()163 164def face_detection(uploaded_img, facemask_dilation_ratio, facecrop_dilation_ratio, width, height):165 face_img = cv2.imread(uploaded_img)166 if face_img is None:167 raise gr.Error("input image should be uploaded or selected.")168 face_mask = np.zeros((face_img.shape[0], face_img.shape[1])).astype('uint8')169 det_bboxes, probs = face_detector.detect(face_img)170 select_bbox = select_face(det_bboxes, probs)171 if select_bbox is None:172 face_mask[:, :] = 255173 else:174 xyxy = select_bbox[:4]175 xyxy = np.round(xyxy).astype('int')176 rb, re, cb, ce = xyxy[1], xyxy[3], xyxy[0], xyxy[2]177 r_pad = int((re - rb) * facemask_dilation_ratio)178 c_pad = int((ce - cb) * facemask_dilation_ratio)179 face_mask[rb - r_pad : re + r_pad, cb - c_pad : ce + c_pad] = 255180 181 r_pad_crop = int((re - rb) * facecrop_dilation_ratio)182 c_pad_crop = int((ce - cb) * facecrop_dilation_ratio)183 crop_rect = [max(0, cb - c_pad_crop), max(0, rb - r_pad_crop), min(ce + c_pad_crop, face_img.shape[1]), min(re + r_pad_crop, face_img.shape[0])]184 face_img = crop_and_pad(face_img, crop_rect)185 face_mask = crop_and_pad(face_mask, crop_rect)186 face_img = cv2.resize(face_img, (width, height))187 face_mask = cv2.resize(face_mask, (width, height))188 189 print('face detect done.')190 return face_img, face_mask191 192@spaces.GPU(duration=200)193def video_pipe(face_img, face_mask, uploaded_audio, width, height, length, context_frames, context_overlap, cfg, steps, sample_rate, fps, device):194 face_mask_tensor = torch.Tensor(face_mask).to(dtype=weight_dtype, device="cuda").unsqueeze(0).unsqueeze(0).unsqueeze(0) / 255.0195 ref_image_pil = Image.fromarray(face_img[:, :, [2, 1, 0]])196 197 video = pipe(198 ref_image_pil,199 uploaded_audio,200 face_mask_tensor,201 width,202 height,203 length,204 steps,205 cfg,206 audio_sample_rate=sample_rate,207 context_frames=context_frames,208 fps=fps,209 context_overlap=context_overlap210 ).videos211 print('video pipe done.')212 213 save_dir = Path("output/tmp")214 save_dir.mkdir(exist_ok=True, parents=True)215 output_video_path = save_dir / "output_video.mp4"216 save_videos_grid(video, str(output_video_path), n_rows=1, fps=fps)217 218 video_clip = VideoFileClip(str(output_video_path))219 audio_clip = AudioFileClip(uploaded_audio)220 final_output_path = save_dir / "output_video_with_audio.mp4"221 video_clip = video_clip.set_audio(audio_clip)222 video_clip.write_videofile(str(final_output_path), codec="libx264", audio_codec="aac")223 224 return final_output_path225 226def process_video(uploaded_img, uploaded_audio, width, height, length, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, device):227 face_img, face_mask = face_detection(uploaded_img, facemask_dilation_ratio, facecrop_dilation_ratio, width, height)228 final_output_path = video_pipe(face_img, face_mask, uploaded_audio, width, height, length, context_frames, context_overlap, cfg, steps, sample_rate, fps, device)229 return final_output_path230 231 232# @spaces.GPU233# def process_video(uploaded_img, uploaded_audio, width, height, length, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, device):234# #### face musk prepare235# face_img = cv2.imread(uploaded_img)236# if face_img is None:237# raise gr.Error("input image should be uploaded or selected.")238# face_mask = np.zeros((face_img.shape[0], face_img.shape[1])).astype('uint8')239# det_bboxes, probs = face_detector.detect(face_img)240# select_bbox = select_face(det_bboxes, probs)241# if select_bbox is None:242# face_mask[:, :] = 255243# else:244# xyxy = select_bbox[:4]245# xyxy = np.round(xyxy).astype('int')246# rb, re, cb, ce = xyxy[1], xyxy[3], xyxy[0], xyxy[2]247# r_pad = int((re - rb) * facemask_dilation_ratio)248# c_pad = int((ce - cb) * facemask_dilation_ratio)249# face_mask[rb - r_pad : re + r_pad, cb - c_pad : ce + c_pad] = 255250 251# #### face crop252# r_pad_crop = int((re - rb) * facecrop_dilation_ratio)253# c_pad_crop = int((ce - cb) * facecrop_dilation_ratio)254# crop_rect = [max(0, cb - c_pad_crop), max(0, rb - r_pad_crop), min(ce + c_pad_crop, face_img.shape[1]), min(re + r_pad_crop, face_img.shape[0])]255# face_img = crop_and_pad(face_img, crop_rect)256# face_mask = crop_and_pad(face_mask, crop_rect)257# face_img = cv2.resize(face_img, (width, height))258# face_mask = cv2.resize(face_mask, (width, height))259# print('face detect done.')260# # ==================== face_locator =====================261# '''262# driver_video = "./assets/driven_videos/c.mp4"263 264# input_frames_cv2 = [cv2.resize(center_crop_cv2(pil_to_cv2(i)), (512, 512)) for i in pils_from_video(driver_video)]265# ref_det = lmk_extractor(face_img)266 267# visualizer = FaceMeshVisualizer(draw_iris=False, draw_mouse=False)268 269# pose_list = []270# sequence_driver_det = []271# try: 272# for frame in input_frames_cv2:273# result = lmk_extractor(frame)274# assert result is not None, "{}, bad video, face not detected".format(driver_video)275# sequence_driver_det.append(result)276# except:277# print("face detection failed")278# exit()279 280# sequence_det_ms = motion_sync(sequence_driver_det, ref_det)281# for p in sequence_det_ms:282# tgt_musk = visualizer.draw_landmarks((width, height), p)283# tgt_musk_pil = Image.fromarray(np.array(tgt_musk).astype(np.uint8)).convert('RGB')284# pose_list.append(torch.Tensor(np.array(tgt_musk_pil)).to(dtype=weight_dtype, device="cuda").permute(2,0,1) / 255.0)285# '''286# # face_mask_tensor = torch.stack(pose_list, dim=1).unsqueeze(0)287# face_mask_tensor = torch.Tensor(face_mask).to(dtype=weight_dtype, device="cuda").unsqueeze(0).unsqueeze(0).unsqueeze(0) / 255.0288 289# ref_image_pil = Image.fromarray(face_img[:, :, [2, 1, 0]])290 291# #del pose_list, sequence_det_ms, sequence_driver_det, input_frames_cv2292 293# video = pipe(294# ref_image_pil,295# uploaded_audio,296# face_mask_tensor,297# width,298# height,299# length,300# steps,301# cfg,302# #generator=generator,303# audio_sample_rate=sample_rate,304# context_frames=context_frames,305# fps=fps,306# context_overlap=context_overlap307# ).videos308# print('video pipe done.')309 310# save_dir = Path("output/tmp")311# save_dir.mkdir(exist_ok=True, parents=True)312# output_video_path = save_dir / "output_video.mp4"313# save_videos_grid(video, str(output_video_path), n_rows=1, fps=fps)314 315# video_clip = VideoFileClip(str(output_video_path))316# audio_clip = AudioFileClip(uploaded_audio)317# final_output_path = save_dir / "output_video_with_audio.mp4"318# video_clip = video_clip.set_audio(audio_clip)319# video_clip.write_videofile(str(final_output_path), codec="libx264", audio_codec="aac")320 321# return final_output_path322 323with gr.Blocks() as demo:324 gr.Markdown('# EchoMimic')325 gr.Markdown('## Lifelike Audio-Driven Portrait Animations through Editable Landmark Conditioning')326 gr.Markdown('Inference time: from ~7mins/240frames to ~50s/240frames on V100 GPU')327 gr.HTML("""328 <div style="display:flex;column-gap:4px;">329 <a href='https://badtobest.github.io/echomimic.html'><img src='https://img.shields.io/badge/Project-Page-blue'></a>330 <a href='https://huggingface.co/BadToBest/EchoMimic'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow'></a>331 <a href='https://arxiv.org/abs/2407.08136'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>332 </div>333 """)334 335 with gr.Row():336 with gr.Column(min_width=250):337 uploaded_img = gr.Image(type="filepath", label="Reference Image")338 with gr.Column(min_width=250):339 uploaded_audio = gr.Audio(type="filepath", label="Input Audio")340 with gr.Accordion(label=advanced_settings_label, open=False):341 with gr.Row():342 width = gr.Slider(label="Width", minimum=128, maximum=1024, value=default_values["width"], interactive=available_property)343 height = gr.Slider(label="Height", minimum=128, maximum=1024, value=default_values["height"], interactive=available_property)344 with gr.Row():345 length = gr.Slider(label="Length", minimum=100, maximum=5000, value=default_values["length"], interactive=available_property)346 seed = gr.Slider(label="Seed", minimum=0, maximum=10000, value=default_values["seed"], interactive=available_property)347 with gr.Row():348 facemask_dilation_ratio = gr.Slider(label="Facemask Dilation Ratio", minimum=0.0, maximum=1.0, step=0.01, value=default_values["facemask_dilation_ratio"], interactive=available_property)349 facecrop_dilation_ratio = gr.Slider(label="Facecrop Dilation Ratio", minimum=0.0, maximum=1.0, step=0.01, value=default_values["facecrop_dilation_ratio"], interactive=available_property)350 with gr.Row():351 context_frames = gr.Slider(label="Context Frames", minimum=0, maximum=50, step=1, value=default_values["context_frames"], interactive=available_property)352 context_overlap = gr.Slider(label="Context Overlap", minimum=0, maximum=10, step=1, value=default_values["context_overlap"], interactive=available_property)353 with gr.Row():354 cfg = gr.Slider(label="CFG", minimum=0.0, maximum=10.0, step=0.1, value=default_values["cfg"], interactive=available_property)355 steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=default_values["steps"], interactive=available_property)356 with gr.Row():357 sample_rate = gr.Slider(label="Sample Rate", minimum=8000, maximum=48000, step=1000, value=default_values["sample_rate"], interactive=available_property)358 fps = gr.Slider(label="FPS", minimum=1, maximum=60, step=1, value=default_values["fps"], interactive=available_property)359 device = gr.Radio(label="Device", choices=["cuda", "cpu"], value=default_values["device"], interactive=available_property)360 361 with gr.Column(min_width=250):362 generate_button = gr.Button("Generate Video")363 output_video = gr.Video()364 with gr.Row():365 366 gr.Examples(367 label = "Portrait examples",368 examples = [369 ['assets/test_imgs/a.png'],370 ['assets/test_imgs/b.png'],371 ['assets/test_imgs/c.png'],372 ['assets/test_imgs/d.png'],373 ['assets/test_imgs/e.png']374 ],375 inputs = [uploaded_img]376 )377 gr.Examples(378 label = "Audio examples",379 examples = [380 ['assets/test_audios/chunnuanhuakai.wav'],381 ['assets/test_audios/chunwang.wav'],382 ['assets/test_audios/echomimic_en_girl.wav'],383 ['assets/test_audios/echomimic_en.wav'],384 ['assets/test_audios/echomimic_girl.wav'],385 ['assets/test_audios/echomimic.wav'],386 ['assets/test_audios/jane.wav'],387 ['assets/test_audios/mei.wav'],388 ['assets/test_audios/walden.wav'],389 ['assets/test_audios/yun.wav'],390 ],391 inputs = [uploaded_audio]392 )393 # gr.HTML("""394 # <div style="display:flex;column-gap:4px;">395 # <a href="https://huggingface.co/spaces/fffiloni/EchoMimic?duplicate=true">396 # <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-xl.svg" alt="Duplicate this Space">397 # </a>398 # <a href="https://huggingface.co/fffiloni">399 # <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-xl-dark.svg" alt="Follow me on HF">400 # </a>401 # </div>402 # """)403 404 # def generate_video(uploaded_img, uploaded_audio, facemask_dilation_ratio=default_values["facemask_dilation_ratio"],405 # facecrop_dilation_ratio=default_values["facecrop_dilation_ratio"],406 # context_frames=default_values["context_frames"],407 # context_overlap=default_values["context_overlap"],408 # cfg=default_values["cfg"],409 # steps=default_values["steps"],410 # sample_rate=default_values["sample_rate"],411 # fps=default_values["fps"],412 # device=default_values["device"],413 # width=default_values["width"],414 # height=default_values["height"],415 # length=default_values["length"] ):416 417 # final_output_path = process_video(418 # uploaded_img, uploaded_audio, width, height, length, seed, facemask_dilation_ratio, facecrop_dilation_ratio, context_frames, context_overlap, cfg, steps, sample_rate, fps, device419 # ) 420 # output_video= final_output_path421 # return final_output_path422 423 # generate_button.click(424 # generate_video,425 # inputs=[426 # uploaded_img,427 # uploaded_audio,428 # # width,429 # # height,430 # # length,431 # # seed,432 # # facemask_dilation_ratio,433 # # facecrop_dilation_ratio,434 # # context_frames,435 # # context_overlap,436 # # cfg,437 # # steps,438 # # sample_rate,439 # # fps,440 # # device441 # ],442 # outputs=output_video,443 # show_api=False444 # )445 def generate_video(uploaded_img, uploaded_audio,446 facemask_dilation_ratio=default_values["facemask_dilation_ratio"],447 facecrop_dilation_ratio=default_values["facecrop_dilation_ratio"],448 context_frames=default_values["context_frames"],449 context_overlap=default_values["context_overlap"],450 cfg=default_values["cfg"],451 steps=default_values["steps"],452 sample_rate=default_values["sample_rate"],453 fps=default_values["fps"],454 device=default_values["device"],455 width=default_values["width"],456 height=default_values["height"],457 length=default_values["length"] ):458 459 final_output_path = process_video(460 uploaded_img, 461 uploaded_audio, width, height, 462 length, facemask_dilation_ratio, 463 facecrop_dilation_ratio, context_frames, 464 context_overlap, cfg, steps, 465 sample_rate, fps, device466 ) 467 output_video = final_output_path468 return final_output_path469 470 generate_button.click(471 generate_video,472 inputs=[473 uploaded_img,474 uploaded_audio475 ],476 outputs=output_video,477 show_progress=True478 )479parser = argparse.ArgumentParser(description='EchoMimic')480parser.add_argument('--server_name', type=str, default='0.0.0.0', help='Server name')481parser.add_argument('--server_port', type=int, default=7680, help='Server port')482args = parser.parse_args()483 484# demo.launch(server_name=args.server_name, server_port=args.server_port, inbrowser=True)485 486if __name__ == '__main__':487 demo.queue(max_size=3).launch(show_api=False, show_error=True)488 #demo.launch(server_name=args.server_name, server_port=args.server_port, inbrowser=True)