Update app.py

app.py

@@ -21,6 +21,36 @@ from lang_list import (
     TEXT_SOURCE_LANGUAGE_NAMES,
 )
 
+from scipy.io import wavfile
+from scipy.io.wavfile import write
+
+from speechbrain.pretrained import SpectralMaskEnhancement
+
+enhance_model = SpectralMaskEnhancement.from_hparams(
+    source="speechbrain/metricgan-plus-voicebank",
+    savedir="pretrained_models/metricgan-plus-voicebank",
+)
+
+knn_vc = torch.hub.load('bshall/knn-vc', 'knn_vc', prematched=True, trust_repo=True, pretrained=True, device='cpu')
+
+def voice_change(audio_in, audio_ref):
+    samplerate1, data1 = wavfile.read(audio_in)
+    samplerate2, data2 = wavfile.read(audio_ref)
+    write("./audio_in.wav", samplerate1, data1)
+    write("./audio_ref.wav", samplerate2, data2)
+
+    query_seq = knn_vc.get_features("./audio_in.wav")
+    matching_set = knn_vc.get_matching_set(["./audio_ref.wav"])
+    out_wav = knn_vc.match(query_seq, matching_set, topk=4)
+    torchaudio.save('output.wav', out_wav[None], 16000)
+    noisy = enhance_model.load_audio(
+        'output.wav'
+    ).unsqueeze(0)
+    enhanced = enhance_model.enhance_batch(noisy, lengths=torch.tensor([1.]))
+    torchaudio.save('enhanced.wav', enhanced.cpu(), 16000)
+    return 'enhanced.wav'
+
+
 CHECKPOINTS_PATH = pathlib.Path(os.getenv("CHECKPOINTS_PATH", "/home/user/app/models"))
 if not CHECKPOINTS_PATH.exists():
     snapshot_download(repo_id="facebook/seamless-m4t-v2-large", repo_type="model", local_dir=CHECKPOINTS_PATH)
@@ -152,27 +182,35 @@ with gr.Blocks() as demo_s2st:
     with gr.Row():
         with gr.Column():
             with gr.Group():
-                input_audio = gr.Audio(label="Input speech", type="filepath")
+                input_audio = gr.Audio(label="请上传一段语音", type="filepath")
                 source_language = gr.Dropdown(
-                    label="Source language",
+                    label="请选择上传语音对应的语言",
                     choices=ASR_TARGET_LANGUAGE_NAMES,
-                    value="English",
+                    value="Mandarin Chinese",
                 )
                 target_language = gr.Dropdown(
-                    label="Target language",
+                    label="请选择翻译后的语言",
                     choices=S2ST_TARGET_LANGUAGE_NAMES,
-                    value=DEFAULT_TARGET_LANGUAGE,
+                    value="English",
                 )
-            btn = gr.Button("Translate")
+            btn = gr.Button("开始AI同声传译之旅吧")
+            btn_vc = gr.Button("恢复原本的音色吧！")
+
         with gr.Column():
             with gr.Group():
                 output_audio = gr.Audio(
-                    label="Translated speech",
+                    label="同声传译后的语音",
+                    autoplay=False,
+                    streaming=False,
+                    type="numpy",
+                )
+                output_text = gr.Textbox(label="翻译后的文本")
+                audio_vc = gr.Audio(
+                    label="相同音色的AI专属语音",
                     autoplay=False,
                     streaming=False,
                     type="numpy",
                 )
-                output_text = gr.Textbox(label="Translated text")
 
     gr.Examples(
         examples=[
@@ -194,25 +232,28 @@ with gr.Blocks() as demo_s2st:
         outputs=[output_audio, output_text],
         api_name="s2st",
     )
+    btn_vc.click(voice_change, [output_audio, input_audio], [audio_vc])
+
 
 with gr.Blocks() as demo_s2tt:
+    
     with gr.Row():
         with gr.Column():
             with gr.Group():
-                input_audio = gr.Audio(label="Input speech", type="filepath")
+                input_audio = gr.Audio(label="请上传一段语音", type="filepath")
                 source_language = gr.Dropdown(
-                    label="Source language",
+                    label="请选择上传语音对应的语言",
                     choices=ASR_TARGET_LANGUAGE_NAMES,
-                    value="English",
+                    value="Mandarin Chinese",
                 )
                 target_language = gr.Dropdown(
-                    label="Target language",
+                    label="请选择翻译后的语言",
                     choices=S2TT_TARGET_LANGUAGE_NAMES,
-                    value=DEFAULT_TARGET_LANGUAGE,
+                    value="English",
                 )
-            btn = gr.Button("Translate")
+            btn = gr.Button("开始AI翻译之旅吧！")
         with gr.Column():
-            output_text = gr.Textbox(label="Translated text")
+            output_text = gr.Textbox(label="翻译后的文本")
 
     gr.Examples(
         examples=[