RabbitRUI/ruispace

1import torch, uuid2import os, sys, shutil3from src.utils.preprocess import CropAndExtract4from src.test_audio2coeff import Audio2Coeff  5from src.facerender.animate import AnimateFromCoeff6from src.generate_batch import get_data7from src.generate_facerender_batch import get_facerender_data8 9from pydub import AudioSegment10 11def mp3_to_wav(mp3_filename,wav_filename,frame_rate):12    mp3_file = AudioSegment.from_file(file=mp3_filename)13    mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav")14 15 16class SadTalker():17 18    def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False):19 20        if torch.cuda.is_available() :21            device = "cuda"22        else:23            device = "cpu"24        25        self.device = device26 27        os.environ['TORCH_HOME']= checkpoint_path28 29        self.checkpoint_path = checkpoint_path30        self.config_path = config_path31 32        self.path_of_lm_croper = os.path.join( checkpoint_path, 'shape_predictor_68_face_landmarks.dat')33        self.path_of_net_recon_model = os.path.join( checkpoint_path, 'epoch_20.pth')34        self.dir_of_BFM_fitting = os.path.join( checkpoint_path, 'BFM_Fitting')35        self.wav2lip_checkpoint = os.path.join( checkpoint_path, 'wav2lip.pth')36 37        self.audio2pose_checkpoint = os.path.join( checkpoint_path, 'auido2pose_00140-model.pth')38        self.audio2pose_yaml_path = os.path.join( config_path, 'auido2pose.yaml')39    40        self.audio2exp_checkpoint = os.path.join( checkpoint_path, 'auido2exp_00300-model.pth')41        self.audio2exp_yaml_path = os.path.join( config_path, 'auido2exp.yaml')42 43        self.free_view_checkpoint = os.path.join( checkpoint_path, 'facevid2vid_00189-model.pth.tar')44 45        self.lazy_load = lazy_load46 47        if not self.lazy_load:48            #init model49            print(self.path_of_lm_croper)50            self.preprocess_model = CropAndExtract(self.path_of_lm_croper, self.path_of_net_recon_model, self.dir_of_BFM_fitting, self.device)51 52            print(self.audio2pose_checkpoint)53            self.audio_to_coeff = Audio2Coeff(self.audio2pose_checkpoint, self.audio2pose_yaml_path, 54                                    self.audio2exp_checkpoint, self.audio2exp_yaml_path, self.wav2lip_checkpoint, self.device)55 56    def test(self, source_image, driven_audio, preprocess='crop', still_mode=False, use_enhancer=False, result_dir='./results/'):57 58        ### crop: only model,59 60        if self.lazy_load:61            #init model62            print(self.path_of_lm_croper)63            self.preprocess_model = CropAndExtract(self.path_of_lm_croper, self.path_of_net_recon_model, self.dir_of_BFM_fitting, self.device)64 65            print(self.audio2pose_checkpoint)66            self.audio_to_coeff = Audio2Coeff(self.audio2pose_checkpoint, self.audio2pose_yaml_path, 67                                    self.audio2exp_checkpoint, self.audio2exp_yaml_path, self.wav2lip_checkpoint, self.device)68        69        if preprocess == 'full': 70            self.mapping_checkpoint = os.path.join(self.checkpoint_path, 'mapping_00109-model.pth.tar')71            self.facerender_yaml_path = os.path.join(self.config_path, 'facerender_still.yaml')72        else:73            self.mapping_checkpoint = os.path.join(self.checkpoint_path, 'mapping_00229-model.pth.tar')74            self.facerender_yaml_path = os.path.join(self.config_path, 'facerender.yaml')75 76        print(self.mapping_checkpoint)77        print(self.free_view_checkpoint)78        self.animate_from_coeff = AnimateFromCoeff(self.free_view_checkpoint, self.mapping_checkpoint, 79                                            self.facerender_yaml_path, self.device)80 81        time_tag = str(uuid.uuid4())82        save_dir = os.path.join(result_dir, time_tag)83        os.makedirs(save_dir, exist_ok=True)84 85        input_dir = os.path.join(save_dir, 'input')86        os.makedirs(input_dir, exist_ok=True)87 88        print(source_image)89        pic_path = os.path.join(input_dir, os.path.basename(source_image)) 90        shutil.move(source_image, input_dir)91 92        if os.path.isfile(driven_audio):93            audio_path = os.path.join(input_dir, os.path.basename(driven_audio))  94 95            #### mp3 to wav96            if '.mp3' in audio_path:97                mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000)98                audio_path = audio_path.replace('.mp3', '.wav')99            else:100                shutil.move(driven_audio, input_dir)101        else:102            raise AttributeError("error audio")103 104 105        os.makedirs(save_dir, exist_ok=True)106        pose_style = 0107        #crop image and extract 3dmm from image108        first_frame_dir = os.path.join(save_dir, 'first_frame_dir')109        os.makedirs(first_frame_dir, exist_ok=True)110        first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir,preprocess)111        112        if first_coeff_path is None:113            raise AttributeError("No face is detected")114 115        #audio2ceoff116        batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=None, still=still_mode) # longer audio?117        coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style)118        #coeff2video119        batch_size = 8120        data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, preprocess=preprocess)121        return_path = self.animate_from_coeff.generate(data, save_dir,  pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess)122        video_name = data['video_name']123        print(f'The generated video is named {video_name} in {save_dir}')124 125        if self.lazy_load:126            del self.preprocess_model127            del self.audio_to_coeff128            del self.animate_from_coeff129 130        torch.cuda.empty_cache()131        torch.cuda.synchronize()132        import gc; gc.collect()133        134        return return_path135 136