tts

app.py

@@ -1,90 +1,88 @@
 import os, sys
 import tempfile
 import gradio as gr
-from modules.text2speech import text2speech 
-from modules.sadtalker_test import SadTalker  
-
-def get_driven_audio(audio):  
-    if os.path.isfile(audio):
-        return audio
-    else:
-        save_path = tempfile.NamedTemporaryFile(
-                delete=False,
-                suffix=("." + "wav"),
-            )
-        gen_audio = text2speech(audio, save_path.name)
-        return gen_audio, gen_audio 
+from src.gradio_demo import SadTalker  
+from src.utils.text2speech import TTSTalker
 
 def get_source_image(image):   
         return image
 
-def sadtalker_demo(result_dir='./tmp/'):
+
+
+def sadtalker_demo():
 
     sad_talker = SadTalker()
+    tts_talker = TTSTalker()
+
     with gr.Blocks(analytics_enabled=False) as sadtalker_interface:
-        gr.Markdown("<div align='center'> <h3> 😭 SadTalker: Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation (CVPR 2023) </h3> \
+        gr.Markdown("<div align='center'> <h2> 😭 SadTalker: Learning Realistic 3D Motion Coefficients for Stylized Audio-Driven Single Image Talking Face Animation (CVPR 2023) </span> </h2> \
                     <a style='font-size:18px;color: #efefef' href='https://arxiv.org/abs/2211.12194'>Arxiv</a> &nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \
                     <a style='font-size:18px;color: #efefef' href='https://sadtalker.github.io'>Homepage</a>  &nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \
-                     <a style='font-size:18px;color: #efefef' href='https://github.com/Winfredy/SadTalker'> Github </a> </div>")
+                     <a style='font-size:18px;color: #efefef' href='https://github.com/Winfredy/SadTalker'> Github </div>")
         
-        with gr.Row():
+        with gr.Row().style(equal_height=False):
             with gr.Column(variant='panel'):
                 with gr.Tabs(elem_id="sadtalker_source_image"):
                     with gr.TabItem('Upload image'):
                         with gr.Row():
-                            source_image = gr.Image(label="Source image", source="upload", type="filepath").style(height=256)
+                            source_image = gr.Image(label="Source image", source="upload", type="filepath").style(height=256,width=256)
  
                 with gr.Tabs(elem_id="sadtalker_driven_audio"):
-                    with gr.TabItem('Upload audio(wav/mp3 only currently)'):
+                    with gr.TabItem('Upload OR TTS'):
                         with gr.Column(variant='panel'):
                             driven_audio = gr.Audio(label="Input audio", source="upload", type="filepath")
+                    
+                        with gr.Column(variant='panel'):
+                            input_text = gr.Textbox(label="Generating audio from text", lines=5, placeholder="Alternatively, you can genreate the audio from text using @Coqui.ai TTS.")
+                            tts = gr.Button('Generate audio',elem_id="sadtalker_audio_generate", variant='primary')
+                            tts.click(fn=tts_talker.test, inputs=[input_text], outputs=[driven_audio])
+                        
 
             with gr.Column(variant='panel'): 
                 with gr.Tabs(elem_id="sadtalker_checkbox"):
                     with gr.TabItem('Settings'):
                         with gr.Column(variant='panel'):
-                            is_still_mode = gr.Checkbox(label="Still Mode (fewer head motion)").style(container=True)
-                            is_resize_mode = gr.Checkbox(label="Resize Mode (⚠️ Resize mode need manually crop the image firstly, can handle larger image crop)").style(container=True)
-                            is_enhance_mode = gr.Checkbox(label="Enhance Mode (better face quality )").style(container=True)
+                            is_still_mode = gr.Checkbox(label="w/ Still Mode (fewer hand motion, works on full body)")
+                            enhancer = gr.Checkbox(label="w/ GFPGAN as Face enhancer")
                             submit = gr.Button('Generate', elem_id="sadtalker_generate", variant='primary')
 
                 with gr.Tabs(elem_id="sadtalker_genearted"):
                         gen_video = gr.Video(label="Generated video", format="mp4").style(width=256)
-                        gen_text = gr.Textbox(visible=False)
-                    
+
         with gr.Row():
             examples = [
                 [
-                    'examples/source_image/art_10.png',
-                    'examples/driven_audio/deyu.wav',
+                    'examples/source_image/full_body_1.png',
+                    'examples/driven_audio/bus_chinese.wav',
                     True,
-                    False,
                     False
                 ],
                 [
-                    'examples/source_image/art_1.png',
-                    'examples/driven_audio/fayu.wav',
+                    'examples/source_image/full_body_2.png',
+                    'examples/driven_audio/itosinger1.wav',
                     True,
+                    False
+                ],
+                [
+                    'examples/source_image/art_13.png',
+                    'examples/driven_audio/fayu.wav',
                     True,
                     False
                 ],
                 [
-                    'examples/source_image/art_9.png',
-                    'examples/driven_audio/itosinger1.wav',
+                    'examples/source_image/art_5.png',
+                    'examples/driven_audio/chinese_news.wav',
                     True,
-                    False,
-                    True
-                ]
+                    False
+                ],
             ]
             gr.Examples(examples=examples,
                         inputs=[
                             source_image,
                             driven_audio,
                             is_still_mode,
-                            is_resize_mode,
-                            is_enhance_mode,
-                            gr.Textbox(value=result_dir, visible=False)], 
-                        outputs=[gen_video, gen_text],
+                            enhancer], 
+                        outputs=[gen_video],
                         fn=sad_talker.test,
                         cache_examples=os.getenv('SYSTEM') == 'spaces')
 
@@ -93,10 +91,8 @@ def sadtalker_demo(result_dir='./tmp/'):
                     inputs=[source_image,
                             driven_audio,
                             is_still_mode,
-                            is_resize_mode,
-                            is_enhance_mode,
-                            gr.Textbox(value=result_dir, visible=False)], 
-                    outputs=[gen_video, gen_text]
+                            enhancer], 
+                    outputs=[gen_video]
                     )
 
     return sadtalker_interface
@@ -104,8 +100,7 @@ def sadtalker_demo(result_dir='./tmp/'):
 
 if __name__ == "__main__":
 
-    sadtalker_result_dir = os.path.join('./', 'results')
-    demo = sadtalker_demo(sadtalker_result_dir)
+    demo = sadtalker_demo()
     demo.launch()
 
 

src/audio2exp_models/audio2exp.py

@@ -22,7 +22,8 @@ class Audio2Exp(nn.Module):
             
             current_mel_input = mel_input[:,i:i+10]
 
-            ref = batch['ref'][:, :, :64].repeat((1,current_mel_input.shape[1],1))           #bs T 64
+            #ref = batch['ref'][:, :, :64].repeat((1,current_mel_input.shape[1],1))           #bs T 64
+            ref = batch['ref'][:, :, :64][:, i:i+10]
             ratio = batch['ratio_gt'][:, i:i+10]                               #bs T
 
             audiox = current_mel_input.view(-1, 1, 80, 16)                  # bs*T 1 80 16

src/audio2pose_models/audio2pose.py

@@ -12,7 +12,7 @@ class Audio2Pose(nn.Module):
         self.latent_dim = cfg.MODEL.CVAE.LATENT_SIZE
         self.device = device
 
-        self.audio_encoder = AudioEncoder(wav2lip_checkpoint)
+        self.audio_encoder = AudioEncoder(wav2lip_checkpoint, device)
         self.audio_encoder.eval()
         for param in self.audio_encoder.parameters():
             param.requires_grad = False
@@ -20,10 +20,6 @@ class Audio2Pose(nn.Module):
         self.netG = CVAE(cfg)
         self.netD_motion = PoseSequenceDiscriminator(cfg)
         
-        self.gan_criterion = nn.MSELoss()
-        self.reg_criterion = nn.L1Loss(reduction='none')
-        self.pair_criterion = nn.PairwiseDistance()
-        self.cosine_loss = nn.CosineSimilarity(dim=1)
         
     def forward(self, x):
 
@@ -81,6 +77,10 @@ class Audio2Pose(nn.Module):
             z = torch.randn(bs, self.latent_dim).to(ref.device)
             batch['z'] = z
             audio_emb = self.audio_encoder(indiv_mels_use[:, -1*self.seq_len:,:,:,:]) #bs seq_len  512
+            if audio_emb.shape[1] != self.seq_len:
+                pad_dim = self.seq_len-audio_emb.shape[1]
+                pad_audio_emb = audio_emb[:, :1].repeat(1, pad_dim, 1) 
+                audio_emb = torch.cat([pad_audio_emb, audio_emb], 1) 
             batch['audio_emb'] = audio_emb
             batch = self.netG.test(batch)
             pose_motion_pred_list.append(batch['pose_motion_pred'][:,-1*re:,:])   

src/audio2pose_models/audio_encoder.py

@@ -19,7 +19,7 @@ class Conv2d(nn.Module):
         return self.act(out)
 
 class AudioEncoder(nn.Module):
-    def __init__(self, wav2lip_checkpoint):
+    def __init__(self, wav2lip_checkpoint, device):
         super(AudioEncoder, self).__init__()
 
         self.audio_encoder = nn.Sequential(
@@ -41,8 +41,8 @@ class AudioEncoder(nn.Module):
             Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
             Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
 
-        #### load the pre-trained audio_encoder\
-        wav2lip_state_dict = torch.load(wav2lip_checkpoint)['state_dict']
+        #### load the pre-trained audio_encoder
+        wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
         state_dict = self.audio_encoder.state_dict()
 
         for k,v in wav2lip_state_dict.items():