Update app.py

app.py

@@ -19,23 +19,48 @@ from models.tts.maskgct.g2p.g2p_generation import g2p, chn_eng_g2p
 
 from transformers import SeamlessM4TFeatureExtractor
 
-# import whisperx
+import whisper
+import langid
 
 processor = SeamlessM4TFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")
 
 device = torch.device("cuda" if torch.cuda.is_available() else "CPU")
 
-# whisper_model = whisperx.load_model("small", "cuda", compute_type="int8")
+whisper_model = whisper.load_model("turbo")
 
-# @torch.no_grad()
-# def get_prompt_text(speech_16k):
-#     asr_result = whisper_model.transcribe(speech_16k)
-#     print("asr_result:", asr_result)
-#     language = asr_result["language"]
-#     #text = asr_result["text"] # whisper asr result
-#     text = asr_result["segments"][0]["text"]
-#     print("prompt text:", text)
-#     return text, language
+def detect_speech_language(speech_file):
+    # load audio and pad/trim it to fit 30 seconds
+    audio = whisper.load_audio(speech_file)
+    audio = whisper.pad_or_trim(audio)
+
+    # make log-Mel spectrogram and move to the same device as the model
+    mel = whisper.log_mel_spectrogram(audio, n_mels=128).to(whisper_model.device)
+
+    # detect the spoken language
+    _, probs = whisper_model.detect_language(mel)
+    return max(probs, key=probs.get)
+
+
+def detect_text_language(text):
+    return langid.classify(text)[0]
+
+@torch.no_grad()
+def get_prompt_text(speech_16k, language):
+    full_prompt_text = ""
+    shot_prompt_text = ""
+    short_prompt_end_ts = 0.0
+
+    asr_result = whisper_model.transcribe(speech_16k, language=language)
+    full_prompt_text = asr_result["text"] # whisper asr result
+    #text = asr_result["segments"][0]["text"] # whisperx asr result
+    shot_prompt_text = ""
+    short_prompt_end_ts = 0.0
+    for segment in asr_result["segments"]:
+        shot_prompt_text = shot_prompt_text +  segment['text']
+        short_prompt_end_ts = segment['end']
+        if short_prompt_end_ts >= 4:
+            break
+    return full_prompt_text, shot_prompt_text, short_prompt_end_ts
 
 
 def g2p_(text, language):
@@ -279,10 +304,7 @@ def load_models():
 @torch.no_grad()
 def maskgct_inference(
     prompt_speech_path,
-    prompt_text,
     target_text,
-    language="en",
-    target_language="en",
     target_len=None,
     n_timesteps=25,
     cfg=2.5,
@@ -295,14 +317,18 @@ def maskgct_inference(
     speech_16k = librosa.load(prompt_speech_path, sr=16000)[0]
     speech = librosa.load(prompt_speech_path, sr=24000)[0]
 
-    # if prompt_text is None:
-    #     prompt_text, language = get_prompt_text(prompt_speech_path)
-    
+    prompt_language = detect_speech_language(prompt_speech_path)
+    full_prompt_text, short_prompt_text, shot_prompt_end_ts = get_prompt_text(prompt_speech_path,
+                                                                              prompt_language)
+    # use the first 4+ seconds wav as the prompt in case the prompt wav is too long
+    speech = speech[0: int(shot_prompt_end_ts * 24000)]
+    speech_16k = speech_16k[0: int(shot_prompt_end_ts*16000)]
+    target_language = detect_text_language(target_text)
     combine_semantic_code, _ = text2semantic(
         device,
         speech_16k,
-        prompt_text,
-        language,
+        short_prompt_text,
+        prompt_language,
         target_text,
         target_language,
         target_len,
@@ -326,21 +352,15 @@ def maskgct_inference(
 @spaces.GPU
 def inference(
     prompt_wav,
-    prompt_text,
     target_text,
     target_len,
     n_timesteps,
-    language,
-    target_language,
 ):
     save_path = "./output/output.wav"
     os.makedirs("./output", exist_ok=True)
     recovered_audio = maskgct_inference(
         prompt_wav,
-        prompt_text,
         target_text,
-        language,
-        target_language,
         target_len=target_len,
         n_timesteps=int(n_timesteps),
         device=device,
@@ -369,7 +389,6 @@ iface = gr.Interface(
     fn=inference,
     inputs=[
         gr.Audio(label="Upload Prompt Wav", type="filepath"),
-        gr.Textbox(label="Prompt Text"),
         gr.Textbox(label="Target Text"),
         gr.Number(
             label="Target Duration (in seconds), if the target duration is less than 0, the system will estimate a duration.", value=-1
@@ -377,8 +396,6 @@ iface = gr.Interface(
         gr.Slider(
             label="Number of Timesteps", minimum=15, maximum=100, value=25, step=1
         ),
-        gr.Dropdown(label="Language", choices=language_list, value="en"),
-        gr.Dropdown(label="Target Language", choices=language_list, value="en"),
     ],
     outputs=gr.Audio(label="Generated Audio"),
     title="MaskGCT TTS Demo",

requirements.txt

@@ -30,4 +30,5 @@ LangSegment
 onnxruntime
 pyopenjtalk
 pykakasi
-openai-whisper
+openai-whisper
+langid