forked from TMElyralab/MuseTalk
-
Notifications
You must be signed in to change notification settings - Fork 0
/
app.py
426 lines (345 loc) · 15.8 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
import os
import time
import pdb
import re
import gradio as gr
import spaces
import numpy as np
import sys
import subprocess
from huggingface_hub import snapshot_download
import requests
import argparse
import os
from omegaconf import OmegaConf
import numpy as np
import cv2
import torch
import glob
import pickle
from tqdm import tqdm
import copy
from argparse import Namespace
import shutil
import gdown
import imageio
import ffmpeg
from moviepy.editor import *
ProjectDir = os.path.abspath(os.path.dirname(__file__))
CheckpointsDir = os.path.join(ProjectDir, "models")
def print_directory_contents(path):
for child in os.listdir(path):
child_path = os.path.join(path, child)
if os.path.isdir(child_path):
print(child_path)
def download_model():
if not os.path.exists(CheckpointsDir):
os.makedirs(CheckpointsDir)
print("Checkpoint Not Downloaded, start downloading...")
tic = time.time()
snapshot_download(
repo_id="TMElyralab/MuseTalk",
local_dir=CheckpointsDir,
max_workers=8,
local_dir_use_symlinks=True,
force_download=True, resume_download=False
)
# weight
os.makedirs(f"{CheckpointsDir}/sd-vae-ft-mse/")
snapshot_download(
repo_id="stabilityai/sd-vae-ft-mse",
local_dir=CheckpointsDir+'/sd-vae-ft-mse',
max_workers=8,
local_dir_use_symlinks=True,
force_download=True, resume_download=False
)
#dwpose
os.makedirs(f"{CheckpointsDir}/dwpose/")
snapshot_download(
repo_id="yzd-v/DWPose",
local_dir=CheckpointsDir+'/dwpose',
max_workers=8,
local_dir_use_symlinks=True,
force_download=True, resume_download=False
)
#vae
url = "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
response = requests.get(url)
# 确保请求成功
if response.status_code == 200:
# 指定文件保存的位置
file_path = f"{CheckpointsDir}/whisper/tiny.pt"
os.makedirs(f"{CheckpointsDir}/whisper/")
# 将文件内容写入指定位置
with open(file_path, "wb") as f:
f.write(response.content)
else:
print(f"请求失败，状态码：{response.status_code}")
#gdown face parse
url = "https://drive.google.com/uc?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812"
os.makedirs(f"{CheckpointsDir}/face-parse-bisent/")
file_path = f"{CheckpointsDir}/face-parse-bisent/79999_iter.pth"
gdown.download(url, file_path, quiet=False)
#resnet
url = "https://download.pytorch.org/models/resnet18-5c106cde.pth"
response = requests.get(url)
# 确保请求成功
if response.status_code == 200:
# 指定文件保存的位置
file_path = f"{CheckpointsDir}/face-parse-bisent/resnet18-5c106cde.pth"
# 将文件内容写入指定位置
with open(file_path, "wb") as f:
f.write(response.content)
else:
print(f"请求失败，状态码：{response.status_code}")
toc = time.time()
print(f"download cost {toc-tic} seconds")
print_directory_contents(CheckpointsDir)
else:
print("Already download the model.")
download_model() # for huggingface deployment.
from musetalk.utils.utils import get_file_type,get_video_fps,datagen
from musetalk.utils.preprocessing import get_landmark_and_bbox,read_imgs,coord_placeholder,get_bbox_range
from musetalk.utils.blending import get_image
from musetalk.utils.utils import load_all_model
@spaces.GPU(duration=600)
@torch.no_grad()
def inference(audio_path,video_path,bbox_shift,progress=gr.Progress(track_tqdm=True)):
args_dict={"result_dir":'./results/output', "fps":25, "batch_size":8, "output_vid_name":'', "use_saved_coord":False}#same with inferenece script
args = Namespace(**args_dict)
input_basename = os.path.basename(video_path).split('.')[0]
audio_basename = os.path.basename(audio_path).split('.')[0]
output_basename = f"{input_basename}_{audio_basename}"
result_img_save_path = os.path.join(args.result_dir, output_basename) # related to video & audio inputs
crop_coord_save_path = os.path.join(result_img_save_path, input_basename+".pkl") # only related to video input
os.makedirs(result_img_save_path,exist_ok =True)
if args.output_vid_name=="":
output_vid_name = os.path.join(args.result_dir, output_basename+".mp4")
else:
output_vid_name = os.path.join(args.result_dir, args.output_vid_name)
############################################## extract frames from source video ##############################################
if get_file_type(video_path)=="video":
save_dir_full = os.path.join(args.result_dir, input_basename)
os.makedirs(save_dir_full,exist_ok = True)
# cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
# os.system(cmd)
# 读取视频
reader = imageio.get_reader(video_path)
# 保存图片
for i, im in enumerate(reader):
imageio.imwrite(f"{save_dir_full}/{i:08d}.png", im)
input_img_list = sorted(glob.glob(os.path.join(save_dir_full, '*.[jpJP][pnPN]*[gG]')))
fps = get_video_fps(video_path)
else: # input img folder
input_img_list = glob.glob(os.path.join(video_path, '*.[jpJP][pnPN]*[gG]'))
input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
fps = args.fps
#print(input_img_list)
############################################## extract audio feature ##############################################
whisper_feature = audio_processor.audio2feat(audio_path)
whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature,fps=fps)
############################################## preprocess input image ##############################################
if os.path.exists(crop_coord_save_path) and args.use_saved_coord:
print("using extracted coordinates")
with open(crop_coord_save_path,'rb') as f:
coord_list = pickle.load(f)
frame_list = read_imgs(input_img_list)
else:
print("extracting landmarks...time consuming")
coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift)
with open(crop_coord_save_path, 'wb') as f:
pickle.dump(coord_list, f)
bbox_shift_text=get_bbox_range(input_img_list, bbox_shift)
i = 0
input_latent_list = []
for bbox, frame in zip(coord_list, frame_list):
if bbox == coord_placeholder:
continue
x1, y1, x2, y2 = bbox
crop_frame = frame[y1:y2, x1:x2]
crop_frame = cv2.resize(crop_frame,(256,256),interpolation = cv2.INTER_LANCZOS4)
latents = vae.get_latents_for_unet(crop_frame)
input_latent_list.append(latents)
# to smooth the first and the last frame
frame_list_cycle = frame_list + frame_list[::-1]
coord_list_cycle = coord_list + coord_list[::-1]
input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
############################################## inference batch by batch ##############################################
print("start inference")
video_num = len(whisper_chunks)
batch_size = args.batch_size
gen = datagen(whisper_chunks,input_latent_list_cycle,batch_size)
res_frame_list = []
for i, (whisper_batch,latent_batch) in enumerate(tqdm(gen,total=int(np.ceil(float(video_num)/batch_size)))):
tensor_list = [torch.FloatTensor(arr) for arr in whisper_batch]
audio_feature_batch = torch.stack(tensor_list).to(unet.device) # torch, B, 5*N,384
audio_feature_batch = pe(audio_feature_batch)
pred_latents = unet.model(latent_batch, timesteps, encoder_hidden_states=audio_feature_batch).sample
recon = vae.decode_latents(pred_latents)
for res_frame in recon:
res_frame_list.append(res_frame)
############################################## pad to full image ##############################################
print("pad talking image to original video")
for i, res_frame in enumerate(tqdm(res_frame_list)):
bbox = coord_list_cycle[i%(len(coord_list_cycle))]
ori_frame = copy.deepcopy(frame_list_cycle[i%(len(frame_list_cycle))])
x1, y1, x2, y2 = bbox
try:
res_frame = cv2.resize(res_frame.astype(np.uint8),(x2-x1,y2-y1))
except:
# print(bbox)
continue
combine_frame = get_image(ori_frame,res_frame,bbox)
cv2.imwrite(f"{result_img_save_path}/{str(i).zfill(8)}.png",combine_frame)
# cmd_img2video = f"ffmpeg -y -v fatal -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=rgb24,scale=out_color_matrix=bt709,format=yuv420p temp.mp4"
# print(cmd_img2video)
# os.system(cmd_img2video)
# 帧率
fps = 25
# 图片路径
# 输出视频路径
output_video = 'temp.mp4'
# 读取图片
def is_valid_image(file):
pattern = re.compile(r'\d{8}\.png')
return pattern.match(file)
images = []
files = [file for file in os.listdir(result_img_save_path) if is_valid_image(file)]
files.sort(key=lambda x: int(x.split('.')[0]))
for file in files:
filename = os.path.join(result_img_save_path, file)
images.append(imageio.imread(filename))
# 保存视频
imageio.mimwrite(output_video, images, 'FFMPEG', fps=fps, codec='libx264', pixelformat='yuv420p')
# cmd_combine_audio = f"ffmpeg -y -v fatal -i {audio_path} -i temp.mp4 {output_vid_name}"
# print(cmd_combine_audio)
# os.system(cmd_combine_audio)
input_video = './temp.mp4'
# Check if the input_video and audio_path exist
if not os.path.exists(input_video):
raise FileNotFoundError(f"Input video file not found: {input_video}")
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
# 读取视频
reader = imageio.get_reader(input_video)
fps = reader.get_meta_data()['fps'] # 获取原视频的帧率
# 将帧存储在列表中
frames = images
# 保存视频并添加音频
# imageio.mimwrite(output_vid_name, frames, 'FFMPEG', fps=fps, codec='libx264', audio_codec='aac', input_params=['-i', audio_path])
# input_video = ffmpeg.input(input_video)
# input_audio = ffmpeg.input(audio_path)
print(len(frames))
# imageio.mimwrite(
# output_video,
# frames,
# 'FFMPEG',
# fps=25,
# codec='libx264',
# audio_codec='aac',
# input_params=['-i', audio_path],
# output_params=['-y'], # Add the '-y' flag to overwrite the output file if it exists
# )
# writer = imageio.get_writer(output_vid_name, fps = 25, codec='libx264', quality=10, pixelformat='yuvj444p')
# for im in frames:
# writer.append_data(im)
# writer.close()
# Load the video
video_clip = VideoFileClip(input_video)
# Load the audio
audio_clip = AudioFileClip(audio_path)
# Set the audio to the video
video_clip = video_clip.set_audio(audio_clip)
# Write the output video
video_clip.write_videofile(output_vid_name, codec='libx264', audio_codec='aac',fps=25)
os.remove("temp.mp4")
#shutil.rmtree(result_img_save_path)
print(f"result is save to {output_vid_name}")
return output_vid_name,bbox_shift_text
# load model weights
audio_processor,vae,unet,pe = load_all_model()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
timesteps = torch.tensor([0], device=device)
def check_video(video):
if not isinstance(video, str):
return video # in case of none type
# Define the output video file name
dir_path, file_name = os.path.split(video)
if file_name.startswith("outputxxx_"):
return video
# Add the output prefix to the file name
output_file_name = "outputxxx_" + file_name
os.makedirs('./results',exist_ok=True)
os.makedirs('./results/output',exist_ok=True)
os.makedirs('./results/input',exist_ok=True)
# Combine the directory path and the new file name
output_video = os.path.join('./results/input', output_file_name)
# # Run the ffmpeg command to change the frame rate to 25fps
# command = f"ffmpeg -i {video} -r 25 -vcodec libx264 -vtag hvc1 -pix_fmt yuv420p crf 18 {output_video} -y"
# read video
reader = imageio.get_reader(video)
fps = reader.get_meta_data()['fps'] # get fps from original video
# conver fps to 25
frames = [im for im in reader]
target_fps = 25
L = len(frames)
L_target = int(L / fps * target_fps)
original_t = [x / fps for x in range(1, L+1)]
t_idx = 0
target_frames = []
for target_t in range(1, L_target+1):
while target_t / target_fps > original_t[t_idx]:
t_idx += 1 # find the first t_idx so that target_t / target_fps <= original_t[t_idx]
if t_idx >= L:
break
target_frames.append(frames[t_idx])
# save video
imageio.mimwrite(output_video, target_frames, 'FFMPEG', fps=25, codec='libx264', quality=9, pixelformat='yuv420p')
return output_video
css = """#input_img {max-width: 1024px !important} #output_vid {max-width: 1024px; max-height: 576px}"""
with gr.Blocks(css=css) as demo:
gr.Markdown(
"<div align='center'> <h1>MuseTalk: Real-Time High Quality Lip Synchronization with Latent Space Inpainting </span> </h1> \
<h2 style='font-weight: 450; font-size: 1rem; margin: 0rem'>\
</br>\
Yue Zhang <sup>\*</sup>,\
Minhao Liu<sup>\*</sup>,\
Zhaokang Chen,\
Bin Wu<sup>†</sup>,\
Yingjie He,\
Chao Zhan,\
Wenjiang Zhou\
(<sup>*</sup>Equal Contribution, <sup>†</sup>Corresponding Author, [email protected])\
Lyra Lab, Tencent Music Entertainment\
</h2> \
<a style='font-size:18px;color: #000000' href='https://github.com/TMElyralab/MuseTalk'>[Github Repo]</a>\
<a style='font-size:18px;color: #000000' href='https://github.com/TMElyralab/MuseTalk'>[Huggingface]</a>\
<a style='font-size:18px;color: #000000' href=''> [Technical report(Coming Soon)] </a>\
<a style='font-size:18px;color: #000000' href=''> [Project Page(Coming Soon)] </a> </div>"
)
with gr.Row():
with gr.Column():
audio = gr.Audio(label="Driven Audio",type="filepath")
video = gr.Video(label="Reference Video",sources=['upload'])
bbox_shift = gr.Number(label="BBox_shift value, px", value=0)
bbox_shift_scale = gr.Textbox(label="BBox_shift recommend value lower bound,The corresponding bbox range is generated after the initial result is generated. \n If the result is not good, it can be adjusted according to this reference value", value="",interactive=False)
btn = gr.Button("Generate")
out1 = gr.Video()
video.change(
fn=check_video, inputs=[video], outputs=[video]
)
btn.click(
fn=inference,
inputs=[
audio,
video,
bbox_shift,
],
outputs=[out1,bbox_shift_scale]
)
# Set the IP and port
ip_address = "0.0.0.0" # Replace with your desired IP address
port_number = 7860 # Replace with your desired port number
demo.queue().launch(
share=False , debug=True, server_name=ip_address, server_port=port_number
)