import os import random from pathlib import Path import numpy as np import torch from diffusers import AutoencoderKL, DDIMScheduler from PIL import Image from src.models.unet_2d_condition import UNet2DConditionModel from src.models.unet_3d_emo import EMOUNet3DConditionModel from src.models.whisper.audio2feature import load_audio_model from src.pipelines.pipeline_echomimicv2 import EchoMimicV2Pipeline from src.utils.util import save_videos_grid from src.models.pose_encoder import PoseEncoder from src.utils.dwpose_util import draw_pose_select_v2 from moviepy.editor import VideoFileClip, AudioFileClip import gradio as gr from datetime import datetime from torchao.quantization import quantize_, int8_weight_only import gc total_vram_in_gb = torch.cuda.get_device_properties(0).total_memory / 1073741824 print(f'\033[32mCUDA版本:{torch.version.cuda}\033[0m') print(f'\033[32mPytorch版本:{torch.__version__}\033[0m') print(f'\033[32m显卡型号:{torch.cuda.get_device_name()}\033[0m') print(f'\033[32m显存大小:{total_vram_in_gb:.2f}GB\033[0m') print(f'\033[32m精度:float16\033[0m') dtype = torch.float16 if torch.cuda.is_available(): device = "cuda" else: print("cuda not available, using cpu") device = "cpu" ffmpeg_path = os.getenv('FFMPEG_PATH') if ffmpeg_path is None: print("please download ffmpeg-static and export to FFMPEG_PATH. \nFor example: export FFMPEG_PATH=./ffmpeg-4.4-amd64-static") elif ffmpeg_path not in os.getenv('PATH'): print("add ffmpeg to path") os.environ["PATH"] = f"{ffmpeg_path}:{os.environ['PATH']}" def generate(image_input, audio_input, pose_input, width, height, length, steps, sample_rate, cfg, fps, context_frames, context_overlap, quantization_input, seed): gc.collect() torch.cuda.empty_cache() torch.cuda.ipc_collect() timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") save_dir = Path("outputs") save_dir.mkdir(exist_ok=True, parents=True) ############# model_init started ############# ## vae init vae = AutoencoderKL.from_pretrained("./pretrained_weights/sd-vae-ft-mse").to(device, dtype=dtype) if quantization_input: quantize_(vae, int8_weight_only()) print("使用int8量化") ## reference net init reference_unet = UNet2DConditionModel.from_pretrained("./pretrained_weights/sd-image-variations-diffusers", subfolder="unet", use_safetensors=False).to(dtype=dtype, device=device) reference_unet.load_state_dict(torch.load("./pretrained_weights/reference_unet.pth", weights_only=True)) if quantization_input: quantize_(reference_unet, int8_weight_only()) ## denoising net init if os.path.exists("./pretrained_weights/motion_module.pth"): print('using motion module') else: exit("motion module not found") ### stage1 + stage2 denoising_unet = EMOUNet3DConditionModel.from_pretrained_2d( "./pretrained_weights/sd-image-variations-diffusers", "./pretrained_weights/motion_module.pth", subfolder="unet", unet_additional_kwargs = { "use_inflated_groupnorm": True, "unet_use_cross_frame_attention": False, "unet_use_temporal_attention": False, "use_motion_module": True, "cross_attention_dim": 384, "motion_module_resolutions": [ 1, 2, 4, 8 ], "motion_module_mid_block": True , "motion_module_decoder_only": False, "motion_module_type": "Vanilla", "motion_module_kwargs":{ "num_attention_heads": 8, "num_transformer_block": 1, "attention_block_types": [ 'Temporal_Self', 'Temporal_Self' ], "temporal_position_encoding": True, "temporal_position_encoding_max_len": 32, "temporal_attention_dim_div": 1, } }, ).to(dtype=dtype, device=device) denoising_unet.load_state_dict(torch.load("./pretrained_weights/denoising_unet.pth", weights_only=True),strict=False) # pose net init pose_net = PoseEncoder(320, conditioning_channels=3, block_out_channels=(16, 32, 96, 256)).to(dtype=dtype, device=device) pose_net.load_state_dict(torch.load("./pretrained_weights/pose_encoder.pth", weights_only=True)) ### load audio processor params audio_processor = load_audio_model(model_path="./pretrained_weights/audio_processor/tiny.pt", device=device) ############# model_init finished ############# sched_kwargs = { "beta_start": 0.00085, "beta_end": 0.012, "beta_schedule": "linear", "clip_sample": False, "steps_offset": 1, "prediction_type": "v_prediction", "rescale_betas_zero_snr": True, "timestep_spacing": "trailing" } scheduler = DDIMScheduler(**sched_kwargs) pipe = EchoMimicV2Pipeline( vae=vae, reference_unet=reference_unet, denoising_unet=denoising_unet, audio_guider=audio_processor, pose_encoder=pose_net, scheduler=scheduler, ) pipe = pipe.to(device, dtype=dtype) if seed is not None and seed > -1: generator = torch.manual_seed(seed) else: seed = random.randint(100, 1000000) generator = torch.manual_seed(seed) inputs_dict = { "refimg": image_input, "audio": audio_input, "pose": pose_input, } print('Pose:', inputs_dict['pose']) print('Reference:', inputs_dict['refimg']) print('Audio:', inputs_dict['audio']) save_name = f"{save_dir}/{timestamp}" ref_image_pil = Image.open(inputs_dict['refimg']).resize((width, height)) audio_clip = AudioFileClip(inputs_dict['audio']) length = min(length, int(audio_clip.duration * fps), len(os.listdir(inputs_dict['pose']))) start_idx = 0 pose_list = [] for index in range(start_idx, start_idx + length): tgt_musk = np.zeros((width, height, 3)).astype('uint8') tgt_musk_path = os.path.join(inputs_dict['pose'], "{}.npy".format(index)) detected_pose = np.load(tgt_musk_path, allow_pickle=True).tolist() imh_new, imw_new, rb, re, cb, ce = detected_pose['draw_pose_params'] im = draw_pose_select_v2(detected_pose, imh_new, imw_new, ref_w=800) im = np.transpose(np.array(im),(1, 2, 0)) tgt_musk[rb:re,cb:ce,:] = im tgt_musk_pil = Image.fromarray(np.array(tgt_musk)).convert('RGB') pose_list.append(torch.Tensor(np.array(tgt_musk_pil)).to(dtype=dtype, device=device).permute(2,0,1) / 255.0) poses_tensor = torch.stack(pose_list, dim=1).unsqueeze(0) audio_clip = AudioFileClip(inputs_dict['audio']) audio_clip = audio_clip.set_duration(length / fps) video = pipe( ref_image_pil, inputs_dict['audio'], poses_tensor[:,:,:length,...], width, height, length, steps, cfg, generator=generator, audio_sample_rate=sample_rate, context_frames=context_frames, fps=fps, context_overlap=context_overlap, start_idx=start_idx, ).videos final_length = min(video.shape[2], poses_tensor.shape[2], length) video_sig = video[:, :, :final_length, :, :] save_videos_grid( video_sig, save_name + "_woa_sig.mp4", n_rows=1, fps=fps, ) video_clip_sig = VideoFileClip(save_name + "_woa_sig.mp4",) video_clip_sig = video_clip_sig.set_audio(audio_clip) video_clip_sig.write_videofile(save_name + "_sig.mp4", codec="libx264", audio_codec="aac", threads=2) video_output = save_name + "_sig.mp4" seed_text = gr.update(visible=True, value=seed) return video_output, seed_text with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("""

EchoMimicV2

🌐 Github | 📜 arXiv
⚠️ 该演示仅供学术研究和体验使用。
""") with gr.Column(): with gr.Row(): with gr.Column(): with gr.Group(): image_input = gr.Image(label="图像输入(自动缩放)", type="filepath") audio_input = gr.Audio(label="音频输入", type="filepath") pose_input = gr.Textbox(label="姿态输入(目录地址)", placeholder="请输入姿态数据的目录地址", value="assets/halfbody_demo/pose/01") with gr.Group(): with gr.Row(): width = gr.Number(label="宽度(16的倍数,推荐768)", value=768) height = gr.Number(label="高度(16的倍数,推荐768)", value=768) length = gr.Number(label="视频长度,推荐240)", value=240) with gr.Row(): steps = gr.Number(label="步骤(推荐30)", value=20) sample_rate = gr.Number(label="采样率(推荐16000)", value=16000) cfg = gr.Number(label="cfg(推荐2.5)", value=2.5, step=0.1) with gr.Row(): fps = gr.Number(label="帧率(推荐24)", value=24) context_frames = gr.Number(label="上下文框架(推荐12)", value=12) context_overlap = gr.Number(label="上下文重叠(推荐3)", value=3) with gr.Row(): quantization_input = gr.Checkbox(label="int8量化(推荐显存12G的用户开启,并使用不超过5秒的音频)", value=False) seed = gr.Number(label="种子(-1为随机)", value=-1) generate_button = gr.Button("🎬 生成视频") with gr.Column(): video_output = gr.Video(label="输出视频") seed_text = gr.Textbox(label="种子", interactive=False, visible=False) gr.Examples( examples=[ ["EMTD_dataset/ref_imgs_by_FLUX/man/0001.png", "assets/halfbody_demo/audio/chinese/echomimicv2_man.wav"], ["EMTD_dataset/ref_imgs_by_FLUX/woman/0077.png", "assets/halfbody_demo/audio/chinese/echomimicv2_woman.wav"], ["EMTD_dataset/ref_imgs_by_FLUX/man/0003.png", "assets/halfbody_demo/audio/chinese/fighting.wav"], ["EMTD_dataset/ref_imgs_by_FLUX/woman/0033.png", "assets/halfbody_demo/audio/chinese/good.wav"], ["EMTD_dataset/ref_imgs_by_FLUX/man/0010.png", "assets/halfbody_demo/audio/chinese/news.wav"], ["EMTD_dataset/ref_imgs_by_FLUX/man/1168.png", "assets/halfbody_demo/audio/chinese/no_smoking.wav"], ["EMTD_dataset/ref_imgs_by_FLUX/woman/0057.png", "assets/halfbody_demo/audio/chinese/ultraman.wav"] ], inputs=[image_input, audio_input], label="预设人物及音频", ) generate_button.click( generate, inputs=[image_input, audio_input, pose_input, width, height, length, steps, sample_rate, cfg, fps, context_frames, context_overlap, quantization_input, seed], outputs=[video_output, seed_text], ) if __name__ == "__main__": demo.queue() demo.launch(inbrowser=True)