diff --git a/README.md b/README.md index 56c6e89..e2ffbc5 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ python app.py --asr_model facebook/hubert-large-ls960-ft ### 3.4 设置背景图片 ``` -python app.py --bg_img bg.jpg +python app.py --bg_img bc.jpg ``` ### 3.5 全身视频拼接 @@ -139,6 +139,34 @@ docker run --rm -it -p 1935:1935 -p 1985:1985 -p 8080:8080 registry.cn-hangzhou. python app.py --transport rtmp --push_url 'rtmp://localhost/live/livestream' ``` 用浏览器打开http://serverip:8010/echo.html + +### 3.9 模型用musetalk +暂不支持rtmp推送 +- 安装依赖库 +```bash +conda install ffmpeg +pip install --no-cache-dir -U openmim +mim install mmengine +mim install "mmcv>=2.0.1" +mim install "mmdet>=3.1.0" +mim install "mmpose>=1.1.0" +``` +- 下载模型 +下载MuseTalk运行需要的模型，提供一个下载地址 https://caiyun.139.com/m/i?2eAjs2nXXnRgr 提取码:qdg2 +解压后，将models下文件拷到本项目的models下 +下载数字人模型，链接: https://caiyun.139.com/m/i?2eAjs8optksop 提取码:3mkt, 解压后将整个文件夹拷到本项目的data/avatars下 +- 运行 +python app.py --model musetalk --transport webrtc +用浏览器打开http://serverip:8010/webrtc.html +可以设置--batch_size 提高显卡利用率，设置--avatar_id 运行不同的数字人 +#### 替换成自己的数字人 +```bash +git clone https://github.com/TMElyralab/MuseTalk.git +cd MuseTalk +修改configs/inference/realtime.yaml，将preparation改为True +python -m scripts.realtime_inference --inference_config configs/inference/realtime.yaml +运行后将results/avatars下文件拷到本项目的data/avatars下 +``` ## 4. Docker Run 不需要第1步的安装，直接运行。 diff --git a/app.py b/app.py index 03e91d1..9a7c80c 100644 --- a/app.py +++ b/app.py @@ -419,9 +419,6 @@ if __name__ == '__main__': # parser.add_argument('--asr_model', type=str, default='facebook/wav2vec2-large-960h-lv60-self') # parser.add_argument('--asr_model', type=str, default='facebook/hubert-large-ls960-ft') - parser.add_argument('--transport', type=str, default='rtcpush') #rtmp webrtc rtcpush - parser.add_argument('--push_url', type=str, default='http://localhost:1985/rtc/v1/whip/?app=live&stream=livestream') #rtmp://localhost/live/livestream - parser.add_argument('--asr_save_feats', action='store_true') # audio FPS parser.add_argument('--fps', type=int, default=50) @@ -437,6 +434,11 @@ if __name__ == '__main__': parser.add_argument('--fullbody_offset_x', type=int, default=0) parser.add_argument('--fullbody_offset_y', type=int, default=0) + #musetalk opt + parser.add_argument('--avatar_id', type=str, default='avator_1') + parser.add_argument('--bbox_shift', type=int, default=5) + parser.add_argument('--batch_size', type=int, default=4) + parser.add_argument('--customvideo', action='store_true', help="custom video") parser.add_argument('--customvideo_img', type=str, default='data/customvideo/img') parser.add_argument('--customvideo_imgnum', type=int, default=1) @@ -447,59 +449,70 @@ if __name__ == '__main__': parser.add_argument('--CHARACTER', type=str, default='test') parser.add_argument('--EMOTION', type=str, default='default') + parser.add_argument('--model', type=str, default='ernerf') #musetalk + + parser.add_argument('--transport', type=str, default='rtcpush') #rtmp webrtc rtcpush + parser.add_argument('--push_url', type=str, default='http://localhost:1985/rtc/v1/whip/?app=live&stream=livestream') #rtmp://localhost/live/livestream + parser.add_argument('--listenport', type=int, default=8010) opt = parser.parse_args() app.config.from_object(opt) - print(app.config) + #print(app.config) tts_type = opt.tts if tts_type == "xtts": print("Computing the latents for a new reference...") gspeaker = get_speaker(opt.REF_FILE, opt.TTS_SERVER) - # assert test mode - opt.test = True - opt.test_train = False - #opt.train_camera =True - # explicit smoothing - opt.smooth_path = True - opt.smooth_lips = True - - assert opt.pose != '', 'Must provide a pose source' - - # if opt.O: - opt.fp16 = True - opt.cuda_ray = True - opt.exp_eye = True - opt.smooth_eye = True - - if opt.torso_imgs=='': #no img,use model output - opt.torso = True - - # assert opt.cuda_ray, "Only support CUDA ray mode." - opt.asr = True - - if opt.patch_size > 1: - # assert opt.patch_size > 16, "patch_size should > 16 to run LPIPS loss." - assert opt.num_rays % (opt.patch_size ** 2) == 0, "patch_size ** 2 should be dividable by num_rays." - seed_everything(opt.seed) - print(opt) - - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - model = NeRFNetwork(opt) - - criterion = torch.nn.MSELoss(reduction='none') - metrics = [] # use no metric in GUI for faster initialization... - print(model) - trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, criterion=criterion, fp16=opt.fp16, metrics=metrics, use_checkpoint=opt.ckpt) - - test_loader = NeRFDataset_Test(opt, device=device).dataloader() - model.aud_features = test_loader._data.auds - model.eye_areas = test_loader._data.eye_area + if opt.model == 'ernerf': + # assert test mode + opt.test = True + opt.test_train = False + #opt.train_camera =True + # explicit smoothing + opt.smooth_path = True + opt.smooth_lips = True + + assert opt.pose != '', 'Must provide a pose source' + + # if opt.O: + opt.fp16 = True + opt.cuda_ray = True + opt.exp_eye = True + opt.smooth_eye = True + + if opt.torso_imgs=='': #no img,use model output + opt.torso = True + + # assert opt.cuda_ray, "Only support CUDA ray mode." + opt.asr = True + + if opt.patch_size > 1: + # assert opt.patch_size > 16, "patch_size should > 16 to run LPIPS loss." + assert opt.num_rays % (opt.patch_size ** 2) == 0, "patch_size ** 2 should be dividable by num_rays." + seed_everything(opt.seed) + print(opt) + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = NeRFNetwork(opt) + + criterion = torch.nn.MSELoss(reduction='none') + metrics = [] # use no metric in GUI for faster initialization... + print(model) + trainer = Trainer('ngp', opt, model, device=device, workspace=opt.workspace, criterion=criterion, fp16=opt.fp16, metrics=metrics, use_checkpoint=opt.ckpt) + + test_loader = NeRFDataset_Test(opt, device=device).dataloader() + model.aud_features = test_loader._data.auds + model.eye_areas = test_loader._data.eye_area + + # we still need test_loader to provide audio features for testing. + nerfreal = NeRFReal(opt, trainer, test_loader) + elif opt.model == 'musetalk': + from musereal import MuseReal + print(opt) + nerfreal = MuseReal(opt) - # we still need test_loader to provide audio features for testing. - nerfreal = NeRFReal(opt, trainer, test_loader) #txt_to_audio('我是中国人,我来自北京') if opt.transport=='rtmp': thread_quit = Event() diff --git a/models/put models here.txt b/models/put models here.txt new file mode 100644 index 0000000..e69de29 diff --git a/museasr.py b/museasr.py new file mode 100644 index 0000000..6d65147 --- /dev/null +++ b/museasr.py @@ -0,0 +1,130 @@ +import time +import torch +import numpy as np +import soundfile as sf +import resampy + +import queue +from queue import Queue +from io import BytesIO + +from musetalk.whisper.audio2feature import Audio2Feature + +class MuseASR: + def __init__(self, opt, audio_processor:Audio2Feature): + self.opt = opt + + self.fps = opt.fps # 20 ms per frame + self.sample_rate = 16000 + self.chunk = self.sample_rate // self.fps # 320 samples per chunk (20ms * 16000 / 1000) + self.queue = Queue() + self.input_stream = BytesIO() + self.output_queue = Queue() + + self.audio_processor = audio_processor + self.batch_size = opt.batch_size + + self.stride_left_size = self.stride_right_size = 6 + self.audio_feats = [] + + self.warm_up() + + def __create_bytes_stream(self,byte_stream): + #byte_stream=BytesIO(buffer) + stream, sample_rate = sf.read(byte_stream) # [T*sample_rate,] float64 + print(f'[INFO]tts audio stream {sample_rate}: {stream.shape}') + stream = stream.astype(np.float32) + + if stream.ndim > 1: + print(f'[WARN] audio has {stream.shape[1]} channels, only use the first.') + stream = stream[:, 0] + + if sample_rate != self.sample_rate and stream.shape[0]>0: + print(f'[WARN] audio sample rate is {sample_rate}, resampling into {self.sample_rate}.') + stream = resampy.resample(x=stream, sr_orig=sample_rate, sr_new=self.sample_rate) + + return stream + + def push_audio(self,buffer): + print(f'[INFO] push_audio {len(buffer)}') + if self.opt.tts == "xtts" or self.opt.tts == "gpt-sovits": + if len(buffer)>0: + stream = np.frombuffer(buffer, dtype=np.int16).astype(np.float32) / 32767 + if self.opt.tts == "xtts": + stream = resampy.resample(x=stream, sr_orig=24000, sr_new=self.sample_rate) + else: + stream = resampy.resample(x=stream, sr_orig=32000, sr_new=self.sample_rate) + #byte_stream=BytesIO(buffer) + #stream = self.__create_bytes_stream(byte_stream) + streamlen = stream.shape[0] + idx=0 + while streamlen >= self.chunk: + self.queue.put(stream[idx:idx+self.chunk]) + streamlen -= self.chunk + idx += self.chunk + # if streamlen>0: #skip last frame(not 20ms) + # self.queue.put(stream[idx:]) + else: #edge tts + self.input_stream.write(buffer) + if len(buffer)<=0: + self.input_stream.seek(0) + stream = self.__create_bytes_stream(self.input_stream) + streamlen = stream.shape[0] + idx=0 + while streamlen >= self.chunk: + self.queue.put(stream[idx:idx+self.chunk]) + streamlen -= self.chunk + idx += self.chunk + #if streamlen>0: #skip last frame(not 20ms) + # self.queue.put(stream[idx:]) + self.input_stream.seek(0) + self.input_stream.truncate() + + def __get_audio_frame(self): + try: + frame = self.queue.get(block=False) + type = 0 + print(f'[INFO] get frame {frame.shape}') + except queue.Empty: + frame = np.zeros(self.chunk, dtype=np.float32) + type = 1 + + return frame,type + + def get_audio_out(self): #get origin audio pcm to nerf + return self.output_queue.get() + + def warm_up(self): + frames = [] + for _ in range(self.stride_left_size + self.stride_right_size): + audio_frame,type=self.__get_audio_frame() + frames.append(audio_frame) + self.output_queue.put((audio_frame,type)) + inputs = np.concatenate(frames) # [N * chunk] + whisper_feature = self.audio_processor.audio2feat(inputs) + for feature in whisper_feature: + self.audio_feats.append(feature) + + for _ in range(self.stride_left_size): + self.output_queue.get() + + def run_step(self): + ############################################## extract audio feature ############################################## + start_time = time.time() + frames = [] + for _ in range(self.batch_size*2): + audio_frame,type=self.__get_audio_frame() + frames.append(audio_frame) + self.output_queue.put((audio_frame,type)) + inputs = np.concatenate(frames) # [N * chunk] + whisper_feature = self.audio_processor.audio2feat(inputs) + for feature in whisper_feature: + self.audio_feats.append(feature) + + #print(f"processing audio costs {(time.time() - start_time) * 1000}ms, inputs shape:{inputs.shape} whisper_feature len:{len(whisper_feature)}") + + def get_next_feat(self): + whisper_chunks = self.audio_processor.feature2chunks(feature_array=self.audio_feats,fps=self.fps/2,batch_size=self.batch_size,start=self.stride_left_size/2 ) + #print(f"whisper_chunks len:{len(whisper_chunks)},self.audio_feats len:{len(self.audio_feats)},self.output_queue len:{self.output_queue.qsize()}") + self.audio_feats = self.audio_feats[-(self.stride_left_size + self.stride_right_size):] + return whisper_chunks \ No newline at end of file diff --git a/musereal.py b/musereal.py new file mode 100644 index 0000000..84eadbb --- /dev/null +++ b/musereal.py @@ -0,0 +1,194 @@ +import math +import torch +import numpy as np + +#from .utils import * +import subprocess +import os +import time +import torch.nn.functional as F +import cv2 +import glob +import pickle +import copy + +import queue +from queue import Queue +from threading import Thread, Event +from io import BytesIO + +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 +from musetalk.utils.blending import get_image,get_image_prepare_material,get_image_blending +from musetalk.utils.utils import load_all_model + +from museasr import MuseASR +import asyncio +from av import AudioFrame, VideoFrame + +class MuseReal: + def __init__(self, opt): + self.opt = opt # shared with the trainer's opt to support in-place modification of rendering parameters. + self.W = opt.W + self.H = opt.H + + self.fps = opt.fps # 20 ms per frame + + #### musetalk + self.avatar_id = opt.avatar_id + self.video_path = '' #video_path + self.bbox_shift = opt.bbox_shift + self.avatar_path = f"./data/avatars/{self.avatar_id}" + self.full_imgs_path = f"{self.avatar_path}/full_imgs" + self.coords_path = f"{self.avatar_path}/coords.pkl" + self.latents_out_path= f"{self.avatar_path}/latents.pt" + self.video_out_path = f"{self.avatar_path}/vid_output/" + self.mask_out_path =f"{self.avatar_path}/mask" + self.mask_coords_path =f"{self.avatar_path}/mask_coords.pkl" + self.avatar_info_path = f"{self.avatar_path}/avator_info.json" + self.avatar_info = { + "avatar_id":self.avatar_id, + "video_path":self.video_path, + "bbox_shift":self.bbox_shift + } + self.batch_size = opt.batch_size + self.idx = 0 + self.res_frame_queue = Queue() + self.__loadmodels() + self.__loadavatar() + + self.asr = MuseASR(opt,self.audio_processor) + + def __loadmodels(self): + # load model weights + self.audio_processor, self.vae, self.unet, self.pe = load_all_model() + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.timesteps = torch.tensor([0], device=device) + self.pe = self.pe.half() + self.vae.vae = self.vae.vae.half() + self.unet.model = self.unet.model.half() + + def __loadavatar(self): + self.input_latent_list_cycle = torch.load(self.latents_out_path) + with open(self.coords_path, 'rb') as f: + self.coord_list_cycle = pickle.load(f) + input_img_list = glob.glob(os.path.join(self.full_imgs_path, '*.[jpJP][pnPN]*[gG]')) + input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0])) + self.frame_list_cycle = read_imgs(input_img_list) + with open(self.mask_coords_path, 'rb') as f: + self.mask_coords_list_cycle = pickle.load(f) + input_mask_list = glob.glob(os.path.join(self.mask_out_path, '*.[jpJP][pnPN]*[gG]')) + input_mask_list = sorted(input_mask_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0])) + self.mask_list_cycle = read_imgs(input_mask_list) + + + def push_audio(self,buffer): + self.asr.push_audio(buffer) + + def __mirror_index(self, index): + size = len(self.coord_list_cycle) + turn = index // size + res = index % size + if turn % 2 == 0: + return res + else: + return size - res - 1 + + def test_step(self,loop=None,audio_track=None,video_track=None): + + # gen = datagen(whisper_chunks, + # self.input_latent_list_cycle, + # self.batch_size) + self.asr.run_step() + whisper_chunks = self.asr.get_next_feat() + whisper_batch = np.stack(whisper_chunks) + latent_batch = [] + for i in range(self.batch_size): + idx = self.__mirror_index(self.idx+i) + latent = self.input_latent_list_cycle[idx] + latent_batch.append(latent) + latent_batch = torch.cat(latent_batch, dim=0) + + # for i, (whisper_batch,latent_batch) in enumerate(gen): + audio_feature_batch = torch.from_numpy(whisper_batch) + audio_feature_batch = audio_feature_batch.to(device=self.unet.device, + dtype=self.unet.model.dtype) + audio_feature_batch = self.pe(audio_feature_batch) + latent_batch = latent_batch.to(dtype=self.unet.model.dtype) + + pred_latents = self.unet.model(latent_batch, + self.timesteps, + encoder_hidden_states=audio_feature_batch).sample + recon = self.vae.decode_latents(pred_latents) + #print('diffusion len=',len(recon)) + for res_frame in recon: + #self.__pushmedia(res_frame,loop,audio_track,video_track) + self.res_frame_queue.put((res_frame,self.__mirror_index(self.idx))) + self.idx = self.idx + 1 + + + def process_frames(self,quit_event,loop=None,audio_track=None,video_track=None): + + while not quit_event.is_set(): + try: + res_frame,idx = self.res_frame_queue.get(block=True, timeout=1) + except queue.Empty: + continue + bbox = self.coord_list_cycle[idx] + ori_frame = copy.deepcopy(self.frame_list_cycle[idx]) + x1, y1, x2, y2 = bbox + try: + res_frame = cv2.resize(res_frame.astype(np.uint8),(x2-x1,y2-y1)) + except: + continue + mask = self.mask_list_cycle[idx] + mask_crop_box = self.mask_coords_list_cycle[idx] + #combine_frame = get_image(ori_frame,res_frame,bbox) + combine_frame = get_image_blending(ori_frame,res_frame,bbox,mask,mask_crop_box) + + image = combine_frame #(outputs['image'] * 255).astype(np.uint8) + new_frame = VideoFrame.from_ndarray(image, format="bgr24") + asyncio.run_coroutine_threadsafe(video_track._queue.put(new_frame), loop) + + audiotype = 0 + for _ in range(2): + frame,type = self.asr.get_audio_out() + audiotype += type + frame = (frame * 32767).astype(np.int16) + new_frame = AudioFrame(format='s16', layout='mono', samples=frame.shape[0]) + new_frame.planes[0].update(frame.tobytes()) + new_frame.sample_rate=16000 + # if audio_track._queue.qsize()>10: + # time.sleep(0.1) + asyncio.run_coroutine_threadsafe(audio_track._queue.put(new_frame), loop) + + def render(self,quit_event,loop=None,audio_track=None,video_track=None): + #if self.opt.asr: + # self.asr.warm_up() + + process_thread = Thread(target=self.process_frames, args=(quit_event,loop,audio_track,video_track)) + process_thread.start() + + count=0 + totaltime=0 + _starttime=time.perf_counter() + #_totalframe=0 + while not quit_event.is_set(): #todo + # update texture every frame + # audio stream thread... + t = time.perf_counter() + self.test_step(loop,audio_track,video_track) + totaltime += (time.perf_counter() - t) + count += self.opt.batch_size + #_totalframe += 1 + if count>=100: + print(f"------actual avg infer fps:{count/totaltime:.4f}") + count=0 + totaltime=0 + if self.res_frame_queue.qsize()>2*self.opt.batch_size: + time.sleep(0.1) + #print('sleep') + # delay = _starttime+_totalframe*0.04-time.perf_counter() #40ms + # if delay > 0: + # time.sleep(delay) + \ No newline at end of file diff --git a/musetalk/whisper/audio2feature.py b/musetalk/whisper/audio2feature.py index 2cfd3a9..645bbc3 100644 --- a/musetalk/whisper/audio2feature.py +++ b/musetalk/whisper/audio2feature.py @@ -59,6 +59,7 @@ class Audio2Feature(): for dt in range(-audio_feat_length[0],audio_feat_length[1]+1): left_idx = int((vid_idx+dt)*50/fps) if left_idx<1 or left_idx>length-1: + print('test-----,left_idx=',left_idx) left_idx = max(0, left_idx) left_idx = min(length-1, left_idx) @@ -78,19 +79,20 @@ class Audio2Feature(): return selected_feature,selected_idx - def feature2chunks(self,feature_array,fps,audio_feat_length = [2,2]): + def feature2chunks(self,feature_array,fps,batch_size,audio_feat_length = [2,2],start=0): whisper_chunks = [] whisper_idx_multiplier = 50./fps i = 0 - print(f"video in {fps} FPS, audio idx in 50FPS") - while 1: - start_idx = int(i * whisper_idx_multiplier) - selected_feature,selected_idx = self.get_sliced_feature(feature_array= feature_array,vid_idx = i,audio_feat_length=audio_feat_length,fps=fps) + #print(f"video in {fps} FPS, audio idx in 50FPS") + for _ in range(batch_size): + # start_idx = int(i * whisper_idx_multiplier) + # if start_idx>=len(feature_array): + # break + selected_feature,selected_idx = self.get_sliced_feature(feature_array= feature_array,vid_idx = i+start,audio_feat_length=audio_feat_length,fps=fps) #print(f"i:{i},selected_idx {selected_idx}") whisper_chunks.append(selected_feature) i += 1 - if start_idx>len(feature_array): - break + return whisper_chunks diff --git a/requirements.txt b/requirements.txt index c3a3f0a..c039bce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -27,9 +27,14 @@ lpips imageio-ffmpeg transformers -edge_tts +edge_tts==6.1.11 flask flask_sockets opencv-python-headless aiortc aiohttp_cors + +ffmpeg-python +omegaconf +diffusers +accelerate