Update src/gradio_demo.py

src/gradio_demo.py

@@ -1,170 +1,171 @@
-import torch, uuid
-import os, sys, shutil, platform
-from src.facerender.pirender_animate import AnimateFromCoeff_PIRender
-from src.utils.preprocess import CropAndExtract
-from src.test_audio2coeff import Audio2Coeff  
-from src.facerender.animate import AnimateFromCoeff
-from src.generate_batch import get_data
-from src.generate_facerender_batch import get_facerender_data
-
-from src.utils.init_path import init_path
-
-from pydub import AudioSegment
-
-
-def mp3_to_wav(mp3_filename,wav_filename,frame_rate):
-    mp3_file = AudioSegment.from_file(file=mp3_filename)
-    mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav")
-
-
-class SadTalker():
-
-    def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False):
-
-        if torch.cuda.is_available():
-            device = "cuda"
-        elif platform.system() == 'Darwin': # macos 
-            device = "mps"
-        else:
-            device = "cpu"
-        
-        self.device = device
-
-        os.environ['TORCH_HOME']= checkpoint_path
-
-        self.checkpoint_path = checkpoint_path
-        self.config_path = config_path
-      
-
-    def test(self, source_image, driven_audio, preprocess='crop', 
-        still_mode=False,  use_enhancer=False, batch_size=1, size=256, 
-        pose_style = 0, 
-        facerender='facevid2vid',
-        exp_scale=1.0, 
-        use_ref_video = False,
-        ref_video = None,
-        ref_info = None,
-        use_idle_mode = False,
-        length_of_audio = 0, use_blink=True,
-        result_dir='./results/'):
-
-        self.sadtalker_paths = init_path(self.checkpoint_path, self.config_path, size, False, preprocess)
-        print(self.sadtalker_paths)
-            
-        self.audio_to_coeff = Audio2Coeff(self.sadtalker_paths, self.device)
-        self.preprocess_model = CropAndExtract(self.sadtalker_paths, self.device)
-        
-        if facerender == 'facevid2vid' and self.device != 'mps':
-            self.animate_from_coeff = AnimateFromCoeff(self.sadtalker_paths, self.device)
-        elif facerender == 'pirender' or self.device == 'mps':
-            self.animate_from_coeff = AnimateFromCoeff_PIRender(self.sadtalker_paths, self.device)
-            facerender = 'pirender'
-        else:
-            raise(RuntimeError('Unknown model: {}'.format(facerender)))
-            
-
-        time_tag = str(uuid.uuid4())
-        save_dir = os.path.join(result_dir, time_tag)
-        os.makedirs(save_dir, exist_ok=True)
-
-        input_dir = os.path.join(save_dir, 'input')
-        os.makedirs(input_dir, exist_ok=True)
-
-        print(source_image)
-        pic_path = os.path.join(input_dir, os.path.basename(source_image)) 
-        shutil.move(source_image, input_dir)
-
-        if driven_audio is not None and os.path.isfile(driven_audio):
-            audio_path = os.path.join(input_dir, os.path.basename(driven_audio))  
-
-            #### mp3 to wav
-            if '.mp3' in audio_path:
-                mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000)
-                audio_path = audio_path.replace('.mp3', '.wav')
-            else:
-                shutil.move(driven_audio, input_dir)
-
-        elif use_idle_mode:
-            audio_path = os.path.join(input_dir, 'idlemode_'+str(length_of_audio)+'.wav') ## generate audio from this new audio_path
-            from pydub import AudioSegment
-            one_sec_segment = AudioSegment.silent(duration=1000*length_of_audio)  #duration in milliseconds
-            one_sec_segment.export(audio_path, format="wav")
-        else:
-            print(use_ref_video, ref_info)
-            assert use_ref_video == True and ref_info == 'all'
-
-        if use_ref_video and ref_info == 'all': # full ref mode
-            ref_video_videoname = os.path.basename(ref_video)
-            audio_path = os.path.join(save_dir, ref_video_videoname+'.wav')
-            print('new audiopath:',audio_path)
-            # if ref_video contains audio, set the audio from ref_video.
-            cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(ref_video, audio_path)
-            os.system(cmd)        
-
-        os.makedirs(save_dir, exist_ok=True)
-        
-        #crop image and extract 3dmm from image
-        first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
-        os.makedirs(first_frame_dir, exist_ok=True)
-        first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir, preprocess, True, size)
-        
-        if first_coeff_path is None:
-            raise AttributeError("No face is detected")
-
-        if use_ref_video:
-            print('using ref video for genreation')
-            ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0]
-            ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname)
-            os.makedirs(ref_video_frame_dir, exist_ok=True)
-            print('3DMM Extraction for the reference video providing pose')
-            ref_video_coeff_path, _, _ =  self.preprocess_model.generate(ref_video, ref_video_frame_dir, preprocess, source_image_flag=False)
-        else:
-            ref_video_coeff_path = None
-
-        if use_ref_video:
-            if ref_info == 'pose':
-                ref_pose_coeff_path = ref_video_coeff_path
-                ref_eyeblink_coeff_path = None
-            elif ref_info == 'blink':
-                ref_pose_coeff_path = None
-                ref_eyeblink_coeff_path = ref_video_coeff_path
-            elif ref_info == 'pose+blink':
-                ref_pose_coeff_path = ref_video_coeff_path
-                ref_eyeblink_coeff_path = ref_video_coeff_path
-            elif ref_info == 'all':            
-                ref_pose_coeff_path = None
-                ref_eyeblink_coeff_path = None
-            else:
-                raise('error in refinfo')
-        else:
-            ref_pose_coeff_path = None
-            ref_eyeblink_coeff_path = None
-
-        #audio2ceoff
-        if use_ref_video and ref_info == 'all':
-            coeff_path = ref_video_coeff_path # self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
-        else:
-            batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path, still=still_mode, \
-                idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink) # longer audio?
-            coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
-
-        #coeff2video
-        data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, \
-            preprocess=preprocess, size=size, expression_scale = exp_scale, facemodel=facerender)
-        return_path = self.animate_from_coeff.generate(data, save_dir,  pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess, img_size=size)
-        video_name = data['video_name']
-        print(f'The generated video is named {video_name} in {save_dir}')
-
-        del self.preprocess_model
-        del self.audio_to_coeff
-        del self.animate_from_coeff
-
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
-            
-        import gc; gc.collect()
-        
-        return return_path
-
+import spaces
+import torch, uuid
+import os, sys, shutil, platform
+from src.facerender.pirender_animate import AnimateFromCoeff_PIRender
+from src.utils.preprocess import CropAndExtract
+from src.test_audio2coeff import Audio2Coeff  
+from src.facerender.animate import AnimateFromCoeff
+from src.generate_batch import get_data
+from src.generate_facerender_batch import get_facerender_data
+
+from src.utils.init_path import init_path
+
+from pydub import AudioSegment
+
+
+def mp3_to_wav(mp3_filename,wav_filename,frame_rate):
+    mp3_file = AudioSegment.from_file(file=mp3_filename)
+    mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav")
+
+
+class SadTalker():
+
+    def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False):
+
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif platform.system() == 'Darwin': # macos 
+            device = "mps"
+        else:
+            device = "cpu"
+        
+        self.device = device
+
+        os.environ['TORCH_HOME']= checkpoint_path
+
+        self.checkpoint_path = checkpoint_path
+        self.config_path = config_path
+      
+    @spaces.GPU
+    def test(self, source_image, driven_audio, preprocess='crop', 
+        still_mode=False,  use_enhancer=False, batch_size=1, size=256, 
+        pose_style = 0, 
+        facerender='facevid2vid',
+        exp_scale=1.0, 
+        use_ref_video = False,
+        ref_video = None,
+        ref_info = None,
+        use_idle_mode = False,
+        length_of_audio = 0, use_blink=True,
+        result_dir='./results/'):
+
+        self.sadtalker_paths = init_path(self.checkpoint_path, self.config_path, size, False, preprocess)
+        print(self.sadtalker_paths)
+            
+        self.audio_to_coeff = Audio2Coeff(self.sadtalker_paths, self.device)
+        self.preprocess_model = CropAndExtract(self.sadtalker_paths, self.device)
+        
+        if facerender == 'facevid2vid' and self.device != 'mps':
+            self.animate_from_coeff = AnimateFromCoeff(self.sadtalker_paths, self.device)
+        elif facerender == 'pirender' or self.device == 'mps':
+            self.animate_from_coeff = AnimateFromCoeff_PIRender(self.sadtalker_paths, self.device)
+            facerender = 'pirender'
+        else:
+            raise(RuntimeError('Unknown model: {}'.format(facerender)))
+            
+
+        time_tag = str(uuid.uuid4())
+        save_dir = os.path.join(result_dir, time_tag)
+        os.makedirs(save_dir, exist_ok=True)
+
+        input_dir = os.path.join(save_dir, 'input')
+        os.makedirs(input_dir, exist_ok=True)
+
+        print(source_image)
+        pic_path = os.path.join(input_dir, os.path.basename(source_image)) 
+        shutil.move(source_image, input_dir)
+
+        if driven_audio is not None and os.path.isfile(driven_audio):
+            audio_path = os.path.join(input_dir, os.path.basename(driven_audio))  
+
+            #### mp3 to wav
+            if '.mp3' in audio_path:
+                mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000)
+                audio_path = audio_path.replace('.mp3', '.wav')
+            else:
+                shutil.move(driven_audio, input_dir)
+
+        elif use_idle_mode:
+            audio_path = os.path.join(input_dir, 'idlemode_'+str(length_of_audio)+'.wav') ## generate audio from this new audio_path
+            from pydub import AudioSegment
+            one_sec_segment = AudioSegment.silent(duration=1000*length_of_audio)  #duration in milliseconds
+            one_sec_segment.export(audio_path, format="wav")
+        else:
+            print(use_ref_video, ref_info)
+            assert use_ref_video == True and ref_info == 'all'
+
+        if use_ref_video and ref_info == 'all': # full ref mode
+            ref_video_videoname = os.path.basename(ref_video)
+            audio_path = os.path.join(save_dir, ref_video_videoname+'.wav')
+            print('new audiopath:',audio_path)
+            # if ref_video contains audio, set the audio from ref_video.
+            cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(ref_video, audio_path)
+            os.system(cmd)        
+
+        os.makedirs(save_dir, exist_ok=True)
+        
+        #crop image and extract 3dmm from image
+        first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
+        os.makedirs(first_frame_dir, exist_ok=True)
+        first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir, preprocess, True, size)
+        
+        if first_coeff_path is None:
+            raise AttributeError("No face is detected")
+
+        if use_ref_video:
+            print('using ref video for genreation')
+            ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0]
+            ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname)
+            os.makedirs(ref_video_frame_dir, exist_ok=True)
+            print('3DMM Extraction for the reference video providing pose')
+            ref_video_coeff_path, _, _ =  self.preprocess_model.generate(ref_video, ref_video_frame_dir, preprocess, source_image_flag=False)
+        else:
+            ref_video_coeff_path = None
+
+        if use_ref_video:
+            if ref_info == 'pose':
+                ref_pose_coeff_path = ref_video_coeff_path
+                ref_eyeblink_coeff_path = None
+            elif ref_info == 'blink':
+                ref_pose_coeff_path = None
+                ref_eyeblink_coeff_path = ref_video_coeff_path
+            elif ref_info == 'pose+blink':
+                ref_pose_coeff_path = ref_video_coeff_path
+                ref_eyeblink_coeff_path = ref_video_coeff_path
+            elif ref_info == 'all':            
+                ref_pose_coeff_path = None
+                ref_eyeblink_coeff_path = None
+            else:
+                raise('error in refinfo')
+        else:
+            ref_pose_coeff_path = None
+            ref_eyeblink_coeff_path = None
+
+        #audio2ceoff
+        if use_ref_video and ref_info == 'all':
+            coeff_path = ref_video_coeff_path # self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
+        else:
+            batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path, still=still_mode, \
+                idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink) # longer audio?
+            coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
+
+        #coeff2video
+        data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, \
+            preprocess=preprocess, size=size, expression_scale = exp_scale, facemodel=facerender)
+        return_path = self.animate_from_coeff.generate(data, save_dir,  pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess, img_size=size)
+        video_name = data['video_name']
+        print(f'The generated video is named {video_name} in {save_dir}')
+
+        del self.preprocess_model
+        del self.audio_to_coeff
+        del self.animate_from_coeff
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+            
+        import gc; gc.collect()
+        
+        return return_path
+