add musetalk

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 全身视频拼接
@@ -140,6 +140,34 @@ 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步的安装，直接运行。
 ```

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,17 +449,23 @@ 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)
 
+    if opt.model == 'ernerf':
         # assert test mode
         opt.test = True
         opt.test_train = False
@@ -500,6 +508,11 @@ if __name__ == '__main__':
 
         # 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)
+
     #txt_to_audio('我是中国人,我来自北京')
     if opt.transport=='rtmp':
         thread_quit = Event()

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

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)
+            

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")
+        #print(f"video in {fps} FPS, audio idx in 50FPS")
-        while 1:
+        for _ in range(batch_size):
-            start_idx = int(i * whisper_idx_multiplier)
+            # 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)
+            # 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
 

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